diff --git a/CMakeLists.txt b/CMakeLists.txt index 6bd3886f5b..e6f51493f8 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -18,7 +18,7 @@ else() endif() project(diskann) -include_directories(${PROJECT_SOURCE_DIR}/include ${PROJECT_SOURCE_DIR}/include/dll) +include_directories(${PROJECT_SOURCE_DIR}/include ${PROJECT_SOURCE_DIR}/include/tsl/include ${PROJECT_SOURCE_DIR}/include/dll ${PROJECT_SOURCE_DIR}/include/ols) #OpenMP find_package(OpenMP) @@ -92,7 +92,7 @@ if(MSVC) set(CMAKE_ARCHIVE_OUTPUT_DIRECTORY_RELEASE ${PROJECT_SOURCE_DIR}/x64/Release) else() set(ENV{TCMALLOC_LARGE_ALLOC_REPORT_THRESHOLD} 500000000000) - # set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -g -DDEBUG -O0 -fsanitize=address -fsanitize=leak -fsanitize=undefined") + #set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -g -DDEBUG -O0 -fsanitize=address -fsanitize=leak -fsanitize=undefined") set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -g -DDEBUG -Wall -Wextra") set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} -Ofast -DNDEBUG -march=native -mtune=native -ftree-vectorize") add_compile_options(-march=native -Wall -fno-builtin-malloc -fno-builtin-calloc -fno-builtin-realloc -fno-builtin-free -fopenmp -fopenmp-simd -funroll-loops -Wfatal-errors -DUSE_ACCELERATED_PQ -DUSE_AVX2) diff --git a/CompilerOptions.cmake b/CompilerOptions.cmake new file mode 100644 index 0000000000..c3266cdb44 --- /dev/null +++ b/CompilerOptions.cmake @@ -0,0 +1,21 @@ +if(MSVC) + #changing default target to X64 + string(REGEX REPLACE "/[M|m][A|a][C|c][H|h][I|i][N|n][E|e]:[X|x]86" "/MACHINE:X64" CMAKE_EXE_LINKER_FLAGS_INIT "${CMAKE_EXE_LINKER_FLAGS_INIT}") + string(REGEX REPLACE "/[M|m][A|a][C|c][H|h][I|i][N|n][E|e]:[X|x]86" "/MACHINE:X64" CMAKE_MODULE_LINKER_FLAGS_INIT "${CMAKE_MODULE_LINKER_FLAGS_INIT}") + string(REGEX REPLACE "/[M|m][A|a][C|c][H|h][I|i][N|n][E|e]:[X|x]86" "/MACHINE:X64" CMAKE_SHARED_LINKER_FLAGS_INIT "${CMAKE_SHARED_LINKER_FLAGS_INIT}") + string(REGEX REPLACE "/[M|m][A|a][C|c][H|h][I|i][N|n][E|e]:[X|x]86" "/MACHINE:X64" CMAKE_STATIC_LINKER_FLAGS_INIT "${CMAKE_STATIC_LINKER_FLAGS_INIT}") + string(REGEX REPLACE "/[M|m][A|a][C|c][H|h][I|i][N|n][E|e]:[X|x]86" "/MACHINE:X64" CMAKE_EXE_LINKER_FLAGS_INIT "${CMAKE_EXE_LINKER_FLAGS_INIT}") + string(REGEX REPLACE "/[M|m][A|a][C|c][H|h][I|i][N|n][E|e]:[X|x]86" "/MACHINE:X64" CMAKE_MODULE_LINKER_FLAGS_INIT "${CMAKE_MODULE_LINKER_FLAGS_INIT}") + string(REGEX REPLACE "/[M|m][A|a][C|c][H|h][I|i][N|n][E|e]:[X|x]86" "/MACHINE:X64" CMAKE_SHARED_LINKER_FLAGS_INIT "${CMAKE_SHARED_LINKER_FLAGS_INIT}") + string(REGEX REPLACE "/[M|m][A|a][C|c][H|h][I|i][N|n][E|e]:[X|x]86" "/MACHINE:X64" CMAKE_STATIC_LINKER_FLAGS_INIT "${CMAKE_STATIC_LINKER_FLAGS_INIT}") + string(REGEX REPLACE "Debug" "Release" CMAKE_BUILD_TYPE_INIT "${CMAKE_BUILD_TYPE_INIT}") +endif() + + +get_cmake_property(_varNames VARIABLES) +list (REMOVE_DUPLICATES _varNames) +list (SORT _varNames) +foreach (_varName ${_varNames}) + message(STATUS "${_varName}=${${_varName}}") +endforeach() + diff --git a/LICENSE b/LICENSE index b7a909e5c0..81b23d8c7d 100644 --- a/LICENSE +++ b/LICENSE @@ -1,3 +1,4 @@ +<<<<<<< HEAD DiskANN MIT License diff --git a/README.md b/README.md index 7fd3007d97..7045ab168a 100644 --- a/README.md +++ b/README.md @@ -1,4 +1,164 @@ # DiskANN +##Linux build: + +Install the following packages through apt-get, and Intel MKL either by downloading the installer or using [apt](https://software.intel.com/en-us/articles/installing-intel-free-libs-and-python-apt-repo) (we tested with build 2019.4-070). +``` +sudo apt install cmake g++ libaio-dev libgoogle-perftools-dev clang-format-4.0 +``` + +Build +``` +mkdir build && cd build && cmake .. && make -j +``` + +##Windows build: + +The Windows version has been tested with the enterprise editions of Visual Studio 2017 and Visual Studio 2019 + +**Prerequisites:** + +* Install CMAKE (v3.15.2 or later) from https://cmake.org +* Install MKL from https://software.intel.com/en-us/mkl +* Download boost (v1.71.0, later versions are not tested) from boost.org + +* Environment variables: + * Set a new System environment variable, called INTEL_ROOT to the "windows" folder under your MKL installation + (For instance, if your install folder is "C:\Program Files (x86)\IntelSWtools", set INTEL_ROOT to "C:\Program Files (x86)\IntelSWtools\compilers_and_libraries\windows") + * Set BOOST_ROOT to your boost download folder + +**Build steps:** +- Open a new developer command prompt +- Create a "build" directory under diskann +- Change to the "build" directory and run +``` +cmake -B. -A x64 .. +``` +**Note: Since VS comes with its own (older) version of cmake, you have to specify the full path to cmake to ensure that the right version is used.** +- This will create a “diskann” solution file. +- Open the "diskann" solution and build the "diskpriority_io" and “nsg_dll” projects in order. +- Then build all the other binaries using the ALL_BUILD project that is part of the solution +- Generated binaries are stored in the diskann/x64/Debug or diskann/x64/Release directories. + +To build from command line, change to the "build" directory and use msbuild to first build the "diskpriority_io" and "nsg_dll" projects. And then build the entire solution, as shown below. +``` +msbuild src\dll\diskpriority_io.vcxproj +msbuild src\dll\nsg_dll.vcxproj +msbuild diskann.sln +``` +Check msbuild docs for additional options including choosing between debug and release builds. + + +##Usage: + +We now detail the main binaries using which one can build and search indices which reside in memory as well as SSD-resident indices. + +**Usage for SSD-based indices** +=============================== + +To generate an SSD-friendly index, use the `tests/build_disk_index` program. +---------------------------------------------------------------------------- + +``` +./tests/build_disk_index [data_type] [data_file.bin] [index_prefix_path] [R] [L] [B] [M] [T]. +``` + +The arguments are as follows: + +(i) data_type: The datatype is the type of dataset you wish to build an index. We support byte indices (signed int8 or unsigned uint8) or float indices. + +(ii) data_file: The input data over which to build an index, in .bin format. The first 4 bytes represent number of points as integer. The next 4 bytes represent the dimension of data as integer. The following n*d*sizeof(T) bytes contain the contents of the data one data point in time. sizeof(T) is 1 for byte indices, and 4 for float indices. This will be read by the program as int8_t for signed indices, uint8_t for unsigned indices or float for float indices. + +(iii) index_prefix_path: the index will generate a few files, all beginning with the specified prefix path. For example, if you provide ~/index_test as the prefix path, build generates files such as ~/index_test_pq_pivots.bin, ~/index_test_pq_compressed.bin, ~/index_test_disk.index, etc. There may be between 8 and 10 files generated with this prefix depending on how we construct the index. + +(iv) R: the degree of our graph index, typically between 60 and 150. Again, larger values will result in bigger indices (with longer indexing times), but better search quality. Try to ensure that the L value is at least the R value unless you need to build indices really quickly, but can somewhat compromise on quality. + +(v) L: the size of search list we maintain during index building. Typical values are between 75 to 200. Larger values will take more time to build but result in indices that provide higher recall for the same search parameters. + +(vi) B: bound on the memory footprint of the index at search time. Once built, the index will use up only the specified RAM limit, the rest will reside on disk. This will dictate how aggressively we compress the data vectors to store in memory. Larger will yield better performance at search time. + +(vii) M: Limit on the memory allowed for building the index. If you specify a value less than what is required to build the index in one pass, the index is built using a divide and conquer approach so that sub-graphs will fit in the RAM budget. The sub-graphs are stitched together to build the overall index. This approach can be upto 1.5 times slower than building the index in one shot. Try to allocate as much memory as possible for index build as your RAM allows. + +(viii) T: number of threads used by the index build process. Since the code is highly parallel, the indexing time improves almost linearly with the number of threads (subject to the cores available on the machine). + +To search the SSD-index, use the `tests/search_disk_index` program. +---------------------------------------------------------------------------- + +``` +./tests/search_disk_index [index_type] [index_prefix_path] [num_nodes_to_cache] [num_threads] [beamwidth (use 0 to optimize internally)] [query_file.bin] [truthset.bin (use "null" for none)] [K] [result_output_prefix] [L1] [L2] etc. +``` + +The arguments are as follows: + +(i) data type: same as (i) above in building index. + +(ii) index_prefix_path: same as (iii) above in building index. + +(iii) num_nodes_to_cache: our program stores the entire graph on disk. For faster search performance, we provide the support to cache a few nodes (which are closest to the starting point) in memory. + +(iv) num_threads: search using specified number of threads in parallel, one thread per query. More will result in more IOs, so find the balance depending on the bandwidth of the SSD. + +(v) beamwidth: maximum number of IO requests each query will issue per iteration of search code. Larger beamwidth williult in fewer IO round-trips per query, but might result in slightly higher number of IO requests to SSD per query. Specifying 0 will optimize the beamwidth depending on the number of threads performing search. + +(vi) query_file.bin: search on these queries, same format as data file (ii) above. The query file must be the same type as specified in (i). + +(vii) truthset.bin file. Must be in the following format, or specify "null": n, the number of queries (4 bytes) followed by d, the number of ground truth elements per query (4 bytes), followed by n*d entries per query representing the d closest IDs per query in integer format, followed by n*d entries representing the corresponding distances (float). Total file size is 8 + 4*n*d + 4*n*d. The groundtruth file, if not available, can be calculated using our program, tests/utils/compute_groundtruth. If you just want to measure the latency numbers of search and output the nearest neighbors without calculating recall, enter "null". + +(viii) K: measure recall@k, meaning the accuracy of retrieving top-k nearest neighbors. + +(ix) result output prefix: search results will be stored in files with specified prefix, in bin format. + +(x, xi, ...) various search_list sizes to perform search with. Larger will result in slower latencies, but higher accuracies. Must be atleast the recall@ value in (vi). + + +**Usage for in-memory indices** +================================ + +To generate index, use the `tests/build_memory_index` program. +-------------------------------------------------------------- + +``` +./tests/build_memory_index [data_type] [data_file.bin] [output_index_file] [R] [L] [alpha] [num_threads_to_use] +``` + +The arguments are as follows: + +(i) data_type: same as (i) above in building disk index. + +(ii) data_file: same as (ii) above in building disk index, the input data file in .bin format of type int8/uint8/float. + +(iii) output_index_file: memory index will be saved here. + +(iv) R: max degree of index: larger is typically better, range (50-150). Preferrably ensure that L is at least R. + +(v) L: candidate_list_size for building index, larger is better (typical range: 75 to 200) + +(vi) alpha: float value which determines how dense our overall graph will be, and diameter will be log of n base alpha (roughly). Typical values are between 1 to 1.5. 1 will yield sparsest graph, 1.5 will yield denser graphs. + +(vii) number of threads to use: indexing uses specified number of threads. + + +To search the generated index, use the `tests/search_memory_index` program: +--------------------------------------------------------------------------- + +``` +./tests/search_memory_index [index_type] [data_file.bin] [memory_index_path] [query_file.bin] [truthset.bin (use "null" for none)] [K] [result_output_prefix] [L1] [L2] etc. +``` + +The arguments are as follows: + +(i) data type: same as (i) above in building index. + +(ii) memory_index_path: enter path of index built (argument (iii) above in building memory index). + +(iii) query_bin: search on these queries, same format as data file (ii) above. The query file must be the same type as specified in (i). + +(iv) Truthset file. Must be in the following format: n, the number of queries (4 bytes) followed by d, the number of ground truth elements per query (4 bytes), followed by n*d entries per query representing the d closest IDs per query in integer format, followed by n*d entries representing the corresponding distances (float). Total file size is 8 + 4*n*d + 4*n*d. The groundtruth file, if not available, can be calculated using our program, tests/utils/compute_groundtruth. + +(v) K: search for recall@k, meaning accuracy of retrieving top-k nearest neighbors. + +(vi) result output prefix: will search and store the computed results in the files with specified prefix in bin format. + +(vii, viii, ...) various search_list sizes to perform search with. Larger will result in slower latencies, but higher accuracies. Must be atleast the recall@ value in (vi). The goal of the project is to build scalable, performant and cost-effective approximate nearest neighbor search algorithms. The initial release has the in-memory version of the [DiskANN paper](https://papers.nips.cc/paper/9527-rand-nsg-fast-accurate-billion-point-nearest-neighbor-search-on-a-single-node.pdf) published in NeurIPS 2019. diff --git a/dependencies/windows/dll/LIBEAY32.dll b/dependencies/windows/dll/LIBEAY32.dll new file mode 100644 index 0000000000..7311c2b3ed Binary files /dev/null and b/dependencies/windows/dll/LIBEAY32.dll differ diff --git a/dependencies/windows/dll/SSLEAY32.dll b/dependencies/windows/dll/SSLEAY32.dll new file mode 100644 index 0000000000..0aa61b0a9f Binary files /dev/null and b/dependencies/windows/dll/SSLEAY32.dll differ diff --git a/dependencies/windows/dll/boost_date_time-vc141-mt-x64-1_70.dll b/dependencies/windows/dll/boost_date_time-vc141-mt-x64-1_70.dll new file mode 100644 index 0000000000..3473c5b70f Binary files /dev/null and b/dependencies/windows/dll/boost_date_time-vc141-mt-x64-1_70.dll differ diff --git a/dependencies/windows/dll/cpprest_2_10.dll b/dependencies/windows/dll/cpprest_2_10.dll new file mode 100644 index 0000000000..a2982e523d Binary files /dev/null and b/dependencies/windows/dll/cpprest_2_10.dll differ diff --git a/dependencies/windows/dll/zlib1.dll b/dependencies/windows/dll/zlib1.dll new file mode 100644 index 0000000000..cc57962a49 Binary files /dev/null and b/dependencies/windows/dll/zlib1.dll differ diff --git a/include/Neighbor_Tag.h b/include/Neighbor_Tag.h new file mode 100644 index 0000000000..dd28576515 --- /dev/null +++ b/include/Neighbor_Tag.h @@ -0,0 +1,33 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include "tsl/robin_set.h" + +#include "parameters.h" + +namespace diskann { + + template + struct Neighbor_Tag { + TagT tag; + float dist; + + Neighbor_Tag() = default; + + Neighbor_Tag(TagT tag, float dist) : tag{tag}, dist{dist} { + } + inline bool operator<(const Neighbor_Tag &other) const { + return (dist < other.dist); + } + inline bool operator==(const Neighbor_Tag &other) const { + return (tag == other.tag); + } + }; +} // namespace diskann diff --git a/include/aligned_file_reader.h b/include/aligned_file_reader.h index c33bacd666..6f18aa2eee 100644 --- a/include/aligned_file_reader.h +++ b/include/aligned_file_reader.h @@ -16,9 +16,10 @@ typedef io_context_t IOContext; #else #include #include +#include #ifndef USE_BING_INFRA -struct IOContext{ +struct IOContext { HANDLE fhandle = NULL; HANDLE iocp = NULL; std::vector reqs; @@ -77,7 +78,7 @@ struct AlignedRead { class AlignedFileReader { protected: tsl::robin_map ctx_map; - std::mutex ctx_mut; + std::mutex ctx_mut; public: // returns the thread-specific context @@ -91,9 +92,11 @@ class AlignedFileReader { // de-register thread-id for a context virtual void deregister_thread() = 0; + virtual void deregister_all_threads() = 0; // Open & close ops // Blocking calls - virtual void open(const std::string& fname) = 0; + virtual void open(const std::string& fname, bool enable_writes, + bool enable_create) = 0; virtual void close() = 0; // process batch of aligned requests in parallel diff --git a/include/ann_exception.h b/include/ann_exception.h index 7939c135a8..bd350dbcac 100644 --- a/include/ann_exception.h +++ b/include/ann_exception.h @@ -19,6 +19,7 @@ namespace diskann { unsigned int lineNum); DISKANN_DLLEXPORT std::string message() const; + DISKANN_DLLEXPORT int errorCode() const; private: int _errorCode; diff --git a/include/aux_utils.h b/include/aux_utils.h index 031a64ba8b..3b94609617 100644 --- a/include/aux_utils.h +++ b/include/aux_utils.h @@ -1,4 +1,4 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. +// Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT license. #pragma once @@ -27,23 +27,41 @@ typedef int FileHandle; #include "cached_io.h" #include "common_includes.h" +#include "tsl/robin_set.h" #include "utils.h" #include "windows_customizations.h" namespace diskann { - const size_t TRAINING_SET_SIZE = 1500000; + + const size_t MAX_PQ_TRAINING_SET_SIZE = 256000; + const size_t MAX_SAMPLE_POINTS_FOR_WARMUP = 1000000; + const double PQ_TRAINING_SET_FRACTION = 0.1; const double SPACE_FOR_CACHED_NODES_IN_GB = 0.25; const double THRESHOLD_FOR_CACHING_IN_GB = 1.0; const uint32_t NUM_NODES_TO_CACHE = 250000; const uint32_t WARMUP_L = 20; - template + template class PQFlashIndex; + DISKANN_DLLEXPORT double get_memory_budget(const std::string &mem_budget_str); + DISKANN_DLLEXPORT double get_memory_budget(double search_ram_budget_in_gb); + DISKANN_DLLEXPORT void add_new_file_to_single_index(std::string index_file, + std::string new_file); + + DISKANN_DLLEXPORT size_t calculate_num_pq_chunks(double final_index_ram_limit, + size_t points_num, + uint32_t dim); + DISKANN_DLLEXPORT double calculate_recall( unsigned num_queries, unsigned *gold_std, float *gs_dist, unsigned dim_gs, unsigned *our_results, unsigned dim_or, unsigned recall_at); + DISKANN_DLLEXPORT double calculate_recall( + unsigned num_queries, unsigned *gold_std, float *gs_dist, unsigned dim_gs, + unsigned *our_results, unsigned dim_or, unsigned recall_at, + const tsl::robin_set &active_tags); + DISKANN_DLLEXPORT void read_idmap(const std::string & fname, std::vector &ivecs); @@ -70,26 +88,29 @@ namespace diskann { template DISKANN_DLLEXPORT int build_merged_vamana_index( - std::string base_file, diskann::Metric _compareMetric, unsigned L, - unsigned R, double sampling_rate, double ram_budget, - std::string mem_index_path, std::string medoids_file, - std::string centroids_file); + std::string base_file, diskann::Metric _compareMetric, + bool single_index_file, unsigned L, unsigned R, double sampling_rate, + double ram_budget, std::string mem_index_path, std::string medoids_file, + std::string centroids_file, const char *tag_file = nullptr); - template + template DISKANN_DLLEXPORT uint32_t optimize_beamwidth( - std::unique_ptr> &_pFlashIndex, T *tuning_sample, - _u64 tuning_sample_num, _u64 tuning_sample_aligned_dim, uint32_t L, - uint32_t nthreads, uint32_t start_bw = 2); + std::unique_ptr> &_pFlashIndex, + T *tuning_sample, _u64 tuning_sample_num, _u64 tuning_sample_aligned_dim, + uint32_t L, uint32_t nthreads, uint32_t start_bw = 2); - template + template DISKANN_DLLEXPORT bool build_disk_index(const char * dataFilePath, const char * indexFilePath, const char * indexBuildParameters, - diskann::Metric _compareMetric); - - template - DISKANN_DLLEXPORT void create_disk_layout(const std::string base_file, - const std::string mem_index_file, - const std::string output_file); + diskann::Metric _compareMetric, + bool single_file_index, + const char * tag_file = nullptr); + template + DISKANN_DLLEXPORT void create_disk_layout( + const std::string &mem_index_file, const std::string &base_file, + const std::string &tag_file, const std::string &pq_pivots_file, + const std::string &pq_compressed_vectors_file, bool single_file_index, + const std::string &output_file); } // namespace diskann diff --git a/include/cached_io.h b/include/cached_io.h index f9d23c4f99..9b8a4cce58 100644 --- a/include/cached_io.h +++ b/include/cached_io.h @@ -9,34 +9,44 @@ #include "logger.h" #include "ann_exception.h" +#include "utils.h" // sequential cached reads class cached_ifstream { public: cached_ifstream() { } - cached_ifstream(const std::string& filename, uint64_t cacheSize) + cached_ifstream(const std::string& filename, uint64_t cacheSize, + uint32_t initial_offset = 0) : cache_size(cacheSize), cur_off(0) { - this->open(filename, cache_size); + this->open(filename, cache_size, initial_offset); } ~cached_ifstream() { - delete[] cache_buf; - reader.close(); + // delete[] cache_buf; + if (reader.is_open()) + reader.close(); } - void open(const std::string& filename, uint64_t cacheSize) { + void close() { + if (reader.is_open()) + reader.close(); + } + + void open(const std::string& filename, uint64_t cacheSize, + size_t initial_offset = 0) { this->cur_off = 0; reader.open(filename, std::ios::binary | std::ios::ate); fsize = reader.tellg(); - reader.seekg(0, std::ios::beg); + reader.seekg(initial_offset, reader.beg); assert(reader.is_open()); assert(cacheSize > 0); - cacheSize = (std::min)(cacheSize, fsize); + cacheSize = (std::min)(cacheSize, fsize - initial_offset); + diskann::cout << "Opened: " << filename.c_str() << ", size: " << fsize + << ", cache_size: " << cacheSize << std::endl; this->cache_size = cacheSize; cache_buf = new char[cacheSize]; + // cache_buf = std::make_unique(cacheSize); reader.read(cache_buf, cacheSize); - diskann::cout << "Opened: " << filename.c_str() << ", size: " << fsize - << ", cache_size: " << cacheSize << std::endl; } size_t get_file_size() { @@ -97,24 +107,42 @@ class cached_ifstream { // sequential cached writes class cached_ofstream { public: - cached_ofstream(const std::string& filename, uint64_t cache_size) + cached_ofstream() { + } + cached_ofstream(const std::string& filename, uint64_t cache_size, + size_t initial_offset = 0) : cache_size(cache_size), cur_off(0) { - writer.open(filename, std::ios::binary); + open(filename, cache_size, initial_offset); + } + void open(const std::string& filename, uint64_t cache_size, + size_t initial_offset = 0) { + open_file_to_write(writer, filename); assert(writer.is_open()); assert(cache_size > 0); + writer.seekp(initial_offset, writer.beg); cache_buf = new char[cache_size]; + fsize = initial_offset; diskann::cout << "Opened: " << filename.c_str() << ", cache_size: " << cache_size << std::endl; } ~cached_ofstream() { + this->close(); + } + + void close() { // dump any remaining data in memory if (cur_off > 0) { this->flush_cache(); } - delete[] cache_buf; - writer.close(); + if (cache_buf != nullptr) { + delete[] cache_buf; + cache_buf = nullptr; + } + + if (writer.is_open()) + writer.close(); diskann::cout << "Finished writing " << fsize << "B" << std::endl; } diff --git a/include/content_buf.h b/include/content_buf.h new file mode 100644 index 0000000000..eafebbc3ba --- /dev/null +++ b/include/content_buf.h @@ -0,0 +1,26 @@ +#pragma once + +#ifdef EXEC_ENV_OLS + +namespace diskann { + class ContentBuf : public std::basic_streambuf { + public: + ContentBuf(char* p, size_t n) { + setg(p, p, p + n); + } + + virtual pos_type seekoff( + off_type off, std::ios_base::seekdir dir, + std::ios_base::openmode which = std::ios_base::in) { + if (dir == std::ios_base::cur) + gbump((int) off); + else if (dir == std::ios_base::end) + setg(eback(), egptr() + off, egptr()); + else if (dir == std::ios_base::beg) + setg(eback(), eback() + off, egptr()); + return gptr() - eback(); + } + }; +} // namespace diskann + +#endif \ No newline at end of file diff --git a/include/cosine_similarity.h b/include/cosine_similarity.h index 375db26591..939f291f66 100644 --- a/include/cosine_similarity.h +++ b/include/cosine_similarity.h @@ -2,17 +2,28 @@ // Licensed under the MIT license. #pragma once + +#include +#include +#include #include #include #include #include +#include +#include +#include + +#include "simd_utils.h" + +extern bool Avx2SupportedCPU; namespace diskann { template inline float compute_l2_norm(const T* vector, uint64_t ndims) { float norm = 0.0f; for (uint64_t i = 0; i < ndims; i++) { - norm += vector[i] * vector[i]; + norm += (float) (vector[i] * vector[i]); } return std::sqrt(norm); } @@ -24,7 +35,7 @@ namespace diskann { float right_norm = compute_l2_norm(right, ndims); float dot = 0.0f; for (uint64_t i = 0; i < ndims; i++) { - dot += left[i] * right[i]; + dot += (float) (left[i] * right[i]); } float cos_sim = dot / (left_norm * right_norm); return cos_sim; diff --git a/include/ctpl_stl.h b/include/ctpl_stl.h new file mode 100644 index 0000000000..7edd6c74fe --- /dev/null +++ b/include/ctpl_stl.h @@ -0,0 +1,272 @@ + +/********************************************************* + * + * Copyright (C) 2014 by Vitaliy Vitsentiy + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + *********************************************************/ + +#ifndef __ctpl_stl_thread_pool_H__ +#define __ctpl_stl_thread_pool_H__ + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// thread pool to run user's functors with signature +// ret func(int id, other_params) +// where id is the index of the thread that runs the functor +// ret is some return type + +namespace ctpl { + + namespace detail { + template + class Queue { + public: + bool push(T const &value) { + std::unique_lock lock(this->mutex); + this->q.push(value); + return true; + } + // deletes the retrieved element, do not use for non integral types + bool pop(T &v) { + std::unique_lock lock(this->mutex); + if (this->q.empty()) + return false; + v = this->q.front(); + this->q.pop(); + return true; + } + bool empty() { + std::unique_lock lock(this->mutex); + return this->q.empty(); + } + + private: + std::queue q; + std::mutex mutex; + }; + } // namespace detail + + class thread_pool { + public: + thread_pool() { + this->init(); + } + thread_pool(int nThreads) { + this->init(); + this->resize(nThreads); + } + + // the destructor waits for all the functions in the queue to be finished + ~thread_pool() { + this->stop(true); + } + + // get the number of running threads in the pool + int size() { + return static_cast(this->threads.size()); + } + + // number of idle threads + int n_idle() { + return this->nWaiting; + } + std::thread &get_thread(int i) { + return *this->threads[i]; + } + + // change the number of threads in the pool + // should be called from one thread, otherwise be careful to not interleave, + // also with this->stop() nThreads must be >= 0 + void resize(int nThreads) { + if (!this->isStop && !this->isDone) { + int oldNThreads = static_cast(this->threads.size()); + if (oldNThreads <= nThreads) { // if the number of threads is increased + this->threads.resize(nThreads); + this->flags.resize(nThreads); + + for (int i = oldNThreads; i < nThreads; ++i) { + this->flags[i] = std::make_shared>(false); + this->set_thread(i); + } + } else { // the number of threads is decreased + for (int i = oldNThreads - 1; i >= nThreads; --i) { + *this->flags[i] = true; // this thread will finish + this->threads[i]->detach(); + } + { + // stop the detached threads that were waiting + std::unique_lock lock(this->mutex); + this->cv.notify_all(); + } + this->threads.resize( + nThreads); // safe to delete because the threads are detached + this->flags.resize( + nThreads); // safe to delete because the threads have copies of + // shared_ptr of the flags, not originals + } + } + } + + // empty the queue + void clear_queue() { + std::function *_f; + while (this->q.pop(_f)) + delete _f; // empty the queue + } + + // pops a functional wrapper to the original function + std::function pop() { + std::function *_f = nullptr; + this->q.pop(_f); + std::unique_ptr> func( + _f); // at return, delete the function even if an exception occurred + std::function f; + if (_f) + f = *_f; + return f; + } + + // wait for all computing threads to finish and stop all threads + // may be called asynchronously to not pause the calling thread while + // waiting if isWait == true, all the functions in the queue are run, + // otherwise the queue is cleared without running the functions + void stop(bool isWait = false) { + if (!isWait) { + if (this->isStop) + return; + this->isStop = true; + for (int i = 0, n = this->size(); i < n; ++i) { + *this->flags[i] = true; // command the threads to stop + } + this->clear_queue(); // empty the queue + } else { + if (this->isDone || this->isStop) + return; + this->isDone = true; // give the waiting threads a command to finish + } + { + std::unique_lock lock(this->mutex); + this->cv.notify_all(); // stop all waiting threads + } + for (int i = 0; i < static_cast(this->threads.size()); + ++i) { // wait for the computing threads to finish + if (this->threads[i]->joinable()) + this->threads[i]->join(); + } + // if there were no threads in the pool but some functors in the queue, + // the functors are not deleted by the threads therefore delete them here + this->clear_queue(); + this->threads.clear(); + this->flags.clear(); + } + + template + auto push(F &&f, Rest &&... rest) -> std::future { + auto pck = + std::make_shared>( + std::bind(std::forward(f), std::placeholders::_1, + std::forward(rest)...)); + auto _f = new std::function([pck](int id) { (*pck)(id); }); + this->q.push(_f); + std::unique_lock lock(this->mutex); + this->cv.notify_one(); + return pck->get_future(); + } + + // run the user's function that excepts argument int - id of the running + // thread. returned value is templatized operator returns std::future, where + // the user can get the result and rethrow the catched exceptins + template + auto push(F &&f) -> std::future { + auto pck = std::make_shared>( + std::forward(f)); + auto _f = new std::function([pck](int id) { (*pck)(id); }); + this->q.push(_f); + std::unique_lock lock(this->mutex); + this->cv.notify_one(); + return pck->get_future(); + } + + private: + // deleted + thread_pool(const thread_pool &); // = delete; + thread_pool(thread_pool &&); // = delete; + thread_pool &operator=(const thread_pool &); // = delete; + thread_pool &operator=(thread_pool &&); // = delete; + + void set_thread(int i) { + std::shared_ptr> flag( + this->flags[i]); // a copy of the shared ptr to the flag + auto f = [this, i, flag /* a copy of the shared ptr to the flag */]() { + std::atomic & _flag = *flag; + std::function *_f; + bool isPop = this->q.pop(_f); + while (true) { + while (isPop) { // if there is anything in the queue + std::unique_ptr> func( + _f); // at return, delete the function even if an exception + // occurred + (*_f)(i); + if (_flag) + return; // the thread is wanted to stop, return even if the queue + // is not empty yet + else + isPop = this->q.pop(_f); + } + // the queue is empty here, wait for the next command + std::unique_lock lock(this->mutex); + ++this->nWaiting; + this->cv.wait(lock, [this, &_f, &isPop, &_flag]() { + isPop = this->q.pop(_f); + return isPop || this->isDone || _flag; + }); + --this->nWaiting; + if (!isPop) + return; // if the queue is empty and this->isDone == true or *flag + // then return + } + }; + this->threads[i].reset( + new std::thread(f)); // compiler may not support std::make_unique() + } + + void init() { + this->nWaiting = 0; + this->isStop = false; + this->isDone = false; + } + + std::vector> threads; + std::vector>> flags; + detail::Queue *> q; + std::atomic isDone; + std::atomic isStop; + std::atomic nWaiting; // how many threads are waiting + + std::mutex mutex; + std::condition_variable cv; + }; + +} // namespace ctpl + +#endif // __ctpl_stl_thread_pool_H__ \ No newline at end of file diff --git a/include/distance.h b/include/distance.h index 3d403d1e56..65ee4ce1f8 100644 --- a/include/distance.h +++ b/include/distance.h @@ -1,6 +1,9 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + #pragma once -#include +#include "windows_customizations.h" #ifdef _WINDOWS #include #include @@ -106,132 +109,56 @@ namespace diskann { } }; - template - class DistanceCosine : public Distance { - float compare(const T *a, const T *b, unsigned length) const { - return diskann::compute_cosine_similarity(a, b, length); - } + class DistanceCosineInt8 : public Distance { + public: + virtual float compare(const int8_t *a, const int8_t *b, + uint32_t length) const; }; - class DistanceL2Int8 : public Distance { + class DistanceCosineFloat : public Distance { public: - float compare(const int8_t *a, const int8_t *b, unsigned size) const { - int32_t result = 0; + virtual float compare(const float *a, const float *b, + uint32_t length) const; + }; -#ifdef _WINDOWS -#ifdef USE_AVX2 - __m256 r = _mm256_setzero_ps(); - char * pX = (char *) a, *pY = (char *) b; - while (size >= 32) { - __m256i r1 = _mm256_subs_epi8(_mm256_loadu_si256((__m256i *) pX), - _mm256_loadu_si256((__m256i *) pY)); - r = _mm256_add_ps(r, _mm256_mul_epi8(r1, r1)); - pX += 32; - pY += 32; - size -= 32; - } - while (size > 0) { - __m128i r2 = _mm_subs_epi8(_mm_loadu_si128((__m128i *) pX), - _mm_loadu_si128((__m128i *) pY)); - r = _mm256_add_ps(r, _mm256_mul32_pi8(r2, r2)); - pX += 4; - pY += 4; - size -= 4; - } - r = _mm256_hadd_ps(_mm256_hadd_ps(r, r), r); - return r.m256_f32[0] + r.m256_f32[4]; -#else -#pragma omp simd reduction(+ : result) aligned(a, b : 8) - for (_s32 i = 0; i < (_s32) size; i++) { - result += ((int32_t)((int16_t) a[i] - (int16_t) b[i])) * - ((int32_t)((int16_t) a[i] - (int16_t) b[i])); - } - return (float) result; -#endif -#else -#pragma omp simd reduction(+ : result) aligned(a, b : 8) - for (_s32 i = 0; i < (_s32) size; i++) { - result += ((int32_t)((int16_t) a[i] - (int16_t) b[i])) * - ((int32_t)((int16_t) a[i] - (int16_t) b[i])); - } - return (float) result; -#endif - } + class SlowDistanceCosineUInt8 : public Distance { + public: + virtual float compare(const uint8_t *a, const uint8_t *b, + uint32_t length) const; + }; + + class DistanceL2Int8 : public Distance { + public: + virtual float compare(const int8_t *a, const int8_t *b, + uint32_t size) const; }; class DistanceL2UInt8 : public Distance { public: - float compare(const uint8_t *a, const uint8_t *b, unsigned size) const { - uint32_t result = 0; -#ifndef _WINDOWS -#pragma omp simd reduction(+ : result) aligned(a, b : 8) -#endif - for (_s32 i = 0; i < (_s32) size; i++) { - result += ((int32_t)((int16_t) a[i] - (int16_t) b[i])) * - ((int32_t)((int16_t) a[i] - (int16_t) b[i])); - } - return (float) result; - } + virtual float compare(const uint8_t *a, const uint8_t *b, + uint32_t size) const; }; class DistanceL2 : public Distance { public: -#ifndef _WINDOWS - float compare(const float *a, const float *b, unsigned size) const - __attribute__((hot)) { - a = (const float *) __builtin_assume_aligned(a, 32); - b = (const float *) __builtin_assume_aligned(b, 32); -#else - float compare(const float *a, const float *b, unsigned size) const { -#endif - - float result = 0; -#ifdef USE_AVX2 - // assume size is divisible by 8 - _u16 niters = size / 8; - __m256 sum = _mm256_setzero_ps(); - for (_u16 j = 0; j < niters; j++) { - // scope is a[8j:8j+7], b[8j:8j+7] - // load a_vec - if (j < (niters - 1)) { - _mm_prefetch((char *) (a + 8 * (j + 1)), _MM_HINT_T0); - _mm_prefetch((char *) (b + 8 * (j + 1)), _MM_HINT_T0); - } - __m256 a_vec = _mm256_load_ps(a + 8 * j); - // load b_vec - __m256 b_vec = _mm256_load_ps(b + 8 * j); - // a_vec - b_vec - __m256 tmp_vec = _mm256_sub_ps(a_vec, b_vec); - /* - // (a_vec - b_vec)**2 - __m256 tmp_vec2 = _mm256_mul_ps(tmp_vec, tmp_vec); - // accumulate sum - sum = _mm256_add_ps(sum, tmp_vec2); - */ - // sum = (tmp_vec**2) + sum - sum = _mm256_fmadd_ps(tmp_vec, tmp_vec, sum); - } - - // horizontal add sum - result = _mm256_reduce_add_ps(sum); +#ifdef _WINDOWS + virtual float compare(const float *a, const float *b, uint32_t size) const; #else -#ifndef _WINDOWS -#pragma omp simd reduction(+ : result) aligned(a, b : 32) -#endif - for (_s32 i = 0; i < (_s32) size; i++) { - result += (a[i] - b[i]) * (a[i] - b[i]); - } + virtual float compare(const float *a, const float *b, uint32_t size) const + __attribute__((hot)); #endif - return result; - } }; - // Slow implementations of the distance functions for machines without AVX2 + // Slow implementations of the distance functions to get diskann to + // work in pre-AVX machines. Performance here is not a concern, so we are + // using the simplest possible implementation. template class SlowDistanceL2Int : public Distance { - virtual float compare(const T *a, const T *b, unsigned length) const { + public: + // Implementing here because this is a template function + virtual float compare(const T *a, const T *b, uint32_t length) const { uint32_t result = 0; - for (_u32 i = 0; i < length; i++) { + for (uint32_t i = 0; i < length; i++) { result += ((int32_t)((int16_t) a[i] - (int16_t) b[i])) * ((int32_t)((int16_t) a[i] - (int16_t) b[i])); } @@ -240,291 +167,22 @@ namespace diskann { }; class SlowDistanceL2Float : public Distance { + public: virtual float compare(const float *a, const float *b, - unsigned length) const { - float result = 0.0f; - for (_u32 i = 0; i < length; i++) { - result += (a[i] - b[i]) * (a[i] - b[i]); - } - return result; - } + uint32_t length) const; }; + // AVX implementations. Borrowed from HNSW code. class AVXDistanceL2Int8 : public Distance { public: virtual float compare(const int8_t *a, const int8_t *b, - unsigned int length) const { -#ifndef _WINDOWS - std::cout << "AVX only supported in Windows build."; - return 0; - } -#else - __m128 r = _mm_setzero_ps(); - __m128i r1; - while (length >= 16) { - r1 = _mm_subs_epi8(_mm_load_si128((__m128i *) a), - _mm_load_si128((__m128i *) b)); - r = _mm_add_ps(r, _mm_mul_epi8(r1)); - a += 16; - b += 16; - length -= 16; - } - r = _mm_hadd_ps(_mm_hadd_ps(r, r), r); - float res = r.m128_f32[0]; - - if (length >= 8) { - __m128 r2 = _mm_setzero_ps(); - __m128i r3 = _mm_subs_epi8(_mm_load_si128((__m128i *) (a - 8)), - _mm_load_si128((__m128i *) (b - 8))); - r2 = _mm_add_ps(r2, _mm_mulhi_epi8(r3)); - a += 8; - b += 8; - length -= 8; - r2 = _mm_hadd_ps(_mm_hadd_ps(r2, r2), r2); - res += r2.m128_f32[0]; - } - - if (length >= 4) { - __m128 r2 = _mm_setzero_ps(); - __m128i r3 = _mm_subs_epi8(_mm_load_si128((__m128i *) (a - 12)), - _mm_load_si128((__m128i *) (b - 12))); - r2 = _mm_add_ps(r2, _mm_mulhi_epi8_shift32(r3)); - res += r2.m128_f32[0] + r2.m128_f32[1]; - } - - return res; - } -#endif + uint32_t length) const; }; class AVXDistanceL2Float : public Distance { public: virtual float compare(const float *a, const float *b, - unsigned int length) const { -#ifndef _WINDOWS - std::cout << "AVX only supported in Windows build."; - return 0; - } -#else - __m128 diff, v1, v2; - __m128 sum = _mm_set1_ps(0); - - while (length >= 4) { - v1 = _mm_loadu_ps(a); - a += 4; - v2 = _mm_loadu_ps(b); - b += 4; - diff = _mm_sub_ps(v1, v2); - sum = _mm_add_ps(sum, _mm_mul_ps(diff, diff)); - length -= 4; - } - - return sum.m128_f32[0] + sum.m128_f32[1] + sum.m128_f32[2] + - sum.m128_f32[3]; - } -#endif + uint32_t length) const; }; - template - class DistanceInnerProduct : public Distance { - public: - float compare(const T *a, const T *b, unsigned size) const { - float result = 0; -#ifdef __GNUC__ -#ifdef __AVX__ -#define AVX_DOT(addr1, addr2, dest, tmp1, tmp2) \ - tmp1 = _mm256_loadu_ps(addr1); \ - tmp2 = _mm256_loadu_ps(addr2); \ - tmp1 = _mm256_mul_ps(tmp1, tmp2); \ - dest = _mm256_add_ps(dest, tmp1); - - __m256 sum; - __m256 l0, l1; - __m256 r0, r1; - unsigned D = (size + 7) & ~7U; - unsigned DR = D % 16; - unsigned DD = D - DR; - const float *l = (float *) a; - const float *r = (float *) b; - const float *e_l = l + DD; - const float *e_r = r + DD; - float unpack[8] __attribute__((aligned(32))) = {0, 0, 0, 0, 0, 0, 0, 0}; - - sum = _mm256_loadu_ps(unpack); - if (DR) { - AVX_DOT(e_l, e_r, sum, l0, r0); - } - - for (unsigned i = 0; i < DD; i += 16, l += 16, r += 16) { - AVX_DOT(l, r, sum, l0, r0); - AVX_DOT(l + 8, r + 8, sum, l1, r1); - } - _mm256_storeu_ps(unpack, sum); - result = unpack[0] + unpack[1] + unpack[2] + unpack[3] + unpack[4] + - unpack[5] + unpack[6] + unpack[7]; - -#else -#ifdef __SSE2__ -#define SSE_DOT(addr1, addr2, dest, tmp1, tmp2) \ - tmp1 = _mm128_loadu_ps(addr1); \ - tmp2 = _mm128_loadu_ps(addr2); \ - tmp1 = _mm128_mul_ps(tmp1, tmp2); \ - dest = _mm128_add_ps(dest, tmp1); - __m128 sum; - __m128 l0, l1, l2, l3; - __m128 r0, r1, r2, r3; - unsigned D = (size + 3) & ~3U; - unsigned DR = D % 16; - unsigned DD = D - DR; - const float *l = a; - const float *r = b; - const float *e_l = l + DD; - const float *e_r = r + DD; - float unpack[4] __attribute__((aligned(16))) = {0, 0, 0, 0}; - - sum = _mm_load_ps(unpack); - switch (DR) { - case 12: - SSE_DOT(e_l + 8, e_r + 8, sum, l2, r2); - case 8: - SSE_DOT(e_l + 4, e_r + 4, sum, l1, r1); - case 4: - SSE_DOT(e_l, e_r, sum, l0, r0); - default: - break; - } - for (unsigned i = 0; i < DD; i += 16, l += 16, r += 16) { - SSE_DOT(l, r, sum, l0, r0); - SSE_DOT(l + 4, r + 4, sum, l1, r1); - SSE_DOT(l + 8, r + 8, sum, l2, r2); - SSE_DOT(l + 12, r + 12, sum, l3, r3); - } - _mm_storeu_ps(unpack, sum); - result += unpack[0] + unpack[1] + unpack[2] + unpack[3]; -#else - - float dot0, dot1, dot2, dot3; - const float *last = a + size; - const float *unroll_group = last - 3; - - /* Process 4 items with each loop for efficiency. */ - while (a < unroll_group) { - dot0 = a[0] * b[0]; - dot1 = a[1] * b[1]; - dot2 = a[2] * b[2]; - dot3 = a[3] * b[3]; - result += dot0 + dot1 + dot2 + dot3; - a += 4; - b += 4; - } - /* Process last 0-3 pixels. Not needed for standard vector lengths. */ - while (a < last) { - result += *a++ * *b++; - } -#endif -#endif -#endif - return result; - } - }; - - template - class DistanceFastL2 : public DistanceInnerProduct { - public: - float norm(const T *a, unsigned size) const { - float result = 0; -#ifdef __GNUC__ -#ifdef __AVX__ -#define AVX_L2NORM(addr, dest, tmp) \ - tmp = _mm256_loadu_ps(addr); \ - tmp = _mm256_mul_ps(tmp, tmp); \ - dest = _mm256_add_ps(dest, tmp); - - __m256 sum; - __m256 l0, l1; - unsigned D = (size + 7) & ~7U; - unsigned DR = D % 16; - unsigned DD = D - DR; - const float *l = (float *) a; - const float *e_l = l + DD; - float unpack[8] __attribute__((aligned(32))) = {0, 0, 0, 0, 0, 0, 0, 0}; - - sum = _mm256_loadu_ps(unpack); - if (DR) { - AVX_L2NORM(e_l, sum, l0); - } - for (unsigned i = 0; i < DD; i += 16, l += 16) { - AVX_L2NORM(l, sum, l0); - AVX_L2NORM(l + 8, sum, l1); - } - _mm256_storeu_ps(unpack, sum); - result = unpack[0] + unpack[1] + unpack[2] + unpack[3] + unpack[4] + - unpack[5] + unpack[6] + unpack[7]; -#else -#ifdef __SSE2__ -#define SSE_L2NORM(addr, dest, tmp) \ - tmp = _mm128_loadu_ps(addr); \ - tmp = _mm128_mul_ps(tmp, tmp); \ - dest = _mm128_add_ps(dest, tmp); - - __m128 sum; - __m128 l0, l1, l2, l3; - unsigned D = (size + 3) & ~3U; - unsigned DR = D % 16; - unsigned DD = D - DR; - const float *l = a; - const float *e_l = l + DD; - float unpack[4] __attribute__((aligned(16))) = {0, 0, 0, 0}; - - sum = _mm_load_ps(unpack); - switch (DR) { - case 12: - SSE_L2NORM(e_l + 8, sum, l2); - case 8: - SSE_L2NORM(e_l + 4, sum, l1); - case 4: - SSE_L2NORM(e_l, sum, l0); - default: - break; - } - for (unsigned i = 0; i < DD; i += 16, l += 16) { - SSE_L2NORM(l, sum, l0); - SSE_L2NORM(l + 4, sum, l1); - SSE_L2NORM(l + 8, sum, l2); - SSE_L2NORM(l + 12, sum, l3); - } - _mm_storeu_ps(unpack, sum); - result += unpack[0] + unpack[1] + unpack[2] + unpack[3]; -#else - float dot0, dot1, dot2, dot3; - const float *last = a + size; - const float *unroll_group = last - 3; - - /* Process 4 items with each loop for efficiency. */ - while (a < unroll_group) { - dot0 = a[0] * a[0]; - dot1 = a[1] * a[1]; - dot2 = a[2] * a[2]; - dot3 = a[3] * a[3]; - result += dot0 + dot1 + dot2 + dot3; - a += 4; - } - /* Process last 0-3 pixels. Not needed for standard vector lengths. */ - while (a < last) { - result += (*a) * (*a); - a++; - } -#endif -#endif -#endif - return result; - } - using DistanceInnerProduct::compare; - float compare(const T *a, const T *b, float norm, - unsigned size) const { // not implement - float result = -2 * DistanceInnerProduct::compare(a, b, size); - result += norm; - return result; - } - }; } // namespace diskann diff --git a/include/exceptions.h b/include/exceptions.h index 0323ac3dcb..eefb0f69cc 100644 --- a/include/exceptions.h +++ b/include/exceptions.h @@ -12,4 +12,4 @@ namespace diskann { : std::logic_error("Function not yet implemented.") { } }; -} +} // namespace diskann diff --git a/include/file_content.h b/include/file_content.h new file mode 100644 index 0000000000..157d988a25 --- /dev/null +++ b/include/file_content.h @@ -0,0 +1,14 @@ +#pragma once + +#ifdef EXEC_ENV_OLS +namespace diskann { + struct FileContent { + public: + FileContent(void* content, size_t size) : _content(content), _size(size) { + } + + void* _content; + size_t _size; + }; +} // namespace diskann +#endif diff --git a/include/index.h b/include/index.h index eedbb1491f..2657ed56c9 100644 --- a/include/index.h +++ b/include/index.h @@ -3,13 +3,16 @@ #pragma once +#include #include #include +#include #include #include #include #include #include "tsl/robin_set.h" +#include "tsl/robin_map.h" #include "distance.h" #include "neighbor.h" @@ -17,98 +20,185 @@ #include "utils.h" #include "windows_customizations.h" +#include "Neighbor_Tag.h" + #define SLACK_FACTOR 1.3 -#define ESTIMATE_RAM_USAGE(size, dim, datasize, degree) \ - (1.30 * (((double) size * dim) * datasize + \ - ((double) size * degree) * sizeof(unsigned) * SLACK_FACTOR)) +#define ESTIMATE_RAM_USAGE(size, dim, datasize, degree) \ + (SLACK_FACTOR * (((double) size * (double) dim) * (double) datasize + \ + ((double) size * (double) degree) * \ + (double) sizeof(unsigned) * SLACK_FACTOR)) namespace diskann { - template + template class Index { public: - DISKANN_DLLEXPORT Index(Metric m, const char *filename, - const size_t max_points = 0, const size_t nd = 0, - const size_t num_frozen_pts = 0, - const bool enable_tags = false, - const bool store_data = true, - const bool support_eager_delete = false); + DISKANN_DLLEXPORT Index(Metric m, const size_t dim, const size_t max_points, + const bool dynamic_index, + const bool save_index_in_one_file, + const bool enable_tags = false, + const bool support_eager_delete = false); + + // DISKANN_DLLEXPORT Index(Index *index); // deep copy DISKANN_DLLEXPORT ~Index(); + // Public Functions for Static Support + // checks if data is consolidated, saves graph, metadata and associated // tags. DISKANN_DLLEXPORT void save(const char *filename); - DISKANN_DLLEXPORT void load(const char *filename, - const bool load_tags = false, - const char *tag_filename = NULL); - // generates one or more frozen points that will never get deleted from the - // graph - DISKANN_DLLEXPORT int generate_random_frozen_points( - const char *filename = NULL); + + DISKANN_DLLEXPORT _u64 save_graph(std::string filename, size_t offset = 0); + DISKANN_DLLEXPORT _u64 save_data(std::string filename, size_t offset = 0); + DISKANN_DLLEXPORT _u64 save_tags(std::string filename, size_t offset = 0); + DISKANN_DLLEXPORT _u64 save_delete_list(const std::string &filename, + size_t offset = 0); + + DISKANN_DLLEXPORT void load(const char *index_file); + + DISKANN_DLLEXPORT size_t load_graph(const std::string filename, + size_t expected_num_points, + size_t offset = 0); + + DISKANN_DLLEXPORT size_t load_data(std::string filename, size_t offset = 0); + + DISKANN_DLLEXPORT size_t load_tags(const std::string tag_file_name, + size_t offset = 0); + DISKANN_DLLEXPORT size_t load_delete_set(const std::string &filename, + size_t offset = 0); + + DISKANN_DLLEXPORT void get_delete_set(tsl::robin_set &del_set); + + DISKANN_DLLEXPORT size_t get_num_points(); + + DISKANN_DLLEXPORT size_t return_max_points(); DISKANN_DLLEXPORT void build( + const char *filename, const size_t num_points_to_load, Parameters & parameters, const std::vector &tags = std::vector()); - // Gopal. Added search overload that takes L as parameter, so that we + DISKANN_DLLEXPORT void build(const char * filename, + const size_t num_points_to_load, + Parameters & parameters, + const char * tag_filename); + // Added search overload that takes L as parameter, so that we // can customize L on a per-query basis without tampering with "Parameters" - DISKANN_DLLEXPORT std::pair search(const T *query, - const size_t K, - const unsigned L, - unsigned *indices); + DISKANN_DLLEXPORT std::pair search( + const T *query, const size_t K, const unsigned L, unsigned *indices, + float *distances = nullptr); DISKANN_DLLEXPORT std::pair search( const T *query, const uint64_t K, const unsigned L, std::vector init_ids, uint64_t *indices, float *distances); - DISKANN_DLLEXPORT std::pair search_with_tags( - const T *query, const size_t K, const unsigned L, TagT *tags, - unsigned frozen_pts, unsigned *indices_buffer = NULL); + DISKANN_DLLEXPORT size_t search_with_tags(const T *query, const uint64_t K, + const unsigned L, TagT *tags, + float * distances, + std::vector &res_vectors); - // repositions frozen points to the end of _data - if they have been moved - // during deletion - DISKANN_DLLEXPORT void readjust_data(unsigned _num_frozen_pts); + DISKANN_DLLEXPORT size_t search_with_tags(const T *query, const size_t K, + const unsigned L, TagT *tags, + float *distances); + + DISKANN_DLLEXPORT std::pair search( + const T *query, const size_t K, const unsigned L, + std::vector> &best_L_tags); + + DISKANN_DLLEXPORT void optimize_graph(); + + DISKANN_DLLEXPORT void search_with_opt_graph(const T *query, size_t K, + size_t L, unsigned *indices); + + DISKANN_DLLEXPORT void clear_index(); + + // Public Functions for Incremental Support /* insertions possible only when id corresponding to tag does not already * exist in the graph */ - DISKANN_DLLEXPORT int insert_point(const T * point, - const Parameters & parameter, - std::vector & pool, - std::vector & tmp, - tsl::robin_set & visited, - std::vector &cut_graph, - const TagT tag); - + DISKANN_DLLEXPORT int insert_point( + const T *point, const Parameters ¶meter, + const TagT tag); // only keep point, tag, parameters // call before triggering deleteions - sets important flags required for // deletion related operations DISKANN_DLLEXPORT int enable_delete(); - // call after all delete requests have been served, checks if deletions were - // executed correctly, rearranges metadata in case of lazy deletes - DISKANN_DLLEXPORT int disable_delete(const Parameters ¶meters, - const bool consolidate = false); - // Record deleted point now and restructure graph later. Return -1 if tag // not found, 0 if OK. Do not call if _eager_delete was called earlier and // data was not consolidated - DISKANN_DLLEXPORT int delete_point(const TagT tag); + DISKANN_DLLEXPORT int lazy_delete(const TagT &tag); + + // Record deleted points now and restructure graph later. Add to failed_tags + // if tag not found. Do not call if _eager_delete was called earlier and + // data was not consolidated. Return -1 if + DISKANN_DLLEXPORT int lazy_delete(const tsl::robin_set &tags, + std::vector & failed_tags); // Delete point from graph and restructure it immediately. Do not call if // _lazy_delete was called earlier and data was not consolidated DISKANN_DLLEXPORT int eager_delete(const TagT tag, - const Parameters ¶meters); + const Parameters ¶meters, + int delete_mode = 1); + // return _data and tag_to_location offset + DISKANN_DLLEXPORT int extract_data( + T *ret_data, std::unordered_map &tag_to_location); + + DISKANN_DLLEXPORT void get_location_to_tag( + std::unordered_map &ret_loc_to_tag); + + DISKANN_DLLEXPORT void prune_all_nbrs(const Parameters ¶meters); + + DISKANN_DLLEXPORT void compact_data_for_insert(); + + DISKANN_DLLEXPORT bool hasIndexBeenSaved(); + // diskv2 API + void iterate_to_fixed_point(const T *node_coords, const unsigned Lindex, + std::vector &expanded_nodes_info, + tsl::robin_map &coord_map, + bool return_frozen_pt = true); + // convenient access to graph + data (aligned) + const std::vector> *get_graph() const { + return &this->_final_graph; + } + T * get_data(); + const std::unordered_map *get_tags() const { + return &this->_location_to_tag; + }; + // repositions frozen points to the end of _data - if they have been moved + // during deletion + DISKANN_DLLEXPORT void reposition_frozen_point_to_end(); + DISKANN_DLLEXPORT void reposition_point(unsigned old_location, + unsigned new_location); - DISKANN_DLLEXPORT void optimize_graph(); + DISKANN_DLLEXPORT void compact_frozen_point(); + DISKANN_DLLEXPORT void compact_data_for_search(); - DISKANN_DLLEXPORT void search_with_opt_graph(const T *query, size_t K, - size_t L, unsigned *indices); + DISKANN_DLLEXPORT void consolidate(Parameters ¶meters); + + // DISKANN_DLLEXPORT void save_index_as_one_file(bool flag); + + DISKANN_DLLEXPORT void get_active_tags(tsl::robin_set &active_tags); + + DISKANN_DLLEXPORT int get_vector_by_tag(TagT &tag, T *vec); + DISKANN_DLLEXPORT const T *get_vector_by_tag(const TagT &tag); + + // TODO: Debugging ONLY + DISKANN_DLLEXPORT void print_status() const; + DISKANN_DLLEXPORT void are_deleted_points_in_graph() const; + DISKANN_DLLEXPORT void print_delete_set() const; + + // This variable MUST be updated if the number of entries in the metadata + // change. + DISKANN_DLLEXPORT static const int METADATA_ROWS = 5; /* Internals of the library */ protected: - typedef std::vector vecNgh; - typedef std::vector> CompactGraph; - CompactGraph _final_graph; - CompactGraph _in_graph; + std::vector> _final_graph; + std::vector> _in_graph; + + // generates one frozen point that will never get deleted from the + // graph + int generate_frozen_point(); // determines navigating node of the graph by calculating medoid of data unsigned calculate_entry_point(); @@ -120,7 +210,7 @@ namespace diskann { const std::vector &init_ids, std::vector & expanded_nodes_info, tsl::robin_set & expanded_nodes_ids, - std::vector & best_L_nodes); + std::vector &best_L_nodes, bool ret_frozen = true); void get_expanded_nodes(const size_t node, const unsigned Lindex, std::vector init_ids, @@ -150,60 +240,85 @@ namespace diskann { void link(Parameters ¶meters); - // WARNING: Do not call reserve_location() without acquiring change_lock_ - unsigned reserve_location(); + // Support for Incremental Indexing + int reserve_location(); + void release_location(); - // get new location corresponding to each undeleted tag after deletions - std::vector get_new_location(unsigned &active); + // Support for resizing the index + // This function must be called ONLY after taking the _change_lock and + // _update_lock. + // Anything else in a MT environment will lead to an inconsistent index. + void resize(uint32_t new_max_points); // renumber nodes, update tag and location maps and compact the graph, mode - // = _consolidated_order in case of lazy deletion and _compacted_order in + // = _compacted_lazy_deletions in case of lazy deletion and + // _compacted_eager_deletions in // case of eager deletion - void compact_data(std::vector new_location, unsigned active, - bool &mode); + void compact_data(); // WARNING: Do not call consolidate_deletes without acquiring change_lock_ // Returns number of live points left after consolidation size_t consolidate_deletes(const Parameters ¶meters); private: - size_t _dim; - size_t _aligned_dim; - T * _data; - size_t _nd; // number of active points i.e. existing in the graph - size_t _max_points; // total number of points in given data set - size_t _num_frozen_pts; - bool _has_built; - Distance *_distance; - unsigned _width; - unsigned _ep; - bool _saturate_graph = false; - std::vector _locks; // Per node lock, cardinality=max_points_ + // DEBUG ONLY + void printTagToLocation(); + + std::shared_timed_mutex _tag_lock; // reader-writer lock on + // _tag_to_location and + std::mutex _change_lock; // Lock taken to synchronously modify _nd + + T *_data = nullptr; // coordinates of all base points + // T *_pq_data = + // nullptr; // coordinates of pq centroid corresponding to every point + Distance * _distance = nullptr; + diskann::Metric _dist_metric; + + size_t _dim; + size_t _aligned_dim; + size_t _nd = 0; // number of active points i.e. existing in the graph + size_t _max_points = 0; // total number of points in given data set + size_t _num_frozen_pts = 0; + unsigned _width = 0; + unsigned _ep = 0; + bool _has_built = false; + bool _saturate_graph = false; + bool _save_as_one_file = false; + bool _dynamic_index = false; + bool _enable_tags = false; char * _opt_graph; size_t _node_size; size_t _data_len; size_t _neighbor_len; - bool _can_delete; - bool _eager_done; // true if eager deletions have been made - bool _lazy_done; // true if lazy deletions have been made - bool _compacted_order; // true if after eager deletions, data has been - // consolidated - bool _enable_tags; - bool _consolidated_order; // true if after lazy deletions, data has been - // consolidated - bool _support_eager_delete; //_support_eager_delete = activates extra data - // structures and functions required for eager - // deletion - bool _store_data; - + // flags for dynamic indexing std::unordered_map _tag_to_location; std::unordered_map _location_to_tag; tsl::robin_set _delete_set; tsl::robin_set _empty_slots; - std::mutex _change_lock; // Allow only 1 thread to insert/delete + bool _support_eager_delete = + false; //_support_eager_delete = activates extra data + // bool _can_delete = false; // only true if deletes can be done (if + // enabled) + bool _eager_done = false; // true if eager deletions have been made + bool _lazy_done = false; // true if lazy deletions have been made + bool _data_compacted = true; // true if data has been consolidated + bool _is_saved = false; // Gopal. Checking if the index is already saved. + + std::vector _locks; // Per node lock, cardinality=max_points_ + std::vector _locks_in; // Per node lock + std::shared_timed_mutex _delete_lock; // Lock on _delete_set and + // _empty_slots when reading and + // writing to them + // _location_to_tag, has a shared lock + // and exclusive lock associated with + // it. + std::shared_timed_mutex _update_lock; // coordinate save() and any change + // being done to the graph. + + const float INDEX_GROWTH_FACTOR = 1.5f; }; } // namespace diskann diff --git a/include/linux_aligned_file_reader.h b/include/linux_aligned_file_reader.h index f9815697cc..d0cb9c9548 100644 --- a/include/linux_aligned_file_reader.h +++ b/include/linux_aligned_file_reader.h @@ -24,15 +24,19 @@ class LinuxAlignedFileReader : public AlignedFileReader { // de-register thread-id for a context void deregister_thread(); + void deregister_all_threads(); + // Open & close ops // Blocking calls - void open(const std::string &fname); + void open(const std::string &fname, bool enable_writes, bool enable_create); void close(); // process batch of aligned requests in parallel // NOTE :: blocking call void read(std::vector &read_reqs, IOContext &ctx, bool async = false); + + void sequential_write(AlignedRead &write_req, IOContext &ctx); }; #endif diff --git a/include/logger_impl.h b/include/logger_impl.h index 035406f50d..794938a287 100644 --- a/include/logger_impl.h +++ b/include/logger_impl.h @@ -66,7 +66,7 @@ namespace diskann { // This implies calling code _must_ either print std::endl or std::flush // to ensure that the message is written immediately. #ifdef EXEC_ENV_OLS - static const int BUFFER_SIZE = 1024; + static const int BUFFER_SIZE = 2048; #else static const int BUFFER_SIZE = 0; #endif diff --git a/include/memory_mapper.h b/include/memory_mapper.h index 4ebe6ec62e..4ccbf6f286 100644 --- a/include/memory_mapper.h +++ b/include/memory_mapper.h @@ -38,4 +38,4 @@ namespace diskann { ~MemoryMapper(); }; -} \ No newline at end of file +} // namespace diskann diff --git a/include/neighbor.h b/include/neighbor.h index 5c37df38e5..7dbf1a3c79 100644 --- a/include/neighbor.h +++ b/include/neighbor.h @@ -6,6 +6,7 @@ #include #include #include +#include #include "utils.h" namespace diskann { diff --git a/include/parameters.h b/include/parameters.h index 42e7d3a463..31a15a1fa0 100644 --- a/include/parameters.h +++ b/include/parameters.h @@ -21,6 +21,10 @@ namespace diskann { // ParamType *ptr = (ParamType *) malloc(sizeof(ParamType)); ParamType *ptr = new ParamType; *ptr = value; + if (params.find(name) != params.end()) { + free(params[name]); + } + params[name] = (void *) ptr; } @@ -54,7 +58,6 @@ namespace diskann { for (auto iter = params.begin(); iter != params.end(); iter++) { if (iter->second != nullptr) free(iter->second); - // delete iter->second; } } diff --git a/include/partition_and_pq.h b/include/partition_and_pq.h index 45b1e26e6d..05dd63a964 100644 --- a/include/partition_and_pq.h +++ b/include/partition_and_pq.h @@ -15,9 +15,18 @@ #include "windows_customizations.h" +const uint32_t NUM_PQ_CENTERS = 256; +const uint32_t NUM_K_MEANS_ITERS = 15; + template void gen_random_slice(const std::string base_file, - const std::string output_prefix, double sampling_rate); + const std::string output_prefix, double sampling_rate, + size_t offset = 0); + +template +void gen_random_slice(const std::string data_file, double p_val, + std::unique_ptr &sampled_data, + size_t &slice_size, size_t &ndims); template void gen_random_slice(const std::string data_file, double p_val, @@ -49,6 +58,12 @@ int partition_with_ram_budget(const std::string data_file, size_t graph_degree, const std::string prefix_path, size_t k_base); +template +DISKANN_DLLEXPORT int generate_pq_pivots( + const std::unique_ptr &passed_train_data, size_t num_train, + unsigned dim, unsigned num_centers, unsigned num_pq_chunks, + unsigned max_k_means_reps, std::string pq_pivots_path); + DISKANN_DLLEXPORT int generate_pq_pivots(const float *train_data, size_t num_train, unsigned dim, unsigned num_centers, @@ -60,4 +75,5 @@ template int generate_pq_data_from_pivots(const std::string data_file, unsigned num_centers, unsigned num_pq_chunks, std::string pq_pivots_path, - std::string pq_compressed_vectors_path); + std::string pq_compressed_vectors_path, + size_t offset = 0); diff --git a/include/percentile_stats.h b/include/percentile_stats.h index 808546c165..8c6be913cc 100644 --- a/include/percentile_stats.h +++ b/include/percentile_stats.h @@ -44,7 +44,7 @@ namespace diskann { vals.begin(), vals.end(), [](const double &left, const double &right) { return left < right; }); - auto retval = vals[(uint64_t)(percentile * len)]; + auto retval = vals[(uint64_t)(percentile * ((float) len))]; vals.clear(); return retval; } @@ -56,6 +56,6 @@ namespace diskann { for (uint64_t i = 0; i < len; i++) { avg += member_fn(stats[i]); } - return avg / len; + return avg / ((double) len); } -} +} // namespace diskann diff --git a/include/pq_flash_index.h b/include/pq_flash_index.h index c7601851c6..23e545dcaa 100644 --- a/include/pq_flash_index.h +++ b/include/pq_flash_index.h @@ -49,13 +49,25 @@ namespace diskann { } }; + template + struct DiskNode { + uint32_t id = 0; + T * coords = nullptr; + uint32_t nnbrs; + uint32_t *nbrs; + + // id : id of node + // sector_buf : sector buf containing `id` data + DiskNode(uint32_t id, T *coords, uint32_t *nhood); + }; + template struct ThreadData { QueryScratch scratch; IOContext ctx; }; - template + template class PQFlashIndex { public: // Gopal. Adapting to the new Bing interface. Since the DiskPriorityIO is @@ -65,26 +77,31 @@ namespace diskann { // created the way we need. Linux will create a simple AlignedFileReader // and pass it. Regular Windows code should create a BingFileReader using // the DiskPriorityIOInterface class, and for running on XTS, create a - // BingFileReader - // using the object passed by the XTS environment. - // Freeing the reader object is now the client's (DiskANNInterface's) - // responsibility. + // BingFileReader using the object passed by the XTS environment. Freeing + // the + // reader object is now the client's (DiskANNInterface's) responsibility. DISKANN_DLLEXPORT PQFlashIndex( - std::shared_ptr &fileReader); + diskann::Metric m, std::shared_ptr &fileReader, + bool single_file_index, bool tags = false); + DISKANN_DLLEXPORT ~PQFlashIndex(); #ifdef EXEC_ENV_OLS - DISKANN_DLLEXPORT int load(diskann::MemoryMappedFiles &files, - uint32_t num_threads, const char *pq_prefix, - const char *disk_index_file); + DISKANN_DLLEXPORT int load(const char *index_prefix, _u32 num_threads, + diskann::MemoryMappedFiles &files, + bool new_index_format = true); #else // load compressed data, and obtains the handle to the disk-resident index - DISKANN_DLLEXPORT int load(uint32_t num_threads, const char *pq_prefix, - const char *disk_index_file); + DISKANN_DLLEXPORT int load(const char *index_prefix, uint32_t num_threads, + bool new_index_format = true); #endif + DISKANN_DLLEXPORT void load_tags(const std::string &tag_file, + size_t offset = 0); DISKANN_DLLEXPORT void load_cache_list(std::vector &node_list); + DISKANN_DLLEXPORT _u64 return_nd(); + #ifdef EXEC_ENV_OLS DISKANN_DLLEXPORT void generate_cache_list_from_sample_queries( MemoryMappedFiles &files, std::string sample_bin, _u64 l_search, @@ -110,18 +127,86 @@ namespace diskann { // setting up thread-specific data // implemented - DISKANN_DLLEXPORT void cached_beam_search( - const T *query, const _u64 k_search, const _u64 l_search, _u64 *res_ids, - float *res_dists, const _u64 beam_width, QueryStats *stats = nullptr, - Distance *output_dist_func = nullptr); - std::shared_ptr &reader; + DISKANN_DLLEXPORT size_t + cached_beam_search(const T *query, const _u64 k_search, const _u64 l_search, + TagT *res_tags, float *res_dists, const _u64 beam_width, + QueryStats *stats = nullptr); - protected: - DISKANN_DLLEXPORT void use_medoids_data_as_centroids(); - DISKANN_DLLEXPORT void setup_thread_data(_u64 nthreads); - DISKANN_DLLEXPORT void destroy_thread_data(); + DISKANN_DLLEXPORT size_t cached_beam_search_ids( + const T *query, const _u64 k_search, const _u64 l_search, + uint64_t *res_ids, float *res_dists, const _u64 beam_width, + QueryStats *stats = nullptr); + + DISKANN_DLLEXPORT size_t cached_beam_search_vectors( + const T *query, const _u64 k_search, const _u64 l_search, + TagT *res_tags, float *res_dists, std::vector &res_vectors, + const _u64 beam_width, QueryStats *stats = nullptr); + + /* diskv2 extra API requirements */ + /* -------------------------------------------------------------------------------------------- + */ + DISKANN_DLLEXPORT void disk_iterate_to_fixed_point( + const T *vec, const uint32_t Lsize, const uint32_t beam_width, + std::vector & expanded_nodes_info, + tsl::robin_map *coord_map = nullptr, + QueryStats *stats = nullptr, ThreadData *passthrough_data = nullptr, + tsl::robin_set *exclude_nodes = nullptr); + std::vector get_init_ids() { + return std::vector(this->medoids, + this->medoids + this->num_medoids); + } + // gives access to backing thread data buf for easy parallelization + std::vector> &get_thread_data() { + return this->thread_data_backing_buf; + } + + // computes PQ dists between src->[ids] into fp_dists (merge, insert) + DISKANN_DLLEXPORT void compute_pq_dists(const _u32 src, const _u32 *ids, + float *fp_dists, const _u32 count, + uint8_t *aligned_scratch = nullptr); + // computes PQ dists between aligned_query->[ids] into fp_dists (search) + DISKANN_DLLEXPORT void compute_pq_dists(const T *query, const _u32 *ids, + float *fp_dists, const _u32 count); + // read/write [start_id:end_id-1] points from disk + // WARNING -- ensure (start_id,end_id) % nnodes_per_sector = 0, + // aligned_scratch is SECTOR_LEN aligned + // WARNING -- ensure aligned_scratch size is >((end_id - + // start_id)/nnodes_per_sector) * SECTOR_LEN bytes + DISKANN_DLLEXPORT _u32 merge_read(std::vector> &disk_nodes, + _u32 &start_id, const _u32 sector_count, + char *scratch); + DISKANN_DLLEXPORT void scan_deleted_nodes( + const tsl::robin_set &delete_set, + std::vector> &deleted_nodes, char *buf, char *backing_buf, + const uint32_t sectors_per_scan); + DISKANN_DLLEXPORT void reload_index(const std::string &new_disk_index_path, + const std::string &new_pq_coords_path, + const std::string &new_tags_file_path); + // DISKANN_DLLEXPORT void passthrough_write(char *buf, const uint64_t + // offset, + // const uint64_t size); + + // deflates `vec` into PQ ids + DISKANN_DLLEXPORT std::vector<_u8> deflate_vector(const T *vec); + std::pair<_u8 *, _u32> get_pq_config() { + return std::make_pair(this->data, (uint32_t) this->n_chunks); + } + DISKANN_DLLEXPORT TagT *get_tags() { + return this->tags; + } + + DISKANN_DLLEXPORT _u64 get_num_frozen_points() { + return this->num_frozen_points; + } + + DISKANN_DLLEXPORT _u64 get_frozen_loc() { + return this->frozen_location; + } + + DISKANN_DLLEXPORT void get_active_tags(tsl::robin_set &active_tags); + + DISKANN_DLLEXPORT int get_vector_by_tag(const TagT &tag, T *vector); - private: // index info // nhood of node `i` is in sector: [i / nnodes_per_sector] // offset in sector: [(i % nnodes_per_sector) * max_node_len] @@ -129,14 +214,24 @@ namespace diskann { // nbrs of node `i`: ((unsigned*)buf) + 1 _u64 max_node_len = 0, nnodes_per_sector = 0, max_degree = 0; + protected: + DISKANN_DLLEXPORT void use_medoids_data_as_centroids(); + DISKANN_DLLEXPORT void setup_thread_data(_u64 nthreads); + DISKANN_DLLEXPORT void destroy_thread_data(); + + private: // data info _u64 num_points = 0; + _u64 num_frozen_points = 0; + _u64 frozen_location = 0; _u64 data_dim = 0; _u64 aligned_dim = 0; - std::string disk_index_file; + std::string _disk_index_file; std::vector> node_visit_counter; + std::shared_ptr &reader; + // PQ data // n_chunks = # of chunks ndims is split into // data: _u8 * n_chunks @@ -148,8 +243,15 @@ namespace diskann { FixedChunkPQTable pq_table; // distance comparator - Distance * dist_cmp = nullptr; - Distance *dist_cmp_float = nullptr; + std::shared_ptr> dist_cmp; + std::shared_ptr> dist_cmp_float; + + // Are we dealing with normalized data? This will be true + // if distance == COSINE and datatype == float. Required + // because going forward, we will normalize vectors when + // asked to search with COSINE similarity. Of course, this + // will be done only for floating point vectors. + bool data_is_normalized = false; // medoid/start info uint32_t *medoids = @@ -162,19 +264,34 @@ namespace diskann { // closest centroid as the starting point of search // nhood_cache - unsigned *nhood_cache_buf = nullptr; + unsigned * nhood_cache_buf = nullptr; tsl::robin_map<_u32, std::pair<_u32, _u32 *>> nhood_cache; // coord_cache - T *coord_cache_buf = nullptr; + T * coord_cache_buf = nullptr; tsl::robin_map<_u32, T *> coord_cache; // thread-specific scratch ConcurrentQueue> thread_data; + std::vector> thread_data_backing_buf; _u64 max_nthreads; bool load_flag = false; bool count_visited_nodes = false; + bool single_index_file; // For fresh + static. Save + // everything in a single file + + // support for tags and dynamic indexing + + bool enable_tags = false; + TagT *tags = nullptr; + + /* diskv2 extra API requirements */ + // ids that don't have disk nhoods, but have in-mem PQ + tsl::robin_set<_u32> invalid_ids; + std::mutex invalid_ids_lock; + +// tags #ifdef EXEC_ENV_OLS // Set to a larger value than the actual header to accommodate // any additions we make to the header. This is an outer limit diff --git a/include/pq_table.h b/include/pq_table.h index 3cac23c15a..cf3059f4d6 100644 --- a/include/pq_table.h +++ b/include/pq_table.h @@ -5,6 +5,9 @@ #include "utils.h" +#define NUM_PQ_CENTROIDS 256 +#define NUM_PQ_OFFSETS 5 + namespace diskann { template class FixedChunkPQTable { @@ -19,6 +22,8 @@ namespace diskann { _u32* rearrangement = nullptr; float* centroid = nullptr; float* tables_T = nullptr; // same as pq_tables, but col-major + float* all_to_all_dists = nullptr; + public: FixedChunkPQTable() { } @@ -35,95 +40,277 @@ namespace diskann { delete[] chunk_offsets; if (centroid != nullptr) delete[] centroid; + if (all_to_all_dists != nullptr) + delete[] all_to_all_dists; #endif } + _u64 get_dim() { + return ndims; + } + #ifdef EXEC_ENV_OLS - void load_pq_centroid_bin(MemoryMappedFiles& files, - const char* pq_table_file, size_t num_chunks){ + void load_pq_pivots_old(MemoryMappedFiles& files, + const std::string& pq_pivots_path, + size_t num_chunks) { +#else + void load_pq_pivots_old(const std::string& pq_pivots_path, + size_t num_chunks) { +#endif + _u64 nr, nc; +// Load the pq pivots. +#ifdef EXEC_ENV_OLS + diskann::load_bin(files, pq_pivots_path.c_str(), tables, nr, nc); +#else + diskann::load_bin(pq_pivots_path.c_str(), tables, nr, nc); +#endif + + if ((nr != NUM_PQ_CENTROIDS)) { + diskann::cout << "Error reading pq_pivots file " << pq_pivots_path + << ". file_num_centers = " << nr << " but expecting " + << NUM_PQ_CENTROIDS << " centers"; + throw diskann::ANNException( + "Error reading pq_pivots file at pivots data.", -1, __FUNCSIG__, + __FILE__, __LINE__); + } + + this->ndims = nc; + +// Load the PQ centroids +//_pq_pivots.bin_centroid.bin +//_chunk_offsets.bin +//_rearrangement_perm.bin +#ifdef EXEC_ENV_OLS + diskann::load_bin(files, pq_pivots_path + "_centroid.bin", + centroid, nr, nc); +#else + diskann::load_bin(pq_pivots_path + "_centroid.bin", centroid, nr, + nc); +#endif + + if ((nr != this->ndims) || (nc != 1)) { + diskann::cerr << "Error reading centroids from pq_pivots file " + << pq_pivots_path << ". file_dim = " << nr + << ", file_cols = " << nc << " but expecting " + << this->ndims << " entries in 1 dimension."; + throw diskann::ANNException( + "Error reading pq_pivots file at centroid data.", -1, __FUNCSIG__, + __FILE__, __LINE__); + } + +#ifdef EXEC_ENV_OLS + diskann::load_bin(files, + pq_pivots_path + "_rearrangement_perm.bin", + rearrangement, nr, nc); #else - void load_pq_centroid_bin(const char* pq_table_file, size_t num_chunks) { + diskann::load_bin(pq_pivots_path + "_rearrangement_perm.bin", + rearrangement, nr, nc); #endif - std::string rearrangement_file = std::string(pq_table_file) + - "_rearrangement_perm.bin"; - std::string chunk_offset_file = - std::string(pq_table_file) + "_chunk_offsets.bin"; - std::string centroid_file = std::string(pq_table_file) + "_centroid.bin"; + if ((nr != this->ndims) || (nc != 1)) { + diskann::cerr << "Error reading re-arrangement data pq_pivots file " + << pq_pivots_path << ". file_dim = " << nr + << ", file_cols = " << nc << " but expecting " + << this->ndims << " entries in 1 dimension."; + throw diskann::ANNException( + "Error reading pq_pivots file at re-arrangement data.", -1, + __FUNCSIG__, __FILE__, __LINE__); + } - // bin structure: [256][ndims][ndims(float)] - uint64_t numr, numc; - size_t npts_u64, ndims_u64; #ifdef EXEC_ENV_OLS - diskann::load_bin(files, pq_table_file, tables, npts_u64, ndims_u64); + diskann::load_bin(files, pq_pivots_path + "_chunk_offsets.bin", + chunk_offsets, nr, nc); #else - diskann::load_bin(pq_table_file, tables, npts_u64, ndims_u64); + diskann::load_bin(pq_pivots_path + "_chunk_offsets.bin", + chunk_offsets, nr, nc); #endif - this->ndims = ndims_u64; - if (file_exists(chunk_offset_file)) { + if (nr != (uint64_t) num_chunks + 1 || nc != 1) { + diskann::cerr + << "Error reading pq_pivots file at chunk offsets; file has nr=" + << nr << ",nc=" << nc << ", expecting nr=" << num_chunks + 1 + << ", nc=1." << std::endl; + throw diskann::ANNException( + "Error reading pq_pivots file at chunk offsets.", -1, __FUNCSIG__, + __FILE__, __LINE__); + } + + this->n_chunks = num_chunks; + diskann::cout << "Loaded PQ Pivots: #ctrs: " << NUM_PQ_CENTROIDS + << ", #dims: " << this->ndims + << ", #chunks: " << this->n_chunks << std::endl; + } + #ifdef EXEC_ENV_OLS - diskann::load_bin<_u32>(files, rearrangement_file, rearrangement, numr, - numc); + void load_pq_pivots_new(MemoryMappedFiles& files, + const std::string& pq_pivots_path, + size_t num_chunks, size_t offset) { #else - diskann::load_bin<_u32>(rearrangement_file, rearrangement, numr, numc); + void load_pq_pivots_new(const std::string& pq_pivots_path, + size_t num_chunks, size_t offset) { #endif - if (numr != ndims_u64 || numc != 1) { - diskann::cerr << "Error loading rearrangement file" << std::endl; - throw diskann::ANNException("Error loading rearrangement file", -1, - __FUNCSIG__, __FILE__, __LINE__); + + _u64 nr, nc; +#ifdef EXEC_ENV_OLS + _u64* file_offset_data; // since load_bin only sets the pointer, no need + // to delete. + diskann::load_bin<_u64>(files, pq_pivots_path, file_offset_data, nr, nc, + offset); +#else + std::unique_ptr<_u64[]> file_offset_data; + diskann::load_bin<_u64>(pq_pivots_path, file_offset_data, nr, nc, offset); +#endif + + if (nr != NUM_PQ_OFFSETS) { + diskann::cout << "Error reading pq_pivots file " << pq_pivots_path + << ". Offsets dont contain correct metadata, # offsets = " + << nr << ", but expecting " << NUM_PQ_OFFSETS; + throw diskann::ANNException( + "Error reading pq_pivots file at offsets data.", -1, __FUNCSIG__, + __FILE__, __LINE__); } + diskann::cout << "Offsets: " << file_offset_data[0] << " " + << file_offset_data[1] << " " << file_offset_data[2] << " " + << file_offset_data[3] << " " << file_offset_data[4] + << std::endl; + #ifdef EXEC_ENV_OLS - diskann::load_bin<_u32>(files, chunk_offset_file, chunk_offsets, numr, - numc); + diskann::load_bin(files, pq_pivots_path.c_str(), tables, nr, nc, + file_offset_data[0] + offset); #else - diskann::load_bin<_u32>(chunk_offset_file, chunk_offsets, numr, numc); -#endif - if (numc != 1 || numr != num_chunks + 1) { - diskann::cerr << "Error loading chunk offsets file. numc: " << numc - << " (should be 1). numr: " << numr << " (should be " - << num_chunks + 1 << ")" << std::endl; - throw diskann::ANNException("Error loading chunk offsets file", -1, - __FUNCSIG__, __FILE__, __LINE__); + diskann::load_bin(pq_pivots_path.c_str(), tables, nr, nc, + file_offset_data[0] + offset); +#endif + + if ((nr != NUM_PQ_CENTROIDS)) { + diskann::cout << "Error reading pq_pivots file " << pq_pivots_path + << ". file_num_centers = " << nr << " but expecting " + << NUM_PQ_CENTROIDS << " centers"; + throw diskann::ANNException( + "Error reading pq_pivots file at pivots data.", -1, __FUNCSIG__, + __FILE__, __LINE__); } - this->n_chunks = numr - 1; + this->ndims = nc; + +#ifdef EXEC_ENV_OLS + diskann::load_bin(files, pq_pivots_path.c_str(), centroid, nr, nc, + file_offset_data[1] + offset); +#else + diskann::load_bin(pq_pivots_path.c_str(), centroid, nr, nc, + file_offset_data[1] + offset); +#endif + + if ((nr != this->ndims) || (nc != 1)) { + diskann::cerr << "Error reading centroids from pq_pivots file " + << pq_pivots_path << ". file_dim = " << nr + << ", file_cols = " << nc << " but expecting " + << this->ndims << " entries in 1 dimension."; + throw diskann::ANNException( + "Error reading pq_pivots file at centroid data.", -1, __FUNCSIG__, + __FILE__, __LINE__); + } #ifdef EXEC_ENV_OLS - diskann::load_bin(files, centroid_file, centroid, numr, numc); + diskann::load_bin(files, pq_pivots_path.c_str(), rearrangement, + nr, nc, file_offset_data[2] + offset); #else - diskann::load_bin(centroid_file, centroid, numr, numc); + diskann::load_bin(pq_pivots_path.c_str(), rearrangement, nr, nc, + file_offset_data[2] + offset); #endif - if (numc != 1 || numr != ndims_u64) { - diskann::cerr << "Error loading centroid file" << std::endl; - throw diskann::ANNException("Error loading centroid file", -1, - __FUNCSIG__, __FILE__, __LINE__); + if ((nr != this->ndims) || (nc != 1)) { + diskann::cerr << "Error reading re-arrangement data pq_pivots file " + << pq_pivots_path << ". file_dim = " << nr + << ", file_cols = " << nc << " but expecting " + << this->ndims << " entries in 1 dimension."; + throw diskann::ANNException( + "Error reading pq_pivots file at re-arrangement data.", -1, + __FUNCSIG__, __FILE__, __LINE__); } - } else { + +#ifdef EXEC_ENV_OLS + diskann::load_bin(files, pq_pivots_path.c_str(), chunk_offsets, + nr, nc, file_offset_data[3] + offset); +#else + diskann::load_bin(pq_pivots_path.c_str(), chunk_offsets, nr, nc, + file_offset_data[3] + offset); +#endif + + if (nr != (uint64_t) num_chunks + 1 || nc != 1) { + diskann::cerr + << "Error reading pq_pivots file at chunk offsets; file has nr=" + << nr << ",nc=" << nc << ", expecting nr=" << num_chunks + 1 + << ", nc=1." << std::endl; + throw diskann::ANNException( + "Error reading pq_pivots file at chunk offsets.", -1, __FUNCSIG__, + __FILE__, __LINE__); + } + this->n_chunks = num_chunks; - rearrangement = new uint32_t[ndims]; - - uint64_t chunk_size = DIV_ROUND_UP(ndims, num_chunks); - for (uint32_t d = 0; d < ndims; d++) - rearrangement[d] = d; - chunk_offsets = new uint32_t[num_chunks + 1]; - for (uint32_t d = 0; d <= num_chunks; d++) - chunk_offsets[d] = (_u32)(std::min)(ndims, d * chunk_size); - centroid = new float[ndims]; - std::memset(centroid, 0, ndims * sizeof(float)); + diskann::cout << "Loaded PQ Pivots: #ctrs: " << NUM_PQ_CENTROIDS + << ", #dims: " << this->ndims + << ", #chunks: " << this->n_chunks << std::endl; + } + +#ifdef EXEC_ENV_OLS + void load_pq_centroid_bin(MemoryMappedFiles& files, + const char* pq_table_file, size_t num_chunks, + size_t offset = 0){ +#else + void load_pq_centroid_bin(const char* pq_table_file, size_t num_chunks, + size_t offset = 0) { +#endif + + std::string pq_pivots_path(pq_table_file); + _u64 nr, nc; + +#ifdef EXEC_ENV_OLS + get_bin_metadata(files, pq_table_file, nr, nc, offset); +#else + get_bin_metadata(pq_table_file, nr, nc, offset); +#endif + + if (nr == NUM_PQ_OFFSETS) { +#ifdef EXEC_ENV_OLS + load_pq_pivots_new(files, pq_table_file, num_chunks, offset); +#else + load_pq_pivots_new(pq_table_file, num_chunks, offset); +#endif + } else if (nr == NUM_PQ_CENTROIDS) { +#ifdef EXEC_ENV_OLS + load_pq_pivots_old(files, pq_table_file, num_chunks); +#else + load_pq_pivots_old(pq_table_file, num_chunks); +#endif } - diskann::cout << "PQ Pivots: #ctrs: " << npts_u64 - << ", #dims: " << ndims_u64 << ", #chunks: " << n_chunks - << std::endl; - // assert((_u64) ndims_u32 == n_chunks * chunk_size); // alloc and compute transpose - tables_T = new float[256 * ndims_u64]; + tables_T = new float[256 * ndims]; for (_u64 i = 0; i < 256; i++) { - for (_u64 j = 0; j < ndims_u64; j++) { - tables_T[j * 256 + i] = tables[i * ndims_u64 + j]; + for (_u64 j = 0; j < ndims; j++) { + tables_T[j * 256 + i] = tables[i * ndims + j]; + } + } + + // added this for easy PQ-PQ squared-distance calculations + // TODO: Create only for StreamingMerger. + all_to_all_dists = new float[256 * 256 * n_chunks]; + std::memset(all_to_all_dists, 0, 256 * 256 * n_chunks * sizeof(float)); + // should perhaps optimize later + for (_u32 i = 0; i < 256; i++) { + for (_u32 j = 0; j < 256; j++) { + for (_u32 c = 0; c < n_chunks; c++) { + for (_u64 d = chunk_offsets[c]; d < chunk_offsets[c + 1]; d++) { + float diff = (tables[i * ndims + d] - tables[j * ndims + d]); + all_to_all_dists[i * 256 * n_chunks + j * n_chunks + c] += + diff * diff; + } + } } } + diskann::cout << "Finished optimizing for PQ-PQ distance compuation " + << std::endl; + } void @@ -137,7 +324,7 @@ namespace diskann { _u64 permuted_dim_in_query = rearrangement[j]; const float* centers_dim_vec = tables_T + (256 * j); for (_u64 idx = 0; idx < 256; idx++) { - // Gopal. Fixing crash in v14 machines. + // Fixing crash in v14 machines. // float diff = centers_dim_vec[idx] - // ((float) query_vec[permuted_dim_in_query] - // centroid[permuted_dim_in_query]); @@ -150,5 +337,45 @@ namespace diskann { } } } -}; + + // computes PQ distance between comp_src and comp_dsts in efficient manner + // comp_src: [nchunks] + // comp_dsts: count * [nchunks] + // dists: [count] + // TODO (perf) :: re-order computation to get better locality + void compute_distances(const _u8* comp_src, const _u8* comp_dsts, + float* dists, const _u32 count) { + std::memset(dists, 0, count * sizeof(float)); + for (_u64 i = 0; i < count; i++) { + for (_u64 c = 0; c < n_chunks; c++) { + dists[i] += + all_to_all_dists[(_u64) comp_src[c] * 256 * n_chunks + + (_u64) comp_dsts[i * n_chunks + c] * n_chunks + c]; + } + } + } + + // fp_vec: [ndims] + // out_pq_vec : [nchunks] + void deflate_vec(const float* fp_vec, _u8* out_pq_vec) { + // permute the vector according to PQ rearrangement, compute all distances + // to 256 centroids and choose the closest (for each chunk) + for (_u32 c = 0; c < n_chunks; c++) { + float closest_dist = std::numeric_limits::max(); + for (_u32 i = 0; i < 256; i++) { + float cur_dist = 0; + for (_u64 d = chunk_offsets[c]; d < chunk_offsets[c + 1]; d++) { + float diff = + (tables[i * ndims + d] - + ((float) fp_vec[rearrangement[d]] - centroid[rearrangement[d]])); + cur_dist += diff * diff; + } + if (cur_dist < closest_dist) { + closest_dist = cur_dist; + out_pq_vec[c] = (_u8) i; + } + } + } + } +}; // namespace diskann } // namespace diskann diff --git a/include/simd_utils.h b/include/simd_utils.h new file mode 100644 index 0000000000..c219485621 --- /dev/null +++ b/include/simd_utils.h @@ -0,0 +1,105 @@ +#pragma once + +#ifdef _WINDOWS +#include +#include +#include +#include +#else +#include +#endif + +namespace diskann { + static inline __m256 _mm256_mul_epi8(__m256i X) { + __m256i zero = _mm256_setzero_si256(); + + __m256i sign_x = _mm256_cmpgt_epi8(zero, X); + + __m256i xlo = _mm256_unpacklo_epi8(X, sign_x); + __m256i xhi = _mm256_unpackhi_epi8(X, sign_x); + + return _mm256_cvtepi32_ps(_mm256_add_epi32(_mm256_madd_epi16(xlo, xlo), + _mm256_madd_epi16(xhi, xhi))); + } + + static inline __m128 _mm_mulhi_epi8(__m128i X) { + __m128i zero = _mm_setzero_si128(); + __m128i sign_x = _mm_cmplt_epi8(X, zero); + __m128i xhi = _mm_unpackhi_epi8(X, sign_x); + + return _mm_cvtepi32_ps( + _mm_add_epi32(_mm_setzero_si128(), _mm_madd_epi16(xhi, xhi))); + } + + static inline __m128 _mm_mulhi_epi8_shift32(__m128i X) { + __m128i zero = _mm_setzero_si128(); + X = _mm_srli_epi64(X, 32); + __m128i sign_x = _mm_cmplt_epi8(X, zero); + __m128i xhi = _mm_unpackhi_epi8(X, sign_x); + + return _mm_cvtepi32_ps( + _mm_add_epi32(_mm_setzero_si128(), _mm_madd_epi16(xhi, xhi))); + } + static inline __m128 _mm_mul_epi8(__m128i X, __m128i Y) { + __m128i zero = _mm_setzero_si128(); + + __m128i sign_x = _mm_cmplt_epi8(X, zero); + __m128i sign_y = _mm_cmplt_epi8(Y, zero); + + __m128i xlo = _mm_unpacklo_epi8(X, sign_x); + __m128i xhi = _mm_unpackhi_epi8(X, sign_x); + __m128i ylo = _mm_unpacklo_epi8(Y, sign_y); + __m128i yhi = _mm_unpackhi_epi8(Y, sign_y); + + return _mm_cvtepi32_ps( + _mm_add_epi32(_mm_madd_epi16(xlo, ylo), _mm_madd_epi16(xhi, yhi))); + } + static inline __m128 _mm_mul_epi8(__m128i X) { + __m128i zero = _mm_setzero_si128(); + __m128i sign_x = _mm_cmplt_epi8(X, zero); + __m128i xlo = _mm_unpacklo_epi8(X, sign_x); + __m128i xhi = _mm_unpackhi_epi8(X, sign_x); + + return _mm_cvtepi32_ps( + _mm_add_epi32(_mm_madd_epi16(xlo, xlo), _mm_madd_epi16(xhi, xhi))); + } + + static inline __m128 _mm_mul32_pi8(__m128i X, __m128i Y) { + __m128i xlo = _mm_cvtepi8_epi16(X), ylo = _mm_cvtepi8_epi16(Y); + return _mm_cvtepi32_ps( + _mm_unpacklo_epi32(_mm_madd_epi16(xlo, ylo), _mm_setzero_si128())); + } + + static inline __m256 _mm256_mul_epi8(__m256i X, __m256i Y) { + __m256i zero = _mm256_setzero_si256(); + + __m256i sign_x = _mm256_cmpgt_epi8(zero, X); + __m256i sign_y = _mm256_cmpgt_epi8(zero, Y); + + __m256i xlo = _mm256_unpacklo_epi8(X, sign_x); + __m256i xhi = _mm256_unpackhi_epi8(X, sign_x); + __m256i ylo = _mm256_unpacklo_epi8(Y, sign_y); + __m256i yhi = _mm256_unpackhi_epi8(Y, sign_y); + + return _mm256_cvtepi32_ps(_mm256_add_epi32(_mm256_madd_epi16(xlo, ylo), + _mm256_madd_epi16(xhi, yhi))); + } + + static inline __m256 _mm256_mul32_pi8(__m128i X, __m128i Y) { + __m256i xlo = _mm256_cvtepi8_epi16(X), ylo = _mm256_cvtepi8_epi16(Y); + return _mm256_blend_ps(_mm256_cvtepi32_ps(_mm256_madd_epi16(xlo, ylo)), + _mm256_setzero_ps(), 252); + } + + static inline float _mm256_reduce_add_ps(__m256 x) { + /* ( x3+x7, x2+x6, x1+x5, x0+x4 ) */ + const __m128 x128 = + _mm_add_ps(_mm256_extractf128_ps(x, 1), _mm256_castps256_ps128(x)); + /* ( -, -, x1+x3+x5+x7, x0+x2+x4+x6 ) */ + const __m128 x64 = _mm_add_ps(x128, _mm_movehl_ps(x128, x128)); + /* ( -, -, -, x0+x1+x2+x3+x4+x5+x6+x7 ) */ + const __m128 x32 = _mm_add_ss(x64, _mm_shuffle_ps(x64, x64, 0x55)); + /* Conversion to float is a no-op on x86-64 */ + return _mm_cvtss_f32(x32); + } +} // namespace diskann diff --git a/include/tcmalloc/LICENSE.txt b/include/tcmalloc/LICENSE.txt new file mode 100644 index 0000000000..26c3f72caa --- /dev/null +++ b/include/tcmalloc/LICENSE.txt @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + https://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. \ No newline at end of file diff --git a/include/tcmalloc/malloc_extension.h b/include/tcmalloc/malloc_extension.h new file mode 100644 index 0000000000..1be138a538 --- /dev/null +++ b/include/tcmalloc/malloc_extension.h @@ -0,0 +1,446 @@ +// -*- Mode: C++; c-basic-offset: 2; indent-tabs-mode: nil -*- +// Copyright (c) 2005, Google Inc. +// All rights reserved. +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions are +// met: +// +// * Redistributions of source code must retain the above copyright +// notice, this list of conditions and the following disclaimer. +// * Redistributions in binary form must reproduce the above +// copyright notice, this list of conditions and the following disclaimer +// in the documentation and/or other materials provided with the +// distribution. +// * Neither the name of Google Inc. nor the names of its +// contributors may be used to endorse or promote products derived from +// this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +// --- +// Author: Sanjay Ghemawat +// +// Extra extensions exported by some malloc implementations. These +// extensions are accessed through a virtual base class so an +// application can link against a malloc that does not implement these +// extensions, and it will get default versions that do nothing. +// +// NOTE FOR C USERS: If you wish to use this functionality from within +// a C program, see malloc_extension_c.h. + +#ifndef BASE_MALLOC_EXTENSION_H_ +#define BASE_MALLOC_EXTENSION_H_ + +#include +// I can't #include config.h in this public API file, but I should +// really use configure (and make malloc_extension.h a .in file) to +// figure out if the system has stdint.h or not. But I'm lazy, so +// for now I'm assuming it's a problem only with MSVC. +#ifndef _MSC_VER +#include +#endif +#include +#include + +// Annoying stuff for windows -- makes sure clients can import these functions +#ifndef PERFTOOLS_DLL_DECL +# ifdef _WIN32 +# define PERFTOOLS_DLL_DECL __declspec(dllimport) +# else +# define PERFTOOLS_DLL_DECL +# endif +#endif + +static const int kMallocHistogramSize = 64; + +// One day, we could support other types of writers (perhaps for C?) +typedef std::string MallocExtensionWriter; + +namespace base { +struct MallocRange; +} + +// Interface to a pluggable system allocator. +class PERFTOOLS_DLL_DECL SysAllocator { + public: + SysAllocator() { + } + virtual ~SysAllocator(); + + // Allocates "size"-byte of memory from system aligned with "alignment". + // Returns NULL if failed. Otherwise, the returned pointer p up to and + // including (p + actual_size -1) have been allocated. + virtual void* Alloc(size_t size, size_t *actual_size, size_t alignment) = 0; +}; + +// The default implementations of the following routines do nothing. +// All implementations should be thread-safe; the current one +// (TCMallocImplementation) is. +class PERFTOOLS_DLL_DECL MallocExtension { + public: + virtual ~MallocExtension(); + + // Call this very early in the program execution -- say, in a global + // constructor -- to set up parameters and state needed by all + // instrumented malloc implemenatations. One example: this routine + // sets environemnt variables to tell STL to use libc's malloc() + // instead of doing its own memory management. This is safe to call + // multiple times, as long as each time is before threads start up. + static void Initialize(); + + // See "verify_memory.h" to see what these routines do + virtual bool VerifyAllMemory(); + virtual bool VerifyNewMemory(const void* p); + virtual bool VerifyArrayNewMemory(const void* p); + virtual bool VerifyMallocMemory(const void* p); + virtual bool MallocMemoryStats(int* blocks, size_t* total, + int histogram[kMallocHistogramSize]); + + // Get a human readable description of the following malloc data structures. + // - Total inuse memory by application. + // - Free memory(thread, central and page heap), + // - Freelist of central cache, each class. + // - Page heap freelist. + // The state is stored as a null-terminated string + // in a prefix of "buffer[0,buffer_length-1]". + // REQUIRES: buffer_length > 0. + virtual void GetStats(char* buffer, int buffer_length); + + // Outputs to "writer" a sample of live objects and the stack traces + // that allocated these objects. The format of the returned output + // is equivalent to the output of the heap profiler and can + // therefore be passed to "pprof". This function is equivalent to + // ReadStackTraces. The main difference is that this function returns + // serialized data appropriately formatted for use by the pprof tool. + // + // Since gperftools 2.8 heap samples are not de-duplicated by the + // library anymore. + // + // NOTE: by default, tcmalloc does not do any heap sampling, and this + // function will always return an empty sample. To get useful + // data from GetHeapSample, you must also set the environment + // variable TCMALLOC_SAMPLE_PARAMETER to a value such as 524288. + virtual void GetHeapSample(MallocExtensionWriter* writer); + + // Outputs to "writer" the stack traces that caused growth in the + // address space size. The format of the returned output is + // equivalent to the output of the heap profiler and can therefore + // be passed to "pprof". This function is equivalent to + // ReadHeapGrowthStackTraces. The main difference is that this function + // returns serialized data appropriately formatted for use by the + // pprof tool. (This does not depend on, or require, + // TCMALLOC_SAMPLE_PARAMETER.) + virtual void GetHeapGrowthStacks(MallocExtensionWriter* writer); + + // Invokes func(arg, range) for every controlled memory + // range. *range is filled in with information about the range. + // + // This is a best-effort interface useful only for performance + // analysis. The implementation may not call func at all. + typedef void (RangeFunction)(void*, const base::MallocRange*); + virtual void Ranges(void* arg, RangeFunction func); + + // ------------------------------------------------------------------- + // Control operations for getting and setting malloc implementation + // specific parameters. Some currently useful properties: + // + // generic + // ------- + // "generic.current_allocated_bytes" + // Number of bytes currently allocated by application + // This property is not writable. + // + // "generic.heap_size" + // Number of bytes in the heap == + // current_allocated_bytes + + // fragmentation + + // freed memory regions + // This property is not writable. + // + // "generic.total_physical_bytes" + // Estimate of total bytes of the physical memory usage by the + // allocator == + // current_allocated_bytes + + // fragmentation + + // metadata + // This property is not writable. + // + // tcmalloc + // -------- + // "tcmalloc.max_total_thread_cache_bytes" + // Upper limit on total number of bytes stored across all + // per-thread caches. Default: 16MB. + // + // "tcmalloc.current_total_thread_cache_bytes" + // Number of bytes used across all thread caches. + // This property is not writable. + // + // "tcmalloc.central_cache_free_bytes" + // Number of free bytes in the central cache that have been + // assigned to size classes. They always count towards virtual + // memory usage, and unless the underlying memory is swapped out + // by the OS, they also count towards physical memory usage. + // This property is not writable. + // + // "tcmalloc.transfer_cache_free_bytes" + // Number of free bytes that are waiting to be transfered between + // the central cache and a thread cache. They always count + // towards virtual memory usage, and unless the underlying memory + // is swapped out by the OS, they also count towards physical + // memory usage. This property is not writable. + // + // "tcmalloc.thread_cache_free_bytes" + // Number of free bytes in thread caches. They always count + // towards virtual memory usage, and unless the underlying memory + // is swapped out by the OS, they also count towards physical + // memory usage. This property is not writable. + // + // "tcmalloc.pageheap_free_bytes" + // Number of bytes in free, mapped pages in page heap. These + // bytes can be used to fulfill allocation requests. They + // always count towards virtual memory usage, and unless the + // underlying memory is swapped out by the OS, they also count + // towards physical memory usage. This property is not writable. + // + // "tcmalloc.pageheap_unmapped_bytes" + // Number of bytes in free, unmapped pages in page heap. + // These are bytes that have been released back to the OS, + // possibly by one of the MallocExtension "Release" calls. + // They can be used to fulfill allocation requests, but + // typically incur a page fault. They always count towards + // virtual memory usage, and depending on the OS, typically + // do not count towards physical memory usage. This property + // is not writable. + // ------------------------------------------------------------------- + + // Get the named "property"'s value. Returns true if the property + // is known. Returns false if the property is not a valid property + // name for the current malloc implementation. + // REQUIRES: property != NULL; value != NULL + virtual bool GetNumericProperty(const char* property, size_t* value); + + // Set the named "property"'s value. Returns true if the property + // is known and writable. Returns false if the property is not a + // valid property name for the current malloc implementation, or + // is not writable. + // REQUIRES: property != NULL + virtual bool SetNumericProperty(const char* property, size_t value); + + // Mark the current thread as "idle". This routine may optionally + // be called by threads as a hint to the malloc implementation that + // any thread-specific resources should be released. Note: this may + // be an expensive routine, so it should not be called too often. + // + // Also, if the code that calls this routine will go to sleep for + // a while, it should take care to not allocate anything between + // the call to this routine and the beginning of the sleep. + // + // Most malloc implementations ignore this routine. + virtual void MarkThreadIdle(); + + // Mark the current thread as "busy". This routine should be + // called after MarkThreadIdle() if the thread will now do more + // work. If this method is not called, performance may suffer. + // + // Most malloc implementations ignore this routine. + virtual void MarkThreadBusy(); + + // Gets the system allocator used by the malloc extension instance. Returns + // NULL for malloc implementations that do not support pluggable system + // allocators. + virtual SysAllocator* GetSystemAllocator(); + + // Sets the system allocator to the specified. + // + // Users could register their own system allocators for malloc implementation + // that supports pluggable system allocators, such as TCMalloc, by doing: + // alloc = new MyOwnSysAllocator(); + // MallocExtension::instance()->SetSystemAllocator(alloc); + // It's up to users whether to fall back (recommended) to the default + // system allocator (use GetSystemAllocator() above) or not. The caller is + // responsible to any necessary locking. + // See tcmalloc/system-alloc.h for the interface and + // tcmalloc/memfs_malloc.cc for the examples. + // + // It's a no-op for malloc implementations that do not support pluggable + // system allocators. + virtual void SetSystemAllocator(SysAllocator *a); + + // Try to release num_bytes of free memory back to the operating + // system for reuse. Use this extension with caution -- to get this + // memory back may require faulting pages back in by the OS, and + // that may be slow. (Currently only implemented in tcmalloc.) + virtual void ReleaseToSystem(size_t num_bytes); + + // Same as ReleaseToSystem() but release as much memory as possible. + virtual void ReleaseFreeMemory(); + + // Sets the rate at which we release unused memory to the system. + // Zero means we never release memory back to the system. Increase + // this flag to return memory faster; decrease it to return memory + // slower. Reasonable rates are in the range [0,10]. (Currently + // only implemented in tcmalloc). + virtual void SetMemoryReleaseRate(double rate); + + // Gets the release rate. Returns a value < 0 if unknown. + virtual double GetMemoryReleaseRate(); + + // Returns the estimated number of bytes that will be allocated for + // a request of "size" bytes. This is an estimate: an allocation of + // SIZE bytes may reserve more bytes, but will never reserve less. + // (Currently only implemented in tcmalloc, other implementations + // always return SIZE.) + // This is equivalent to malloc_good_size() in OS X. + virtual size_t GetEstimatedAllocatedSize(size_t size); + + // Returns the actual number N of bytes reserved by tcmalloc for the + // pointer p. The client is allowed to use the range of bytes + // [p, p+N) in any way it wishes (i.e. N is the "usable size" of this + // allocation). This number may be equal to or greater than the number + // of bytes requested when p was allocated. + // p must have been allocated by this malloc implementation, + // must not be an interior pointer -- that is, must be exactly + // the pointer returned to by malloc() et al., not some offset + // from that -- and should not have been freed yet. p may be NULL. + // (Currently only implemented in tcmalloc; other implementations + // will return 0.) + // This is equivalent to malloc_size() in OS X, malloc_usable_size() + // in glibc, and _msize() for windows. + virtual size_t GetAllocatedSize(const void* p); + + // Returns kOwned if this malloc implementation allocated the memory + // pointed to by p, or kNotOwned if some other malloc implementation + // allocated it or p is NULL. May also return kUnknownOwnership if + // the malloc implementation does not keep track of ownership. + // REQUIRES: p must be a value returned from a previous call to + // malloc(), calloc(), realloc(), memalign(), posix_memalign(), + // valloc(), pvalloc(), new, or new[], and must refer to memory that + // is currently allocated (so, for instance, you should not pass in + // a pointer after having called free() on it). + enum Ownership { + // NOTE: Enum values MUST be kept in sync with the version in + // malloc_extension_c.h + kUnknownOwnership = 0, + kOwned, + kNotOwned + }; + virtual Ownership GetOwnership(const void* p); + + // The current malloc implementation. Always non-NULL. + static MallocExtension* instance(); + + // Change the malloc implementation. Typically called by the + // malloc implementation during initialization. + static void Register(MallocExtension* implementation); + + // Returns detailed information about malloc's freelists. For each list, + // return a FreeListInfo: + struct FreeListInfo { + size_t min_object_size; + size_t max_object_size; + size_t total_bytes_free; + const char* type; + }; + // Each item in the vector refers to a different freelist. The lists + // are identified by the range of allocations that objects in the + // list can satisfy ([min_object_size, max_object_size]) and the + // type of freelist (see below). The current size of the list is + // returned in total_bytes_free (which count against a processes + // resident and virtual size). + // + // Currently supported types are: + // + // "tcmalloc.page{_unmapped}" - tcmalloc's page heap. An entry for each size + // class in the page heap is returned. Bytes in "page_unmapped" + // are no longer backed by physical memory and do not count against + // the resident size of a process. + // + // "tcmalloc.large{_unmapped}" - tcmalloc's list of objects larger + // than the largest page heap size class. Only one "large" + // entry is returned. There is no upper-bound on the size + // of objects in the large free list; this call returns + // kint64max for max_object_size. Bytes in + // "large_unmapped" are no longer backed by physical memory + // and do not count against the resident size of a process. + // + // "tcmalloc.central" - tcmalloc's central free-list. One entry per + // size-class is returned. Never unmapped. + // + // "debug.free_queue" - free objects queued by the debug allocator + // and not returned to tcmalloc. + // + // "tcmalloc.thread" - tcmalloc's per-thread caches. Never unmapped. + virtual void GetFreeListSizes(std::vector* v); + + // Get a list of stack traces of sampled allocation points. Returns + // a pointer to a "new[]-ed" result array, and stores the sample + // period in "sample_period". + // + // The state is stored as a sequence of adjacent entries + // in the returned array. Each entry has the following form: + // uintptr_t count; // Number of objects with following trace + // uintptr_t size; // Total size of objects with following trace + // uintptr_t depth; // Number of PC values in stack trace + // void* stack[depth]; // PC values that form the stack trace + // + // The list of entries is terminated by a "count" of 0. + // + // It is the responsibility of the caller to "delete[]" the returned array. + // + // May return NULL to indicate no results. + // + // This is an internal extension. Callers should use the more + // convenient "GetHeapSample(string*)" method defined above. + virtual void** ReadStackTraces(int* sample_period); + + // Like ReadStackTraces(), but returns stack traces that caused growth + // in the address space size. + virtual void** ReadHeapGrowthStackTraces(); + + // Returns the size in bytes of the calling threads cache. + virtual size_t GetThreadCacheSize(); + + // Like MarkThreadIdle, but does not destroy the internal data + // structures of the thread cache. When the thread resumes, it wil + // have an empty cache but will not need to pay to reconstruct the + // cache data structures. + virtual void MarkThreadTemporarilyIdle(); +}; + +namespace base { + +// Information passed per range. More fields may be added later. +struct MallocRange { + enum Type { + INUSE, // Application is using this range + FREE, // Range is currently free + UNMAPPED, // Backing physical memory has been returned to the OS + UNKNOWN + // More enum values may be added in the future + }; + + uintptr_t address; // Address of range + size_t length; // Byte length of range + Type type; // Type of this range + double fraction; // Fraction of range that is being used (0 if !INUSE) + + // Perhaps add the following: + // - stack trace if this range was sampled + // - heap growth stack trace if applicable to this range + // - age when allocated (for inuse) or freed (if not in use) +}; + +} // namespace base + +#endif // BASE_MALLOC_EXTENSION_H_ diff --git a/include/threadpool.h b/include/threadpool.h new file mode 100644 index 0000000000..bfaaeffe1f --- /dev/null +++ b/include/threadpool.h @@ -0,0 +1,115 @@ +/***************************************************************************** + +Copyright (c) 2012 Jakob Progsch, Vclav Zeman + +This software is provided 'as-is', without any express or implied +warranty. In no event will the authors be held liable for any damages +arising from the use of this software. + +Permission is granted to anyone to use this software for any purpose, +including commercial applications, and to alter it and redistribute it +freely, subject to the following restrictions: + + 1. The origin of this software must not be misrepresented; you must not + claim that you wrote the original software. If you use this software + in a product, an acknowledgment in the product documentation would be + appreciated but is not required. + + 2. Altered source versions must be plainly marked as such, and must not be + misrepresented as being the original software. + + 3. This notice may not be removed or altered from any source + distribution. + + *******************************************************************************/ + +#ifndef THREAD_POOL_H +#define THREAD_POOL_H + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +class ThreadPool { + public: + ThreadPool(size_t); + template + auto enqueue(F&& f, Args&&... args) + -> std::future::type>; + ~ThreadPool(); + + private: + // need to keep track of threads so we can join them + std::vector workers; + // the task queue + std::queue> tasks; + + // synchronization + std::mutex queue_mutex; + std::condition_variable condition; + bool stop; +}; + +// the constructor just launches some amount of workers +inline ThreadPool::ThreadPool(size_t threads) : stop(false) { + for (size_t i = 0; i < threads; ++i) + workers.emplace_back([this] { + for (;;) { + std::function task; + + { + std::unique_lock lock(this->queue_mutex); + this->condition.wait( + lock, [this] { return this->stop || !this->tasks.empty(); }); + if (this->stop && this->tasks.empty()) + return; + task = std::move(this->tasks.front()); + this->tasks.pop(); + } + + task(); + } + }); +} + +// add new work item to the pool +template +auto ThreadPool::enqueue(F&& f, Args&&... args) + -> std::future::type> { + using return_type = typename std::result_of::type; + + auto task = std::make_shared>( + std::bind(std::forward(f), std::forward(args)...)); + + std::future res = task->get_future(); + { + std::unique_lock lock(queue_mutex); + + // don't allow enqueueing after stopping the pool + if (stop) + throw std::runtime_error("enqueue on stopped ThreadPool"); + + tasks.emplace([task]() { (*task)(); }); + } + condition.notify_one(); + return res; +} + +// the destructor joins all threads +inline ThreadPool::~ThreadPool() { + { + std::unique_lock lock(queue_mutex); + stop = true; + } + condition.notify_all(); + for (std::thread& worker : workers) + worker.join(); +} + +#endif \ No newline at end of file diff --git a/include/timer.h b/include/timer.h index 4671c33beb..bf52ed8835 100644 --- a/include/timer.h +++ b/include/timer.h @@ -22,4 +22,4 @@ namespace diskann { .count(); } }; -} +} // namespace diskann diff --git a/include/tsl/CMakeLists.txt b/include/tsl/CMakeLists.txt new file mode 100644 index 0000000000..924f38f022 --- /dev/null +++ b/include/tsl/CMakeLists.txt @@ -0,0 +1,78 @@ +cmake_minimum_required(VERSION 3.1) +include(GNUInstallDirs) + + +project(tsl-robin-map VERSION 0.6.3) + +add_library(robin_map INTERFACE) +# Use tsl::robin_map as target, more consistent with other libraries conventions (Boost, Qt, ...) +add_library(tsl::robin_map ALIAS robin_map) + +target_include_directories(robin_map INTERFACE + "$" + "$") + +list(APPEND headers "${CMAKE_CURRENT_SOURCE_DIR}/include/tsl/robin_growth_policy.h" + "${CMAKE_CURRENT_SOURCE_DIR}/include/tsl/robin_hash.h" + "${CMAKE_CURRENT_SOURCE_DIR}/include/tsl/robin_map.h" + "${CMAKE_CURRENT_SOURCE_DIR}/include/tsl/robin_set.h") +target_sources(robin_map INTERFACE "$") + +if(MSVC) + target_sources(robin_map INTERFACE + "$" + "$") +endif() + + + + +# Installation (only compatible with CMake version >= 3.3) +if(${CMAKE_VERSION} VERSION_GREATER "3.2") + include(CMakePackageConfigHelpers) + + ## Install include directory and potential natvis file + install(DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/include/tsl" + DESTINATION "${CMAKE_INSTALL_INCLUDEDIR}") + + if(MSVC) + install(FILES "${CMAKE_CURRENT_SOURCE_DIR}/tsl-robin-map.natvis" + DESTINATION "${CMAKE_INSTALL_DATAROOTDIR}") + endif() + + + + ## Create and install tsl-robin-mapConfig.cmake + configure_package_config_file("${CMAKE_CURRENT_SOURCE_DIR}/cmake/tsl-robin-mapConfig.cmake.in" + "${CMAKE_CURRENT_BINARY_DIR}/tsl-robin-mapConfig.cmake" + INSTALL_DESTINATION "${CMAKE_INSTALL_DATAROOTDIR}/cmake/tsl-robin-map") + + install(FILES "${CMAKE_CURRENT_BINARY_DIR}/tsl-robin-mapConfig.cmake" + DESTINATION "${CMAKE_INSTALL_DATAROOTDIR}/cmake/tsl-robin-map") + + + + ## Create and install tsl-robin-mapTargets.cmake + install(TARGETS robin_map + EXPORT tsl-robin-mapTargets) + + install(EXPORT tsl-robin-mapTargets + NAMESPACE tsl:: + DESTINATION "${CMAKE_INSTALL_DATAROOTDIR}/cmake/tsl-robin-map") + + + + ## Create and install tsl-robin-mapConfigVersion.cmake + # tsl-robin-map is header-only and does not depend on the architecture. + # Remove CMAKE_SIZEOF_VOID_P from tsl-robin-mapConfigVersion.cmake so that a + # tsl-robin-mapConfig.cmake generated for a 64 bit target can be used for 32 bit + # targets and vice versa. + set(CMAKE_SIZEOF_VOID_P_BACKUP ${CMAKE_SIZEOF_VOID_P}) + unset(CMAKE_SIZEOF_VOID_P) + write_basic_package_version_file("${CMAKE_CURRENT_BINARY_DIR}/tsl-robin-mapConfigVersion.cmake" + COMPATIBILITY SameMajorVersion) + set(CMAKE_SIZEOF_VOID_P ${CMAKE_SIZEOF_VOID_P_BACKUP}) + + install(FILES "${CMAKE_CURRENT_BINARY_DIR}/tsl-robin-mapConfigVersion.cmake" + DESTINATION "${CMAKE_INSTALL_DATAROOTDIR}/cmake/tsl-robin-map") +endif() diff --git a/include/tsl/LICENSE b/include/tsl/LICENSE new file mode 100644 index 0000000000..e9c5ae95f3 --- /dev/null +++ b/include/tsl/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2017 Thibaut Goetghebuer-Planchon + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/include/tsl/README.md b/include/tsl/README.md new file mode 100644 index 0000000000..612b277b13 --- /dev/null +++ b/include/tsl/README.md @@ -0,0 +1,483 @@ +[![Build Status](https://travis-ci.org/Tessil/robin-map.svg?branch=master)](https://travis-ci.org/Tessil/robin-map) [![Build status](https://ci.appveyor.com/api/projects/status/lo79n4ya4nta79q4/branch/master?svg=true)](https://ci.appveyor.com/project/Tessil/robin-map/branch/master) + +## A C++ implementation of a fast hash map and hash set using robin hood hashing + +The robin-map library is a C++ implementation of a fast hash map and hash set using open-addressing and linear robin hood hashing with backward shift deletion to resolve collisions. + +Four classes are provided: `tsl::robin_map`, `tsl::robin_set`, `tsl::robin_pg_map` and `tsl::robin_pg_set`. The first two are faster and use a power of two growth policy, the last two use a prime growth policy instead and are able to cope better with a poor hash function. Use the prime version if there is a chance of repeating patterns in the lower bits of your hash (e.g. you are storing pointers with an identity hash function). See [GrowthPolicy](#growth-policy) for details. + +A **benchmark** of `tsl::robin_map` against other hash maps may be found [here](https://tessil.github.io/2016/08/29/benchmark-hopscotch-map.html). This page also gives some advices on which hash table structure you should try for your use case (useful if you are a bit lost with the multiple hash tables implementations in the `tsl` namespace). + +### Key features + +- Header-only library, just add the [include](include/) directory to your include path and you are ready to go. If you use CMake, you can also use the `tsl::robin_map` exported target from the [CMakeLists.txt](CMakeLists.txt). +- Fast hash table, check the [benchmark](https://tessil.github.io/2016/08/29/benchmark-hopscotch-map.html) for some numbers. +- Support for move-only and non-default constructible key/value. +- Support for heterogeneous lookups allowing the usage of `find` with a type different than `Key` (e.g. if you have a map that uses `std::unique_ptr` as key, you can use a `foo*` or a `std::uintptr_t` as key parameter to `find` without constructing a `std::unique_ptr`, see [example](#heterogeneous-lookups)). +- No need to reserve any sentinel value from the keys. +- Possibility to store the hash value alongside the stored key-value for faster rehash and lookup if the hash or the key equal functions are expensive to compute. Note that hash may be stored even if not asked explicitly when the library can detect that it will have no impact on the size of the structure in memory due to alignment. See the [StoreHash](https://tessil.github.io/robin-map/classtsl_1_1robin__map.html#details) template parameter for details. +- If the hash is known before a lookup, it is possible to pass it as parameter to speed-up the lookup (see `precalculated_hash` parameter in [API](https://tessil.github.io/robin-map/classtsl_1_1robin__map.html#a35021b11aabb61820236692a54b3a0f8)). +- Support for efficient serialization and deserialization (see [example](#serialization) and the `serialize/deserialize` methods in the [API](https://tessil.github.io/robin-map/classtsl_1_1robin__map.html) for details). +- The library can be used with exceptions disabled (through `-fno-exceptions` option on Clang and GCC, without an `/EH` option on MSVC or simply by defining `TSL_NO_EXCEPTIONS`). `std::terminate` is used in replacement of the `throw` instruction when exceptions are disabled. +- API closely similar to `std::unordered_map` and `std::unordered_set`. + +### Differences compared to `std::unordered_map` + +`tsl::robin_map` tries to have an interface similar to `std::unordered_map`, but some differences exist. +- The **strong exception guarantee only holds** if the following statement is true `std::is_nothrow_swappable::value && std::is_nothrow_move_constructible::value` (where `value_type` is `Key` for `tsl::robin_set` and `std::pair` for `tsl::robin_map`). Otherwise if an exception is thrown during the swap or the move, the structure may end up in a undefined state. Note that per the standard, a `value_type` with a noexcept copy constructor and no move constructor also satisfies this condition and will thus guarantee the strong exception guarantee for the structure (see [API](https://tessil.github.io/robin-map/classtsl_1_1robin__map.html#details) for details). +- The type `Key`, and also `T` in case of map, must be swappable. They must also be copy and/or move constructible. +- Iterator invalidation doesn't behave in the same way, any operation modifying the hash table invalidate them (see [API](https://tessil.github.io/robin-map/classtsl_1_1robin__map.html#details) for details). +- References and pointers to keys or values in the map are invalidated in the same way as iterators to these keys-values. +- For iterators of `tsl::robin_map`, `operator*()` and `operator->()` return a reference and a pointer to `const std::pair` instead of `std::pair` making the value `T` not modifiable. To modify the value you have to call the `value()` method of the iterator to get a mutable reference. Example: +```c++ +tsl::robin_map map = {{1, 1}, {2, 1}, {3, 1}}; +for(auto it = map.begin(); it != map.end(); ++it) { + //it->second = 2; // Illegal + it.value() = 2; // Ok +} +``` +- No support for some buckets related methods (like `bucket_size`, `bucket`, ...). + +These differences also apply between `std::unordered_set` and `tsl::robin_set`. + +Thread-safety guarantees are the same as `std::unordered_map/set` (i.e. possible to have multiple readers with no writer). + +### Growth policy + +The library supports multiple growth policies through the `GrowthPolicy` template parameter. Three policies are provided by the library but you can easily implement your own if needed. + +* **[tsl::rh::power_of_two_growth_policy.](https://tessil.github.io/robin-map/classtsl_1_1rh_1_1power__of__two__growth__policy.html)** Default policy used by `tsl::robin_map/set`. This policy keeps the size of the bucket array of the hash table to a power of two. This constraint allows the policy to avoid the usage of the slow modulo operation to map a hash to a bucket, instead of hash % 2n, it uses hash & (2n - 1) (see [fast modulo](https://en.wikipedia.org/wiki/Modulo_operation#Performance_issues)). Fast but this may cause a lot of collisions with a poor hash function as the modulo with a power of two only masks the most significant bits in the end. +* **[tsl::rh::prime_growth_policy.](https://tessil.github.io/robin-map/classtsl_1_1rh_1_1prime__growth__policy.html)** Default policy used by `tsl::robin_pg_map/set`. The policy keeps the size of the bucket array of the hash table to a prime number. When mapping a hash to a bucket, using a prime number as modulo will result in a better distribution of the hash across the buckets even with a poor hash function. To allow the compiler to optimize the modulo operation, the policy use a lookup table with constant primes modulos (see [API](https://tessil.github.io/robin-map/classtsl_1_1rh_1_1prime__growth__policy.html#details) for details). Slower than `tsl::rh::power_of_two_growth_policy` but more secure. +* **[tsl::rh::mod_growth_policy.](https://tessil.github.io/robin-map/classtsl_1_1rh_1_1mod__growth__policy.html)** The policy grows the map by a customizable growth factor passed in parameter. It then just use the modulo operator to map a hash to a bucket. Slower but more flexible. + + +To implement your own policy, you have to implement the following interface. + +```c++ +struct custom_policy { + // Called on hash table construction and rehash, min_bucket_count_in_out is the minimum buckets + // that the hash table needs. The policy can change it to a higher number of buckets if needed + // and the hash table will use this value as bucket count. If 0 bucket is asked, then the value + // must stay at 0. + explicit custom_policy(std::size_t& min_bucket_count_in_out); + + // Return the bucket [0, bucket_count()) to which the hash belongs. + // If bucket_count() is 0, it must always return 0. + std::size_t bucket_for_hash(std::size_t hash) const noexcept; + + // Return the number of buckets that should be used on next growth + std::size_t next_bucket_count() const; + + // Maximum number of buckets supported by the policy + std::size_t max_bucket_count() const; + + // Reset the growth policy as if the policy was created with a bucket count of 0. + // After a clear, the policy must always return 0 when bucket_for_hash() is called. + void clear() noexcept; +} +``` + +### Installation + +To use robin-map, just add the [include](include/) directory to your include path. It is a **header-only** library. + +If you use CMake, you can also use the `tsl::robin_map` exported target from the [CMakeLists.txt](CMakeLists.txt) with `target_link_libraries`. +```cmake +# Example where the robin-map project is stored in a third-party directory +add_subdirectory(third-party/robin-map) +target_link_libraries(your_target PRIVATE tsl::robin_map) +``` + +If the project has been installed through `make install`, you can also use `find_package(tsl-robin-map REQUIRED)` instead of `add_subdirectory`. + +The library is available in [vcpkg](https://github.com/Microsoft/vcpkg/tree/master/ports/robin-map) and [conan](https://bintray.com/tessil/tsl/tsl-robin-map%3Atessil). It's also present in [Debian](https://packages.debian.org/buster/robin-map-dev), [Ubuntu](https://packages.ubuntu.com/disco/robin-map-dev) and [Fedora](https://apps.fedoraproject.org/packages/robin-map-devel) package repositories. + +The code should work with any C++11 standard-compliant compiler and has been tested with GCC 4.8.4, Clang 3.5.0 and Visual Studio 2015. + +To run the tests you will need the Boost Test library and CMake. + +```bash +git clone https://github.com/Tessil/robin-map.git +cd robin-map/tests +mkdir build +cd build +cmake .. +cmake --build . +./tsl_robin_map_tests +``` + +### Usage + +The API can be found [here](https://tessil.github.io/robin-map/). + +All methods are not documented yet, but they replicate the behavior of the ones in `std::unordered_map` and `std::unordered_set`, except if specified otherwise. + + +### Example + +```c++ +#include +#include +#include +#include +#include + +int main() { + tsl::robin_map map = {{"a", 1}, {"b", 2}}; + map["c"] = 3; + map["d"] = 4; + + map.insert({"e", 5}); + map.erase("b"); + + for(auto it = map.begin(); it != map.end(); ++it) { + //it->second += 2; // Not valid. + it.value() += 2; + } + + // {d, 6} {a, 3} {e, 7} {c, 5} + for(const auto& key_value : map) { + std::cout << "{" << key_value.first << ", " << key_value.second << "}" << std::endl; + } + + + if(map.find("a") != map.end()) { + std::cout << "Found \"a\"." << std::endl; + } + + const std::size_t precalculated_hash = std::hash()("a"); + // If we already know the hash beforehand, we can pass it in parameter to speed-up lookups. + if(map.find("a", precalculated_hash) != map.end()) { + std::cout << "Found \"a\" with hash " << precalculated_hash << "." << std::endl; + } + + + /* + * Calculating the hash and comparing two std::string may be slow. + * We can store the hash of each std::string in the hash map to make + * the inserts and lookups faster by setting StoreHash to true. + */ + tsl::robin_map, + std::equal_to, + std::allocator>, + true> map2; + + map2["a"] = 1; + map2["b"] = 2; + + // {a, 1} {b, 2} + for(const auto& key_value : map2) { + std::cout << "{" << key_value.first << ", " << key_value.second << "}" << std::endl; + } + + + + + tsl::robin_set set; + set.insert({1, 9, 0}); + set.insert({2, -1, 9}); + + // {0} {1} {2} {9} {-1} + for(const auto& key : set) { + std::cout << "{" << key << "}" << std::endl; + } +} +``` + +#### Heterogeneous lookups + +Heterogeneous overloads allow the usage of other types than `Key` for lookup and erase operations as long as the used types are hashable and comparable to `Key`. + +To activate the heterogeneous overloads in `tsl::robin_map/set`, the qualified-id `KeyEqual::is_transparent` must be valid. It works the same way as for [`std::map::find`](http://en.cppreference.com/w/cpp/container/map/find). You can either use [`std::equal_to<>`](http://en.cppreference.com/w/cpp/utility/functional/equal_to_void) or define your own function object. + +Both `KeyEqual` and `Hash` will need to be able to deal with the different types. + +```c++ +#include +#include +#include +#include + + +struct employee { + employee(int id, std::string name) : m_id(id), m_name(std::move(name)) { + } + + // Either we include the comparators in the class and we use `std::equal_to<>`... + friend bool operator==(const employee& empl, int empl_id) { + return empl.m_id == empl_id; + } + + friend bool operator==(int empl_id, const employee& empl) { + return empl_id == empl.m_id; + } + + friend bool operator==(const employee& empl1, const employee& empl2) { + return empl1.m_id == empl2.m_id; + } + + + int m_id; + std::string m_name; +}; + +// ... or we implement a separate class to compare employees. +struct equal_employee { + using is_transparent = void; + + bool operator()(const employee& empl, int empl_id) const { + return empl.m_id == empl_id; + } + + bool operator()(int empl_id, const employee& empl) const { + return empl_id == empl.m_id; + } + + bool operator()(const employee& empl1, const employee& empl2) const { + return empl1.m_id == empl2.m_id; + } +}; + +struct hash_employee { + std::size_t operator()(const employee& empl) const { + return std::hash()(empl.m_id); + } + + std::size_t operator()(int id) const { + return std::hash()(id); + } +}; + + +int main() { + // Use std::equal_to<> which will automatically deduce and forward the parameters + tsl::robin_map> map; + map.insert({employee(1, "John Doe"), 2001}); + map.insert({employee(2, "Jane Doe"), 2002}); + map.insert({employee(3, "John Smith"), 2003}); + + // John Smith 2003 + auto it = map.find(3); + if(it != map.end()) { + std::cout << it->first.m_name << " " << it->second << std::endl; + } + + map.erase(1); + + + + // Use a custom KeyEqual which has an is_transparent member type + tsl::robin_map map2; + map2.insert({employee(4, "Johnny Doe"), 2004}); + + // 2004 + std::cout << map2.at(4) << std::endl; +} +``` + + +#### Serialization + +The library provides an efficient way to serialize and deserialize a map or a set so that it can be saved to a file or send through the network. +To do so, it requires the user to provide a function object for both serialization and deserialization. + +```c++ +struct serializer { + // Must support the following types for U: std::int16_t, std::uint32_t, + // std::uint64_t, float and std::pair if a map is used or Key for + // a set. + template + void operator()(const U& value); +}; +``` + +```c++ +struct deserializer { + // Must support the following types for U: std::int16_t, std::uint32_t, + // std::uint64_t, float and std::pair if a map is used or Key for + // a set. + template + U operator()(); +}; +``` + +Note that the implementation leaves binary compatibility (endianness, float binary representation, size of int, ...) of the types it serializes/deserializes in the hands of the provided function objects if compatibility is required. + +More details regarding the `serialize` and `deserialize` methods can be found in the [API](https://tessil.github.io/robin-map/classtsl_1_1robin__map.html). + +```c++ +#include +#include +#include +#include +#include + + +class serializer { +public: + serializer(const char* file_name) { + m_ostream.exceptions(m_ostream.badbit | m_ostream.failbit); + m_ostream.open(file_name, std::ios::binary); + } + + template::value>::type* = nullptr> + void operator()(const T& value) { + m_ostream.write(reinterpret_cast(&value), sizeof(T)); + } + + void operator()(const std::pair& value) { + (*this)(value.first); + (*this)(value.second); + } + +private: + std::ofstream m_ostream; +}; + +class deserializer { +public: + deserializer(const char* file_name) { + m_istream.exceptions(m_istream.badbit | m_istream.failbit | m_istream.eofbit); + m_istream.open(file_name, std::ios::binary); + } + + template + T operator()() { + T value; + deserialize(value); + + return value; + } + +private: + template::value>::type* = nullptr> + void deserialize(T& value) { + m_istream.read(reinterpret_cast(&value), sizeof(T)); + } + + void deserialize(std::pair& value) { + deserialize(value.first); + deserialize(value.second); + } + +private: + std::ifstream m_istream; +}; + + +int main() { + const tsl::robin_map map = {{1, -1}, {2, -2}, {3, -3}, {4, -4}}; + + + const char* file_name = "robin_map.data"; + { + serializer serial(file_name); + map.serialize(serial); + } + + { + deserializer dserial(file_name); + auto map_deserialized = tsl::robin_map::deserialize(dserial); + + assert(map == map_deserialized); + } + + { + deserializer dserial(file_name); + + /** + * If the serialized and deserialized map are hash compatibles (see conditions in API), + * setting the argument to true speed-up the deserialization process as we don't have + * to recalculate the hash of each key. We also know how much space each bucket needs. + */ + const bool hash_compatible = true; + auto map_deserialized = + tsl::robin_map::deserialize(dserial, hash_compatible); + + assert(map == map_deserialized); + } +} +``` + +##### Serialization with Boost Serialization and compression with zlib + +It is possible to use a serialization library to avoid the boilerplate. + +The following example uses Boost Serialization with the Boost zlib compression stream to reduce the size of the resulting serialized file. The example requires C++20 due to the usage of the template parameter list syntax in lambdas, but it can be adapted to less recent versions. + +```c++ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + + +namespace boost { namespace serialization { + template + void serialize(Archive & ar, tsl::robin_map& map, const unsigned int version) { + split_free(ar, map, version); + } + + template + void save(Archive & ar, const tsl::robin_map& map, const unsigned int /*version*/) { + auto serializer = [&ar](const auto& v) { ar & v; }; + map.serialize(serializer); + } + + template + void load(Archive & ar, tsl::robin_map& map, const unsigned int /*version*/) { + auto deserializer = [&ar]() { U u; ar & u; return u; }; + map = tsl::robin_map::deserialize(deserializer); + } +}} + + +int main() { + tsl::robin_map map = {{1, -1}, {2, -2}, {3, -3}, {4, -4}}; + + + const char* file_name = "robin_map.data"; + { + std::ofstream ofs; + ofs.exceptions(ofs.badbit | ofs.failbit); + ofs.open(file_name, std::ios::binary); + + boost::iostreams::filtering_ostream fo; + fo.push(boost::iostreams::zlib_compressor()); + fo.push(ofs); + + boost::archive::binary_oarchive oa(fo); + + oa << map; + } + + { + std::ifstream ifs; + ifs.exceptions(ifs.badbit | ifs.failbit | ifs.eofbit); + ifs.open(file_name, std::ios::binary); + + boost::iostreams::filtering_istream fi; + fi.push(boost::iostreams::zlib_decompressor()); + fi.push(ifs); + + boost::archive::binary_iarchive ia(fi); + + tsl::robin_map map_deserialized; + ia >> map_deserialized; + + assert(map == map_deserialized); + } +} +``` + +### License + +The code is licensed under the MIT license, see the [LICENSE file](LICENSE) for details. diff --git a/include/tsl/include/tsl/robin_growth_policy.h b/include/tsl/include/tsl/robin_growth_policy.h new file mode 100644 index 0000000000..ed17c840c4 --- /dev/null +++ b/include/tsl/include/tsl/robin_growth_policy.h @@ -0,0 +1,348 @@ +/** + * MIT License + * + * Copyright (c) 2017 Thibaut Goetghebuer-Planchon + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in all + * copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + */ +#ifndef TSL_ROBIN_GROWTH_POLICY_H +#define TSL_ROBIN_GROWTH_POLICY_H + + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + + +#ifdef TSL_DEBUG +# define tsl_rh_assert(expr) assert(expr) +#else +# define tsl_rh_assert(expr) (static_cast(0)) +#endif + + +/** + * If exceptions are enabled, throw the exception passed in parameter, otherwise call std::terminate. + */ +#if (defined(__cpp_exceptions) || defined(__EXCEPTIONS) || (defined (_MSC_VER) && defined (_CPPUNWIND))) && !defined(TSL_NO_EXCEPTIONS) +# define TSL_RH_THROW_OR_TERMINATE(ex, msg) throw ex(msg) +#else +# define TSL_RH_NO_EXCEPTIONS +# ifdef NDEBUG +# define TSL_RH_THROW_OR_TERMINATE(ex, msg) std::terminate() +# else +# include +# define TSL_RH_THROW_OR_TERMINATE(ex, msg) do { std::cerr << msg << std::endl; std::terminate(); } while(0) +# endif +#endif + + +#if defined(__GNUC__) || defined(__clang__) +# define TSL_RH_LIKELY(exp) (__builtin_expect(!!(exp), true)) +#else +# define TSL_RH_LIKELY(exp) (exp) +#endif + + +namespace tsl { +namespace rh { + +/** + * Grow the hash table by a factor of GrowthFactor keeping the bucket count to a power of two. It allows + * the table to use a mask operation instead of a modulo operation to map a hash to a bucket. + * + * GrowthFactor must be a power of two >= 2. + */ +template +class power_of_two_growth_policy { +public: + /** + * Called on the hash table creation and on rehash. The number of buckets for the table is passed in parameter. + * This number is a minimum, the policy may update this value with a higher value if needed (but not lower). + * + * If 0 is given, min_bucket_count_in_out must still be 0 after the policy creation and + * bucket_for_hash must always return 0 in this case. + */ + explicit power_of_two_growth_policy(std::size_t& min_bucket_count_in_out) { + if(min_bucket_count_in_out > max_bucket_count()) { + TSL_RH_THROW_OR_TERMINATE(std::length_error, "The hash table exceeds its maximum size."); + } + + if(min_bucket_count_in_out > 0) { + min_bucket_count_in_out = round_up_to_power_of_two(min_bucket_count_in_out); + m_mask = min_bucket_count_in_out - 1; + } + else { + m_mask = 0; + } + } + + /** + * Return the bucket [0, bucket_count()) to which the hash belongs. + * If bucket_count() is 0, it must always return 0. + */ + std::size_t bucket_for_hash(std::size_t hash) const noexcept { + return hash & m_mask; + } + + /** + * Return the number of buckets that should be used on next growth. + */ + std::size_t next_bucket_count() const { + if((m_mask + 1) > max_bucket_count() / GrowthFactor) { + TSL_RH_THROW_OR_TERMINATE(std::length_error, "The hash table exceeds its maximum size."); + } + + return (m_mask + 1) * GrowthFactor; + } + + /** + * Return the maximum number of buckets supported by the policy. + */ + std::size_t max_bucket_count() const { + // Largest power of two. + return (std::numeric_limits::max() / 2) + 1; + } + + /** + * Reset the growth policy as if it was created with a bucket count of 0. + * After a clear, the policy must always return 0 when bucket_for_hash is called. + */ + void clear() noexcept { + m_mask = 0; + } + +private: + static std::size_t round_up_to_power_of_two(std::size_t value) { + if(is_power_of_two(value)) { + return value; + } + + if(value == 0) { + return 1; + } + + --value; + for(std::size_t i = 1; i < sizeof(std::size_t) * CHAR_BIT; i *= 2) { + value |= value >> i; + } + + return value + 1; + } + + static constexpr bool is_power_of_two(std::size_t value) { + return value != 0 && (value & (value - 1)) == 0; + } + +protected: + static_assert(is_power_of_two(GrowthFactor) && GrowthFactor >= 2, "GrowthFactor must be a power of two >= 2."); + + std::size_t m_mask; +}; + + +/** + * Grow the hash table by GrowthFactor::num / GrowthFactor::den and use a modulo to map a hash + * to a bucket. Slower but it can be useful if you want a slower growth. + */ +template> +class mod_growth_policy { +public: + explicit mod_growth_policy(std::size_t& min_bucket_count_in_out) { + if(min_bucket_count_in_out > max_bucket_count()) { + TSL_RH_THROW_OR_TERMINATE(std::length_error, "The hash table exceeds its maximum size."); + } + + if(min_bucket_count_in_out > 0) { + m_mod = min_bucket_count_in_out; + } + else { + m_mod = 1; + } + } + + std::size_t bucket_for_hash(std::size_t hash) const noexcept { + return hash % m_mod; + } + + std::size_t next_bucket_count() const { + if(m_mod == max_bucket_count()) { + TSL_RH_THROW_OR_TERMINATE(std::length_error, "The hash table exceeds its maximum size."); + } + + const double next_bucket_count = std::ceil(double(m_mod) * REHASH_SIZE_MULTIPLICATION_FACTOR); + if(!std::isnormal(next_bucket_count)) { + TSL_RH_THROW_OR_TERMINATE(std::length_error, "The hash table exceeds its maximum size."); + } + + if(next_bucket_count > double(max_bucket_count())) { + return max_bucket_count(); + } + else { + return std::size_t(next_bucket_count); + } + } + + std::size_t max_bucket_count() const { + return MAX_BUCKET_COUNT; + } + + void clear() noexcept { + m_mod = 1; + } + +private: + static constexpr double REHASH_SIZE_MULTIPLICATION_FACTOR = 1.0 * GrowthFactor::num / GrowthFactor::den; + static const std::size_t MAX_BUCKET_COUNT = + std::size_t(double( + std::numeric_limits::max() / REHASH_SIZE_MULTIPLICATION_FACTOR + )); + + static_assert(REHASH_SIZE_MULTIPLICATION_FACTOR >= 1.1, "Growth factor should be >= 1.1."); + + std::size_t m_mod; +}; + + + +namespace detail { + +#if SIZE_MAX >= ULLONG_MAX +#define TSL_RH_NB_PRIMES 51 +#elif SIZE_MAX >= ULONG_MAX +#define TSL_RH_NB_PRIMES 40 +#else +#define TSL_RH_NB_PRIMES 23 +#endif + +static constexpr const std::array PRIMES = {{ + 1u, 5u, 17u, 29u, 37u, 53u, 67u, 79u, 97u, 131u, 193u, 257u, 389u, 521u, 769u, 1031u, + 1543u, 2053u, 3079u, 6151u, 12289u, 24593u, 49157u, +#if SIZE_MAX >= ULONG_MAX + 98317ul, 196613ul, 393241ul, 786433ul, 1572869ul, 3145739ul, 6291469ul, 12582917ul, + 25165843ul, 50331653ul, 100663319ul, 201326611ul, 402653189ul, 805306457ul, 1610612741ul, + 3221225473ul, 4294967291ul, +#endif +#if SIZE_MAX >= ULLONG_MAX + 6442450939ull, 12884901893ull, 25769803751ull, 51539607551ull, 103079215111ull, 206158430209ull, + 412316860441ull, 824633720831ull, 1649267441651ull, 3298534883309ull, 6597069766657ull, +#endif +}}; + +template +static constexpr std::size_t mod(std::size_t hash) { return hash % PRIMES[IPrime]; } + +// MOD_PRIME[iprime](hash) returns hash % PRIMES[iprime]. This table allows for faster modulo as the +// compiler can optimize the modulo code better with a constant known at the compilation. +static constexpr const std::array MOD_PRIME = {{ + &mod<0>, &mod<1>, &mod<2>, &mod<3>, &mod<4>, &mod<5>, &mod<6>, &mod<7>, &mod<8>, &mod<9>, &mod<10>, + &mod<11>, &mod<12>, &mod<13>, &mod<14>, &mod<15>, &mod<16>, &mod<17>, &mod<18>, &mod<19>, &mod<20>, + &mod<21>, &mod<22>, +#if SIZE_MAX >= ULONG_MAX + &mod<23>, &mod<24>, &mod<25>, &mod<26>, &mod<27>, &mod<28>, &mod<29>, &mod<30>, &mod<31>, &mod<32>, + &mod<33>, &mod<34>, &mod<35>, &mod<36>, &mod<37> , &mod<38>, &mod<39>, +#endif +#if SIZE_MAX >= ULLONG_MAX + &mod<40>, &mod<41>, &mod<42>, &mod<43>, &mod<44>, &mod<45>, &mod<46>, &mod<47>, &mod<48>, &mod<49>, + &mod<50>, +#endif +}}; + +} + +/** + * Grow the hash table by using prime numbers as bucket count. Slower than tsl::rh::power_of_two_growth_policy in + * general but will probably distribute the values around better in the buckets with a poor hash function. + * + * To allow the compiler to optimize the modulo operation, a lookup table is used with constant primes numbers. + * + * With a switch the code would look like: + * \code + * switch(iprime) { // iprime is the current prime of the hash table + * case 0: hash % 5ul; + * break; + * case 1: hash % 17ul; + * break; + * case 2: hash % 29ul; + * break; + * ... + * } + * \endcode + * + * Due to the constant variable in the modulo the compiler is able to optimize the operation + * by a series of multiplications, substractions and shifts. + * + * The 'hash % 5' could become something like 'hash - (hash * 0xCCCCCCCD) >> 34) * 5' in a 64 bits environment. + */ +class prime_growth_policy { +public: + explicit prime_growth_policy(std::size_t& min_bucket_count_in_out) { + auto it_prime = std::lower_bound(detail::PRIMES.begin(), + detail::PRIMES.end(), min_bucket_count_in_out); + if(it_prime == detail::PRIMES.end()) { + TSL_RH_THROW_OR_TERMINATE(std::length_error, "The hash table exceeds its maximum size."); + } + + m_iprime = static_cast(std::distance(detail::PRIMES.begin(), it_prime)); + if(min_bucket_count_in_out > 0) { + min_bucket_count_in_out = *it_prime; + } + else { + min_bucket_count_in_out = 0; + } + } + + std::size_t bucket_for_hash(std::size_t hash) const noexcept { + return detail::MOD_PRIME[m_iprime](hash); + } + + std::size_t next_bucket_count() const { + if(m_iprime + 1 >= detail::PRIMES.size()) { + TSL_RH_THROW_OR_TERMINATE(std::length_error, "The hash table exceeds its maximum size."); + } + + return detail::PRIMES[m_iprime + 1]; + } + + std::size_t max_bucket_count() const { + return detail::PRIMES.back(); + } + + void clear() noexcept { + m_iprime = 0; + } + +private: + unsigned int m_iprime; + + static_assert(std::numeric_limits::max() >= detail::PRIMES.size(), + "The type of m_iprime is not big enough."); +}; + +} +} + +#endif diff --git a/include/tsl/include/tsl/robin_hash.h b/include/tsl/include/tsl/robin_hash.h new file mode 100644 index 0000000000..e8bc5e68fb --- /dev/null +++ b/include/tsl/include/tsl/robin_hash.h @@ -0,0 +1,1451 @@ +/** + * MIT License + * + * Copyright (c) 2017 Thibaut Goetghebuer-Planchon + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in all + * copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + */ +#ifndef TSL_ROBIN_HASH_H +#define TSL_ROBIN_HASH_H + + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "robin_growth_policy.h" + + +namespace tsl { + +namespace detail_robin_hash { + +template +struct make_void { + using type = void; +}; + +template +struct has_is_transparent: std::false_type { +}; + +template +struct has_is_transparent::type>: std::true_type { +}; + +template +struct is_power_of_two_policy: std::false_type { +}; + +template +struct is_power_of_two_policy>: std::true_type { +}; + +// Only available in C++17, we need to be compatible with C++11 +template +const T& clamp( const T& v, const T& lo, const T& hi) { + return std::min(hi, std::max(lo, v)); +} + +template +static T numeric_cast(U value, const char* error_message = "numeric_cast() failed.") { + T ret = static_cast(value); + if(static_cast(ret) != value) { + TSL_RH_THROW_OR_TERMINATE(std::runtime_error, error_message); + } + + const bool is_same_signedness = (std::is_unsigned::value && std::is_unsigned::value) || + (std::is_signed::value && std::is_signed::value); + if(!is_same_signedness && (ret < T{}) != (value < U{})) { + TSL_RH_THROW_OR_TERMINATE(std::runtime_error, error_message); + } + + return ret; +} + + +using truncated_hash_type = std::uint_least32_t; + +/** + * Helper class that stores a truncated hash if StoreHash is true and nothing otherwise. + */ +template +class bucket_entry_hash { +public: + bool bucket_hash_equal(std::size_t /*hash*/) const noexcept { + return true; + } + + truncated_hash_type truncated_hash() const noexcept { + return 0; + } + +protected: + void set_hash(truncated_hash_type /*hash*/) noexcept { + } +}; + +template<> +class bucket_entry_hash { +public: + bool bucket_hash_equal(std::size_t hash) const noexcept { + return m_hash == truncated_hash_type(hash); + } + + truncated_hash_type truncated_hash() const noexcept { + return m_hash; + } + +protected: + void set_hash(truncated_hash_type hash) noexcept { + m_hash = truncated_hash_type(hash); + } + +private: + truncated_hash_type m_hash; +}; + + +/** + * Each bucket entry has: + * - A value of type `ValueType`. + * - An integer to store how far the value of the bucket, if any, is from its ideal bucket + * (ex: if the current bucket 5 has the value 'foo' and `hash('foo') % nb_buckets` == 3, + * `dist_from_ideal_bucket()` will return 2 as the current value of the bucket is two + * buckets away from its ideal bucket) + * If there is no value in the bucket (i.e. `empty()` is true) `dist_from_ideal_bucket()` will be < 0. + * - A marker which tells us if the bucket is the last bucket of the bucket array (useful for the + * iterator of the hash table). + * - If `StoreHash` is true, 32 bits of the hash of the value, if any, are also stored in the bucket. + * If the size of the hash is more than 32 bits, it is truncated. We don't store the full hash + * as storing the hash is a potential opportunity to use the unused space due to the alignment + * of the bucket_entry structure. We can thus potentially store the hash without any extra space + * (which would not be possible with 64 bits of the hash). + */ +template +class bucket_entry: public bucket_entry_hash { + using bucket_hash = bucket_entry_hash; + +public: + using value_type = ValueType; + using distance_type = std::int_least16_t; + + + bucket_entry() noexcept: bucket_hash(), m_dist_from_ideal_bucket(EMPTY_MARKER_DIST_FROM_IDEAL_BUCKET), + m_last_bucket(false) + { + tsl_rh_assert(empty()); + } + + bucket_entry(bool last_bucket) noexcept: bucket_hash(), m_dist_from_ideal_bucket(EMPTY_MARKER_DIST_FROM_IDEAL_BUCKET), + m_last_bucket(last_bucket) + { + tsl_rh_assert(empty()); + } + + bucket_entry(const bucket_entry& other) noexcept(std::is_nothrow_copy_constructible::value): + bucket_hash(other), + m_dist_from_ideal_bucket(EMPTY_MARKER_DIST_FROM_IDEAL_BUCKET), + m_last_bucket(other.m_last_bucket) + { + if(!other.empty()) { + ::new (static_cast(std::addressof(m_value))) value_type(other.value()); + m_dist_from_ideal_bucket = other.m_dist_from_ideal_bucket; + } + } + + /** + * Never really used, but still necessary as we must call resize on an empty `std::vector`. + * and we need to support move-only types. See robin_hash constructor for details. + */ + bucket_entry(bucket_entry&& other) noexcept(std::is_nothrow_move_constructible::value): + bucket_hash(std::move(other)), + m_dist_from_ideal_bucket(EMPTY_MARKER_DIST_FROM_IDEAL_BUCKET), + m_last_bucket(other.m_last_bucket) + { + if(!other.empty()) { + ::new (static_cast(std::addressof(m_value))) value_type(std::move(other.value())); + m_dist_from_ideal_bucket = other.m_dist_from_ideal_bucket; + } + } + + bucket_entry& operator=(const bucket_entry& other) + noexcept(std::is_nothrow_copy_constructible::value) + { + if(this != &other) { + clear(); + + bucket_hash::operator=(other); + if(!other.empty()) { + ::new (static_cast(std::addressof(m_value))) value_type(other.value()); + } + + m_dist_from_ideal_bucket = other.m_dist_from_ideal_bucket; + m_last_bucket = other.m_last_bucket; + } + + return *this; + } + + bucket_entry& operator=(bucket_entry&& ) = delete; + + ~bucket_entry() noexcept { + clear(); + } + + void clear() noexcept { + if(!empty()) { + destroy_value(); + m_dist_from_ideal_bucket = EMPTY_MARKER_DIST_FROM_IDEAL_BUCKET; + } + } + + bool empty() const noexcept { + return m_dist_from_ideal_bucket == EMPTY_MARKER_DIST_FROM_IDEAL_BUCKET; + } + + value_type& value() noexcept { + tsl_rh_assert(!empty()); + return *reinterpret_cast(std::addressof(m_value)); + } + + const value_type& value() const noexcept { + tsl_rh_assert(!empty()); + return *reinterpret_cast(std::addressof(m_value)); + } + + distance_type dist_from_ideal_bucket() const noexcept { + return m_dist_from_ideal_bucket; + } + + bool last_bucket() const noexcept { + return m_last_bucket; + } + + void set_as_last_bucket() noexcept { + m_last_bucket = true; + } + + template + void set_value_of_empty_bucket(distance_type dist_from_ideal_bucket, + truncated_hash_type hash, Args&&... value_type_args) + { + tsl_rh_assert(dist_from_ideal_bucket >= 0); + tsl_rh_assert(empty()); + + ::new (static_cast(std::addressof(m_value))) value_type(std::forward(value_type_args)...); + this->set_hash(hash); + m_dist_from_ideal_bucket = dist_from_ideal_bucket; + + tsl_rh_assert(!empty()); + } + + void swap_with_value_in_bucket(distance_type& dist_from_ideal_bucket, + truncated_hash_type& hash, value_type& value) + { + tsl_rh_assert(!empty()); + + using std::swap; + swap(value, this->value()); + swap(dist_from_ideal_bucket, m_dist_from_ideal_bucket); + + // Avoid warning of unused variable if StoreHash is false + (void) hash; + if(StoreHash) { + const truncated_hash_type tmp_hash = this->truncated_hash(); + this->set_hash(hash); + hash = tmp_hash; + } + } + + static truncated_hash_type truncate_hash(std::size_t hash) noexcept { + return truncated_hash_type(hash); + } + +private: + void destroy_value() noexcept { + tsl_rh_assert(!empty()); + value().~value_type(); + } + +public: + static const distance_type DIST_FROM_IDEAL_BUCKET_LIMIT = 4096; + static_assert(DIST_FROM_IDEAL_BUCKET_LIMIT <= std::numeric_limits::max() - 1, + "DIST_FROM_IDEAL_BUCKET_LIMIT must be <= std::numeric_limits::max() - 1."); + +private: + using storage = typename std::aligned_storage::type; + + static const distance_type EMPTY_MARKER_DIST_FROM_IDEAL_BUCKET = -1; + + distance_type m_dist_from_ideal_bucket; + bool m_last_bucket; + storage m_value; +}; + + + +/** + * Internal common class used by `robin_map` and `robin_set`. + * + * ValueType is what will be stored by `robin_hash` (usually `std::pair` for map and `Key` for set). + * + * `KeySelect` should be a `FunctionObject` which takes a `ValueType` in parameter and returns a + * reference to the key. + * + * `ValueSelect` should be a `FunctionObject` which takes a `ValueType` in parameter and returns a + * reference to the value. `ValueSelect` should be void if there is no value (in a set for example). + * + * The strong exception guarantee only holds if the expression + * `std::is_nothrow_swappable::value && std::is_nothrow_move_constructible::value` is true. + * + * Behaviour is undefined if the destructor of `ValueType` throws. + */ +template +class robin_hash: private Hash, private KeyEqual, private GrowthPolicy { +private: + template + using has_mapped_type = typename std::integral_constant::value>; + + static_assert(noexcept(std::declval().bucket_for_hash(std::size_t(0))), "GrowthPolicy::bucket_for_hash must be noexcept."); + static_assert(noexcept(std::declval().clear()), "GrowthPolicy::clear must be noexcept."); + +public: + template + class robin_iterator; + + using key_type = typename KeySelect::key_type; + using value_type = ValueType; + using size_type = std::size_t; + using difference_type = std::ptrdiff_t; + using hasher = Hash; + using key_equal = KeyEqual; + using allocator_type = Allocator; + using reference = value_type&; + using const_reference = const value_type&; + using pointer = value_type*; + using const_pointer = const value_type*; + using iterator = robin_iterator; + using const_iterator = robin_iterator; + + +private: + /** + * Either store the hash because we are asked by the `StoreHash` template parameter + * or store the hash because it doesn't cost us anything in size and can be used to speed up rehash. + */ + static constexpr bool STORE_HASH = StoreHash || + ( + (sizeof(tsl::detail_robin_hash::bucket_entry) == + sizeof(tsl::detail_robin_hash::bucket_entry)) + && + (sizeof(std::size_t) == sizeof(truncated_hash_type) || + is_power_of_two_policy::value) + && + // Don't store the hash for primitive types with default hash. + (!std::is_arithmetic::value || + !std::is_same>::value) + ); + + /** + * Only use the stored hash on lookup if we are explicitly asked. We are not sure how slow + * the KeyEqual operation is. An extra comparison may slow things down with a fast KeyEqual. + */ + static constexpr bool USE_STORED_HASH_ON_LOOKUP = StoreHash; + + /** + * We can only use the hash on rehash if the size of the hash type is the same as the stored one or + * if we use a power of two modulo. In the case of the power of two modulo, we just mask + * the least significant bytes, we just have to check that the truncated_hash_type didn't truncated + * more bytes. + */ + static bool USE_STORED_HASH_ON_REHASH(size_type bucket_count) { + (void) bucket_count; + if(STORE_HASH && sizeof(std::size_t) == sizeof(truncated_hash_type)) { + return true; + } + else if(STORE_HASH && is_power_of_two_policy::value) { + tsl_rh_assert(bucket_count > 0); + return (bucket_count - 1) <= std::numeric_limits::max(); + } + else { + return false; + } + } + + using bucket_entry = tsl::detail_robin_hash::bucket_entry; + using distance_type = typename bucket_entry::distance_type; + + using buckets_allocator = typename std::allocator_traits::template rebind_alloc; + using buckets_container_type = std::vector; + + +public: + /** + * The 'operator*()' and 'operator->()' methods return a const reference and const pointer respectively to the + * stored value type. + * + * In case of a map, to get a mutable reference to the value associated to a key (the '.second' in the + * stored pair), you have to call 'value()'. + * + * The main reason for this is that if we returned a `std::pair&` instead + * of a `const std::pair&`, the user may modify the key which will put the map in a undefined state. + */ + template + class robin_iterator { + friend class robin_hash; + + private: + using bucket_entry_ptr = typename std::conditional::type; + + + robin_iterator(bucket_entry_ptr bucket) noexcept: m_bucket(bucket) { + } + + public: + using iterator_category = std::forward_iterator_tag; + using value_type = const typename robin_hash::value_type; + using difference_type = std::ptrdiff_t; + using reference = value_type&; + using pointer = value_type*; + + + robin_iterator() noexcept { + } + + // Copy constructor from iterator to const_iterator. + template::type* = nullptr> + robin_iterator(const robin_iterator& other) noexcept: m_bucket(other.m_bucket) { + } + + robin_iterator(const robin_iterator& other) = default; + robin_iterator(robin_iterator&& other) = default; + robin_iterator& operator=(const robin_iterator& other) = default; + robin_iterator& operator=(robin_iterator&& other) = default; + + const typename robin_hash::key_type& key() const { + return KeySelect()(m_bucket->value()); + } + + template::value && IsConst>::type* = nullptr> + const typename U::value_type& value() const { + return U()(m_bucket->value()); + } + + template::value && !IsConst>::type* = nullptr> + typename U::value_type& value() const { + return U()(m_bucket->value()); + } + + reference operator*() const { + return m_bucket->value(); + } + + pointer operator->() const { + return std::addressof(m_bucket->value()); + } + + robin_iterator& operator++() { + while(true) { + if(m_bucket->last_bucket()) { + ++m_bucket; + return *this; + } + + ++m_bucket; + if(!m_bucket->empty()) { + return *this; + } + } + } + + robin_iterator operator++(int) { + robin_iterator tmp(*this); + ++*this; + + return tmp; + } + + friend bool operator==(const robin_iterator& lhs, const robin_iterator& rhs) { + return lhs.m_bucket == rhs.m_bucket; + } + + friend bool operator!=(const robin_iterator& lhs, const robin_iterator& rhs) { + return !(lhs == rhs); + } + + private: + bucket_entry_ptr m_bucket; + }; + + +public: +#if defined(__cplusplus) && __cplusplus >= 201402L + robin_hash(size_type bucket_count, + const Hash& hash, + const KeyEqual& equal, + const Allocator& alloc, + float min_load_factor = DEFAULT_MIN_LOAD_FACTOR, + float max_load_factor = DEFAULT_MAX_LOAD_FACTOR): + Hash(hash), + KeyEqual(equal), + GrowthPolicy(bucket_count), + m_buckets_data( + [&]() { + if(bucket_count > max_bucket_count()) { + TSL_RH_THROW_OR_TERMINATE(std::length_error, + "The map exceeds its maximum bucket count."); + } + + return bucket_count; + }(), alloc + ), + m_buckets(m_buckets_data.empty()?static_empty_bucket_ptr():m_buckets_data.data()), + m_bucket_count(bucket_count), + m_nb_elements(0), + m_grow_on_next_insert(false), + m_try_shrink_on_next_insert(false) + { + if(m_bucket_count > 0) { + tsl_rh_assert(!m_buckets_data.empty()); + m_buckets_data.back().set_as_last_bucket(); + } + + this->min_load_factor(min_load_factor); + this->max_load_factor(max_load_factor); + } +#else + /** + * C++11 doesn't support the creation of a std::vector with a custom allocator and 'count' default-inserted elements. + * The needed contructor `explicit vector(size_type count, const Allocator& alloc = Allocator());` is only + * available in C++14 and later. We thus must resize after using the `vector(const Allocator& alloc)` constructor. + * + * We can't use `vector(size_type count, const T& value, const Allocator& alloc)` as it requires the + * value T to be copyable. + */ + robin_hash(size_type bucket_count, + const Hash& hash, + const KeyEqual& equal, + const Allocator& alloc, + float min_load_factor = DEFAULT_MIN_LOAD_FACTOR, + float max_load_factor = DEFAULT_MAX_LOAD_FACTOR): + Hash(hash), + KeyEqual(equal), + GrowthPolicy(bucket_count), + m_buckets_data(alloc), + m_buckets(static_empty_bucket_ptr()), + m_bucket_count(bucket_count), + m_nb_elements(0), + m_grow_on_next_insert(false), + m_try_shrink_on_next_insert(false) + { + if(bucket_count > max_bucket_count()) { + TSL_RH_THROW_OR_TERMINATE(std::length_error, "The map exceeds its maximum bucket count."); + } + + if(m_bucket_count > 0) { + m_buckets_data.resize(m_bucket_count); + m_buckets = m_buckets_data.data(); + + tsl_rh_assert(!m_buckets_data.empty()); + m_buckets_data.back().set_as_last_bucket(); + } + + this->min_load_factor(min_load_factor); + this->max_load_factor(max_load_factor); + } +#endif + + robin_hash(const robin_hash& other): Hash(other), + KeyEqual(other), + GrowthPolicy(other), + m_buckets_data(other.m_buckets_data), + m_buckets(m_buckets_data.empty()?static_empty_bucket_ptr():m_buckets_data.data()), + m_bucket_count(other.m_bucket_count), + m_nb_elements(other.m_nb_elements), + m_load_threshold(other.m_load_threshold), + m_min_load_factor(other.m_min_load_factor), + m_max_load_factor(other.m_max_load_factor), + m_grow_on_next_insert(other.m_grow_on_next_insert), + m_try_shrink_on_next_insert(other.m_try_shrink_on_next_insert) + { + } + + robin_hash(robin_hash&& other) noexcept(std::is_nothrow_move_constructible::value && + std::is_nothrow_move_constructible::value && + std::is_nothrow_move_constructible::value && + std::is_nothrow_move_constructible::value) + : Hash(std::move(static_cast(other))), + KeyEqual(std::move(static_cast(other))), + GrowthPolicy(std::move(static_cast(other))), + m_buckets_data(std::move(other.m_buckets_data)), + m_buckets(m_buckets_data.empty()?static_empty_bucket_ptr():m_buckets_data.data()), + m_bucket_count(other.m_bucket_count), + m_nb_elements(other.m_nb_elements), + m_load_threshold(other.m_load_threshold), + m_min_load_factor(other.m_min_load_factor), + m_max_load_factor(other.m_max_load_factor), + m_grow_on_next_insert(other.m_grow_on_next_insert), + m_try_shrink_on_next_insert(other.m_try_shrink_on_next_insert) + { + other.clear_and_shrink(); + } + + robin_hash& operator=(const robin_hash& other) { + if(&other != this) { + Hash::operator=(other); + KeyEqual::operator=(other); + GrowthPolicy::operator=(other); + + m_buckets_data = other.m_buckets_data; + m_buckets = m_buckets_data.empty()?static_empty_bucket_ptr(): + m_buckets_data.data(); + m_bucket_count = other.m_bucket_count; + m_nb_elements = other.m_nb_elements; + + m_load_threshold = other.m_load_threshold; + m_min_load_factor = other.m_min_load_factor; + m_max_load_factor = other.m_max_load_factor; + + m_grow_on_next_insert = other.m_grow_on_next_insert; + m_try_shrink_on_next_insert = other.m_try_shrink_on_next_insert; + } + + return *this; + } + + robin_hash& operator=(robin_hash&& other) { + other.swap(*this); + other.clear(); + + return *this; + } + + allocator_type get_allocator() const { + return m_buckets_data.get_allocator(); + } + + + /* + * Iterators + */ + iterator begin() noexcept { + std::size_t i = 0; + while(i < m_bucket_count && m_buckets[i].empty()) { + i++; + } + + return iterator(m_buckets + i); + } + + const_iterator begin() const noexcept { + return cbegin(); + } + + const_iterator cbegin() const noexcept { + std::size_t i = 0; + while(i < m_bucket_count && m_buckets[i].empty()) { + i++; + } + + return const_iterator(m_buckets + i); + } + + iterator end() noexcept { + return iterator(m_buckets + m_bucket_count); + } + + const_iterator end() const noexcept { + return cend(); + } + + const_iterator cend() const noexcept { + return const_iterator(m_buckets + m_bucket_count); + } + + + /* + * Capacity + */ + bool empty() const noexcept { + return m_nb_elements == 0; + } + + size_type size() const noexcept { + return m_nb_elements; + } + + size_type max_size() const noexcept { + return m_buckets_data.max_size(); + } + + /* + * Modifiers + */ + void clear() noexcept { + if(m_min_load_factor > 0.0f) { + clear_and_shrink(); + } + else { + for(auto& bucket: m_buckets_data) { + bucket.clear(); + } + + m_nb_elements = 0; + m_grow_on_next_insert = false; + } + } + + + + template + std::pair insert(P&& value) { + return insert_impl(KeySelect()(value), std::forward

(value)); + } + + template + iterator insert_hint(const_iterator hint, P&& value) { + if(hint != cend() && compare_keys(KeySelect()(*hint), KeySelect()(value))) { + return mutable_iterator(hint); + } + + return insert(std::forward

(value)).first; + } + + template + void insert(InputIt first, InputIt last) { + if(std::is_base_of::iterator_category>::value) + { + const auto nb_elements_insert = std::distance(first, last); + const size_type nb_free_buckets = m_load_threshold - size(); + tsl_rh_assert(m_load_threshold >= size()); + + if(nb_elements_insert > 0 && nb_free_buckets < size_type(nb_elements_insert)) { + reserve(size() + size_type(nb_elements_insert)); + } + } + + for(; first != last; ++first) { + insert(*first); + } + } + + + + template + std::pair insert_or_assign(K&& key, M&& obj) { + auto it = try_emplace(std::forward(key), std::forward(obj)); + if(!it.second) { + it.first.value() = std::forward(obj); + } + + return it; + } + + template + iterator insert_or_assign(const_iterator hint, K&& key, M&& obj) { + if(hint != cend() && compare_keys(KeySelect()(*hint), key)) { + auto it = mutable_iterator(hint); + it.value() = std::forward(obj); + + return it; + } + + return insert_or_assign(std::forward(key), std::forward(obj)).first; + } + + + template + std::pair emplace(Args&&... args) { + return insert(value_type(std::forward(args)...)); + } + + template + iterator emplace_hint(const_iterator hint, Args&&... args) { + return insert_hint(hint, value_type(std::forward(args)...)); + } + + + + template + std::pair try_emplace(K&& key, Args&&... args) { + return insert_impl(key, std::piecewise_construct, + std::forward_as_tuple(std::forward(key)), + std::forward_as_tuple(std::forward(args)...)); + } + + template + iterator try_emplace_hint(const_iterator hint, K&& key, Args&&... args) { + if(hint != cend() && compare_keys(KeySelect()(*hint), key)) { + return mutable_iterator(hint); + } + + return try_emplace(std::forward(key), std::forward(args)...).first; + } + + /** + * Here to avoid `template size_type erase(const K& key)` being used when + * we use an `iterator` instead of a `const_iterator`. + */ + iterator erase(iterator pos) { + erase_from_bucket(pos); + + /** + * Erase bucket used a backward shift after clearing the bucket. + * Check if there is a new value in the bucket, if not get the next non-empty. + */ + if(pos.m_bucket->empty()) { + ++pos; + } + + m_try_shrink_on_next_insert = true; + + return pos; + } + + iterator erase(const_iterator pos) { + return erase(mutable_iterator(pos)); + } + + iterator erase(const_iterator first, const_iterator last) { + if(first == last) { + return mutable_iterator(first); + } + + auto first_mutable = mutable_iterator(first); + auto last_mutable = mutable_iterator(last); + for(auto it = first_mutable.m_bucket; it != last_mutable.m_bucket; ++it) { + if(!it->empty()) { + it->clear(); + m_nb_elements--; + } + } + + if(last_mutable == end()) { + m_try_shrink_on_next_insert = true; + return end(); + } + + + /* + * Backward shift on the values which come after the deleted values. + * We try to move the values closer to their ideal bucket. + */ + std::size_t icloser_bucket = static_cast(first_mutable.m_bucket - m_buckets); + std::size_t ito_move_closer_value = static_cast(last_mutable.m_bucket - m_buckets); + tsl_rh_assert(ito_move_closer_value > icloser_bucket); + + const std::size_t ireturn_bucket = ito_move_closer_value - + std::min(ito_move_closer_value - icloser_bucket, + std::size_t(m_buckets[ito_move_closer_value].dist_from_ideal_bucket())); + + while(ito_move_closer_value < m_bucket_count && m_buckets[ito_move_closer_value].dist_from_ideal_bucket() > 0) { + icloser_bucket = ito_move_closer_value - + std::min(ito_move_closer_value - icloser_bucket, + std::size_t(m_buckets[ito_move_closer_value].dist_from_ideal_bucket())); + + + tsl_rh_assert(m_buckets[icloser_bucket].empty()); + const distance_type new_distance = distance_type(m_buckets[ito_move_closer_value].dist_from_ideal_bucket() - + (ito_move_closer_value - icloser_bucket)); + m_buckets[icloser_bucket].set_value_of_empty_bucket(new_distance, + m_buckets[ito_move_closer_value].truncated_hash(), + std::move(m_buckets[ito_move_closer_value].value())); + m_buckets[ito_move_closer_value].clear(); + + + ++icloser_bucket; + ++ito_move_closer_value; + } + + m_try_shrink_on_next_insert = true; + + return iterator(m_buckets + ireturn_bucket); + } + + + template + size_type erase(const K& key) { + return erase(key, hash_key(key)); + } + + template + size_type erase(const K& key, std::size_t hash) { + auto it = find(key, hash); + if(it != end()) { + erase_from_bucket(it); + m_try_shrink_on_next_insert = true; + + return 1; + } + else { + return 0; + } + } + + + + + + void swap(robin_hash& other) { + using std::swap; + + swap(static_cast(*this), static_cast(other)); + swap(static_cast(*this), static_cast(other)); + swap(static_cast(*this), static_cast(other)); + swap(m_buckets_data, other.m_buckets_data); + swap(m_buckets, other.m_buckets); + swap(m_bucket_count, other.m_bucket_count); + swap(m_nb_elements, other.m_nb_elements); + swap(m_load_threshold, other.m_load_threshold); + swap(m_min_load_factor, other.m_min_load_factor); + swap(m_max_load_factor, other.m_max_load_factor); + swap(m_grow_on_next_insert, other.m_grow_on_next_insert); + swap(m_try_shrink_on_next_insert, other.m_try_shrink_on_next_insert); + } + + + /* + * Lookup + */ + template::value>::type* = nullptr> + typename U::value_type& at(const K& key) { + return at(key, hash_key(key)); + } + + template::value>::type* = nullptr> + typename U::value_type& at(const K& key, std::size_t hash) { + return const_cast(static_cast(this)->at(key, hash)); + } + + + template::value>::type* = nullptr> + const typename U::value_type& at(const K& key) const { + return at(key, hash_key(key)); + } + + template::value>::type* = nullptr> + const typename U::value_type& at(const K& key, std::size_t hash) const { + auto it = find(key, hash); + if(it != cend()) { + return it.value(); + } + else { + TSL_RH_THROW_OR_TERMINATE(std::out_of_range, "Couldn't find key."); + } + } + + template::value>::type* = nullptr> + typename U::value_type& operator[](K&& key) { + return try_emplace(std::forward(key)).first.value(); + } + + + template + size_type count(const K& key) const { + return count(key, hash_key(key)); + } + + template + size_type count(const K& key, std::size_t hash) const { + if(find(key, hash) != cend()) { + return 1; + } + else { + return 0; + } + } + + + template + iterator find(const K& key) { + return find_impl(key, hash_key(key)); + } + + template + iterator find(const K& key, std::size_t hash) { + return find_impl(key, hash); + } + + + template + const_iterator find(const K& key) const { + return find_impl(key, hash_key(key)); + } + + template + const_iterator find(const K& key, std::size_t hash) const { + return find_impl(key, hash); + } + + + template + bool contains(const K& key) const { + return contains(key, hash_key(key)); + } + + template + bool contains(const K& key, std::size_t hash) const { + return count(key, hash) != 0; + } + + + template + std::pair equal_range(const K& key) { + return equal_range(key, hash_key(key)); + } + + template + std::pair equal_range(const K& key, std::size_t hash) { + iterator it = find(key, hash); + return std::make_pair(it, (it == end())?it:std::next(it)); + } + + + template + std::pair equal_range(const K& key) const { + return equal_range(key, hash_key(key)); + } + + template + std::pair equal_range(const K& key, std::size_t hash) const { + const_iterator it = find(key, hash); + return std::make_pair(it, (it == cend())?it:std::next(it)); + } + + /* + * Bucket interface + */ + size_type bucket_count() const { + return m_bucket_count; + } + + size_type max_bucket_count() const { + return std::min(GrowthPolicy::max_bucket_count(), m_buckets_data.max_size()); + } + + /* + * Hash policy + */ + float load_factor() const { + if(bucket_count() == 0) { + return 0; + } + + return float(m_nb_elements)/float(bucket_count()); + } + + float min_load_factor() const { + return m_min_load_factor; + } + + float max_load_factor() const { + return m_max_load_factor; + } + + void min_load_factor(float ml) { + m_min_load_factor = clamp(ml, float(MINIMUM_MIN_LOAD_FACTOR), + float(MAXIMUM_MIN_LOAD_FACTOR)); + } + + void max_load_factor(float ml) { + m_max_load_factor = clamp(ml, float(MINIMUM_MAX_LOAD_FACTOR), + float(MAXIMUM_MAX_LOAD_FACTOR)); + m_load_threshold = size_type(float(bucket_count())*m_max_load_factor); + } + + void rehash(size_type count) { + count = std::max(count, size_type(std::ceil(float(size())/max_load_factor()))); + rehash_impl(count); + } + + void reserve(size_type count) { + rehash(size_type(std::ceil(float(count)/max_load_factor()))); + } + + /* + * Observers + */ + hasher hash_function() const { + return static_cast(*this); + } + + key_equal key_eq() const { + return static_cast(*this); + } + + + /* + * Other + */ + iterator mutable_iterator(const_iterator pos) { + return iterator(const_cast(pos.m_bucket)); + } + +private: + template + std::size_t hash_key(const K& key) const { + return Hash::operator()(key); + } + + template + bool compare_keys(const K1& key1, const K2& key2) const { + return KeyEqual::operator()(key1, key2); + } + + std::size_t bucket_for_hash(std::size_t hash) const { + const std::size_t bucket = GrowthPolicy::bucket_for_hash(hash); + tsl_rh_assert(bucket < m_bucket_count || (bucket == 0 && m_bucket_count == 0)); + + return bucket; + } + + template::value>::type* = nullptr> + std::size_t next_bucket(std::size_t index) const noexcept { + tsl_rh_assert(index < bucket_count()); + + return (index + 1) & this->m_mask; + } + + template::value>::type* = nullptr> + std::size_t next_bucket(std::size_t index) const noexcept { + tsl_rh_assert(index < bucket_count()); + + index++; + return (index != bucket_count())?index:0; + } + + + + template + iterator find_impl(const K& key, std::size_t hash) { + return mutable_iterator(static_cast(this)->find(key, hash)); + } + + template + const_iterator find_impl(const K& key, std::size_t hash) const { + std::size_t ibucket = bucket_for_hash(hash); + distance_type dist_from_ideal_bucket = 0; + + while(dist_from_ideal_bucket <= m_buckets[ibucket].dist_from_ideal_bucket()) { + if(TSL_RH_LIKELY((!USE_STORED_HASH_ON_LOOKUP || m_buckets[ibucket].bucket_hash_equal(hash)) && + compare_keys(KeySelect()(m_buckets[ibucket].value()), key))) + { + return const_iterator(m_buckets + ibucket); + } + + ibucket = next_bucket(ibucket); + dist_from_ideal_bucket++; + } + + return cend(); + } + + void erase_from_bucket(iterator pos) { + pos.m_bucket->clear(); + m_nb_elements--; + + /** + * Backward shift, swap the empty bucket, previous_ibucket, with the values on its right, ibucket, + * until we cross another empty bucket or if the other bucket has a distance_from_ideal_bucket == 0. + * + * We try to move the values closer to their ideal bucket. + */ + std::size_t previous_ibucket = static_cast(pos.m_bucket - m_buckets); + std::size_t ibucket = next_bucket(previous_ibucket); + + while(m_buckets[ibucket].dist_from_ideal_bucket() > 0) { + tsl_rh_assert(m_buckets[previous_ibucket].empty()); + + const distance_type new_distance = distance_type(m_buckets[ibucket].dist_from_ideal_bucket() - 1); + m_buckets[previous_ibucket].set_value_of_empty_bucket(new_distance, m_buckets[ibucket].truncated_hash(), + std::move(m_buckets[ibucket].value())); + m_buckets[ibucket].clear(); + + previous_ibucket = ibucket; + ibucket = next_bucket(ibucket); + } + } + + template + std::pair insert_impl(const K& key, Args&&... value_type_args) { + const std::size_t hash = hash_key(key); + + std::size_t ibucket = bucket_for_hash(hash); + distance_type dist_from_ideal_bucket = 0; + + while(dist_from_ideal_bucket <= m_buckets[ibucket].dist_from_ideal_bucket()) { + if((!USE_STORED_HASH_ON_LOOKUP || m_buckets[ibucket].bucket_hash_equal(hash)) && + compare_keys(KeySelect()(m_buckets[ibucket].value()), key)) + { + return std::make_pair(iterator(m_buckets + ibucket), false); + } + + ibucket = next_bucket(ibucket); + dist_from_ideal_bucket++; + } + + if(rehash_on_extreme_load()) { + ibucket = bucket_for_hash(hash); + dist_from_ideal_bucket = 0; + + while(dist_from_ideal_bucket <= m_buckets[ibucket].dist_from_ideal_bucket()) { + ibucket = next_bucket(ibucket); + dist_from_ideal_bucket++; + } + } + + + if(m_buckets[ibucket].empty()) { + m_buckets[ibucket].set_value_of_empty_bucket(dist_from_ideal_bucket, bucket_entry::truncate_hash(hash), + std::forward(value_type_args)...); + } + else { + insert_value(ibucket, dist_from_ideal_bucket, bucket_entry::truncate_hash(hash), + std::forward(value_type_args)...); + } + + + m_nb_elements++; + /* + * The value will be inserted in ibucket in any case, either because it was + * empty or by stealing the bucket (robin hood). + */ + return std::make_pair(iterator(m_buckets + ibucket), true); + } + + + template + void insert_value(std::size_t ibucket, distance_type dist_from_ideal_bucket, + truncated_hash_type hash, Args&&... value_type_args) + { + value_type value(std::forward(value_type_args)...); + insert_value_impl(ibucket, dist_from_ideal_bucket, hash, value); + } + + void insert_value(std::size_t ibucket, distance_type dist_from_ideal_bucket, + truncated_hash_type hash, value_type&& value) + { + insert_value_impl(ibucket, dist_from_ideal_bucket, hash, value); + } + + /* + * We don't use `value_type&& value` as last argument due to a bug in MSVC when `value_type` is a pointer, + * The compiler is not able to see the difference between `std::string*` and `std::string*&&` resulting in + * a compilation error. + * + * The `value` will be in a moved state at the end of the function. + */ + void insert_value_impl(std::size_t ibucket, distance_type dist_from_ideal_bucket, + truncated_hash_type hash, value_type& value) + { + m_buckets[ibucket].swap_with_value_in_bucket(dist_from_ideal_bucket, hash, value); + ibucket = next_bucket(ibucket); + dist_from_ideal_bucket++; + + while(!m_buckets[ibucket].empty()) { + if(dist_from_ideal_bucket > m_buckets[ibucket].dist_from_ideal_bucket()) { + if(dist_from_ideal_bucket >= bucket_entry::DIST_FROM_IDEAL_BUCKET_LIMIT) { + /** + * The number of probes is really high, rehash the map on the next insert. + * Difficult to do now as rehash may throw an exception. + */ + m_grow_on_next_insert = true; + } + + m_buckets[ibucket].swap_with_value_in_bucket(dist_from_ideal_bucket, hash, value); + } + + ibucket = next_bucket(ibucket); + dist_from_ideal_bucket++; + } + + m_buckets[ibucket].set_value_of_empty_bucket(dist_from_ideal_bucket, hash, std::move(value)); + } + + + void rehash_impl(size_type count) { + robin_hash new_table(count, static_cast(*this), static_cast(*this), + get_allocator(), m_min_load_factor, m_max_load_factor); + + const bool use_stored_hash = USE_STORED_HASH_ON_REHASH(new_table.bucket_count()); + for(auto& bucket: m_buckets_data) { + if(bucket.empty()) { + continue; + } + + const std::size_t hash = use_stored_hash?bucket.truncated_hash(): + new_table.hash_key(KeySelect()(bucket.value())); + + new_table.insert_value_on_rehash(new_table.bucket_for_hash(hash), 0, + bucket_entry::truncate_hash(hash), std::move(bucket.value())); + } + + new_table.m_nb_elements = m_nb_elements; + new_table.swap(*this); + } + + void clear_and_shrink() noexcept { + GrowthPolicy::clear(); + m_buckets_data.clear(); + m_buckets = static_empty_bucket_ptr(); + m_bucket_count = 0; + m_nb_elements = 0; + m_load_threshold = 0; + m_grow_on_next_insert = false; + m_try_shrink_on_next_insert = false; + } + + void insert_value_on_rehash(std::size_t ibucket, distance_type dist_from_ideal_bucket, + truncated_hash_type hash, value_type&& value) + { + while(true) { + if(dist_from_ideal_bucket > m_buckets[ibucket].dist_from_ideal_bucket()) { + if(m_buckets[ibucket].empty()) { + m_buckets[ibucket].set_value_of_empty_bucket(dist_from_ideal_bucket, hash, std::move(value)); + return; + } + else { + m_buckets[ibucket].swap_with_value_in_bucket(dist_from_ideal_bucket, hash, value); + } + } + + dist_from_ideal_bucket++; + ibucket = next_bucket(ibucket); + } + } + + + + /** + * Grow the table if m_grow_on_next_insert is true or we reached the max_load_factor. + * Shrink the table if m_try_shrink_on_next_insert is true (an erase occurred) and + * we're below the min_load_factor. + * + * Return true if the table has been rehashed. + */ + bool rehash_on_extreme_load() { + if(m_grow_on_next_insert || size() >= m_load_threshold) { + rehash_impl(GrowthPolicy::next_bucket_count()); + m_grow_on_next_insert = false; + + return true; + } + + if(m_try_shrink_on_next_insert) { + m_try_shrink_on_next_insert = false; + if(m_min_load_factor != 0.0f && load_factor() < m_min_load_factor) { + reserve(size() + 1); + + return true; + } + } + + return false; + } + + +public: + static const size_type DEFAULT_INIT_BUCKETS_SIZE = 0; + + static constexpr float DEFAULT_MAX_LOAD_FACTOR = 0.5f; + static constexpr float MINIMUM_MAX_LOAD_FACTOR = 0.2f; + static constexpr float MAXIMUM_MAX_LOAD_FACTOR = 0.95f; + + static constexpr float DEFAULT_MIN_LOAD_FACTOR = 0.0f; + static constexpr float MINIMUM_MIN_LOAD_FACTOR = 0.0f; + static constexpr float MAXIMUM_MIN_LOAD_FACTOR = 0.15f; + + static_assert(MINIMUM_MAX_LOAD_FACTOR < MAXIMUM_MAX_LOAD_FACTOR, + "MINIMUM_MAX_LOAD_FACTOR should be < MAXIMUM_MAX_LOAD_FACTOR"); + static_assert(MINIMUM_MIN_LOAD_FACTOR < MAXIMUM_MIN_LOAD_FACTOR, + "MINIMUM_MIN_LOAD_FACTOR should be < MAXIMUM_MIN_LOAD_FACTOR"); + static_assert(MAXIMUM_MIN_LOAD_FACTOR < MINIMUM_MAX_LOAD_FACTOR, + "MAXIMUM_MIN_LOAD_FACTOR should be < MINIMUM_MAX_LOAD_FACTOR"); + +private: + /** + * Return an always valid pointer to an static empty bucket_entry with last_bucket() == true. + */ + bucket_entry* static_empty_bucket_ptr() noexcept { + static bucket_entry empty_bucket(true); + return &empty_bucket; + } + +private: + buckets_container_type m_buckets_data; + + /** + * Points to m_buckets_data.data() if !m_buckets_data.empty() otherwise points to static_empty_bucket_ptr. + * This variable is useful to avoid the cost of checking if m_buckets_data is empty when trying + * to find an element. + * + * TODO Remove m_buckets_data and only use a pointer instead of a pointer+vector to save some space in the robin_hash object. + * Manage the Allocator manually. + */ + bucket_entry* m_buckets; + + /** + * Used a lot in find, avoid the call to m_buckets_data.size() which is a bit slower. + */ + size_type m_bucket_count; + + size_type m_nb_elements; + + size_type m_load_threshold; + + float m_min_load_factor; + float m_max_load_factor; + + bool m_grow_on_next_insert; + + /** + * We can't shrink down the map on erase operations as the erase methods need to return the next iterator. + * Shrinking the map would invalidate all the iterators and we could not return the next iterator in a meaningful way, + * On erase, we thus just indicate on erase that we should try to shrink the hash table on the next insert + * if we go below the min_load_factor. + */ + bool m_try_shrink_on_next_insert; +}; + +} + +} + +#endif diff --git a/include/tsl/include/tsl/robin_map.h b/include/tsl/include/tsl/robin_map.h new file mode 100644 index 0000000000..86761f5308 --- /dev/null +++ b/include/tsl/include/tsl/robin_map.h @@ -0,0 +1,715 @@ +/** + * MIT License + * + * Copyright (c) 2017 Thibaut Goetghebuer-Planchon + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in all + * copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + */ +#ifndef TSL_ROBIN_MAP_H +#define TSL_ROBIN_MAP_H + + +#include +#include +#include +#include +#include +#include +#include "robin_hash.h" + + +namespace tsl { + + +/** + * Implementation of a hash map using open-addressing and the robin hood hashing algorithm with backward shift deletion. + * + * For operations modifying the hash map (insert, erase, rehash, ...), the strong exception guarantee + * is only guaranteed when the expression `std::is_nothrow_swappable>::value && + * std::is_nothrow_move_constructible>::value` is true, otherwise if an exception + * is thrown during the swap or the move, the hash map may end up in a undefined state. Per the standard + * a `Key` or `T` with a noexcept copy constructor and no move constructor also satisfies the + * `std::is_nothrow_move_constructible>::value` criterion (and will thus guarantee the + * strong exception for the map). + * + * When `StoreHash` is true, 32 bits of the hash are stored alongside the values. It can improve + * the performance during lookups if the `KeyEqual` function takes time (if it engenders a cache-miss for example) + * as we then compare the stored hashes before comparing the keys. When `tsl::rh::power_of_two_growth_policy` is used + * as `GrowthPolicy`, it may also speed-up the rehash process as we can avoid to recalculate the hash. + * When it is detected that storing the hash will not incur any memory penalty due to alignment (i.e. + * `sizeof(tsl::detail_robin_hash::bucket_entry) == + * sizeof(tsl::detail_robin_hash::bucket_entry)`) and `tsl::rh::power_of_two_growth_policy` is + * used, the hash will be stored even if `StoreHash` is false so that we can speed-up the rehash (but it will + * not be used on lookups unless `StoreHash` is true). + * + * `GrowthPolicy` defines how the map grows and consequently how a hash value is mapped to a bucket. + * By default the map uses `tsl::rh::power_of_two_growth_policy`. This policy keeps the number of buckets + * to a power of two and uses a mask to map the hash to a bucket instead of the slow modulo. + * Other growth policies are available and you may define your own growth policy, + * check `tsl::rh::power_of_two_growth_policy` for the interface. + * + * `std::pair` must be swappable. + * + * `Key` and `T` must be copy and/or move constructible. + * + * If the destructor of `Key` or `T` throws an exception, the behaviour of the class is undefined. + * + * Iterators invalidation: + * - clear, operator=, reserve, rehash: always invalidate the iterators. + * - insert, emplace, emplace_hint, operator[]: if there is an effective insert, invalidate the iterators. + * - erase: always invalidate the iterators. + */ +template, + class KeyEqual = std::equal_to, + class Allocator = std::allocator>, + bool StoreHash = false, + class GrowthPolicy = tsl::rh::power_of_two_growth_policy<2>> +class robin_map { +private: + template + using has_is_transparent = tsl::detail_robin_hash::has_is_transparent; + + class KeySelect { + public: + using key_type = Key; + + const key_type& operator()(const std::pair& key_value) const noexcept { + return key_value.first; + } + + key_type& operator()(std::pair& key_value) noexcept { + return key_value.first; + } + }; + + class ValueSelect { + public: + using value_type = T; + + const value_type& operator()(const std::pair& key_value) const noexcept { + return key_value.second; + } + + value_type& operator()(std::pair& key_value) noexcept { + return key_value.second; + } + }; + + using ht = detail_robin_hash::robin_hash, KeySelect, ValueSelect, + Hash, KeyEqual, Allocator, StoreHash, GrowthPolicy>; + +public: + using key_type = typename ht::key_type; + using mapped_type = T; + using value_type = typename ht::value_type; + using size_type = typename ht::size_type; + using difference_type = typename ht::difference_type; + using hasher = typename ht::hasher; + using key_equal = typename ht::key_equal; + using allocator_type = typename ht::allocator_type; + using reference = typename ht::reference; + using const_reference = typename ht::const_reference; + using pointer = typename ht::pointer; + using const_pointer = typename ht::const_pointer; + using iterator = typename ht::iterator; + using const_iterator = typename ht::const_iterator; + + +public: + /* + * Constructors + */ + robin_map(): robin_map(ht::DEFAULT_INIT_BUCKETS_SIZE) { + } + + explicit robin_map(size_type bucket_count, + const Hash& hash = Hash(), + const KeyEqual& equal = KeyEqual(), + const Allocator& alloc = Allocator()): + m_ht(bucket_count, hash, equal, alloc) + { + } + + robin_map(size_type bucket_count, + const Allocator& alloc): robin_map(bucket_count, Hash(), KeyEqual(), alloc) + { + } + + robin_map(size_type bucket_count, + const Hash& hash, + const Allocator& alloc): robin_map(bucket_count, hash, KeyEqual(), alloc) + { + } + + explicit robin_map(const Allocator& alloc): robin_map(ht::DEFAULT_INIT_BUCKETS_SIZE, alloc) { + } + + template + robin_map(InputIt first, InputIt last, + size_type bucket_count = ht::DEFAULT_INIT_BUCKETS_SIZE, + const Hash& hash = Hash(), + const KeyEqual& equal = KeyEqual(), + const Allocator& alloc = Allocator()): robin_map(bucket_count, hash, equal, alloc) + { + insert(first, last); + } + + template + robin_map(InputIt first, InputIt last, + size_type bucket_count, + const Allocator& alloc): robin_map(first, last, bucket_count, Hash(), KeyEqual(), alloc) + { + } + + template + robin_map(InputIt first, InputIt last, + size_type bucket_count, + const Hash& hash, + const Allocator& alloc): robin_map(first, last, bucket_count, hash, KeyEqual(), alloc) + { + } + + robin_map(std::initializer_list init, + size_type bucket_count = ht::DEFAULT_INIT_BUCKETS_SIZE, + const Hash& hash = Hash(), + const KeyEqual& equal = KeyEqual(), + const Allocator& alloc = Allocator()): + robin_map(init.begin(), init.end(), bucket_count, hash, equal, alloc) + { + } + + robin_map(std::initializer_list init, + size_type bucket_count, + const Allocator& alloc): + robin_map(init.begin(), init.end(), bucket_count, Hash(), KeyEqual(), alloc) + { + } + + robin_map(std::initializer_list init, + size_type bucket_count, + const Hash& hash, + const Allocator& alloc): + robin_map(init.begin(), init.end(), bucket_count, hash, KeyEqual(), alloc) + { + } + + robin_map& operator=(std::initializer_list ilist) { + m_ht.clear(); + + m_ht.reserve(ilist.size()); + m_ht.insert(ilist.begin(), ilist.end()); + + return *this; + } + + allocator_type get_allocator() const { return m_ht.get_allocator(); } + + + /* + * Iterators + */ + iterator begin() noexcept { return m_ht.begin(); } + const_iterator begin() const noexcept { return m_ht.begin(); } + const_iterator cbegin() const noexcept { return m_ht.cbegin(); } + + iterator end() noexcept { return m_ht.end(); } + const_iterator end() const noexcept { return m_ht.end(); } + const_iterator cend() const noexcept { return m_ht.cend(); } + + + /* + * Capacity + */ + bool empty() const noexcept { return m_ht.empty(); } + size_type size() const noexcept { return m_ht.size(); } + size_type max_size() const noexcept { return m_ht.max_size(); } + + /* + * Modifiers + */ + void clear() noexcept { m_ht.clear(); } + + + + std::pair insert(const value_type& value) { + return m_ht.insert(value); + } + + template::value>::type* = nullptr> + std::pair insert(P&& value) { + return m_ht.emplace(std::forward

(value)); + } + + std::pair insert(value_type&& value) { + return m_ht.insert(std::move(value)); + } + + + iterator insert(const_iterator hint, const value_type& value) { + return m_ht.insert_hint(hint, value); + } + + template::value>::type* = nullptr> + iterator insert(const_iterator hint, P&& value) { + return m_ht.emplace_hint(hint, std::forward

(value)); + } + + iterator insert(const_iterator hint, value_type&& value) { + return m_ht.insert_hint(hint, std::move(value)); + } + + + template + void insert(InputIt first, InputIt last) { + m_ht.insert(first, last); + } + + void insert(std::initializer_list ilist) { + m_ht.insert(ilist.begin(), ilist.end()); + } + + + + + template + std::pair insert_or_assign(const key_type& k, M&& obj) { + return m_ht.insert_or_assign(k, std::forward(obj)); + } + + template + std::pair insert_or_assign(key_type&& k, M&& obj) { + return m_ht.insert_or_assign(std::move(k), std::forward(obj)); + } + + template + iterator insert_or_assign(const_iterator hint, const key_type& k, M&& obj) { + return m_ht.insert_or_assign(hint, k, std::forward(obj)); + } + + template + iterator insert_or_assign(const_iterator hint, key_type&& k, M&& obj) { + return m_ht.insert_or_assign(hint, std::move(k), std::forward(obj)); + } + + + + /** + * Due to the way elements are stored, emplace will need to move or copy the key-value once. + * The method is equivalent to insert(value_type(std::forward(args)...)); + * + * Mainly here for compatibility with the std::unordered_map interface. + */ + template + std::pair emplace(Args&&... args) { + return m_ht.emplace(std::forward(args)...); + } + + + + /** + * Due to the way elements are stored, emplace_hint will need to move or copy the key-value once. + * The method is equivalent to insert(hint, value_type(std::forward(args)...)); + * + * Mainly here for compatibility with the std::unordered_map interface. + */ + template + iterator emplace_hint(const_iterator hint, Args&&... args) { + return m_ht.emplace_hint(hint, std::forward(args)...); + } + + + + + template + std::pair try_emplace(const key_type& k, Args&&... args) { + return m_ht.try_emplace(k, std::forward(args)...); + } + + template + std::pair try_emplace(key_type&& k, Args&&... args) { + return m_ht.try_emplace(std::move(k), std::forward(args)...); + } + + template + iterator try_emplace(const_iterator hint, const key_type& k, Args&&... args) { + return m_ht.try_emplace_hint(hint, k, std::forward(args)...); + } + + template + iterator try_emplace(const_iterator hint, key_type&& k, Args&&... args) { + return m_ht.try_emplace_hint(hint, std::move(k), std::forward(args)...); + } + + + + + iterator erase(iterator pos) { return m_ht.erase(pos); } + iterator erase(const_iterator pos) { return m_ht.erase(pos); } + iterator erase(const_iterator first, const_iterator last) { return m_ht.erase(first, last); } + size_type erase(const key_type& key) { return m_ht.erase(key); } + + /** + * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same + * as hash_function()(key). Useful to speed-up the lookup to the value if you already have the hash. + */ + size_type erase(const key_type& key, std::size_t precalculated_hash) { + return m_ht.erase(key, precalculated_hash); + } + + /** + * This overload only participates in the overload resolution if the typedef KeyEqual::is_transparent exists. + * If so, K must be hashable and comparable to Key. + */ + template::value>::type* = nullptr> + size_type erase(const K& key) { return m_ht.erase(key); } + + /** + * @copydoc erase(const K& key) + * + * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same + * as hash_function()(key). Useful to speed-up the lookup to the value if you already have the hash. + */ + template::value>::type* = nullptr> + size_type erase(const K& key, std::size_t precalculated_hash) { + return m_ht.erase(key, precalculated_hash); + } + + + + void swap(robin_map& other) { other.m_ht.swap(m_ht); } + + + + /* + * Lookup + */ + T& at(const Key& key) { return m_ht.at(key); } + + /** + * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same + * as hash_function()(key). Useful to speed-up the lookup if you already have the hash. + */ + T& at(const Key& key, std::size_t precalculated_hash) { return m_ht.at(key, precalculated_hash); } + + + const T& at(const Key& key) const { return m_ht.at(key); } + + /** + * @copydoc at(const Key& key, std::size_t precalculated_hash) + */ + const T& at(const Key& key, std::size_t precalculated_hash) const { return m_ht.at(key, precalculated_hash); } + + + /** + * This overload only participates in the overload resolution if the typedef KeyEqual::is_transparent exists. + * If so, K must be hashable and comparable to Key. + */ + template::value>::type* = nullptr> + T& at(const K& key) { return m_ht.at(key); } + + /** + * @copydoc at(const K& key) + * + * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same + * as hash_function()(key). Useful to speed-up the lookup if you already have the hash. + */ + template::value>::type* = nullptr> + T& at(const K& key, std::size_t precalculated_hash) { return m_ht.at(key, precalculated_hash); } + + + /** + * @copydoc at(const K& key) + */ + template::value>::type* = nullptr> + const T& at(const K& key) const { return m_ht.at(key); } + + /** + * @copydoc at(const K& key, std::size_t precalculated_hash) + */ + template::value>::type* = nullptr> + const T& at(const K& key, std::size_t precalculated_hash) const { return m_ht.at(key, precalculated_hash); } + + + + + T& operator[](const Key& key) { return m_ht[key]; } + T& operator[](Key&& key) { return m_ht[std::move(key)]; } + + + + + size_type count(const Key& key) const { return m_ht.count(key); } + + /** + * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same + * as hash_function()(key). Useful to speed-up the lookup if you already have the hash. + */ + size_type count(const Key& key, std::size_t precalculated_hash) const { + return m_ht.count(key, precalculated_hash); + } + + /** + * This overload only participates in the overload resolution if the typedef KeyEqual::is_transparent exists. + * If so, K must be hashable and comparable to Key. + */ + template::value>::type* = nullptr> + size_type count(const K& key) const { return m_ht.count(key); } + + /** + * @copydoc count(const K& key) const + * + * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same + * as hash_function()(key). Useful to speed-up the lookup if you already have the hash. + */ + template::value>::type* = nullptr> + size_type count(const K& key, std::size_t precalculated_hash) const { return m_ht.count(key, precalculated_hash); } + + + + + iterator find(const Key& key) { return m_ht.find(key); } + + /** + * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same + * as hash_function()(key). Useful to speed-up the lookup if you already have the hash. + */ + iterator find(const Key& key, std::size_t precalculated_hash) { return m_ht.find(key, precalculated_hash); } + + const_iterator find(const Key& key) const { return m_ht.find(key); } + + /** + * @copydoc find(const Key& key, std::size_t precalculated_hash) + */ + const_iterator find(const Key& key, std::size_t precalculated_hash) const { + return m_ht.find(key, precalculated_hash); + } + + /** + * This overload only participates in the overload resolution if the typedef KeyEqual::is_transparent exists. + * If so, K must be hashable and comparable to Key. + */ + template::value>::type* = nullptr> + iterator find(const K& key) { return m_ht.find(key); } + + /** + * @copydoc find(const K& key) + * + * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same + * as hash_function()(key). Useful to speed-up the lookup if you already have the hash. + */ + template::value>::type* = nullptr> + iterator find(const K& key, std::size_t precalculated_hash) { return m_ht.find(key, precalculated_hash); } + + /** + * @copydoc find(const K& key) + */ + template::value>::type* = nullptr> + const_iterator find(const K& key) const { return m_ht.find(key); } + + /** + * @copydoc find(const K& key) + * + * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same + * as hash_function()(key). Useful to speed-up the lookup if you already have the hash. + */ + template::value>::type* = nullptr> + const_iterator find(const K& key, std::size_t precalculated_hash) const { + return m_ht.find(key, precalculated_hash); + } + + + + + bool contains(const Key& key) const { return m_ht.contains(key); } + + /** + * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same + * as hash_function()(key). Useful to speed-up the lookup if you already have the hash. + */ + bool contains(const Key& key, std::size_t precalculated_hash) const { + return m_ht.contains(key, precalculated_hash); + } + + /** + * This overload only participates in the overload resolution if the typedef KeyEqual::is_transparent exists. + * If so, K must be hashable and comparable to Key. + */ + template::value>::type* = nullptr> + bool contains(const K& key) const { return m_ht.contains(key); } + + /** + * @copydoc contains(const K& key) const + * + * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same + * as hash_function()(key). Useful to speed-up the lookup if you already have the hash. + */ + template::value>::type* = nullptr> + bool contains(const K& key, std::size_t precalculated_hash) const { + return m_ht.contains(key, precalculated_hash); + } + + + + + std::pair equal_range(const Key& key) { return m_ht.equal_range(key); } + + /** + * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same + * as hash_function()(key). Useful to speed-up the lookup if you already have the hash. + */ + std::pair equal_range(const Key& key, std::size_t precalculated_hash) { + return m_ht.equal_range(key, precalculated_hash); + } + + std::pair equal_range(const Key& key) const { return m_ht.equal_range(key); } + + /** + * @copydoc equal_range(const Key& key, std::size_t precalculated_hash) + */ + std::pair equal_range(const Key& key, std::size_t precalculated_hash) const { + return m_ht.equal_range(key, precalculated_hash); + } + + /** + * This overload only participates in the overload resolution if the typedef KeyEqual::is_transparent exists. + * If so, K must be hashable and comparable to Key. + */ + template::value>::type* = nullptr> + std::pair equal_range(const K& key) { return m_ht.equal_range(key); } + + + /** + * @copydoc equal_range(const K& key) + * + * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same + * as hash_function()(key). Useful to speed-up the lookup if you already have the hash. + */ + template::value>::type* = nullptr> + std::pair equal_range(const K& key, std::size_t precalculated_hash) { + return m_ht.equal_range(key, precalculated_hash); + } + + /** + * @copydoc equal_range(const K& key) + */ + template::value>::type* = nullptr> + std::pair equal_range(const K& key) const { return m_ht.equal_range(key); } + + /** + * @copydoc equal_range(const K& key, std::size_t precalculated_hash) + */ + template::value>::type* = nullptr> + std::pair equal_range(const K& key, std::size_t precalculated_hash) const { + return m_ht.equal_range(key, precalculated_hash); + } + + + + + /* + * Bucket interface + */ + size_type bucket_count() const { return m_ht.bucket_count(); } + size_type max_bucket_count() const { return m_ht.max_bucket_count(); } + + + /* + * Hash policy + */ + float load_factor() const { return m_ht.load_factor(); } + + float min_load_factor() const { return m_ht.min_load_factor(); } + float max_load_factor() const { return m_ht.max_load_factor(); } + + /** + * Set the `min_load_factor` to `ml`. When the `load_factor` of the map goes + * below `min_load_factor` after some erase operations, the map will be + * shrunk when an insertion occurs. The erase method itself never shrinks + * the map. + * + * The default value of `min_load_factor` is 0.0f, the map never shrinks by default. + */ + void min_load_factor(float ml) { m_ht.min_load_factor(ml); } + void max_load_factor(float ml) { m_ht.max_load_factor(ml); } + + void rehash(size_type count) { m_ht.rehash(count); } + void reserve(size_type count) { m_ht.reserve(count); } + + + /* + * Observers + */ + hasher hash_function() const { return m_ht.hash_function(); } + key_equal key_eq() const { return m_ht.key_eq(); } + + /* + * Other + */ + + /** + * Convert a const_iterator to an iterator. + */ + iterator mutable_iterator(const_iterator pos) { + return m_ht.mutable_iterator(pos); + } + + friend bool operator==(const robin_map& lhs, const robin_map& rhs) { + if(lhs.size() != rhs.size()) { + return false; + } + + for(const auto& element_lhs: lhs) { + const auto it_element_rhs = rhs.find(element_lhs.first); + if(it_element_rhs == rhs.cend() || element_lhs.second != it_element_rhs->second) { + return false; + } + } + + return true; + } + + friend bool operator!=(const robin_map& lhs, const robin_map& rhs) { + return !operator==(lhs, rhs); + } + + friend void swap(robin_map& lhs, robin_map& rhs) { + lhs.swap(rhs); + } + +private: + ht m_ht; +}; + + +/** + * Same as `tsl::robin_map`. + */ +template, + class KeyEqual = std::equal_to, + class Allocator = std::allocator>, + bool StoreHash = false> +using robin_pg_map = robin_map; + +} // end namespace tsl + +#endif diff --git a/include/tsl/include/tsl/robin_set.h b/include/tsl/include/tsl/robin_set.h new file mode 100644 index 0000000000..8f1b9082f8 --- /dev/null +++ b/include/tsl/include/tsl/robin_set.h @@ -0,0 +1,582 @@ +/** + * MIT License + * + * Copyright (c) 2017 Thibaut Goetghebuer-Planchon + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in all + * copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + */ +#ifndef TSL_ROBIN_SET_H +#define TSL_ROBIN_SET_H + + +#include +#include +#include +#include +#include +#include +#include "robin_hash.h" + + +namespace tsl { + + +/** + * Implementation of a hash set using open-addressing and the robin hood hashing algorithm with backward shift deletion. + * + * For operations modifying the hash set (insert, erase, rehash, ...), the strong exception guarantee + * is only guaranteed when the expression `std::is_nothrow_swappable::value && + * std::is_nothrow_move_constructible::value` is true, otherwise if an exception + * is thrown during the swap or the move, the hash set may end up in a undefined state. Per the standard + * a `Key` with a noexcept copy constructor and no move constructor also satisfies the + * `std::is_nothrow_move_constructible::value` criterion (and will thus guarantee the + * strong exception for the set). + * + * When `StoreHash` is true, 32 bits of the hash are stored alongside the values. It can improve + * the performance during lookups if the `KeyEqual` function takes time (or engenders a cache-miss for example) + * as we then compare the stored hashes before comparing the keys. When `tsl::rh::power_of_two_growth_policy` is used + * as `GrowthPolicy`, it may also speed-up the rehash process as we can avoid to recalculate the hash. + * When it is detected that storing the hash will not incur any memory penalty due to alignment (i.e. + * `sizeof(tsl::detail_robin_hash::bucket_entry) == + * sizeof(tsl::detail_robin_hash::bucket_entry)`) and `tsl::rh::power_of_two_growth_policy` is + * used, the hash will be stored even if `StoreHash` is false so that we can speed-up the rehash (but it will + * not be used on lookups unless `StoreHash` is true). + * + * `GrowthPolicy` defines how the set grows and consequently how a hash value is mapped to a bucket. + * By default the set uses `tsl::rh::power_of_two_growth_policy`. This policy keeps the number of buckets + * to a power of two and uses a mask to set the hash to a bucket instead of the slow modulo. + * Other growth policies are available and you may define your own growth policy, + * check `tsl::rh::power_of_two_growth_policy` for the interface. + * + * `Key` must be swappable. + * + * `Key` must be copy and/or move constructible. + * + * If the destructor of `Key` throws an exception, the behaviour of the class is undefined. + * + * Iterators invalidation: + * - clear, operator=, reserve, rehash: always invalidate the iterators. + * - insert, emplace, emplace_hint, operator[]: if there is an effective insert, invalidate the iterators. + * - erase: always invalidate the iterators. + */ +template, + class KeyEqual = std::equal_to, + class Allocator = std::allocator, + bool StoreHash = false, + class GrowthPolicy = tsl::rh::power_of_two_growth_policy<2>> +class robin_set { +private: + template + using has_is_transparent = tsl::detail_robin_hash::has_is_transparent; + + class KeySelect { + public: + using key_type = Key; + + const key_type& operator()(const Key& key) const noexcept { + return key; + } + + key_type& operator()(Key& key) noexcept { + return key; + } + }; + + using ht = detail_robin_hash::robin_hash; + +public: + using key_type = typename ht::key_type; + using value_type = typename ht::value_type; + using size_type = typename ht::size_type; + using difference_type = typename ht::difference_type; + using hasher = typename ht::hasher; + using key_equal = typename ht::key_equal; + using allocator_type = typename ht::allocator_type; + using reference = typename ht::reference; + using const_reference = typename ht::const_reference; + using pointer = typename ht::pointer; + using const_pointer = typename ht::const_pointer; + using iterator = typename ht::iterator; + using const_iterator = typename ht::const_iterator; + + + /* + * Constructors + */ + robin_set(): robin_set(ht::DEFAULT_INIT_BUCKETS_SIZE) { + } + + explicit robin_set(size_type bucket_count, + const Hash& hash = Hash(), + const KeyEqual& equal = KeyEqual(), + const Allocator& alloc = Allocator()): + m_ht(bucket_count, hash, equal, alloc) + { + } + + robin_set(size_type bucket_count, + const Allocator& alloc): robin_set(bucket_count, Hash(), KeyEqual(), alloc) + { + } + + robin_set(size_type bucket_count, + const Hash& hash, + const Allocator& alloc): robin_set(bucket_count, hash, KeyEqual(), alloc) + { + } + + explicit robin_set(const Allocator& alloc): robin_set(ht::DEFAULT_INIT_BUCKETS_SIZE, alloc) { + } + + template + robin_set(InputIt first, InputIt last, + size_type bucket_count = ht::DEFAULT_INIT_BUCKETS_SIZE, + const Hash& hash = Hash(), + const KeyEqual& equal = KeyEqual(), + const Allocator& alloc = Allocator()): robin_set(bucket_count, hash, equal, alloc) + { + insert(first, last); + } + + template + robin_set(InputIt first, InputIt last, + size_type bucket_count, + const Allocator& alloc): robin_set(first, last, bucket_count, Hash(), KeyEqual(), alloc) + { + } + + template + robin_set(InputIt first, InputIt last, + size_type bucket_count, + const Hash& hash, + const Allocator& alloc): robin_set(first, last, bucket_count, hash, KeyEqual(), alloc) + { + } + + robin_set(std::initializer_list init, + size_type bucket_count = ht::DEFAULT_INIT_BUCKETS_SIZE, + const Hash& hash = Hash(), + const KeyEqual& equal = KeyEqual(), + const Allocator& alloc = Allocator()): + robin_set(init.begin(), init.end(), bucket_count, hash, equal, alloc) + { + } + + robin_set(std::initializer_list init, + size_type bucket_count, + const Allocator& alloc): + robin_set(init.begin(), init.end(), bucket_count, Hash(), KeyEqual(), alloc) + { + } + + robin_set(std::initializer_list init, + size_type bucket_count, + const Hash& hash, + const Allocator& alloc): + robin_set(init.begin(), init.end(), bucket_count, hash, KeyEqual(), alloc) + { + } + + + robin_set& operator=(std::initializer_list ilist) { + m_ht.clear(); + + m_ht.reserve(ilist.size()); + m_ht.insert(ilist.begin(), ilist.end()); + + return *this; + } + + allocator_type get_allocator() const { return m_ht.get_allocator(); } + + + /* + * Iterators + */ + iterator begin() noexcept { return m_ht.begin(); } + const_iterator begin() const noexcept { return m_ht.begin(); } + const_iterator cbegin() const noexcept { return m_ht.cbegin(); } + + iterator end() noexcept { return m_ht.end(); } + const_iterator end() const noexcept { return m_ht.end(); } + const_iterator cend() const noexcept { return m_ht.cend(); } + + + /* + * Capacity + */ + bool empty() const noexcept { return m_ht.empty(); } + size_type size() const noexcept { return m_ht.size(); } + size_type max_size() const noexcept { return m_ht.max_size(); } + + /* + * Modifiers + */ + void clear() noexcept { m_ht.clear(); } + + + + + std::pair insert(const value_type& value) { + return m_ht.insert(value); + } + + std::pair insert(value_type&& value) { + return m_ht.insert(std::move(value)); + } + + iterator insert(const_iterator hint, const value_type& value) { + return m_ht.insert_hint(hint, value); + } + + iterator insert(const_iterator hint, value_type&& value) { + return m_ht.insert_hint(hint, std::move(value)); + } + + template + void insert(InputIt first, InputIt last) { + m_ht.insert(first, last); + } + + void insert(std::initializer_list ilist) { + m_ht.insert(ilist.begin(), ilist.end()); + } + + + + + /** + * Due to the way elements are stored, emplace will need to move or copy the key-value once. + * The method is equivalent to insert(value_type(std::forward(args)...)); + * + * Mainly here for compatibility with the std::unordered_map interface. + */ + template + std::pair emplace(Args&&... args) { + return m_ht.emplace(std::forward(args)...); + } + + + + /** + * Due to the way elements are stored, emplace_hint will need to move or copy the key-value once. + * The method is equivalent to insert(hint, value_type(std::forward(args)...)); + * + * Mainly here for compatibility with the std::unordered_map interface. + */ + template + iterator emplace_hint(const_iterator hint, Args&&... args) { + return m_ht.emplace_hint(hint, std::forward(args)...); + } + + + + iterator erase(iterator pos) { return m_ht.erase(pos); } + iterator erase(const_iterator pos) { return m_ht.erase(pos); } + iterator erase(const_iterator first, const_iterator last) { return m_ht.erase(first, last); } + size_type erase(const key_type& key) { return m_ht.erase(key); } + + /** + * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same + * as hash_function()(key). Useful to speed-up the lookup to the value if you already have the hash. + */ + size_type erase(const key_type& key, std::size_t precalculated_hash) { + return m_ht.erase(key, precalculated_hash); + } + + /** + * This overload only participates in the overload resolution if the typedef KeyEqual::is_transparent exists. + * If so, K must be hashable and comparable to Key. + */ + template::value>::type* = nullptr> + size_type erase(const K& key) { return m_ht.erase(key); } + + /** + * @copydoc erase(const K& key) + * + * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same + * as hash_function()(key). Useful to speed-up the lookup to the value if you already have the hash. + */ + template::value>::type* = nullptr> + size_type erase(const K& key, std::size_t precalculated_hash) { + return m_ht.erase(key, precalculated_hash); + } + + + + void swap(robin_set& other) { other.m_ht.swap(m_ht); } + + + + /* + * Lookup + */ + size_type count(const Key& key) const { return m_ht.count(key); } + + /** + * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same + * as hash_function()(key). Useful to speed-up the lookup if you already have the hash. + */ + size_type count(const Key& key, std::size_t precalculated_hash) const { return m_ht.count(key, precalculated_hash); } + + /** + * This overload only participates in the overload resolution if the typedef KeyEqual::is_transparent exists. + * If so, K must be hashable and comparable to Key. + */ + template::value>::type* = nullptr> + size_type count(const K& key) const { return m_ht.count(key); } + + /** + * @copydoc count(const K& key) const + * + * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same + * as hash_function()(key). Useful to speed-up the lookup if you already have the hash. + */ + template::value>::type* = nullptr> + size_type count(const K& key, std::size_t precalculated_hash) const { return m_ht.count(key, precalculated_hash); } + + + + + iterator find(const Key& key) { return m_ht.find(key); } + + /** + * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same + * as hash_function()(key). Useful to speed-up the lookup if you already have the hash. + */ + iterator find(const Key& key, std::size_t precalculated_hash) { return m_ht.find(key, precalculated_hash); } + + const_iterator find(const Key& key) const { return m_ht.find(key); } + + /** + * @copydoc find(const Key& key, std::size_t precalculated_hash) + */ + const_iterator find(const Key& key, std::size_t precalculated_hash) const { return m_ht.find(key, precalculated_hash); } + + /** + * This overload only participates in the overload resolution if the typedef KeyEqual::is_transparent exists. + * If so, K must be hashable and comparable to Key. + */ + template::value>::type* = nullptr> + iterator find(const K& key) { return m_ht.find(key); } + + /** + * @copydoc find(const K& key) + * + * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same + * as hash_function()(key). Useful to speed-up the lookup if you already have the hash. + */ + template::value>::type* = nullptr> + iterator find(const K& key, std::size_t precalculated_hash) { return m_ht.find(key, precalculated_hash); } + + /** + * @copydoc find(const K& key) + */ + template::value>::type* = nullptr> + const_iterator find(const K& key) const { return m_ht.find(key); } + + /** + * @copydoc find(const K& key) + * + * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same + * as hash_function()(key). Useful to speed-up the lookup if you already have the hash. + */ + template::value>::type* = nullptr> + const_iterator find(const K& key, std::size_t precalculated_hash) const { return m_ht.find(key, precalculated_hash); } + + + + + bool contains(const Key& key) const { return m_ht.contains(key); } + + /** + * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same + * as hash_function()(key). Useful to speed-up the lookup if you already have the hash. + */ + bool contains(const Key& key, std::size_t precalculated_hash) const { + return m_ht.contains(key, precalculated_hash); + } + + /** + * This overload only participates in the overload resolution if the typedef KeyEqual::is_transparent exists. + * If so, K must be hashable and comparable to Key. + */ + template::value>::type* = nullptr> + bool contains(const K& key) const { return m_ht.contains(key); } + + /** + * @copydoc contains(const K& key) const + * + * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same + * as hash_function()(key). Useful to speed-up the lookup if you already have the hash. + */ + template::value>::type* = nullptr> + bool contains(const K& key, std::size_t precalculated_hash) const { + return m_ht.contains(key, precalculated_hash); + } + + + + + std::pair equal_range(const Key& key) { return m_ht.equal_range(key); } + + /** + * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same + * as hash_function()(key). Useful to speed-up the lookup if you already have the hash. + */ + std::pair equal_range(const Key& key, std::size_t precalculated_hash) { + return m_ht.equal_range(key, precalculated_hash); + } + + std::pair equal_range(const Key& key) const { return m_ht.equal_range(key); } + + /** + * @copydoc equal_range(const Key& key, std::size_t precalculated_hash) + */ + std::pair equal_range(const Key& key, std::size_t precalculated_hash) const { + return m_ht.equal_range(key, precalculated_hash); + } + + /** + * This overload only participates in the overload resolution if the typedef KeyEqual::is_transparent exists. + * If so, K must be hashable and comparable to Key. + */ + template::value>::type* = nullptr> + std::pair equal_range(const K& key) { return m_ht.equal_range(key); } + + /** + * @copydoc equal_range(const K& key) + * + * Use the hash value 'precalculated_hash' instead of hashing the key. The hash value should be the same + * as hash_function()(key). Useful to speed-up the lookup if you already have the hash. + */ + template::value>::type* = nullptr> + std::pair equal_range(const K& key, std::size_t precalculated_hash) { + return m_ht.equal_range(key, precalculated_hash); + } + + /** + * @copydoc equal_range(const K& key) + */ + template::value>::type* = nullptr> + std::pair equal_range(const K& key) const { return m_ht.equal_range(key); } + + /** + * @copydoc equal_range(const K& key, std::size_t precalculated_hash) + */ + template::value>::type* = nullptr> + std::pair equal_range(const K& key, std::size_t precalculated_hash) const { + return m_ht.equal_range(key, precalculated_hash); + } + + + + + /* + * Bucket interface + */ + size_type bucket_count() const { return m_ht.bucket_count(); } + size_type max_bucket_count() const { return m_ht.max_bucket_count(); } + + + /* + * Hash policy + */ + float load_factor() const { return m_ht.load_factor(); } + + float min_load_factor() const { return m_ht.min_load_factor(); } + float max_load_factor() const { return m_ht.max_load_factor(); } + + /** + * Set the `min_load_factor` to `ml`. When the `load_factor` of the set goes + * below `min_load_factor` after some erase operations, the set will be + * shrunk when an insertion occurs. The erase method itself never shrinks + * the set. + * + * The default value of `min_load_factor` is 0.0f, the set never shrinks by default. + */ + void min_load_factor(float ml) { m_ht.min_load_factor(ml); } + void max_load_factor(float ml) { m_ht.max_load_factor(ml); } + + void rehash(size_type count) { m_ht.rehash(count); } + void reserve(size_type count) { m_ht.reserve(count); } + + + /* + * Observers + */ + hasher hash_function() const { return m_ht.hash_function(); } + key_equal key_eq() const { return m_ht.key_eq(); } + + + /* + * Other + */ + + /** + * Convert a const_iterator to an iterator. + */ + iterator mutable_iterator(const_iterator pos) { + return m_ht.mutable_iterator(pos); + } + + friend bool operator==(const robin_set& lhs, const robin_set& rhs) { + if(lhs.size() != rhs.size()) { + return false; + } + + for(const auto& element_lhs: lhs) { + const auto it_element_rhs = rhs.find(element_lhs); + if(it_element_rhs == rhs.cend()) { + return false; + } + } + + return true; + } + + friend bool operator!=(const robin_set& lhs, const robin_set& rhs) { + return !operator==(lhs, rhs); + } + + friend void swap(robin_set& lhs, robin_set& rhs) { + lhs.swap(rhs); + } + +private: + ht m_ht; +}; + + +/** + * Same as `tsl::robin_set`. + */ +template, + class KeyEqual = std::equal_to, + class Allocator = std::allocator, + bool StoreHash = false> +using robin_pg_set = robin_set; + +} // end namespace tsl + +#endif + diff --git a/include/utils.h b/include/utils.h index 6b9db5bf62..76f3d26631 100644 --- a/include/utils.h +++ b/include/utils.h @@ -1,4 +1,4 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. +// Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT license. #pragma once @@ -13,6 +13,7 @@ #include #include #include +#include #ifdef __APPLE__ #else #include @@ -26,8 +27,9 @@ typedef HANDLE FileHandle; typedef int FileHandle; #endif +#include "distance.h" #include "logger.h" -#include "cached_io.h" +#include "ann_exception.h" #include "common_includes.h" #include "windows_customizations.h" @@ -51,7 +53,9 @@ typedef int FileHandle; #define IS_ALIGNED(X, Y) ((uint64_t)(X) % (uint64_t)(Y) == 0) #define IS_512_ALIGNED(X) IS_ALIGNED(X, 512) #define IS_4096_ALIGNED(X) IS_ALIGNED(X, 4096) - +#define METADATA_SIZE \ + 4096 // all metadata of individual sub-component files is written in first + // 4KB for unified files typedef uint64_t _u64; typedef int64_t _s64; typedef uint32_t _u32; @@ -61,10 +65,130 @@ typedef int16_t _s16; typedef uint8_t _u8; typedef int8_t _s8; +inline bool file_exists(const std::string& name, bool dirCheck = false) { + int val; +#ifndef _WINDOWS + struct stat buffer; + val = stat(name.c_str(), &buffer); +#else + struct _stat64 buffer; + val = _stat64(name.c_str(), &buffer); +#endif + + diskann::cout << " Stat(" << name.c_str() << ") returned: " << val + << std::endl; + if (val != 0) { + switch (errno) { + case EINVAL: + diskann::cout << "Invalid argument passed to stat()" << std::endl; + break; + case ENOENT: + diskann::cout << "File " << name.c_str() << " does not exist" + << std::endl; + break; + default: + diskann::cout << "Unexpected error in stat():" << errno << std::endl; + break; + } + return false; + } else { + // the file entry exists. If reqd, check if this is a directory. + return dirCheck ? buffer.st_mode & S_IFDIR : true; + } +} + +inline std::string getTempFilePath(const std::string& workingDir, + const std::string& suffix) { + std::string retFile; +#ifdef _WINDOWS + char temp[MAX_PATH]; + + do { + if (!tmpnam_s(temp, MAX_PATH) == 0) { + throw diskann::ANNException("Could not create temporary name.", -1); + } + + std::string tempFile(temp); + memset(temp, 0, MAX_PATH); + // GetTempPath returns number of chars in path incl a trailing '\' + int numCharsInTempPath = 0; + if ((numCharsInTempPath = GetTempPathA(MAX_PATH, temp)) == 0) { + throw diskann::ANNException("GetTempPathA failed with error code: ", + GetLastError()); + } + + tempFile.erase(0, numCharsInTempPath); + retFile = workingDir + "\\" + tempFile + "_" + suffix; + } while (file_exists( + retFile)); // To handle the rare case that the file may exist already. +#else + int i = 0; + std::string temp = "temp"; + do { + retFile = workingDir + temp + std::to_string(i) + "_" + suffix; + i++; + } while (file_exists(retFile)); +#endif + return retFile; +} + +inline void open_file_to_write(std::ofstream& writer, + const std::string& filename) { + writer.exceptions(std::ofstream::failbit | std::ofstream::badbit); + if (!file_exists(filename)) + writer.open(filename, std::ios::binary | std::ios::out); + else + writer.open(filename, std::ios::binary | std::ios::in | std::ios::out); + + if (writer.fail()) { + diskann::cerr << std::string("Failed to open file") + filename + + " for write because " + << std::strerror(errno) << std::endl; + throw diskann::ANNException( + std::string("Failed to open file ") + filename + + " for write because: " + std::strerror(errno), + -1); + } +} + +inline _u64 get_file_size(const std::string& fname) { + std::ifstream reader(fname, std::ios::binary | std::ios::ate); + if (!reader.fail() && reader.is_open()) { + _u64 end_pos = reader.tellg(); + reader.close(); + return end_pos; + } else { + diskann::cerr << "Could not open file: " << fname << std::endl; + return 0; + } +} + +inline int delete_file(const std::string& fileName) { + if (file_exists(fileName)) { + auto rc = ::remove(fileName.c_str()); + if (rc != 0) { + diskann::cerr + << "Could not delete file: " << fileName + << " even though it exists. This might indicate a permissions issue. " + "If you see this message, please contact the diskann team." + << std::endl; + } + return rc; + } else { + return 0; + } +} + namespace diskann { static const size_t MAX_SIZE_OF_STREAMBUF = 2LL * 1024 * 1024 * 1024; - enum Metric { L2 = 0, INNER_PRODUCT = 1, FAST_L2 = 2, PQ = 3 }; + enum Metric { L2 = 0, INNER_PRODUCT = 1, FAST_L2 = 2, PQ = 3, COSINE = 4 }; + + DISKANN_DLLEXPORT float calc_recall_set_tags( + unsigned num_queries, unsigned* gold_std, unsigned dim_gs, + unsigned* our_results_tags, unsigned dim_or, unsigned recall_at, + unsigned subset_size, std::string gt_tag_filename, + std::string current_tag_filename); inline void alloc_aligned(void** ptr, size_t size, size_t align) { *ptr = nullptr; @@ -77,9 +201,33 @@ namespace diskann { assert(*ptr != nullptr); } + inline void realloc_aligned(void** ptr, size_t size, size_t align) { + assert(IS_ALIGNED(size, align)); +#ifdef _WINDOWS + *ptr = ::_aligned_realloc(*ptr, size, align); +#endif + assert(*ptr != nullptr); + } + + inline void realloc_aligned(void** ptr, void** ptr_new, size_t old_size, + size_t new_size, size_t align) { + assert(IS_ALIGNED(new_size, align)); +#ifndef _WINDOWS + alloc_aligned((void**) &ptr_new, new_size, align); + memcpy(*ptr_new, *ptr, old_size); + ::free(*ptr); + *ptr = *ptr_new; +#endif + assert(*ptr != nullptr); + } + + inline void check_stop(std::string arnd) { + int brnd; + diskann::cout << arnd << std::endl; + std::cin >> brnd; + } + inline void aligned_free(void* ptr) { - // Gopal. Must have a check here if the pointer was actually allocated by - // _alloc_aligned if (ptr == nullptr) { return; } @@ -110,8 +258,10 @@ namespace diskann { // get_bin_metadata functions START inline void get_bin_metadata_impl(std::basic_istream& reader, - size_t& nrows, size_t& ncols) { + size_t& nrows, size_t& ncols, + size_t offset = 0) { int nrows_32, ncols_32; + reader.seekg(offset, reader.beg); reader.read((char*) &nrows_32, sizeof(int)); reader.read((char*) &ncols_32, sizeof(int)); nrows = nrows_32; @@ -121,19 +271,22 @@ namespace diskann { #ifdef EXEC_ENV_OLS inline void get_bin_metadata(MemoryMappedFiles& files, const std::string& bin_file, size_t& nrows, - size_t& ncols) { + size_t& ncols, size_t offset = 0) { diskann::cout << "Getting metadata for file: " << bin_file << std::endl; - auto fc = files.getContent(bin_file); - auto cb = ContentBuf((char*) fc._content, fc._size); - std::basic_istream reader(&cb); - get_bin_metadata_impl(reader, nrows, ncols); + auto fc = files.getContent(bin_file); + int nrows_32, ncols_32; + int32_t* metadata_ptr = (int32_t*) ((char*) fc._content + offset); + nrows_32 = *metadata_ptr; + ncols_32 = *(metadata_ptr + 1); + nrows = nrows_32; + ncols = ncols_32; } #endif inline void get_bin_metadata(const std::string& bin_file, size_t& nrows, - size_t& ncols) { + size_t& ncols, size_t offset = 0) { std::ifstream reader(bin_file.c_str(), std::ios::binary); - get_bin_metadata_impl(reader, nrows, ncols); + get_bin_metadata_impl(reader, nrows, ncols, offset); } // get_bin_metadata functions END @@ -151,10 +304,11 @@ namespace diskann { // load_bin functions START template - inline void load_bin_impl(std::basic_istream& reader, - size_t actual_file_size, T*& data, size_t& npts, - size_t& dim) { + inline void load_bin_impl(std::basic_istream& reader, T*& data, + size_t& npts, size_t& dim, size_t file_offset = 0) { int npts_i32, dim_i32; + + reader.seekg(file_offset, reader.beg); reader.read((char*) &npts_i32, sizeof(int)); reader.read((char*) &dim_i32, sizeof(int)); npts = (unsigned) npts_i32; @@ -163,64 +317,34 @@ namespace diskann { diskann::cout << "Metadata: #pts = " << npts << ", #dims = " << dim << "..." << std::endl; - size_t expected_actual_file_size = - npts * dim * sizeof(T) + 2 * sizeof(uint32_t); - if (actual_file_size != expected_actual_file_size) { - std::stringstream stream; - stream << "Error. File size mismatch. Actual size is " << actual_file_size - << " while expected size is " << expected_actual_file_size - << " npts = " << npts << " dim = " << dim - << " size of = " << sizeof(T) << std::endl; - diskann::cout << stream.str(); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); - } - data = new T[npts * dim]; reader.read((char*) data, npts * dim * sizeof(T)); - - // diskann::cout << "Last bytes: " - // << getValues(data + (npts - 2) * dim, dim); - // diskann::cout << "Finished reading bin file." << std::endl; } #ifdef EXEC_ENV_OLS template inline void load_bin(MemoryMappedFiles& files, const std::string& bin_file, - T*& data, size_t& npts, size_t& dim) { - diskann::cout << "Reading bin file " << bin_file.c_str() << " ..." - << std::endl; - + T*& data, size_t& npts, size_t& dim, size_t offset = 0) { + diskann::cout << "Reading bin file " << bin_file.c_str() + << " at offset: " << offset << "..." << std::endl; auto fc = files.getContent(bin_file); uint32_t t_npts, t_dim; - uint32_t* contentAsIntPtr = (uint32_t*) (fc._content); + uint32_t* contentAsIntPtr = (uint32_t*) ((char*) fc._content + offset); t_npts = *(contentAsIntPtr); t_dim = *(contentAsIntPtr + 1); npts = t_npts; dim = t_dim; - auto actual_file_size = npts * dim * sizeof(T) + 2 * sizeof(uint32_t); - if (actual_file_size != fc._size) { - std::stringstream stream; - stream << "Error. File size mismatch. Actual size is " << fc._size - << " while expected size is " << actual_file_size - << " npts = " << npts << " dim = " << dim - << " size of = " << sizeof(T) << std::endl; - diskann::cout << stream.str(); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); - } - - data = - (T*) ((char*) fc._content + 2 * sizeof(uint32_t)); // No need to copy! + data = (T*) ((char*) fc._content + offset + + 2 * sizeof(uint32_t)); // No need to copy! } #endif template inline void load_bin(const std::string& bin_file, T*& data, size_t& npts, - size_t& dim) { + size_t& dim, size_t offset = 0) { // OLS //_u64 read_blk_size = 64 * 1024 * 1024; // cached_ifstream reader(bin_file, read_blk_size); @@ -229,20 +353,20 @@ namespace diskann { diskann::cout << "Reading bin file " << bin_file.c_str() << " ..." << std::endl; std::ifstream reader(bin_file, std::ios::binary | std::ios::ate); - uint64_t fsize = reader.tellg(); + // uint64_t fsize = reader.tellg(); reader.seekg(0); - load_bin_impl(reader, fsize, data, npts, dim); + load_bin_impl(reader, data, npts, dim, offset); } // load_bin functions END inline void load_truthset(const std::string& bin_file, uint32_t*& ids, - float*& dists, size_t& npts, size_t& dim) { - _u64 read_blk_size = 64 * 1024 * 1024; - cached_ifstream reader(bin_file, read_blk_size); - diskann::cout << "Reading truthset file " << bin_file.c_str() << " ..." + float*& dists, size_t& npts, size_t& dim, + uint32_t** tags = nullptr) { + std::ifstream reader(bin_file, std::ios::binary); + diskann::cout << "Reading truthset file " << bin_file.c_str() << "..." << std::endl; - size_t actual_file_size = reader.get_file_size(); + size_t actual_file_size = get_file_size(bin_file); int npts_i32, dim_i32; reader.read((char*) &npts_i32, sizeof(int)); @@ -264,9 +388,15 @@ namespace diskann { size_t expected_file_size_just_ids = npts * dim * sizeof(uint32_t) + 2 * sizeof(uint32_t); + size_t with_tags_actual_file_size = + 3 * npts * dim * sizeof(uint32_t) + 2 * sizeof(uint32_t); + if (actual_file_size == expected_file_size_just_ids) truthset_type = 2; + if (actual_file_size == with_tags_actual_file_size) + truthset_type = 3; + if (truthset_type == -1) { std::stringstream stream; stream << "Error. File size mismatch. File should have bin format, with " @@ -283,120 +413,141 @@ namespace diskann { ids = new uint32_t[npts * dim]; reader.read((char*) ids, npts * dim * sizeof(uint32_t)); - if (truthset_type == 1) { + if ((truthset_type == 1) || (truthset_type == 3)) { dists = new float[npts * dim]; reader.read((char*) dists, npts * dim * sizeof(float)); } + if (truthset_type == 3) { + *tags = new uint32_t[npts * dim]; + reader.read((char*) *tags, npts * dim * sizeof(uint32_t)); + } } #ifdef EXEC_ENV_OLS template inline void load_bin(MemoryMappedFiles& files, const std::string& bin_file, - std::unique_ptr& data, size_t& npts, size_t& dim) { + std::unique_ptr& data, size_t& npts, size_t& dim, + size_t offset = 0) { T* ptr; - load_bin(files, bin_file, ptr, npts, dim); + load_bin(files, bin_file, ptr, npts, dim, offset); data.reset(ptr); } + #endif template inline void load_bin(const std::string& bin_file, std::unique_ptr& data, - size_t& npts, size_t& dim) { + size_t& npts, size_t& dim, size_t offset = 0) { T* ptr; - load_bin(bin_file, ptr, npts, dim); + load_bin(bin_file, ptr, npts, dim, offset); data.reset(ptr); } template - inline void save_bin(const std::string& filename, T* data, size_t npts, - size_t ndims) { - std::ofstream writer(filename, std::ios::binary | std::ios::out); + inline uint64_t save_bin(const std::string& filename, T* data, size_t npts, + size_t ndims, size_t offset = 0) { + std::ofstream writer; + open_file_to_write(writer, filename); + diskann::cout << "Writing bin: " << filename.c_str() << std::endl; - int npts_i32 = (int) npts, ndims_i32 = (int) ndims; + writer.seekp(offset, writer.beg); + int npts_i32 = (int) npts, ndims_i32 = (int) ndims; + size_t bytes_written = npts * ndims * sizeof(T) + 2 * sizeof(uint32_t); writer.write((char*) &npts_i32, sizeof(int)); writer.write((char*) &ndims_i32, sizeof(int)); diskann::cout << "bin: #pts = " << npts << ", #dims = " << ndims - << ", size = " << npts * ndims * sizeof(T) + 2 * sizeof(int) - << "B" << std::endl; + << ", size = " << bytes_written << "B" << std::endl; - // data = new T[npts_u64 * ndims_u64]; writer.write((char*) data, npts * ndims * sizeof(T)); writer.close(); diskann::cout << "Finished writing bin." << std::endl; + return bytes_written; } // load_aligned_bin functions START template - inline void load_aligned_bin_impl(std::basic_istream& reader, - size_t actual_file_size, T*& data, + inline void load_aligned_bin_impl(std::basic_istream& reader, T*& data, size_t& npts, size_t& dim, - size_t& rounded_dim) { + size_t& rounded_dim, size_t offset = 0) { int npts_i32, dim_i32; + reader.seekg(offset, reader.beg); reader.read((char*) &npts_i32, sizeof(int)); reader.read((char*) &dim_i32, sizeof(int)); + npts = (unsigned) npts_i32; dim = (unsigned) dim_i32; - - size_t expected_actual_file_size = - npts * dim * sizeof(T) + 2 * sizeof(uint32_t); - if (actual_file_size != expected_actual_file_size) { - std::stringstream stream; - stream << "Error. File size mismatch. Actual size is " << actual_file_size - << " while expected size is " << expected_actual_file_size - << " npts = " << npts << " dim = " << dim - << " size of = " << sizeof(T) << std::endl; - diskann::cout << stream.str() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); - } rounded_dim = ROUND_UP(dim, 8); diskann::cout << "Metadata: #pts = " << npts << ", #dims = " << dim << ", aligned_dim = " << rounded_dim << "..." << std::flush; size_t allocSize = npts * rounded_dim * sizeof(T); - diskann::cout << "allocating aligned memory, " << allocSize << " bytes..." - << std::flush; alloc_aligned(((void**) &data), allocSize, 8 * sizeof(T)); - diskann::cout << "done. Copying data..." << std::flush; for (size_t i = 0; i < npts; i++) { reader.read((char*) (data + i * rounded_dim), dim * sizeof(T)); memset(data + i * rounded_dim + dim, 0, (rounded_dim - dim) * sizeof(T)); } - diskann::cout << " done." << std::endl; + diskann::cout << " Allocated " << allocSize << "bytes and copied data " + << std::endl; } #ifdef EXEC_ENV_OLS template inline void load_aligned_bin(MemoryMappedFiles& files, const std::string& bin_file, T*& data, - size_t& npts, size_t& dim, size_t& rounded_dim) { - diskann::cout << "Reading bin file " << bin_file << " ..." << std::flush; - FileContent fc = files.getContent(bin_file); - ContentBuf buf((char*) fc._content, fc._size); - std::basic_istream reader(&buf); + size_t& npts, size_t& dim, size_t& rounded_dim, + size_t offset = 0) { + diskann::cout << "Reading bin file " << bin_file << " at offset: " << offset + << "..." << std::flush; + FileContent fc = files.getContent(bin_file); + // ContentBuf buf((char*) fc._content, fc._size); + + char* read_addr = (((char*) fc._content) + offset); + + int npts_32 = *((int*) read_addr); + int ndim_32 = *((int*) (read_addr + sizeof(int))); + + npts = (uint32_t) npts_32; + dim = (uint32_t) ndim_32; + + char* data_start = ((char*) fc._content) + offset + 2 * sizeof(int); + rounded_dim = ROUND_UP(dim, 8); + diskann::cout << "Metadata: #pts = " << npts << ", #dims = " << dim + << ", aligned_dim = " << rounded_dim << "..." << std::flush; + + size_t allocSize = npts * rounded_dim * sizeof(T); + + diskann::cout << "allocating aligned memory, " << allocSize << " bytes..." + << std::flush; + + alloc_aligned(((void**) &data), allocSize, 8 * sizeof(T)); + diskann::cout << "done. Copying data..." << std::flush; + + for (size_t i = 0; i < npts; i++) { + memcpy((data + i * rounded_dim), data_start, dim * sizeof(T)); + memset(data + i * rounded_dim + dim, 0, (rounded_dim - dim) * sizeof(T)); + data_start += dim * sizeof(T); + } - size_t actual_file_size = fc._size; - load_aligned_bin_impl(reader, actual_file_size, data, npts, dim, - rounded_dim); + if (data_start - read_addr != dim * sizeof(T) * npts + 2 * sizeof(int)) { + diskann::cerr << "Read " << data_start - read_addr + << " bytes of data instead of: " << dim * sizeof(T) * npts + << std::endl; + } + diskann::cout << " done." << std::endl; } #endif template inline void load_aligned_bin(const std::string& bin_file, T*& data, - size_t& npts, size_t& dim, size_t& rounded_dim) { - diskann::cout << "Reading bin file " << bin_file << " ..." << std::flush; - // START OLS - //_u64 read_blk_size = 64 * 1024 * 1024; - // cached_ifstream reader(bin_file, read_blk_size); - // size_t actual_file_size = reader.get_file_size(); - // END OLS - + size_t& npts, size_t& dim, size_t& rounded_dim, + size_t offset = 0) { + diskann::cout << "Reading bin file " << bin_file << " at offset " << offset + << "..." << std::flush; std::ifstream reader(bin_file, std::ios::binary | std::ios::ate); - uint64_t fsize = reader.tellg(); reader.seekg(0); - load_aligned_bin_impl(reader, fsize, data, npts, dim, rounded_dim); + load_aligned_bin_impl(reader, data, npts, dim, rounded_dim, offset); } template @@ -410,24 +561,58 @@ namespace diskann { } } - // plain saves data as npts X ndims array into filename template - void save_Tvecs(const char* filename, T* data, size_t npts, size_t ndims) { - std::string fname(filename); + inline void load_aligned_bin(const std::string& bin_file, + std::unique_ptr& data, size_t& npts, + size_t& dim, size_t& rounded_dim, + size_t offset = 0) { + T* ptr; + load_aligned_bin(bin_file, ptr, npts, dim, rounded_dim, offset); + data.reset(ptr); + } - // create cached ofstream with 64MB cache - cached_ofstream writer(fname, 64 * 1048576); + template + inline uint64_t save_data_in_base_dimensions(const std::string& filename, + T* data, size_t npts, + size_t ndims, size_t aligned_dim, + size_t offset = 0) { + std::ofstream writer; //(filename, std::ios::binary | std::ios::out); + open_file_to_write(writer, filename); + int npts_i32 = (int) npts, ndims_i32 = (int) ndims; + _u64 bytes_written = 2 * sizeof(uint32_t) + npts * ndims * sizeof(T); + writer.seekp(offset, writer.beg); + writer.write((char*) &npts_i32, sizeof(int)); + writer.write((char*) &ndims_i32, sizeof(int)); + for (size_t i = 0; i < npts; i++) { + writer.write((char*) (data + i * aligned_dim), ndims * sizeof(T)); + } + writer.close(); + return bytes_written; + } - unsigned dims_u32 = (unsigned) ndims; + template + inline void copy_aligned_data_from_file(const std::string bin_file, T*& data, + size_t& npts, size_t& dim, + const size_t& rounded_dim, + size_t offset = 0) { + if (data == nullptr) { + diskann::cout << "Memory was not allocated for " << data + << " before calling the load function. Exiting..." + << std::endl; + exit(-1); + } + std::ifstream reader(bin_file, std::ios::binary); + reader.seekg(offset, reader.beg); - // start writing - for (uint64_t i = 0; i < npts; i++) { - // write dims in u32 - writer.write((char*) &dims_u32, sizeof(unsigned)); + int npts_i32, dim_i32; + reader.read((char*) &npts_i32, sizeof(int)); + reader.read((char*) &dim_i32, sizeof(int)); + npts = (unsigned) npts_i32; + dim = (unsigned) dim_i32; - // get cur point in data - T* cur_pt = data + i * ndims; - writer.write((char*) cur_pt, ndims * sizeof(T)); + for (size_t i = 0; i < npts; i++) { + reader.read((char*) (data + i * rounded_dim), dim * sizeof(T)); + memset(data + i * rounded_dim + dim, 0, (rounded_dim - dim) * sizeof(T)); } } @@ -444,7 +629,17 @@ namespace diskann { for (size_t d = 0; d < max_prefetch_size; d += 64) _mm_prefetch((const char*) vec + d, _MM_HINT_T1); } -}; // namespace diskann + + // NOTE: Implementation in utils.cpp. + void block_convert(std::ofstream& writr, std::ifstream& readr, + float* read_buf, _u64 npts, _u64 ndims); + + DISKANN_DLLEXPORT void normalize_data_file(const std::string& inFileName, + const std::string& outFileName); + + template + Distance* get_distance_function(Metric m); +} // namespace diskann struct PivotContainer { PivotContainer() = default; @@ -465,28 +660,6 @@ struct PivotContainer { float piv_dist; }; -inline bool file_exists(const std::string& name) { - struct stat buffer; - auto val = stat(name.c_str(), &buffer); - diskann::cout << " Stat(" << name.c_str() << ") returned: " << val - << std::endl; - return (val == 0); -} - -inline _u64 get_file_size(const std::string& fname) { - std::ifstream reader(fname, std::ios::binary | std::ios::ate); - if (!reader.fail() && reader.is_open()) { - _u64 end_pos = reader.tellg(); - diskann::cout << " Tellg: " << reader.tellg() << " as u64: " << end_pos - << std::endl; - reader.close(); - return end_pos; - } else { - diskann::cout << "Could not open file: " << fname << std::endl; - return 0; - } -} - inline bool validate_file_size(const std::string& name) { std::ifstream in(std::string(name), std::ios::binary); in.seekg(0, in.end); @@ -495,11 +668,12 @@ inline bool validate_file_size(const std::string& name) { size_t expected_file_size; in.read((char*) &expected_file_size, sizeof(uint64_t)); if (actual_file_size != expected_file_size) { - diskann::cout << "Error loading" << name << ". Expected " - "size (metadata): " + diskann::cerr << "Error loading" << name + << ". Expected " + "size (metadata): " << expected_file_size - << ", actual file size : " << actual_file_size - << ". Exitting." << std::endl; + << ", actual file size : " << actual_file_size << ". Exiting." + << std::endl; in.close(); return false; } @@ -507,18 +681,50 @@ inline bool validate_file_size(const std::string& name) { return true; } +template +diskann::Distance* get_distance_function(diskann::Metric m); + +extern bool AvxSupportedCPU; +extern bool Avx2SupportedCPU; + #ifdef _WINDOWS #include #include +inline size_t getMemoryUsage() { + PROCESS_MEMORY_COUNTERS_EX pmc; + GetProcessMemoryInfo(GetCurrentProcess(), (PROCESS_MEMORY_COUNTERS*) &pmc, + sizeof(pmc)); + return pmc.PrivateUsage; +} + +inline std::string getWindowsErrorMessage(DWORD lastError) { + char* errorText; + FormatMessageA( + // use system message tables to retrieve error text + FORMAT_MESSAGE_FROM_SYSTEM + // allocate buffer on local heap for error text + | FORMAT_MESSAGE_ALLOCATE_BUFFER + // Important! will fail otherwise, since we're not + // (and CANNOT) pass insertion parameters + | FORMAT_MESSAGE_IGNORE_INSERTS, + NULL, // unused with FORMAT_MESSAGE_FROM_SYSTEM + lastError, MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), + (LPSTR) &errorText, // output + 0, // minimum size for output buffer + NULL); // arguments - see note + + return errorText != nullptr ? std::string(errorText) : std::string(); +} + inline void printProcessMemory(const char* message) { PROCESS_MEMORY_COUNTERS counters; HANDLE h = GetCurrentProcess(); GetProcessMemoryInfo(h, &counters, sizeof(counters)); diskann::cout << message << " [Peaking Working Set size: " - << counters.PeakWorkingSetSize * 1.0 / (1024 * 1024 * 1024) + << counters.PeakWorkingSetSize * 1.0 / (1024.0 * 1024 * 1024) << "GB Working set size: " - << counters.WorkingSetSize * 1.0 / (1024 * 1024 * 1024) + << counters.WorkingSetSize * 1.0 / (1024.0 * 1024 * 1024) << "GB Private bytes " << counters.PagefileUsage * 1.0 / (1024 * 1024 * 1024) << "GB]" << std::endl; @@ -529,11 +735,6 @@ inline void printProcessMemory(const char* message) { inline bool avx2Supported() { return true; } - -inline void printProcessMemory(const char* message) { - diskann::cout << message << std::endl; +inline void printProcessMemory(const char*) { } #endif - -extern bool AvxSupportedCPU; -extern bool Avx2SupportedCPU; diff --git a/include/v2/aux_dist.h b/include/v2/aux_dist.h new file mode 100644 index 0000000000..986977159f --- /dev/null +++ b/include/v2/aux_dist.h @@ -0,0 +1,5 @@ +#pragma once + +namespace diskann { + class FloatFloat +} // namespace diskann \ No newline at end of file diff --git a/include/v2/delete_set.h b/include/v2/delete_set.h new file mode 100644 index 0000000000..654b43dd18 --- /dev/null +++ b/include/v2/delete_set.h @@ -0,0 +1,37 @@ +#pragma once + +#include "v2/graph_delta.h" +#include "tsl/robin_map.h" +#include "tsl/robin_set.h" +#include +#include +#include +#include +#include + +namespace diskann { + class DeleteSet { + public: + // max # track_merge calls before `id` becomes free + DeleteSet(uint32_t max_merges); + ~DeleteSet(); + + // adds `id` to deleted set + void add_delete(uint32_t id); + + // checks if `id` is in delete set + bool is_dead(uint32_t id); + void batch_is_dead(const uint32_t *ids, bool* dead, const uint32_t count); + + // track merge + release merged nodes + void merge_start(); + + // returns nodes + std::vector track_merge(); + private: + tsl::robin_map *primary = nullptr; + tsl::robin_map *secondary = nullptr; + uint32_t max_merges; + std::mutex lock; + }; +} // namespace diskann \ No newline at end of file diff --git a/include/v2/fs_allocator.h b/include/v2/fs_allocator.h new file mode 100644 index 0000000000..d360652d36 --- /dev/null +++ b/include/v2/fs_allocator.h @@ -0,0 +1,30 @@ +#pragma once + +#include "tsl/robin_map.h" +#include "tsl/robin_set.h" +#include +#include +#include +#include +#include + +namespace diskann { + // cached allocator for fast aligned mallocs + template + class FixedSizeAlignedAllocator { + public: + // create aligned buffer with at least max_count * ndims elements + FixedSizeAlignedAllocator(const uint32_t ndims, const uint32_t max_count); + // destruct allocator, free mem + ~FixedSizeAlignedAllocator(); + // allocate ndims buffer + T* allocate(); + // deallocate ndims elements + void deallocate(T* ptr); + private: + std::mutex lock; + T* buf = nullptr; + tsl::robin_set free_set; + uint32_t count; + }; +} // namespace diskann \ No newline at end of file diff --git a/include/v2/graph_delta.h b/include/v2/graph_delta.h new file mode 100644 index 0000000000..be395cddd3 --- /dev/null +++ b/include/v2/graph_delta.h @@ -0,0 +1,36 @@ +#pragma once + +#include +#include +#include +#include +#include + +namespace diskann { + class GraphDelta { + public: + GraphDelta(const uint32_t offset, const uint32_t max_nodes); + // inserts node `id` into graph with `nhood` as neighbors + // SUCCEEDS ONLY IF `id` belongs to the range [offset, offset + max_nodes] + void insert_vector(const uint32_t id, const uint32_t*nhood, const uint32_t nnbrs); + + // adds required back-edges from `srcs` to `dest` + void inter_insert(const uint32_t dest, const uint32_t* srcs, const uint32_t src_count); + + // get nhood for single ID + const std::vector get_nhood(const uint32_t id); + + void rename_edges(const tsl::robin_map& rename_map); + void rename_edges(const std::function &rename_func); + private: + bool is_relevant(const uint32_t id); + // in-memory graph + std::vector> graph; + // locks to access nodes in graph + std::unique_ptr locks; + // max nodes + uint32_t offset; + uint32_t max_nodes; + // id 'n' nhood located at graph[n - offset] if offset <= n <= offset + max_nodes + }; +}; diff --git a/include/v2/index_merger.h b/include/v2/index_merger.h new file mode 100644 index 0000000000..c3323a8759 --- /dev/null +++ b/include/v2/index_merger.h @@ -0,0 +1,184 @@ +#pragma once + +#include "v2/graph_delta.h" +#include "tsl/robin_map.h" +#include "tsl/robin_set.h" +#include "pq_flash_index.h" +#include "linux_aligned_file_reader.h" +#include "index.h" +#include +#include +#include +#include +#include +#include "windows_customizations.h" + +namespace diskann { + template + class StreamingMerger { + public: + // constructor to read a constructed index, allocated IDs + // disk_in : SSD-DiskANN index to merge into + // mem_in : list of mem-DiskANN indices to merge into disk_in + // disk_out : SSD-DiskANN index to write out + // delete_list : list of IDs to delete from disk_in + // ndims : dimensionality of full-prec vectors + // dist : distance comparator -- WARNING :: assumed to be L2 + // beam_width : BW for search on disk_in + // range : max out-degree + // l_index : L param for indexing + // maxc : max num of candidates to consider while pruning + /* StreamingMerger(const char* disk_in, const std::vector + &mem_in, const char* disk_out, const char* deleted_tags, const uint32_t + ndims, Distance* dist, const uint32_t beam_width, const uint32_t + range, const uint32_t l_index, const float alpha, const uint32_t maxc, + bool single_file_index);*/ + DISKANN_DLLEXPORT StreamingMerger(const uint32_t ndims, Distance *dist, + diskann::Metric dist_metric, + const uint32_t beam_width, + const uint32_t range, + const uint32_t l_index, const float alpha, + const uint32_t maxc, + bool single_file_index); + + DISKANN_DLLEXPORT ~StreamingMerger(); + + DISKANN_DLLEXPORT void merge(const char * disk_in, + const std::vector &mem_in, + const char * disk_out, + std::vector*> &deleted_tags, + std::string &working_folder); + // merge all memory indices into the disk index and write out new disk index + void mergeImpl(); + + private: + /* insert related funcs */ + void process_inserts(); + void process_inserts_pq(); + void insert_mem_vec(const T *vec, const uint32_t offset_id); + void offset_iterate_to_fixed_point( + const T *vec, const uint32_t Lsize, + std::vector & expanded_nodes_info, + tsl::robin_map &coord_map); + // used to prune insert() edges + void prune_neighbors(const tsl::robin_map &coord_map, + std::vector & pool, + std::vector & pruned_list); + // used to prune inter-insert() edges + void prune_neighbors_pq(std::vector &pool, + std::vector &pruned_list, + uint8_t * scratch = nullptr); + void occlude_list(std::vector & pool, + const tsl::robin_map &coord_map, + std::vector & result, + std::vector & occlude_factor); + void occlude_list_pq(std::vector &pool, + std::vector &result, + std::vector & occlude_factor, + uint8_t * scratch = nullptr); + + void dump_to_disk(const uint32_t start_id, const char *buf, + const uint32_t n_sector, std::ofstream& output_writer); + + /* delete related funcs */ + // converts tags into deleted IDs + void compute_deleted_ids(); + // process all deletes + void process_deletes(); + // reads nhoods of all deleted nods + void populate_deleted_nhoods(); + // eliminates references to deleted nodes in id_nhoods + void consolidate_deletes(DiskNode &disk_node, + uint8_t * scratch = nullptr); + // whether the specific node is deleted / node id not in use + bool is_deleted(const DiskNode &disk_node); + + /* rename related funcs */ + // assign smallest free IDs to new inserts + void compute_rename_map(); + void rename(DiskNode &node) const; + void rename(std::vector &nhood) const; + // returns uint32_t::max() upon failure + uint32_t rename(uint32_t id) const; + // returns uint32_t::max() upon failure + uint32_t rename_inverse(uint32_t renamed_id) const; + // returns ID of mem index offset_id belongs to; uint32_t::max() otherwise + uint32_t get_index_id(const uint32_t offset_id) const; + std::vector get_edge_list(const uint32_t offset_id); + const T * get_mem_data(const uint32_t offset_id); + + /* merge related funcs */ + void write_tag_file(const std::string &tag_out_filename, + const uint32_t npts); + void process_merges(); + + // deletes + tsl::robin_set deleted_tags; + tsl::robin_map> disk_deleted_nhoods; + tsl::robin_set disk_deleted_ids; + std::vector> mem_deleted_ids; + char *delete_backing_buf = nullptr; + + // rename stuff + tsl::robin_map rename_map; + tsl::robin_map inverse_map; + std::vector> rename_list; + std::vector> inverse_list; + + // disk index + GraphDelta * disk_delta; + PQFlashIndex *disk_index; + std::vector init_ids; + uint8_t * pq_data = nullptr; + TagT * disk_tags = nullptr; + uint32_t pq_nchunks; + uint32_t max_node_len, nnodes_per_sector, disk_npts; + std::string disk_index_out_path, disk_index_in_path, pq_coords_file; + std::string temp_disk_index_path, temp_pq_coords_path, temp_tags_path; + std::string final_index_file, final_pq_coords_file, final_tags_file; + //std::fstream output_writer; + std::vector> disk_thread_data; + + // mem-index + std::vector mem_deltas; + //std::vector *> mem_indices; + //std::vector> *> mem_graphs; + std::vector mem_data; + std::vector> mem_tags; + std::vector offset_ids; + std::vector mem_npts; + Distance * dist_cmp; + diskann::Metric dist_metric; + //T * _data_load; + std::vector> latter_deleted_tags; + + + + + // allocators + // FixedSizeAlignedAllocator *fp_alloc = nullptr; + // FixedSizeAlignedAllocator *pq_alloc = nullptr; + + // book keeping + std::vector free_ids; + uint8_t * thread_pq_scratch = nullptr; + std::vector thread_bufs; + // vector info + uint32_t ndims, aligned_ndims; + // search + index params + uint32_t beam_width; + uint32_t l_index, range, maxc; + float alpha; + + uint64_t disk_index_num_frozen = 0; // assuming that every merge + // folds a single mem index into the PQFlashIndex + uint64_t disk_index_frozen_loc = 0; + // timing stuff + std::vector insert_times, delta_times, search_times, prune_times; + std::mutex _print_lock; + + bool _single_file_index = false; + + std::string TMP_FOLDER; + }; +}; // namespace diskann diff --git a/include/v2/merge_insert.h b/include/v2/merge_insert.h new file mode 100644 index 0000000000..bc8c5851b4 --- /dev/null +++ b/include/v2/merge_insert.h @@ -0,0 +1,156 @@ +#pragma once + +#include "v2/graph_delta.h" +#include "v2/fs_allocator.h" +#include "v2/index_merger.h" +#include "tsl/robin_map.h" +#include "tsl/robin_set.h" +#include "pq_flash_index.h" +#include "linux_aligned_file_reader.h" +#include "index.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "parameters.h" +#include "Neighbor_Tag.h" +#include "threadpool.h" + +#include "windows_customizations.h" + +#define MERGE_TH 18000000 + +namespace diskann { + + template + class MergeInsert { + public: + // constructor + DISKANN_DLLEXPORT MergeInsert( + Parameters& parameters, size_t dim, const std::string mem_prefix, + const std::string disk_prefix_in, const std::string disk_prefix_out, + Distance* dist, diskann::Metric disk_metric, bool single_file_index, + std::string working_folder); + + DISKANN_DLLEXPORT ~MergeInsert(); + + // insertion function - insert into short_term_index + DISKANN_DLLEXPORT int insert(const T* point, const TagT& tag); + + DISKANN_DLLEXPORT void lazy_delete(const TagT& tag); + //DISKANN_DLLEXPORT void lazy_delete(tsl::robin_set& delete_list); + + // search function - search both short_term_index and long_term_index and + // return with top L candidate tags of the shard + DISKANN_DLLEXPORT void search_sync(const T* query, const uint64_t K, + const uint64_t search_L, TagT* tags, + float* distances, QueryStats* stats); + // void return_active_tags(tsl::robin_set& active_tags); + + // continuously runs in background to check if mem index size has exceeded + // its threshold - triggers index switch and merge + DISKANN_DLLEXPORT int trigger_merge(); + + DISKANN_DLLEXPORT void final_merge(); + + DISKANN_DLLEXPORT std::string ret_merge_prefix(); + + protected: + // call constructor to StreamingMerger object + void construct_index_merger(); + + // call StreamingMerger destructor to explicitly de-register threads + void destruct_index_merger(); + + //_active_index flag will be modified only inside this function + void switch_index(); // function to atomically switch btw indices, makes + // older index inactive(read-only), saves it, makes new + // index active (r/w) + + // save currently active mem_index and make it inactive + int save(); + + // make a local copy of _deletion_set and save it to a _deleted_tags_file + void save_del_set(); + + // call merge on a StreamingMerger object, only if index switching and + // saving is successful + void merge(); + + private: + size_t _merge_th = 0; + size_t _mem_points = 0; // reflects number of points in active mem index + size_t _index_points = 0; + size_t _dim; + _u32 _num_nodes_to_cache; + _u32 _num_search_threads; + uint64_t _beamwidth; + + std::unordered_map curr_location_to_tag; + + std::shared_ptr> _mem_index_0 = nullptr; + std::shared_ptr> _mem_index_1 = nullptr; + std::shared_ptr reader = nullptr; + PQFlashIndex* _disk_index = nullptr; + StreamingMerger * _merger = nullptr; + std::string TMP_FOLDER; + + diskann::Metric _dist_metric; + Distance* _dist_comp; + + diskann::Parameters _paras_mem; + diskann::Parameters _paras_disk; + + tsl::robin_set _deletion_set_0; + tsl::robin_set _deletion_set_1; + + std::vector*> _deleted_tags_vector; + + int _active_index = 0; // reflects value of writable index + int _active_delete_set = 0; // reflects active _deletion_set + std::atomic_bool _active_0; // true except when merging + std::atomic_bool _active_1; // true except when merging + std::atomic_bool _active_del_0; // true except when being saved + std::atomic_bool _active_del_1; // true except when being saved + std::atomic_bool _clearing_index_0; // don't search mem_index if true + std::atomic_bool _clearing_index_1; // don't search mem_index if true + std::atomic_bool _switching_disk_prefixes = + false; // wait if true, search when false + std::atomic_bool _check_switch_index = + false; // true when switch_index acquires _index_lock in writer mode, + // insert threads wait till it turns back to false + std::atomic_bool _check_switch_delete = + false; // true when switching between _deletion_sets, _delete_lock + // acquired in write mode, delete thread waits till it turns back to false + + bool _single_file_index = false; + + std::shared_timed_mutex _delete_lock; // lock to access _deletion_set + std::shared_timed_mutex _index_lock; // mutex to switch between mem indices + std::shared_timed_mutex _change_lock; // mutex to switch increment _mem_pts + std::shared_timed_mutex _disk_lock; // mutex to switch between disk indices + std::shared_timed_mutex + _clear_lock_0; // lock to prevent an index from being cleared when it + // is being searched and vice versa + std::shared_timed_mutex + _clear_lock_1; // lock to prevent an index from being cleared when it + // is being searched and vice versa + + ThreadPool* _search_tpool; + + std::string _mem_index_prefix; + std::string _disk_index_prefix_in; + std::string _disk_index_prefix_out; + std::string _deleted_tags_file; + }; +}; // namespace diskann diff --git a/include/windows_aligned_file_reader.h b/include/windows_aligned_file_reader.h index 8fec3d4f0e..433d3c0bf7 100644 --- a/include/windows_aligned_file_reader.h +++ b/include/windows_aligned_file_reader.h @@ -31,7 +31,7 @@ class WindowsAlignedFileReader : public AlignedFileReader { // Open & close ops // Blocking calls DISKANN_DLLEXPORT virtual void open(const std::string &fname); - DISKANN_DLLEXPORT virtual void close(); + DISKANN_DLLEXPORT virtual void close(); DISKANN_DLLEXPORT virtual void register_thread(); DISKANN_DLLEXPORT virtual void deregister_thread() { @@ -41,8 +41,7 @@ class WindowsAlignedFileReader : public AlignedFileReader { // process batch of aligned requests in parallel // NOTE :: blocking call for the calling thread, but can thread-safe DISKANN_DLLEXPORT virtual void read(std::vector &read_reqs, - IOContext &ctx, - bool async); + IOContext &ctx, bool async); }; #endif // USE_BING_INFRA #endif //_WINDOWS diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 9cd8a86bc4..fe49a71d2a 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -9,7 +9,9 @@ else() #file(GLOB CPP_SOURCES *.cpp) set(CPP_SOURCES ann_exception.cpp aux_utils.cpp index.cpp linux_aligned_file_reader.cpp math_utils.cpp memory_mapper.cpp - partition_and_pq.cpp pq_flash_index.cpp logger.cpp utils.cpp) + partition_and_pq.cpp pq_flash_index.cpp logger.cpp distance.cpp + utils.cpp v2/graph_delta.cpp v2/index_merger.cpp + v2/merge_insert.cpp) add_library(${PROJECT_NAME} ${CPP_SOURCES}) add_library(${PROJECT_NAME}_s STATIC ${CPP_SOURCES}) endif() diff --git a/src/ann_exception.cpp b/src/ann_exception.cpp index 23f3db9937..c24985eaae 100644 --- a/src/ann_exception.cpp +++ b/src/ann_exception.cpp @@ -33,4 +33,8 @@ namespace diskann { return sstream.str(); } -} // namespace diskann \ No newline at end of file + int ANNException::errorCode() const { + return _errorCode; + } + +} // namespace diskann diff --git a/src/aux_utils.cpp b/src/aux_utils.cpp index 6a2990ba66..18a7a27b27 100644 --- a/src/aux_utils.cpp +++ b/src/aux_utils.cpp @@ -21,19 +21,98 @@ #include "pq_flash_index.h" #include "utils.h" +#include "pq_flash_index.h" +#include "tsl/robin_set.h" +#include "utils.h" + +#define NUM_KMEANS 15 + namespace diskann { - double get_memory_budget(const std::string &mem_budget_str) { - double mem_ram_budget = atof(mem_budget_str.c_str()); - double final_index_ram_limit = mem_ram_budget; - if (mem_ram_budget - SPACE_FOR_CACHED_NODES_IN_GB > + void add_new_file_to_single_index(std::string index_file, + std::string new_file) { + std::unique_ptr<_u64[]> metadata; + _u64 nr, nc; + diskann::load_bin<_u64>(index_file, metadata, nr, nc, 0); + if (nc != 1) { + std::stringstream stream; + stream << "Error, index file specified does not have correct metadata. " + << std::endl; + throw diskann::ANNException(stream.str(), -1); + } + size_t index_ending_offset = metadata[nr - 1]; + _u64 read_blk_size = 64 * 1024 * 1024; + cached_ofstream writer(index_file, read_blk_size, index_ending_offset); + _u64 check_file_size = get_file_size(index_file); + if (check_file_size != index_ending_offset) { + std::stringstream stream; + stream << "Error, index file specified does not have correct metadata " + "(last entry must match the filesize). " + << std::endl; + throw diskann::ANNException(stream.str(), -1); + } + + cached_ifstream reader(new_file, read_blk_size); + size_t fsize = reader.get_file_size(); + if (fsize == 0) { + std::stringstream stream; + stream << "Error, new file specified is empty. Not appending. " + << std::endl; + throw diskann::ANNException(stream.str(), -1); + } + + size_t num_blocks = DIV_ROUND_UP(fsize, read_blk_size); + char * dump = new char[read_blk_size]; + for (_u64 i = 0; i < num_blocks; i++) { + size_t cur_block_size = read_blk_size > fsize - (i * read_blk_size) + ? fsize - (i * read_blk_size) + : read_blk_size; + reader.read(dump, cur_block_size); + writer.write(dump, cur_block_size); + } + reader.close(); + writer.close(); + + delete[] dump; + std::vector<_u64> new_meta; + for (_u64 i = 0; i < nr; i++) + new_meta.push_back(metadata[i]); + new_meta.push_back(metadata[nr - 1] + fsize); + + diskann::save_bin<_u64>(index_file, new_meta.data(), new_meta.size(), 1, 0); + } + + double get_memory_budget(double search_ram_budget) { + double final_index_ram_limit = search_ram_budget; + if (search_ram_budget - SPACE_FOR_CACHED_NODES_IN_GB > THRESHOLD_FOR_CACHING_IN_GB) { // slack for space used by cached // nodes - final_index_ram_limit = mem_ram_budget - SPACE_FOR_CACHED_NODES_IN_GB; + final_index_ram_limit = search_ram_budget - SPACE_FOR_CACHED_NODES_IN_GB; } return final_index_ram_limit * 1024 * 1024 * 1024; } + double get_memory_budget(const std::string &mem_budget_str) { + double search_ram_budget = atof(mem_budget_str.c_str()); + return get_memory_budget(search_ram_budget); + } + + size_t calculate_num_pq_chunks(double final_index_ram_limit, + size_t points_num, uint32_t dim) { + size_t num_pq_chunks = + (size_t)(std::floor)(_u64(final_index_ram_limit / (double) points_num)); + + diskann::cout << "Calculated num_pq_chunks :" << num_pq_chunks << std::endl; + num_pq_chunks = num_pq_chunks <= 0 ? 1 : num_pq_chunks; + num_pq_chunks = num_pq_chunks > dim ? dim : num_pq_chunks; + num_pq_chunks = + num_pq_chunks > MAX_PQ_CHUNKS ? MAX_PQ_CHUNKS : num_pq_chunks; + + diskann::cout << "Compressing " << dim << "-dimensional data into " + << num_pq_chunks << " bytes per vector." << std::endl; + return num_pq_chunks; + } + double calculate_recall(unsigned num_queries, unsigned *gold_std, float *gs_dist, unsigned dim_gs, unsigned *our_results, unsigned dim_or, @@ -48,8 +127,8 @@ namespace diskann { unsigned *res_vec = our_results + dim_or * i; size_t tie_breaker = recall_at; if (gs_dist != nullptr) { - tie_breaker = recall_at - 1; float *gt_dist_vec = gs_dist + dim_gs * i; + tie_breaker = recall_at - 1; while (tie_breaker < dim_gs && gt_dist_vec[tie_breaker] == gt_dist_vec[recall_at - 1]) tie_breaker++; @@ -57,9 +136,10 @@ namespace diskann { gt.insert(gt_vec, gt_vec + tie_breaker); res.insert(res_vec, res_vec + recall_at); + unsigned cur_recall = 0; - for (auto &v : gt) { - if (res.find(v) != res.end()) { + for (auto &v : res) { + if (gt.find(v) != gt.end()) { cur_recall++; } } @@ -68,6 +148,63 @@ namespace diskann { return total_recall / (num_queries) * (100.0 / recall_at); } + double calculate_recall(unsigned num_queries, unsigned *gold_std, + float *gs_dist, unsigned dim_gs, + unsigned *our_results, unsigned dim_or, + unsigned recall_at, + const tsl::robin_set &active_tags) { + double total_recall = 0; + std::set gt, res; + bool printed = false; + for (size_t i = 0; i < num_queries; i++) { + gt.clear(); + res.clear(); + unsigned *gt_vec = gold_std + dim_gs * i; + unsigned *res_vec = our_results + dim_or * i; + size_t tie_breaker = recall_at; + unsigned active_points_count = 0; + unsigned cur_counter = 0; + while (active_points_count < recall_at && cur_counter < dim_gs) { + if (active_tags.find(*(gt_vec + cur_counter)) != active_tags.end()) { + active_points_count++; + } + cur_counter++; + } + if (active_tags.empty()) + cur_counter = recall_at; + + if ((active_points_count < recall_at && !active_tags.empty()) && + !printed) { + diskann::cout << "Warning: Couldn't find enough closest neighbors " + << active_points_count << "/" << recall_at + << " from " + "truthset for query # " + << i << ". Will result in under-reported value of recall." + << std::endl; + printed = true; + } + if (gs_dist != nullptr) { + tie_breaker = cur_counter - 1; + float *gt_dist_vec = gs_dist + dim_gs * i; + while (tie_breaker < dim_gs && + gt_dist_vec[tie_breaker] == gt_dist_vec[cur_counter - 1]) + tie_breaker++; + } + + gt.insert(gt_vec, gt_vec + tie_breaker); + res.insert(res_vec, res_vec + recall_at); + unsigned cur_recall = 0; + for (auto &v : res) { + if (gt.find(v) != gt.end()) { + cur_recall++; + } + } + total_recall += cur_recall; + } + return ((double) (total_recall / (num_queries))) * + ((double) (100.0 / recall_at)); + } + template T *generateRandomWarmup(uint64_t warmup_num, uint64_t warmup_dim, uint64_t warmup_aligned_dim) { @@ -102,6 +239,12 @@ namespace diskann { if (files.fileExists(cache_warmup_file)) { diskann::load_aligned_bin(files, cache_warmup_file, warmup, warmup_num, file_dim, file_aligned_dim); + diskann::cout << "In the warmup file: " << cache_warmup_file + << " File dim: " << file_dim + << " File aligned dim: " << file_aligned_dim + << " Expected dim: " << warmup_dim + << " Expected aligned dim: " << warmup_aligned_dim + << std::endl; if (file_dim != warmup_dim || file_aligned_dim != warmup_aligned_dim) { std::stringstream stream; stream << "Mismatched dimensions in sample file. file_dim = " @@ -152,9 +295,8 @@ namespace diskann { std::ifstream reader(fname.c_str(), std::ios::binary); reader.read((char *) &npts32, sizeof(uint32_t)); reader.read((char *) &dim, sizeof(uint32_t)); - if (dim != 1 || - actual_file_size != - ((size_t) npts32) * sizeof(uint32_t) + 2 * sizeof(uint32_t)) { + if (dim != 1 || actual_file_size != ((size_t) npts32) * sizeof(uint32_t) + + 2 * sizeof(uint32_t)) { std::stringstream stream; stream << "Error reading idmap file. Check if the file is bin file with " "1 dimensional data. Actual: " @@ -209,32 +351,23 @@ namespace diskann { node_shard.push_back(std::make_pair((_u32) node_id, (_u32) shard)); } } - std::sort(node_shard.begin(), node_shard.end(), [](const auto &left, - const auto &right) { - return left.first < right.first || - (left.first == right.first && left.second < right.second); - }); + std::sort(node_shard.begin(), node_shard.end(), + [](const auto &left, const auto &right) { + return left.first < right.first || (left.first == right.first && + left.second < right.second); + }); diskann::cout << "Finished computing node -> shards map" << std::endl; // create cached vamana readers std::vector vamana_readers(nshards); for (_u64 i = 0; i < nshards; i++) { vamana_readers[i].open(vamana_names[i], 1024 * 1048576); - size_t actual_file_size = get_file_size(vamana_names[i]); size_t expected_file_size; vamana_readers[i].read((char *) &expected_file_size, sizeof(uint64_t)); - if (actual_file_size != expected_file_size) { - std::stringstream stream; - stream << "Error in Vamana Index file " << vamana_names[i] - << " Actual file size: " << actual_file_size - << " does not match expected file size: " << expected_file_size - << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); - } } - size_t merged_index_size = 16; + size_t merged_index_size = 24; + size_t merged_index_frozen = 0; // create cached vamana writers cached_ofstream diskann_writer(output_vamana, 1024 * 1048576); diskann_writer.write((char *) &merged_index_size, sizeof(uint64_t)); @@ -259,10 +392,13 @@ namespace diskann { medoid_writer.write((char *) &nshards_u32, sizeof(uint32_t)); medoid_writer.write((char *) &one_val, sizeof(uint32_t)); + _u64 vamana_index_frozen = 0; for (_u64 shard = 0; shard < nshards; shard++) { unsigned medoid; // read medoid vamana_readers[shard].read((char *) &medoid, sizeof(unsigned)); + vamana_readers[shard].read((char *) &vamana_index_frozen, sizeof(_u64)); + assert(vamana_index_frozen == false); // rename medoid medoid = idmaps[shard][medoid]; @@ -271,11 +407,12 @@ namespace diskann { if (shard == (nshards - 1)) //--> uncomment if running hierarchical diskann_writer.write((char *) &medoid, sizeof(unsigned)); } + diskann_writer.write((char *) &merged_index_frozen, sizeof(_u64)); medoid_writer.close(); diskann::cout << "Starting merge" << std::endl; - // Gopal. random_shuffle() is deprecated. + // random_shuffle() is deprecated. std::random_device rng; std::mt19937 urng(rng()); @@ -288,7 +425,7 @@ namespace diskann { unsigned node_id = id_shard.first; unsigned shard_id = id_shard.second; if (cur_id < node_id) { - // Gopal. random_shuffle() is deprecated. + // random_shuffle() is deprecated. std::shuffle(final_nhood.begin(), final_nhood.end(), urng); nnbrs = (unsigned) (std::min)(final_nhood.size(), (uint64_t) max_degree); @@ -321,7 +458,7 @@ namespace diskann { } } - // Gopal. random_shuffle() is deprecated. + // random_shuffle() is deprecated. std::shuffle(final_nhood.begin(), final_nhood.end(), urng); nnbrs = (unsigned) (std::min)(final_nhood.size(), (uint64_t) max_degree); // write into merged ofstream @@ -342,12 +479,11 @@ namespace diskann { } template - int build_merged_vamana_index(std::string base_file, - diskann::Metric _compareMetric, unsigned L, - unsigned R, double sampling_rate, - double ram_budget, std::string mem_index_path, - std::string medoids_file, - std::string centroids_file) { + int build_merged_vamana_index( + std::string base_file, diskann::Metric _compareMetric, + bool single_file_index, unsigned L, unsigned R, double sampling_rate, + double ram_budget, std::string mem_index_path, std::string medoids_file, + std::string centroids_file, const char *tag_file) { size_t base_num, base_dim; diskann::get_bin_metadata(base_file, base_num, base_dim); @@ -360,20 +496,40 @@ namespace diskann { paras.Set("L", (unsigned) L); paras.Set("R", (unsigned) R); paras.Set("C", 750); - paras.Set("alpha", 2.0f); + paras.Set("alpha", 1.2f); paras.Set("num_rnds", 2); - paras.Set("saturate_graph", 1); + paras.Set("saturate_graph", 1); // was 0 earlier. paras.Set("save_path", mem_index_path); + bool tags_enabled; + if (tag_file == nullptr) + tags_enabled = false; + else + tags_enabled = true; + std::unique_ptr> _pvamanaIndex = std::unique_ptr>( - new diskann::Index(_compareMetric, base_file.c_str())); - _pvamanaIndex->build(paras); + new diskann::Index(_compareMetric, base_dim, base_num, false, + single_file_index, tags_enabled)); + if (tags_enabled) + _pvamanaIndex->build(base_file.c_str(), base_num, paras, tag_file); + else + _pvamanaIndex->build(base_file.c_str(), base_num, paras); + _pvamanaIndex->save(mem_index_path.c_str()); std::remove(medoids_file.c_str()); std::remove(centroids_file.c_str()); return 0; } + + if (single_file_index || tag_file != nullptr) { + diskann::cout << "Cannot build merged index if single_file_index is " + "required or if tags are specified. Please contact " + "rakri@microsoft.com if this is required" + << std::endl; + return 1; + } + std::string merged_index_prefix = mem_index_path + "_tempFiles"; int num_parts = partition_with_ram_budget(base_file, sampling_rate, ram_budget, @@ -392,15 +548,18 @@ namespace diskann { paras.Set("L", L); paras.Set("R", (2 * (R / 3))); paras.Set("C", 750); - paras.Set("alpha", 2.0f); + paras.Set("alpha", 1.2f); paras.Set("num_rnds", 2); - paras.Set("saturate_graph", 1); + paras.Set("saturate_graph", 0); paras.Set("save_path", shard_index_file); + _u64 shard_base_dim, shard_base_pts; + get_bin_metadata(shard_base_file, shard_base_pts, shard_base_dim); std::unique_ptr> _pvamanaIndex = - std::unique_ptr>( - new diskann::Index(_compareMetric, shard_base_file.c_str())); - _pvamanaIndex->build(paras); + std::unique_ptr>(new diskann::Index( + _compareMetric, shard_base_dim, shard_base_pts, false, + single_file_index)); // TODO: Single? + _pvamanaIndex->build(shard_base_file.c_str(), shard_base_pts, paras); _pvamanaIndex->save(shard_index_file.c_str()); } @@ -416,9 +575,13 @@ namespace diskann { std::to_string(p) + "_ids_uint32.bin"; std::string shard_index_file = merged_index_prefix + "_subshard-" + std::to_string(p) + "_mem.index"; + // Required if Index.cpp thinks we are building a multi-file index. + std::string shard_index_file_data = shard_index_file + ".data"; + std::remove(shard_base_file.c_str()); std::remove(shard_id_file.c_str()); std::remove(shard_index_file.c_str()); + std::remove(shard_index_file_data.c_str()); } return 0; } @@ -429,71 +592,103 @@ namespace diskann { // optimizes the beamwidth to maximize QPS for a given L_search subject to // 99.9 latency not blowing up - template - uint32_t optimize_beamwidth( - std::unique_ptr> &pFlashIndex, T - *tuning_sample, - _u64 tuning_sample_num, _u64 tuning_sample_aligned_dim, uint32_t L, - uint32_t nthreads, uint32_t start_bw) { - uint32_t cur_bw = start_bw; - double max_qps = 0; - uint32_t best_bw = start_bw; - bool stop_flag = false; - - while (!stop_flag) { - std::vector tuning_sample_result_ids_64(tuning_sample_num, - 0); - std::vector tuning_sample_result_dists(tuning_sample_num, - 0); - diskann::QueryStats * stats = new - diskann::QueryStats[tuning_sample_num]; - - auto s = std::chrono::high_resolution_clock::now(); - #pragma omp parallel for schedule(dynamic, 1) num_threads(nthreads) - for (_s64 i = 0; i < (int64_t) tuning_sample_num; i++) { - pFlashIndex->cached_beam_search( - tuning_sample + (i * tuning_sample_aligned_dim), 1, L, - tuning_sample_result_ids_64.data() + (i * 1), - tuning_sample_result_dists.data() + (i * 1), cur_bw, stats + - i); - } - auto e = std::chrono::high_resolution_clock::now(); - std::chrono::duration diff = e - s; - double qps = (1.0f * tuning_sample_num) / (1.0f * diff.count()); - - double lat_999 = diskann::get_percentile_stats( - stats, tuning_sample_num, 0.999, - [](const diskann::QueryStats &stats) { return stats.total_us; }); - - double mean_latency = diskann::get_mean_stats( - stats, tuning_sample_num, - [](const diskann::QueryStats &stats) { return stats.total_us; }); - - if (qps > max_qps && lat_999 < (15000) + mean_latency * 2) { - max_qps = qps; - best_bw = cur_bw; - cur_bw = (uint32_t)(std::ceil)((float) cur_bw * 1.1); - } else { - stop_flag = true; - } - if (cur_bw > 64) - stop_flag = true; - - delete[] stats; + template + uint32_t optimize_beamwidth( + std::unique_ptr> &pFlashIndex, + T *tuning_sample, _u64 tuning_sample_num, _u64 tuning_sample_aligned_dim, + uint32_t L, uint32_t nthreads, uint32_t start_bw) { + uint32_t cur_bw = start_bw; + double max_qps = 0; + uint32_t best_bw = start_bw; + bool stop_flag = false; + + while (!stop_flag) { + std::vector tuning_sample_result_tags(tuning_sample_num, 0); + std::vector tuning_sample_result_dists(tuning_sample_num, 0); + diskann::QueryStats *stats = new diskann::QueryStats[tuning_sample_num]; + + auto s = std::chrono::high_resolution_clock::now(); +#pragma omp parallel for schedule(dynamic, 1) num_threads(nthreads) + for (_s64 i = 0; i < (int64_t) tuning_sample_num; i++) { + pFlashIndex->cached_beam_search( + tuning_sample + (i * tuning_sample_aligned_dim), 1, L, + tuning_sample_result_tags.data() + (i * 1), + tuning_sample_result_dists.data() + (i * 1), cur_bw, stats + i); } - return best_bw; + auto e = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = e - s; + double qps = + (1.0f * (float) tuning_sample_num) / (1.0f * (float) diff.count()); + + double lat_999 = diskann::get_percentile_stats( + stats, tuning_sample_num, 0.999f, + [](const diskann::QueryStats &stats) { return stats.total_us; }); + + double mean_latency = diskann::get_mean_stats( + stats, tuning_sample_num, + [](const diskann::QueryStats &stats) { return stats.total_us; }); + + if (qps > max_qps && lat_999 < (15000) + mean_latency * 2) { + max_qps = qps; + best_bw = cur_bw; + cur_bw = (uint32_t)(std::ceil)((float) cur_bw * 1.1f); + } else { + stop_flag = true; + } + if (cur_bw > 64) + stop_flag = true; + + delete[] stats; } + return best_bw; + } - template - void create_disk_layout(const std::string base_file, - const std::string mem_index_file, - const std::string output_file) { + // if single_index format is true, we assume that the entire mem index is in + // mem_index_file, and the entire disk index will be in output_file. + template + void create_disk_layout(const std::string &mem_index_file, + const std::string &base_file, + const std::string &tag_file, + const std::string &pq_pivots_file, + const std::string &pq_vectors_file, + bool single_file_index, + const std::string &output_file) { unsigned npts, ndims; // amount to read or write in one shot _u64 read_blk_size = 64 * 1024 * 1024; _u64 write_blk_size = read_blk_size; - cached_ifstream base_reader(base_file, read_blk_size); + cached_ifstream base_reader; + std::ifstream vamana_reader; + _u64 base_offset = 0, vamana_offset = 0, tags_offset = 0; + bool tags_enabled = false; + + if (single_file_index) { + _u64 nr, nc; + std::unique_ptr<_u64[]> offsets; + diskann::load_bin<_u64>(mem_index_file, offsets, nr, nc); + if (nr != Index::METADATA_ROWS && nc != 1) { + std::stringstream stream; + stream + << "Vamana Single Index file size does not meet meta-data criteria." + << std::endl; + + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } + vamana_offset = offsets[0]; + base_offset = offsets[1]; + tags_offset = offsets[2]; + tags_enabled = tags_offset != offsets[3]; + vamana_reader.open(mem_index_file, std::ios::binary); + vamana_reader.seekg(vamana_offset, vamana_reader.beg); + base_reader.open(mem_index_file, read_blk_size, base_offset); + } else { + base_reader.open(base_file, read_blk_size); + vamana_reader.open(mem_index_file, std::ios::binary); + tags_enabled = tag_file != ""; + } + base_reader.read((char *) &npts, sizeof(uint32_t)); base_reader.read((char *) &ndims, sizeof(uint32_t)); @@ -502,33 +697,27 @@ namespace diskann { ndims_64 = ndims; // create cached reader + writer - size_t actual_file_size = get_file_size(mem_index_file); - cached_ifstream vamana_reader(mem_index_file, read_blk_size); - cached_ofstream diskann_writer(output_file, write_blk_size); + // size_t actual_file_size = get_file_size(mem_index_file); + std::remove(output_file.c_str()); + cached_ofstream diskann_writer; + diskann_writer.open(output_file, write_blk_size); // metadata: width, medoid unsigned width_u32, medoid_u32; size_t index_file_size; vamana_reader.read((char *) &index_file_size, sizeof(uint64_t)); - if (index_file_size != actual_file_size) { - std::stringstream stream; - stream << "Vamana Index file size does not match expected size per " - "meta-data." - << " file size from file: " << index_file_size - << " actual file size: " << actual_file_size << std::endl; - - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); - } + _u64 vamana_frozen_num = false, vamana_frozen_loc = 0; vamana_reader.read((char *) &width_u32, sizeof(unsigned)); vamana_reader.read((char *) &medoid_u32, sizeof(unsigned)); - + vamana_reader.read((char *) &vamana_frozen_num, sizeof(_u64)); // compute _u64 medoid, max_node_len, nnodes_per_sector; npts_64 = (_u64) npts; medoid = (_u64) medoid_u32; + if (vamana_frozen_num == 1) + vamana_frozen_loc = medoid; max_node_len = (((_u64) width_u32 + 1) * sizeof(unsigned)) + (ndims_64 * sizeof(T)); nnodes_per_sector = SECTOR_LEN / max_node_len; @@ -539,26 +728,31 @@ namespace diskann { << std::endl; // SECTOR_LEN buffer for each sector - std::unique_ptr sector_buf = - std::make_unique(SECTOR_LEN); - std::unique_ptr node_buf = - std::make_unique(max_node_len); + std::unique_ptr sector_buf = std::make_unique(SECTOR_LEN); + std::unique_ptr node_buf = std::make_unique(max_node_len); unsigned &nnbrs = *(unsigned *) (node_buf.get() + ndims_64 * sizeof(T)); unsigned *nhood_buf = (unsigned *) (node_buf.get() + (ndims_64 * sizeof(T)) + sizeof(unsigned)); // number of sectors (1 for meta data) - _u64 n_sectors = ROUND_UP(npts_64, nnodes_per_sector) / - nnodes_per_sector; + _u64 n_sectors = ROUND_UP(npts_64, nnodes_per_sector) / nnodes_per_sector; _u64 disk_index_file_size = (n_sectors + 1) * SECTOR_LEN; - // write first sector with metadata - *(_u64 *) (sector_buf.get() + 0 * sizeof(_u64)) = disk_index_file_size; - *(_u64 *) (sector_buf.get() + 1 * sizeof(_u64)) = npts_64; - *(_u64 *) (sector_buf.get() + 2 * sizeof(_u64)) = medoid; - *(_u64 *) (sector_buf.get() + 3 * sizeof(_u64)) = max_node_len; - *(_u64 *) (sector_buf.get() + 4 * sizeof(_u64)) = nnodes_per_sector; - diskann_writer.write(sector_buf.get(), SECTOR_LEN); + + std::vector<_u64> output_file_meta; + output_file_meta.push_back(npts_64); + output_file_meta.push_back(ndims_64); + output_file_meta.push_back(medoid); + output_file_meta.push_back(max_node_len); + output_file_meta.push_back(nnodes_per_sector); + output_file_meta.push_back(vamana_frozen_num); + output_file_meta.push_back(vamana_frozen_loc); + output_file_meta.push_back(disk_index_file_size); + + diskann_writer.write(sector_buf.get(), SECTOR_LEN); // write out the empty + // first sector, will + // be populated at the + // end. std::unique_ptr cur_node_coords = std::make_unique(ndims_64); diskann::cout << "# sectors: " << n_sectors << std::endl; @@ -576,24 +770,32 @@ namespace diskann { vamana_reader.read((char *) &nnbrs, sizeof(unsigned)); // sanity checks on nnbrs - assert(nnbrs > 0); - assert(nnbrs <= width_u32); + if (nnbrs == 0) { + diskann::cout << "ERROR. Found point with no out-neighbors; Point#: " + << cur_node_id << std::endl; + exit(-1); + } // read node's nhood - vamana_reader.read((char *) nhood_buf, nnbrs * sizeof(unsigned)); + vamana_reader.read((char *) nhood_buf, + (std::min)(nnbrs, width_u32) * sizeof(unsigned)); + if (nnbrs > width_u32) { + vamana_reader.seekg((nnbrs - width_u32) * sizeof(unsigned), + vamana_reader.cur); + } // write coords of node first // T *node_coords = data + ((_u64) ndims_64 * cur_node_id); - base_reader.read((char *) cur_node_coords.get(), sizeof(T) * - ndims_64); + base_reader.read((char *) cur_node_coords.get(), sizeof(T) * ndims_64); memcpy(node_buf.get(), cur_node_coords.get(), ndims_64 * sizeof(T)); // write nnbrs - *(unsigned *) (node_buf.get() + ndims_64 * sizeof(T)) = nnbrs; + *(unsigned *) (node_buf.get() + ndims_64 * sizeof(T)) = + (std::min)(nnbrs, width_u32); // write nhood next memcpy(node_buf.get() + ndims_64 * sizeof(T) + sizeof(unsigned), - nhood_buf, nnbrs * sizeof(unsigned)); + nhood_buf, (std::min)(nnbrs, width_u32) * sizeof(unsigned)); // get offset into sector_buf char *sector_node_buf = @@ -606,13 +808,93 @@ namespace diskann { // flush sector to disk diskann_writer.write(sector_buf.get(), SECTOR_LEN); } + diskann_writer.close(); + size_t tag_bytes_written = 0; + + // frozen point implies dynamic index which must have tags + if (vamana_frozen_num > 0) { + std::unique_ptr mem_index_tags; + size_t nr, nc; + if (single_file_index) + diskann::load_bin(mem_index_file, mem_index_tags, nr, nc, + tags_offset); + else + diskann::load_bin(tag_file, mem_index_tags, nr, nc, tags_offset); + + if (nr != npts_64 && nc != 1) { + std::stringstream stream; + stream << "Error loading tags file. File dims are " << nr << ", " << nc + << ", but expecting " << npts_64 + << " tags in 1 dimension (bin format)." << std::endl; + + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } + + if (single_file_index) { + tag_bytes_written = diskann::save_bin( + output_file, mem_index_tags.get(), nr, nc, + output_file_meta[output_file_meta.size() - 1]); + } else { + diskann::save_bin(output_file + std::string(".tags"), + mem_index_tags.get(), nr, nc); + } + } else { + if (tags_enabled) { + std::unique_ptr mem_index_tags; + size_t nr, nc; + if (single_file_index) { + diskann::load_bin(mem_index_file, mem_index_tags, nr, nc, + tags_offset); + } else { + if (!file_exists(tag_file)) { + diskann::cout << "Static vamana index, tag file " << tag_file + << "does not exist. Exiting...." << std::endl; + exit(-1); + } + + diskann::load_bin(tag_file, mem_index_tags, nr, nc, + tags_offset); + } + + if (nr != npts_64 && nc != 1) { + std::stringstream stream; + stream << "Error loading tags file. File dims are " << nr << ", " + << nc << ", but expecting " << npts_64 + << " tags in 1 dimension (bin format)." << std::endl; + + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } + + if (single_file_index) { + tag_bytes_written = diskann::save_bin( + output_file, mem_index_tags.get(), nr, nc, + output_file_meta[output_file_meta.size() - 1]); + } else { + diskann::save_bin(output_file + std::string(".tags"), + mem_index_tags.get(), nr, nc); + } + } + } + + output_file_meta.push_back(output_file_meta[output_file_meta.size() - 1] + + tag_bytes_written); + diskann::save_bin<_u64>(output_file, output_file_meta.data(), + output_file_meta.size(), 1, 0); + + if (single_file_index) { + add_new_file_to_single_index(output_file, pq_pivots_file); + add_new_file_to_single_index(output_file, pq_vectors_file); + } diskann::cout << "Output file written." << std::endl; } - template - bool build_disk_index(const char *dataFilePath, const char *indexFilePath, + template + bool build_disk_index(const char *dataPath, const char *indexFilePath, const char * indexBuildParameters, - diskann::Metric _compareMetric) { + diskann::Metric _compareMetric, bool single_file_index, + const char *tag_file) { std::stringstream parser; parser << std::string(indexBuildParameters); std::string cur_param; @@ -620,17 +902,19 @@ namespace diskann { while (parser >> cur_param) param_list.push_back(cur_param); - if (param_list.size() != 5) { + if (param_list.size() != 5 && param_list.size() != 6) { diskann::cout - << "Correct usage of parameters is R (max degree) " - "L (indexing list size, better if >= R) B (RAM limit of final " - "index in " - "GB) M (memory limit while indexing) T (number of threads for " - "indexing)" + << "Correct usage of parameters is: R (max degree)" + " L (indexing list size, should be >= R) " + " B (RAM limit of final index in GB) " + " M (memory limit while indexing in GB)" + " T (number of threads for indexing) " + " [C (compression ratio for PQ. Overrides parameter value B)] " << std::endl; return false; } + std::string dataFilePath(dataPath); std::string index_prefix_path(indexFilePath); std::string pq_pivots_path = index_prefix_path + "_pq_pivots.bin"; std::string pq_compressed_vectors_path = @@ -646,15 +930,15 @@ namespace diskann { double final_index_ram_limit = get_memory_budget(param_list[2]); if (final_index_ram_limit <= 0) { - std::cerr << "Insufficient memory budget (or string was not in right " - "format). Should be > 0." - << std::endl; + diskann::cerr << "Insufficient memory budget (or string was not in right " + "format). Should be > 0." + << std::endl; return false; } double indexing_ram_budget = (float) atof(param_list[3].c_str()); if (indexing_ram_budget <= 0) { - std::cerr << "Not building index. Please provide more RAM budget" - << std::endl; + diskann::cerr << "Not building index. Please provide more RAM budget" + << std::endl; return false; } _u32 num_threads = (_u32) atoi(param_list[4].c_str()); @@ -666,80 +950,160 @@ namespace diskann { diskann::cout << "Starting index build: R=" << R << " L=" << L << " Query RAM budget: " << final_index_ram_limit - << " Indexing ram budget: " << indexing_ram_budget - << " T: " << num_threads << std::endl; + << " Indexing RAM budget: " << indexing_ram_budget + << " T: " << num_threads << " Final index will be in " + << (single_file_index ? "single file" : "multiple files") + << std::endl; + + std::string normalized_file_path = dataFilePath; + if (_compareMetric == diskann::Metric::COSINE) { + if (std::is_floating_point::value) { + diskann::cout << "Cosine metric chosen. Normalizing vectors and " + "changing distance to L2 to boost accuracy." + << std::endl; + + normalized_file_path = + std::string(indexFilePath) + "_data.normalized.bin"; + normalize_data_file(dataFilePath, normalized_file_path); + _compareMetric = diskann::Metric::L2; + } else { + diskann::cerr << "WARNING: Cannot normalize integral data types." + << " Using cosine distance with integer data types may " + "result in poor recall." + << " Consider using L2 distance with integral data types." + << std::endl; + } + } auto s = std::chrono::high_resolution_clock::now(); size_t points_num, dim; - diskann::get_bin_metadata(dataFilePath, points_num, dim); + diskann::get_bin_metadata(normalized_file_path, points_num, dim); + auto training_set_size = + PQ_TRAINING_SET_FRACTION * points_num > MAX_PQ_TRAINING_SET_SIZE + ? MAX_PQ_TRAINING_SET_SIZE + : (_u32) std::round(PQ_TRAINING_SET_FRACTION * points_num); + training_set_size = (training_set_size == 0) ? 1 : training_set_size; + diskann::cout << "(Normalized, if required) file : " << normalized_file_path + << " has: " << points_num + << " points. Changing training set size to " + << training_set_size << " points" << std::endl; size_t num_pq_chunks = - (size_t)(std::floor)(_u64(final_index_ram_limit / points_num)); - - num_pq_chunks = num_pq_chunks <= 0 ? 1 : num_pq_chunks; - num_pq_chunks = num_pq_chunks > dim ? dim : num_pq_chunks; - num_pq_chunks = - num_pq_chunks > MAX_PQ_CHUNKS ? MAX_PQ_CHUNKS : num_pq_chunks; - - diskann::cout << "Compressing " << dim << "-dimensional data into " - << num_pq_chunks << " bytes per vector." << std::endl; + calculate_num_pq_chunks(final_index_ram_limit, points_num, dim); size_t train_size, train_dim; float *train_data; - double p_val = ((double) TRAINING_SET_SIZE / (double) points_num); - // generates random sample and sets it to train_data and updates - // train_size - gen_random_slice(dataFilePath, p_val, train_data, train_size, - train_dim); - - diskann::cout << "Training data loaded of size " << train_size << - std::endl; + auto start = std::chrono::high_resolution_clock::now(); + double p_val = ((double) training_set_size / (double) points_num); + // generates random sample and sets it to train_data and updates train_size + gen_random_slice(normalized_file_path, p_val, train_data, train_size, + train_dim); + diskann::cout << "Generating PQ pivots with training data of size: " + << train_size << " num PQ chunks: " << num_pq_chunks + << std::endl; generate_pq_pivots(train_data, train_size, (uint32_t) dim, 256, - (uint32_t) num_pq_chunks, 15, pq_pivots_path); - generate_pq_data_from_pivots(dataFilePath, 256, (uint32_t) - num_pq_chunks, - pq_pivots_path, - pq_compressed_vectors_path); + (uint32_t) num_pq_chunks, NUM_KMEANS, pq_pivots_path); + auto end = std::chrono::high_resolution_clock::now(); + diskann::cout << "Pivots generated in " + << std::chrono::duration(end - start).count() << "s." + << std::endl; + start = std::chrono::high_resolution_clock::now(); + generate_pq_data_from_pivots(normalized_file_path, 256, + (uint32_t) num_pq_chunks, pq_pivots_path, + pq_compressed_vectors_path); delete[] train_data; - train_data = nullptr; - + end = std::chrono::high_resolution_clock::now(); + diskann::cout << "Compressed data generated and written in: " + << std::chrono::duration(end - start).count() << "s." + << std::endl; + start = std::chrono::high_resolution_clock::now(); diskann::build_merged_vamana_index( - dataFilePath, _compareMetric, L, R, p_val, indexing_ram_budget, - mem_index_path, medoids_path, centroids_path); - - diskann::create_disk_layout(dataFilePath, mem_index_path, - disk_index_path); + normalized_file_path, _compareMetric, single_file_index, L, R, p_val, + indexing_ram_budget, mem_index_path, medoids_path, centroids_path, + tag_file); + end = std::chrono::high_resolution_clock::now(); + diskann::cout << "Vamana index built in: " + << std::chrono::duration(end - start).count() << "s." + << std::endl; - double sample_sampling_rate = (150000.0 / points_num); - gen_random_slice(dataFilePath, sample_base_prefix, - sample_sampling_rate); + if (tag_file == nullptr) { + diskann::create_disk_layout( + mem_index_path, normalized_file_path, "", pq_pivots_path, + pq_compressed_vectors_path, single_file_index, disk_index_path); + } else { + std::string tag_filename = std::string(tag_file); + diskann::create_disk_layout( + mem_index_path, normalized_file_path, tag_filename, pq_pivots_path, + pq_compressed_vectors_path, single_file_index, disk_index_path); + } + double ten_percent_points = std::ceil(points_num * 0.1); + double num_sample_points = ten_percent_points > MAX_SAMPLE_POINTS_FOR_WARMUP + ? MAX_SAMPLE_POINTS_FOR_WARMUP + : ten_percent_points; + double sample_sampling_rate = num_sample_points / points_num; + diskann::cout << "Generating warmup file with " << num_sample_points + << " points using a sampling rate of: " + << sample_sampling_rate << std::endl; + gen_random_slice(normalized_file_path, sample_base_prefix, + sample_sampling_rate); + + diskann::cout << "Deleting memory index file: " << mem_index_path + << std::endl; std::remove(mem_index_path.c_str()); + // TODO: This is poor design. The decision to add the ".data" prefix + // is taken by build_vamana_index. So, we shouldn't repeate it here. + // Checking to see if we can merge the data and index into one file. + std::remove((mem_index_path + ".data").c_str()); + if (normalized_file_path != dataFilePath) { + // then we created a normalized vector file. Delete it. + diskann::cout << "Deleting normalized vector file: " + << normalized_file_path << std::endl; + std::remove(normalized_file_path.c_str()); + } - auto e = - std::chrono::high_resolution_clock::now(); + auto e = std::chrono::high_resolution_clock::now(); std::chrono::duration diff = e - s; diskann::cout << "Indexing time: " << diff.count() << std::endl; - return true; } - template DISKANN_DLLEXPORT void create_disk_layout( - const std::string base_file, const std::string mem_index_file, - const std::string output_file); - - template DISKANN_DLLEXPORT void create_disk_layout( - const std::string base_file, const std::string mem_index_file, - const std::string output_file); - template DISKANN_DLLEXPORT void create_disk_layout( - const std::string base_file, const std::string mem_index_file, - const std::string output_file); + template DISKANN_DLLEXPORT void create_disk_layout( + const std::string &mem_index_file, const std::string &base_file, + const std::string &tag_file, const std::string &pq_pivots_file, + const std::string &pq_compressed_vectors_file, bool single_file_index, + const std::string &output_file); + template DISKANN_DLLEXPORT void create_disk_layout( + const std::string &mem_index_file, const std::string &base_file, + const std::string &tag_file, const std::string &pq_pivots_file, + const std::string &pq_compressed_vectors_file, bool single_file_index, + const std::string &output_file); + template DISKANN_DLLEXPORT void create_disk_layout( + const std::string &mem_index_file, const std::string &base_file, + const std::string &tag_file, const std::string &pq_pivots_file, + const std::string &pq_compressed_vectors_file, bool single_file_index, + const std::string &output_file); + template DISKANN_DLLEXPORT void create_disk_layout( + const std::string &mem_index_file, const std::string &base_file, + const std::string &tag_file, const std::string &pq_pivots_file, + const std::string &pq_compressed_vectors_file, bool single_file_index, + const std::string &output_file); + template DISKANN_DLLEXPORT void create_disk_layout( + const std::string &mem_index_file, const std::string &base_file, + const std::string &tag_file, const std::string &pq_pivots_file, + const std::string &pq_compressed_vectors_file, bool single_file_index, + const std::string &output_file); + template DISKANN_DLLEXPORT void create_disk_layout( + const std::string &mem_index_file, const std::string &base_file, + const std::string &tag_file, const std::string &pq_pivots_file, + const std::string &pq_compressed_vectors_file, bool single_file_index, + const std::string &output_file); template DISKANN_DLLEXPORT int8_t *load_warmup( const std::string &cache_warmup_file, uint64_t &warmup_num, @@ -763,45 +1127,76 @@ namespace diskann { uint64_t &warmup_num, uint64_t warmup_dim, uint64_t warmup_aligned_dim); #endif - template DISKANN_DLLEXPORT uint32_t optimize_beamwidth( - std::unique_ptr> &pFlashIndex, + template DISKANN_DLLEXPORT uint32_t optimize_beamwidth( + std::unique_ptr> &pFlashIndex, int8_t *tuning_sample, _u64 tuning_sample_num, _u64 tuning_sample_aligned_dim, uint32_t L, uint32_t nthreads, uint32_t start_bw); - template DISKANN_DLLEXPORT uint32_t optimize_beamwidth( - std::unique_ptr> &pFlashIndex, + template DISKANN_DLLEXPORT uint32_t optimize_beamwidth( + std::unique_ptr> &pFlashIndex, uint8_t *tuning_sample, _u64 tuning_sample_num, _u64 tuning_sample_aligned_dim, uint32_t L, uint32_t nthreads, uint32_t start_bw); - template DISKANN_DLLEXPORT uint32_t optimize_beamwidth( - std::unique_ptr> &pFlashIndex, + template DISKANN_DLLEXPORT uint32_t optimize_beamwidth( + std::unique_ptr> &pFlashIndex, float *tuning_sample, _u64 tuning_sample_num, _u64 tuning_sample_aligned_dim, uint32_t L, uint32_t nthreads, uint32_t start_bw); - template DISKANN_DLLEXPORT bool build_disk_index( + template DISKANN_DLLEXPORT uint32_t optimize_beamwidth( + std::unique_ptr> &pFlashIndex, + int8_t *tuning_sample, _u64 tuning_sample_num, + _u64 tuning_sample_aligned_dim, uint32_t L, uint32_t nthreads, + uint32_t start_bw); + template DISKANN_DLLEXPORT uint32_t optimize_beamwidth( + std::unique_ptr> &pFlashIndex, + uint8_t *tuning_sample, _u64 tuning_sample_num, + _u64 tuning_sample_aligned_dim, uint32_t L, uint32_t nthreads, + uint32_t start_bw); + template DISKANN_DLLEXPORT uint32_t optimize_beamwidth( + std::unique_ptr> &pFlashIndex, + float *tuning_sample, _u64 tuning_sample_num, + _u64 tuning_sample_aligned_dim, uint32_t L, uint32_t nthreads, + uint32_t start_bw); + + template DISKANN_DLLEXPORT bool build_disk_index( + const char *dataFilePath, const char *indexFilePath, + const char *indexBuildParameters, diskann::Metric _compareMetric, + bool singleFileIndex, const char *tag_file); + template DISKANN_DLLEXPORT bool build_disk_index( + const char *dataFilePath, const char *indexFilePath, + const char *indexBuildParameters, diskann::Metric _compareMetric, + bool singleFileIndex, const char *tag_file); + template DISKANN_DLLEXPORT bool build_disk_index( + const char *dataFilePath, const char *indexFilePath, + const char *indexBuildParameters, diskann::Metric _compareMetric, + bool singleFileIndex, const char *tag_file); + template DISKANN_DLLEXPORT bool build_disk_index( const char *dataFilePath, const char *indexFilePath, - const char *indexBuildParameters, diskann::Metric _compareMetric); - template DISKANN_DLLEXPORT bool build_disk_index( + const char *indexBuildParameters, diskann::Metric _compareMetric, + bool singleFileIndex, const char *tag_file); + template DISKANN_DLLEXPORT bool build_disk_index( const char *dataFilePath, const char *indexFilePath, - const char *indexBuildParameters, diskann::Metric _compareMetric); - template DISKANN_DLLEXPORT bool build_disk_index( + const char *indexBuildParameters, diskann::Metric _compareMetric, + bool singleFileIndex, const char *tag_file); + template DISKANN_DLLEXPORT bool build_disk_index( const char *dataFilePath, const char *indexFilePath, - const char *indexBuildParameters, diskann::Metric _compareMetric); + const char *indexBuildParameters, diskann::Metric _compareMetric, + bool singleFileIndex, const char *tag_file); template DISKANN_DLLEXPORT int build_merged_vamana_index( - std::string base_file, diskann::Metric _compareMetric, unsigned L, - unsigned R, double sampling_rate, double ram_budget, - std::string mem_index_path, std::string medoids_path, - std::string centroids_file); + std::string base_file, diskann::Metric _compareMetric, + bool single_file_index, unsigned L, unsigned R, double sampling_rate, + double ram_budget, std::string mem_index_path, std::string medoids_path, + std::string centroids_file, const char *tag_file); template DISKANN_DLLEXPORT int build_merged_vamana_index( - std::string base_file, diskann::Metric _compareMetric, unsigned L, - unsigned R, double sampling_rate, double ram_budget, - std::string mem_index_path, std::string medoids_path, - std::string centroids_file); + std::string base_file, diskann::Metric _compareMetric, + bool single_file_index, unsigned L, unsigned R, double sampling_rate, + double ram_budget, std::string mem_index_path, std::string medoids_path, + std::string centroids_file, const char *tag_file); template DISKANN_DLLEXPORT int build_merged_vamana_index( - std::string base_file, diskann::Metric _compareMetric, unsigned L, - unsigned R, double sampling_rate, double ram_budget, - std::string mem_index_path, std::string medoids_path, - std::string centroids_file); + std::string base_file, diskann::Metric _compareMetric, + bool single_file_index, unsigned L, unsigned R, double sampling_rate, + double ram_budget, std::string mem_index_path, std::string medoids_path, + std::string centroids_file, const char *tag_file); }; // namespace diskann diff --git a/src/convert/convert.cpp b/src/convert/convert.cpp new file mode 100644 index 0000000000..1d445d1b2c --- /dev/null +++ b/src/convert/convert.cpp @@ -0,0 +1,72 @@ +// convert.cpp : This file contains the 'main' function. Program execution +// begins and ends there. +// + +#include +#include "util.h" + +template +int aux_main(int argc, char** argv) { + source_type* sourceVec = nullptr; + size_t num_points, num_dims; + diskann::cout << "Loading source vector from file " << argv[2] << std::endl; + diskann::load_bin(argv[2], sourceVec, num_points, num_dims); + + dest_type* destVec = new dest_type[num_points * num_dims]; + diskann::convert_types(sourceVec, destVec, num_points, + num_dims); + diskann::cout << "Converted vector from type " << argv[1] << "to type " << argv[3] + << std::endl; + diskann::save_bin(argv[4], destVec, num_points, num_dims); + diskann::cout << "Saved converted vector to file " << argv[4] << std::endl; + return 0; +} + +int main(int argc, char** argv) { + if (argc < 5) { + diskann::cout + << "Usage: convert " + "where type is one of [float/uint8/int8], case-sensitive. " + << std::endl; + } + + std::string sourceType(argv[1]); + std::string destType(argv[3]); + + if (sourceType == destType) { + diskann::cout << "Nothing to do. source and destination types are the same." + << std::endl; + return 0; + } + + if (sourceType == "int8" && destType == "uint8") { + return aux_main(argc, argv); + } else if (sourceType == "int8" && destType == "float") { + return aux_main(argc, argv); + } else if (sourceType == "uint8" && destType == "int8") { + return aux_main(argc, argv); + } else if (sourceType == "uint8" && destType == "float") { + return aux_main(argc, argv); + } else if (sourceType == "float" && destType == "int8") { + return aux_main(argc, argv); + } else if (sourceType == "float" && destType == "uint8") { + return aux_main(argc, argv); + } else { + diskann::cout << "One of the source or destination types is unknown: " + << argv[0] << "," << argv[2] << std::endl; + return 0; + } +} + +// Run program: Ctrl + F5 or Debug > Start Without Debugging menu +// Debug program: F5 or Debug > Start Debugging menu + +// Tips for Getting Started: +// 1. Use the Solution Explorer window to add/manage files +// 2. Use the Team Explorer window to connect to source control +// 3. Use the Output window to see build output and other messages +// 4. Use the Error List window to view errors +// 5. Go to Project > Add New Item to create new code files, or Project > Add +// Existing Item to add existing code files to the project +// 6. In the future, to open this project again, go to File > Open > Project +// and select the .sln file diff --git a/src/distance.cpp b/src/distance.cpp new file mode 100644 index 0000000000..cc8ec79ab6 --- /dev/null +++ b/src/distance.cpp @@ -0,0 +1,233 @@ +#ifdef _WINDOWS +#include +#include +#include +#include +#else +#include +#endif + +#include "simd_utils.h" +#include +#include + +#include "distance.h" + +namespace diskann { + + // Cosine similarity. + float DistanceCosineInt8::compare(const int8_t *a, const int8_t *b, + uint32_t length) const { +#ifdef _WINDOWS + return diskann::CosineSimilarity2(a, b, length); +#else + return diskann::compute_cosine_similarity(a, b, length); +#endif + } + + float DistanceCosineFloat::compare(const float *a, const float *b, + uint32_t length) const { +#ifdef _WINDOWS + return diskann::CosineSimilarity2(a, b, length); +#else + return diskann::compute_cosine_similarity(a, b, length); +#endif + } + + float SlowDistanceCosineUInt8::compare(const uint8_t *a, const uint8_t *b, + uint32_t length) const { + int magA = 0, magB = 0, scalarProduct = 0; + for (uint32_t i = 0; i < length; i++) { + magA += ((uint32_t) a[i]) * ((uint32_t) a[i]); + magB += ((uint32_t) b[i]) * ((uint32_t) b[i]); + scalarProduct += ((uint32_t) a[i]) * ((uint32_t) b[i]); + } + // similarity == 1-cosine distance + return 1.0f - (float) (scalarProduct / (sqrt(magA) * sqrt(magB))); + } + + // L2 distance functions. + float DistanceL2Int8::compare(const int8_t *a, const int8_t *b, + uint32_t size) const { + int32_t result = 0; + +#ifdef _WINDOWS +#ifdef USE_AVX2 + __m256 r = _mm256_setzero_ps(); + char * pX = (char *) a, *pY = (char *) b; + while (size >= 32) { + __m256i r1 = _mm256_subs_epi8(_mm256_loadu_si256((__m256i *) pX), + _mm256_loadu_si256((__m256i *) pY)); + r = _mm256_add_ps(r, _mm256_mul_epi8(r1, r1)); + pX += 32; + pY += 32; + size -= 32; + } + while (size > 0) { + __m128i r2 = _mm_subs_epi8(_mm_loadu_si128((__m128i *) pX), + _mm_loadu_si128((__m128i *) pY)); + r = _mm256_add_ps(r, _mm256_mul32_pi8(r2, r2)); + pX += 4; + pY += 4; + size -= 4; + } + r = _mm256_hadd_ps(_mm256_hadd_ps(r, r), r); + return r.m256_f32[0] + r.m256_f32[4]; +#else +#pragma omp simd reduction(+ : result) aligned(a, b : 8) + for (_s32 i = 0; i < (_s32) size; i++) { + result += ((int32_t)((int16_t) a[i] - (int16_t) b[i])) * + ((int32_t)((int16_t) a[i] - (int16_t) b[i])); + } + return (float) result; +#endif +#else +#pragma omp simd reduction(+ : result) aligned(a, b : 8) + for (int32_t i = 0; i < (int32_t) size; i++) { + result += ((int32_t)((int16_t) a[i] - (int16_t) b[i])) * + ((int32_t)((int16_t) a[i] - (int16_t) b[i])); + } + return (float) result; +#endif + } + + float DistanceL2UInt8::compare(const uint8_t *a, const uint8_t *b, + uint32_t size) const { + uint32_t result = 0; +#ifndef _WINDOWS +#pragma omp simd reduction(+ : result) aligned(a, b : 8) +#endif + for (int32_t i = 0; i < (int32_t) size; i++) { + result += ((int32_t)((int16_t) a[i] - (int16_t) b[i])) * + ((int32_t)((int16_t) a[i] - (int16_t) b[i])); + } + return (float) result; + } + +#ifndef _WINDOWS + float DistanceL2::compare(const float *a, const float *b, + uint32_t size) const { + a = (const float *) __builtin_assume_aligned(a, 32); + b = (const float *) __builtin_assume_aligned(b, 32); +#else + float DistanceL2::compare(const float *a, const float *b, + uint32_t size) const { +#endif + + float result = 0; +#ifdef USE_AVX2 + // assume size is divisible by 8 + uint16_t niters = (uint16_t)(size / 8); + __m256 sum = _mm256_setzero_ps(); + for (uint16_t j = 0; j < niters; j++) { + // scope is a[8j:8j+7], b[8j:8j+7] + // load a_vec + if (j < (niters - 1)) { + _mm_prefetch((char *) (a + 8 * (j + 1)), _MM_HINT_T0); + _mm_prefetch((char *) (b + 8 * (j + 1)), _MM_HINT_T0); + } + __m256 a_vec = _mm256_load_ps(a + 8 * j); + // load b_vec + __m256 b_vec = _mm256_load_ps(b + 8 * j); + // a_vec - b_vec + __m256 tmp_vec = _mm256_sub_ps(a_vec, b_vec); + /* + // (a_vec - b_vec)**2 + __m256 tmp_vec2 = _mm256_mul_ps(tmp_vec, tmp_vec); + // accumulate sum + sum = _mm256_add_ps(sum, tmp_vec2); + */ + // sum = (tmp_vec**2) + sum + sum = _mm256_fmadd_ps(tmp_vec, tmp_vec, sum); + } + + // horizontal add sum + result = _mm256_reduce_add_ps(sum); +#else +#ifndef _WINDOWS +#pragma omp simd reduction(+ : result) aligned(a, b : 32) +#endif + for (_s32 i = 0; i < (_s32) size; i++) { + result += (a[i] - b[i]) * (a[i] - b[i]); + } +#endif + return result; + } + + float SlowDistanceL2Float::compare(const float *a, const float *b, + uint32_t length) const { + float result = 0.0f; + for (uint32_t i = 0; i < length; i++) { + result += (a[i] - b[i]) * (a[i] - b[i]); + } + return result; + } + +#ifdef _WINDOWS + float AVXDistanceL2Int8::compare(const int8_t *a, const int8_t *b, + uint32_t length) const { + __m128 r = _mm_setzero_ps(); + __m128i r1; + while (length >= 16) { + r1 = _mm_subs_epi8(_mm_load_si128((__m128i *) a), + _mm_load_si128((__m128i *) b)); + r = _mm_add_ps(r, _mm_mul_epi8(r1)); + a += 16; + b += 16; + length -= 16; + } + r = _mm_hadd_ps(_mm_hadd_ps(r, r), r); + float res = r.m128_f32[0]; + + if (length >= 8) { + __m128 r2 = _mm_setzero_ps(); + __m128i r3 = _mm_subs_epi8(_mm_load_si128((__m128i *) (a - 8)), + _mm_load_si128((__m128i *) (b - 8))); + r2 = _mm_add_ps(r2, _mm_mulhi_epi8(r3)); + a += 8; + b += 8; + length -= 8; + r2 = _mm_hadd_ps(_mm_hadd_ps(r2, r2), r2); + res += r2.m128_f32[0]; + } + + if (length >= 4) { + __m128 r2 = _mm_setzero_ps(); + __m128i r3 = _mm_subs_epi8(_mm_load_si128((__m128i *) (a - 12)), + _mm_load_si128((__m128i *) (b - 12))); + r2 = _mm_add_ps(r2, _mm_mulhi_epi8_shift32(r3)); + res += r2.m128_f32[0] + r2.m128_f32[1]; + } + + return res; + } + + float AVXDistanceL2Float::compare(const float *a, const float *b, + uint32_t length) const { + __m128 diff, v1, v2; + __m128 sum = _mm_set1_ps(0); + + while (length >= 4) { + v1 = _mm_loadu_ps(a); + a += 4; + v2 = _mm_loadu_ps(b); + b += 4; + diff = _mm_sub_ps(v1, v2); + sum = _mm_add_ps(sum, _mm_mul_ps(diff, diff)); + length -= 4; + } + + return sum.m128_f32[0] + sum.m128_f32[1] + sum.m128_f32[2] + + sum.m128_f32[3]; + } +#else + float AVXDistanceL2Int8::compare(const int8_t *, const int8_t *, + uint32_t) const { + return 0; + } + float AVXDistanceL2Float::compare(const float *, const float *, + uint32_t) const { + return 0; + } +#endif +} // namespace diskann diff --git a/src/dll/CMakeLists.txt b/src/dll/CMakeLists.txt index 75ee6ea7fa..d76bee441e 100644 --- a/src/dll/CMakeLists.txt +++ b/src/dll/CMakeLists.txt @@ -1,8 +1,11 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT license. -add_library(diskann_dll SHARED dllmain.cpp ../partition_and_pq.cpp ../pq_flash_index.cpp ../logger.cpp ../utils.cpp - ../windows_aligned_file_reader.cpp ../memory_mapper.cpp ../index.cpp ../math_utils.cpp ../aux_utils.cpp ../ann_exception.cpp) +add_library(diskann_dll SHARED dllmain.cpp + diskann_interface.cpp diskann_incr_index.cpp ../partition_and_pq.cpp ../pq_flash_index.cpp + ../windows_aligned_file_reader.cpp bing_aligned_file_reader.cpp ../logger.cpp ../utils.cpp + ../memory_mapper.cpp ../distance.cpp ../index.cpp ../math_utils.cpp ../aux_utils.cpp ../v2/graph_delta.cpp + ../v2/index_merger.cpp ../v2/merge_insert.cpp ../ann_exception.cpp ) if (MSVC) add_definitions(-D_USRDLL -D_WINDLL -DDISKANN_DLL) add_compile_options(/MD) @@ -11,8 +14,8 @@ if (MSVC) target_link_options(diskann_dll PRIVATE $<$:/IMPLIB:${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}/diskann_dll.lib> $<$:/IMPLIB:${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}/diskann_dll.lib> ) - target_link_libraries(diskann_dll debug ${PROJECT_SOURCE_DIR}/dependencies/windows/tcmalloc/libtcmalloc_minimal.lib) - target_link_libraries(diskann_dll optimized ${PROJECT_SOURCE_DIR}/dependencies/windows/tcmalloc/libtcmalloc_minimal.lib) + target_link_libraries(diskann_dll debug ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}/diskpriority_io.lib ${PROJECT_SOURCE_DIR}/dependencies/windows/tcmalloc/libtcmalloc_minimal.lib) + target_link_libraries(diskann_dll optimized ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}/diskpriority_io.lib ${PROJECT_SOURCE_DIR}/dependencies/windows/tcmalloc/libtcmalloc_minimal.lib) add_custom_command(TARGET @@ -47,10 +50,36 @@ if (MSVC) diskann_dll POST_BUILD COMMAND ${CMAKE_COMMAND} -E copy ${INTEL_ROOT}redist/intel64/mkl/mkl_rt.dll "$<$:\"${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}\">$<$:\"${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}\">" ) + add_custom_command(TARGET + diskann_dll + POST_BUILD + COMMAND ${CMAKE_COMMAND} -E copy ${PROJECT_SOURCE_DIR}/dependencies/windows/dll/boost_date_time-vc141-mt-x64-1_70.dll "$<$:\"${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}\">$<$:\"${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}\">" ) + add_custom_command(TARGET + diskann_dll + POST_BUILD + COMMAND ${CMAKE_COMMAND} -E copy ${PROJECT_SOURCE_DIR}/dependencies/windows/dll/cpprest_2_10.dll "$<$:\"${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}\">$<$:\"${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}\">" ) + add_custom_command(TARGET + diskann_dll + POST_BUILD + COMMAND ${CMAKE_COMMAND} -E copy ${PROJECT_SOURCE_DIR}/dependencies/windows/dll/LIBEAY32.dll "$<$:\"${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}\">$<$:\"${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}\">" ) + add_custom_command(TARGET + diskann_dll + POST_BUILD + COMMAND ${CMAKE_COMMAND} -E copy ${PROJECT_SOURCE_DIR}/dependencies/windows/dll/SSLEAY32.dll "$<$:\"${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}\">$<$:\"${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}\">" ) + add_custom_command(TARGET + diskann_dll + POST_BUILD + COMMAND ${CMAKE_COMMAND} -E copy ${PROJECT_SOURCE_DIR}/dependencies/windows/dll/zlib1.dll "$<$:\"${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}\">$<$:\"${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}\">" ) + add_custom_command(TARGET diskann_dll POST_BUILD COMMAND ${CMAKE_COMMAND} -E copy ${PROJECT_SOURCE_DIR}/dependencies/windows/tcmalloc/libtcmalloc_minimal.dll "$<$:\"${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}\">$<$:\"${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}\">" ) + add_custom_command(TARGET + diskann_dll + POST_BUILD + COMMAND ${CMAKE_COMMAND} -E copy ${PROJECT_SOURCE_DIR}/config/SharedFileStoreTemplate.ini "$<$:\"${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}\">$<$:\"${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}\">" ) + endif() diff --git a/src/index.cpp b/src/index.cpp index 2405878669..4cbddb34c2 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -2,6 +2,7 @@ // Licensed under the MIT license. #include +#include #include #include #include @@ -14,12 +15,17 @@ #include #include #include +#include +#include #include #include #include +#include #include #include #include "tsl/robin_set.h" +#include "tsl/robin_map.h" +#include #include #include @@ -35,318 +41,507 @@ #include "timer.h" #include "utils.h" #include "windows_customizations.h" +#include "ann_exception.h" +#include "tcmalloc/malloc_extension.h" #ifdef _WINDOWS #include #endif +#include "Neighbor_Tag.h" // only L2 implemented. Need to implement inner product search -namespace { - template - diskann::Distance *get_distance_function(diskann::Metric m); - - template<> - diskann::Distance *get_distance_function(diskann::Metric m) { - if (m == diskann::Metric::FAST_L2) { - std::cout << "Here" << std::endl; - return new diskann::DistanceFastL2(); - } else if (m == diskann::Metric::L2) { - if (Avx2SupportedCPU) { - std::cout << "Using AVX2 distance computation" << std::endl; - return new diskann::DistanceL2(); - } else if (AvxSupportedCPU) { - std::cout << "AVX2 not supported. Using AVX distance computation" + +namespace diskann { + // Initialize an index with metric m, load the data of type T with filename + // (bin), and initialize max_points + template + + Index::Index(Metric m, const size_t dim, const size_t max_points, + const bool dynamic_index, + const bool save_index_in_one_file, + const bool enable_tags, const bool support_eager_delete) + : _dist_metric(m), _dim(dim), _max_points(max_points), + _save_as_one_file(save_index_in_one_file), + _dynamic_index(dynamic_index), _enable_tags(enable_tags), + _support_eager_delete(support_eager_delete) { + if (dynamic_index && !enable_tags) { + throw diskann::ANNException( + "ERROR: Eager Deletes must have Dynamic Indexing enabled.", -1, + __FUNCSIG__, __FILE__, __LINE__); + diskann::cerr + << "WARNING: Dynamic Indices must have tags enabled. Auto-enabling." + << std::endl; + _enable_tags = true; + } + if (support_eager_delete && !dynamic_index) { + diskann::cout << "ERROR: Eager Deletes must have Dynamic Indexing " + "enabled. Exitting." + << std::endl; + exit(-1); + } + // data is stored to _nd * aligned_dim matrix with necessary + // zero-padding + _aligned_dim = ROUND_UP(_dim, 8); + + if (dynamic_index) + _num_frozen_pts = 1; + + alloc_aligned(((void **) &_data), + (_max_points + _num_frozen_pts) * _aligned_dim * sizeof(T), + 8 * sizeof(T)); + std::memset(_data, 0, + (_max_points + _num_frozen_pts) * _aligned_dim * sizeof(T)); + + _ep = (unsigned) _max_points; + + _final_graph.reserve(_max_points + _num_frozen_pts); + _final_graph.resize(_max_points + _num_frozen_pts); + + for (size_t i = 0; i < _max_points + _num_frozen_pts; i++) + _final_graph[i].clear(); + + if (_support_eager_delete) { + _in_graph.reserve(_max_points + _num_frozen_pts); + _in_graph.resize(_max_points + _num_frozen_pts); + } + + diskann::cout << "Getting distance function for metric: " + << (m == diskann::Metric::COSINE ? "cosine" : "l2") << std::endl; - return new diskann::AVXDistanceL2Float(); + this->_distance = get_distance_function(m); + _locks = std::vector(_max_points + _num_frozen_pts); + + if (_support_eager_delete) + _locks_in = std::vector(_max_points + _num_frozen_pts); + + _width = 0; + } + + template + Index::~Index() { + delete this->_distance; + aligned_free(_data); + } + + template + void Index::clear_index() { + memset(_data, 0, + _aligned_dim * (_max_points + _num_frozen_pts) * sizeof(T)); + _nd = 0; + for (size_t i = 0; i < _final_graph.size(); i++) + _final_graph[i].clear(); + + _tag_to_location.clear(); + _location_to_tag.clear(); + + _delete_set.clear(); + _empty_slots.clear(); + } + + template + _u64 Index::save_tags(std::string tags_file, size_t offset) { + if (!_enable_tags) { + diskann::cout << "Not saving tags as they are not enabled." << std::endl; + return 0; + } + size_t tag_bytes_written; + TagT * tag_data = new TagT[_nd + _num_frozen_pts]; + for (_u32 i = 0; i < _nd; i++) { + if (_location_to_tag.find(i) != _location_to_tag.end()) { + tag_data[i] = _location_to_tag[i]; } else { - std::cout << "Older CPU. Using slow distance computation" << std::endl; - return new diskann::SlowDistanceL2Float(); + // catering to future when tagT can be any type. + std::memset((char *) &tag_data[i], 0, sizeof(TagT)); } - } else { - std::stringstream stream; - stream << "Only L2 metric supported as of now. Email " - "gopalsr@microsoft.com if you need cosine similarity or inner " - "product." - << std::endl; - std::cerr << stream.str() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); } + if (_num_frozen_pts > 0) { + std::memset((char *) &tag_data[_ep], 0, sizeof(TagT)); + } + tag_bytes_written = + save_bin(tags_file, tag_data, _nd + _num_frozen_pts, 1, offset); + delete[] tag_data; + return tag_bytes_written; } - template<> - diskann::Distance *get_distance_function(diskann::Metric m) { - if (m == diskann::Metric::L2) { - if (Avx2SupportedCPU) { - std::cout << "Using AVX2 distance computation" << std::endl; - return new diskann::DistanceL2Int8(); - } else if (AvxSupportedCPU) { - std::cout << "AVX2 not supported. Using AVX distance computation" - << std::endl; - return new diskann::AVXDistanceL2Int8(); - } else { - std::cout << "Older CPU. Using slow distance computation" << std::endl; - return new diskann::SlowDistanceL2Int(); - } + template + _u64 Index::save_data(std::string data_file, size_t offset) { + return save_data_in_base_dimensions(data_file, _data, _nd + _num_frozen_pts, + _dim, _aligned_dim, offset); + } + + // save the graph index on a file as an adjacency list. For each point, + // first store the number of neighbors, and then the neighbor list (each as + // 4 byte unsigned) + template + _u64 Index::save_graph(std::string graph_file, size_t offset) { + std::ofstream out; + open_file_to_write(out, graph_file); + + out.seekp(offset, out.beg); + _u64 index_size = 24; + _u32 max_degree = 0; + out.write((char *) &index_size, sizeof(uint64_t)); + out.write((char *) &_width, sizeof(unsigned)); + unsigned ep_u32 = _ep; + out.write((char *) &ep_u32, sizeof(unsigned)); + out.write((char *) &_num_frozen_pts, sizeof(_u64)); + for (unsigned i = 0; i < _nd + _num_frozen_pts; i++) { + unsigned GK = (unsigned) _final_graph[i].size(); + out.write((char *) &GK, sizeof(unsigned)); + out.write((char *) _final_graph[i].data(), GK * sizeof(unsigned)); + max_degree = _final_graph[i].size() > max_degree + ? (_u32) _final_graph[i].size() + : max_degree; + index_size += (_u64)(sizeof(unsigned) * (GK + 1)); + } + out.seekp(offset, out.beg); + out.write((char *) &index_size, sizeof(uint64_t)); + out.write((char *) &max_degree, sizeof(_u32)); + out.close(); + return index_size; // number of bytes written + } + + template + _u64 Index::save_delete_list(const std::string &filename, + _u64 file_offset) { + if (_delete_set.size() == 0) { + return 0; + } + std::unique_ptr<_u32[]> delete_list = + std::make_unique<_u32[]>(_delete_set.size()); + _u32 i = 0; + for (auto &del : _delete_set) { + delete_list[i++] = del; + } + return save_bin<_u32>(filename, delete_list.get(), _delete_set.size(), 1, + file_offset); + } + + template + void Index::save(const char *filename) { + // first check if no thread is inserting + auto start = std::chrono::high_resolution_clock::now(); + std::unique_lock lock(_update_lock); + _change_lock.lock(); + + // compact_data(); + compact_frozen_point(); + if (!_save_as_one_file) { + std::string graph_file = std::string(filename); + std::string tags_file = std::string(filename) + ".tags"; + std::string data_file = std::string(filename) + ".data"; + std::string delete_list_file = std::string(filename) + ".del"; + + // Because the save_* functions use append mode, ensure that + // the files are deleted before save. Ideally, we should check + // the error code for delete_file, but will ignore now because + // delete should succeed if save will succeed. + delete_file(graph_file); + save_graph(graph_file); + delete_file(data_file); + save_data(data_file); + delete_file(tags_file); + save_tags(tags_file); + delete_file(delete_list_file); + save_delete_list(delete_list_file); } else { + delete_file(filename); + std::vector cumul_bytes(5, 0); + cumul_bytes[0] = METADATA_SIZE; + cumul_bytes[1] = + cumul_bytes[0] + save_graph(std::string(filename), cumul_bytes[0]); + cumul_bytes[2] = + cumul_bytes[1] + save_data(std::string(filename), cumul_bytes[1]); + cumul_bytes[3] = + cumul_bytes[2] + save_tags(std::string(filename), cumul_bytes[2]); + cumul_bytes[4] = + cumul_bytes[3] + save_delete_list(filename, cumul_bytes[3]); + diskann::save_bin<_u64>(filename, cumul_bytes.data(), cumul_bytes.size(), + 1, 0); + + diskann::cout << "Saved index as one file to " << filename << " of size " + << cumul_bytes[cumul_bytes.size() - 1] << "B." << std::endl; + } + + reposition_frozen_point_to_end(); + + _change_lock.unlock(); + auto stop = std::chrono::high_resolution_clock::now(); + auto timespan = + std::chrono::duration_cast>(stop - start); + diskann::cout << "Time taken for save: " << timespan.count() << "s." + << std::endl; + } + + template + size_t Index::load_tags(const std::string tag_filename, + size_t offset) { + if (_enable_tags && !file_exists(tag_filename)) { + diskann::cerr << "Tag file provided does not exist!" << std::endl; + throw diskann::ANNException("Tag file provided does not exist!", -1, + __FUNCSIG__, __FILE__, __LINE__); + } + + if (!_enable_tags) { + diskann::cout << "Tags not loaded as tags not enabled." << std::endl; + return 0; + } + + size_t file_dim, file_num_points; + TagT * tag_data; + load_bin(std::string(tag_filename), tag_data, file_num_points, + file_dim, offset); + + if (file_dim != 1) { std::stringstream stream; - stream << "Only L2 metric supported as of now. Email " - "gopalsr@microsoft.com if you need cosine similarity or inner " - "product." - << std::endl; + stream << "ERROR: Loading " << file_dim << " dimensions for tags," + << "but tag file must have 1 dimension." << std::endl; std::cerr << stream.str() << std::endl; + delete[] tag_data; throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } + + size_t num_data_points = + _num_frozen_pts > 0 ? file_num_points - 1 : file_num_points; + for (_u32 i = 0; i < (_u32) num_data_points; i++) { + TagT tag = *(tag_data + i); + if (_delete_set.find(i) == _delete_set.end()) { + _location_to_tag[i] = tag; + _tag_to_location[tag] = (_u32) i; + } + } + diskann::cout << "Tags loaded." << std::endl; + delete[] tag_data; + return file_num_points; } - template<> - diskann::Distance *get_distance_function(diskann::Metric m) { - if (m == diskann::Metric::L2) { - std::cout << "AVX/AVX2 distance function not defined for Uint8. Using " - "slow version. " - "Contact gopalsr@microsoft.com if you need AVX/AVX2 support." - << std::endl; - return new diskann::DistanceL2UInt8(); - } else { + template + size_t Index::load_data(std::string filename, size_t offset) { + if (!file_exists(filename)) { std::stringstream stream; - stream << "Only L2 metric supported as of now. Email " - "gopalsr@microsoft.com if you need cosine similarity or inner " - "product." + stream << "ERROR: data file " << filename << " does not exist." << std::endl; std::cerr << stream.str() << std::endl; + aligned_free(_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } - } -} // namespace - -namespace diskann { - - // Initialize an index with metric m, load the data of type T with filename - // (bin), and initialize max_points - template - Index::Index(Metric m, const char *filename, const size_t max_points, - const size_t nd, const size_t num_frozen_pts, - const bool enable_tags, const bool store_data, - const bool support_eager_delete) - : _num_frozen_pts(num_frozen_pts), _has_built(false), _width(0), - _can_delete(false), _eager_done(true), _lazy_done(true), - _compacted_order(true), _enable_tags(enable_tags), - _consolidated_order(true), _support_eager_delete(support_eager_delete), - _store_data(store_data) { - // data is stored to _nd * aligned_dim matrix with necessary - // zero-padding - diskann::cout << "Number of frozen points = " << _num_frozen_pts - << std::endl; - load_aligned_bin(std::string(filename), _data, _nd, _dim, _aligned_dim); + size_t file_dim, file_num_points; + diskann::get_bin_metadata(filename, file_num_points, file_dim, offset); - if (nd > 0) { - if (_nd >= nd) - _nd = nd; // Consider the first _nd points and ignore the rest. - else { - std::stringstream stream; - stream << "ERROR: Driver requests loading " << _nd << " points," - << "but file has fewer (" << nd << ") points" << std::endl; - diskann::cerr << stream.str() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); - } - } + // since we are loading a new dataset, _empty_slots must be cleared + _empty_slots.clear(); - _max_points = (max_points > 0) ? max_points : _nd; - if (_max_points < _nd) { + if (file_dim != _dim) { std::stringstream stream; - stream << "ERROR: max_points must be >= data size; max_points: " - << _max_points << " n: " << _nd << std::endl; - diskann::cerr << stream.str() << std::endl; + stream << "ERROR: Driver requests loading " << _dim << " dimension," + << "but file has " << file_dim << " dimension." << std::endl; + std::cerr << stream.str() << std::endl; + aligned_free(_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } - // Allocate space for max points and frozen points, - // and add frozen points at the end of the array - if (_num_frozen_pts > 0) { - auto temp = _data; - _data = (T *) realloc( - _data, (_max_points + _num_frozen_pts) * _aligned_dim * sizeof(T)); - if (_data == nullptr) { - free(temp); - diskann::cout << "Realloc failed, killing programme" << std::endl; - throw diskann::ANNException("Realloc failed", -1, __FUNCSIG__, __FILE__, - __LINE__); - } - } - - this->_distance = ::get_distance_function(m); - _locks = std::vector(_max_points + _num_frozen_pts); + if (file_num_points > _max_points + _num_frozen_pts) { + //_change_lock is already locked in load() + std::unique_lock tl(_tag_lock); + std::unique_lock growth_lock(_update_lock); - _width = 0; - } + resize(file_num_points); + } - template<> - Index::~Index() { - delete this->_distance; - aligned_free(_data); + copy_aligned_data_from_file(std::string(filename), _data, + file_num_points, file_dim, _aligned_dim, + offset); + return file_num_points; } - template<> - Index<_s8>::~Index() { - delete this->_distance; - aligned_free(_data); + template + void Index::get_delete_set(tsl::robin_set &del_set) { + del_set = _delete_set; } - template<> - Index<_u8>::~Index() { - delete this->_distance; - aligned_free(_data); + template + size_t Index::load_delete_set(const std::string &filename, + size_t offset) { + std::unique_ptr<_u32[]> delete_list; + _u64 npts, ndim; + load_bin<_u32>(filename, delete_list, npts, ndim, offset); + assert(ndim == 1); + for (size_t i = 0; i < npts; i++) { + _delete_set.insert(delete_list[i]); + } + return npts; } - // save the graph index on a file as an adjacency list. For each point, - // first store the number of neighbors, and then the neighbor list (each as - // 4 byte unsigned) + // load the index from file and update the width (max_degree), ep (navigating + // node id), and _final_graph (adjacency list) template - void Index::save(const char *filename) { - long long total_gr_edges = 0; - size_t index_size = 0; - std::ofstream out(std::string(filename), std::ios::binary | std::ios::out); - - if (_support_eager_delete) - if (_eager_done && (!_compacted_order)) { - if (_nd < _max_points) { - assert(_final_graph.size() == _max_points + _num_frozen_pts); - unsigned active = 0; - std::vector new_location = get_new_location(active); - diskann::cout << "Size of new_location = " << new_location.size() - << std::endl; - for (unsigned i = 0; i < new_location.size(); i++) - if ((_delete_set.find(i) == _delete_set.end()) && - (new_location[i] >= _max_points + _num_frozen_pts)) - diskann::cout << "Wrong new_location assigned to " << i - << std::endl; - else { - if ((_delete_set.find(i) != _delete_set.end()) && - (new_location[i] < _max_points + _num_frozen_pts)) - diskann::cout << "Wrong location assigned to " - "delete point " - << i << std::endl; - } - compact_data(new_location, active, _compacted_order); - - if (_support_eager_delete) - update_in_graph(); + void Index::load(const char *filename) { + _change_lock.lock(); + + size_t tags_file_num_pts = 0, graph_num_pts = 0, data_file_num_pts = 0; + + if (!_save_as_one_file) { + std::string data_file = std::string(filename) + ".data"; + std::string tags_file = std::string(filename) + ".tags"; + std::string delete_set_file = std::string(filename) + ".del"; + std::string graph_file = std::string(filename); + data_file_num_pts = load_data(data_file); + if (file_exists(delete_set_file)) { + load_delete_set(delete_set_file); + } + if (_enable_tags) { + tags_file_num_pts = load_tags(tags_file); + } + graph_num_pts = load_graph(graph_file, data_file_num_pts); - } else { - assert(_final_graph.size() == _max_points + _num_frozen_pts); - if (_enable_tags) { - _change_lock.lock(); - if (_can_delete) { - std::cerr << "Disable deletes and consolidate " - "index before saving." - << std::endl; - throw diskann::ANNException( - "Disable deletes and consolidate index before " - "saving.", - -1, __FUNCSIG__, __FILE__, __LINE__); - } - } - } + } else { + _u64 nr, nc; + std::unique_ptr<_u64[]> file_offset_data; + + std::string index_file(filename); + + diskann::load_bin<_u64>(index_file, file_offset_data, nr, nc, 0); + // Loading data first so that we know how many points to expect. + data_file_num_pts = load_data(index_file, file_offset_data[1]); + graph_num_pts = + load_graph(index_file, data_file_num_pts, file_offset_data[0]); + if (file_offset_data[3] != file_offset_data[4]) { + load_delete_set(index_file, file_offset_data[3]); } - if (_lazy_done) { - assert(_final_graph.size() == _max_points + _num_frozen_pts); if (_enable_tags) { - _change_lock.lock(); - if (_can_delete || (!_consolidated_order)) { - diskann::cout << "Disable deletes and consolidate index before " - "saving." - << std::endl; - throw diskann::ANNException( - "Disable deletes and consolidate index before saving.", -1, - __FUNCSIG__, __FILE__, __LINE__); - } + tags_file_num_pts = load_tags(index_file, file_offset_data[2]); } } - out.write((char *) &index_size, sizeof(uint64_t)); - out.write((char *) &_width, sizeof(unsigned)); - out.write((char *) &_ep, sizeof(unsigned)); - for (unsigned i = 0; i < _nd + _num_frozen_pts; i++) { - unsigned GK = (unsigned) _final_graph[i].size(); - out.write((char *) &GK, sizeof(unsigned)); - out.write((char *) _final_graph[i].data(), GK * sizeof(unsigned)); - total_gr_edges += GK; + + if (data_file_num_pts != graph_num_pts || + (data_file_num_pts != tags_file_num_pts && _enable_tags)) { + std::stringstream stream; + stream << "ERROR: When loading index, loaded " << data_file_num_pts + << " points from datafile, " << graph_num_pts + << " from graph, and " << tags_file_num_pts + << " tags, with num_frozen_pts being set to " << _num_frozen_pts + << " in constructor." << std::endl; + std::cerr << stream.str() << std::endl; + aligned_free(_data); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); } - index_size = out.tellp(); - out.seekp(0, std::ios::beg); - out.write((char *) &index_size, sizeof(uint64_t)); - out.close(); - diskann::cout << "Avg degree: " - << ((float) total_gr_edges) / - ((float) (_nd + _num_frozen_pts)) - << std::endl; + _nd = data_file_num_pts - _num_frozen_pts; + _empty_slots.clear(); + for (_u32 i = _nd; i < _max_points; i++) { + _empty_slots.insert(i); + } + + _lazy_done = _delete_set.size() != 0; + + reposition_frozen_point_to_end(); + diskann::cout << "Num frozen points:" << _num_frozen_pts << " _nd: " << _nd + << " _ep: " << _ep + << " size(_location_to_tag): " << _location_to_tag.size() + << " size(_tag_to_location):" << _tag_to_location.size() + << " Max points: " << _max_points << std::endl; + + _change_lock.unlock(); } - // load the index from file and update the width (max_degree), ep - // (navigating node id), and _final_graph (adjacency list) template - void Index::load(const char *filename, const bool load_tags, - const char *tag_filename) { - if (!validate_file_size(filename)) { - return; - } + size_t Index::load_graph(std::string filename, + size_t expected_num_points, size_t offset) { std::ifstream in(filename, std::ios::binary); - size_t expected_file_size; + in.seekg(offset, in.beg); + size_t expected_file_size; + _u64 file_frozen_pts; in.read((char *) &expected_file_size, sizeof(_u64)); in.read((char *) &_width, sizeof(unsigned)); in.read((char *) &_ep, sizeof(unsigned)); + in.read((char *) &file_frozen_pts, sizeof(_u64)); + + if (file_frozen_pts != _num_frozen_pts) { + std::stringstream stream; + if (file_frozen_pts == 1) + stream << "ERROR: When loading index, detected dynamic index, but " + "constructor asks for static index. Exitting." + << std::endl; + else + stream << "ERROR: When loading index, detected static index, but " + "constructor asks for dynamic index. Exitting." + << std::endl; + std::cerr << stream.str() << std::endl; + aligned_free(_data); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } diskann::cout << "Loading vamana index " << filename << "..." << std::flush; + // Sanity check. In case the user gave us fewer points as max_points than + // the number + // of points in the dataset, resize the _final_graph to the larger size. + if (_max_points < (expected_num_points - _num_frozen_pts)) { + diskann::cout << "Number of points in data: " << expected_num_points + << " is more than max_points argument: " + << _final_graph.size() + << " Setting max points to: " << expected_num_points + << std::endl; + _final_graph.resize(expected_num_points); + _max_points = expected_num_points - _num_frozen_pts; + // changed expected_num to expected_num - frozen_num + } + + size_t bytes_read = 24; size_t cc = 0; unsigned nodes = 0; - while (!in.eof()) { + while (bytes_read != expected_file_size) { unsigned k; in.read((char *) &k, sizeof(unsigned)); - if (in.eof()) - break; + if (k == 0) { + diskann::cerr << "ERROR: Point found with no out-neighbors, point#" + << nodes << std::endl; + } + // if (in.eof()) + // break; cc += k; ++nodes; std::vector tmp(k); + tmp.reserve(k); in.read((char *) tmp.data(), k * sizeof(unsigned)); - - _final_graph.emplace_back(tmp); + _final_graph[nodes - 1].swap(tmp); + bytes_read += sizeof(uint32_t) * ((_u64) k + 1); if (nodes % 10000000 == 0) diskann::cout << "." << std::flush; } - if (_final_graph.size() != _nd) { - diskann::cout << "ERROR. mismatch in number of points. Graph has " - << _final_graph.size() << " points and loaded dataset has " - << _nd << " points. " << std::endl; - return; - } - diskann::cout << "..done. Index has " << nodes << " nodes and " << cc - << " out-edges" << std::endl; + diskann::cout << "done. Index has " << nodes << " nodes and " << cc + << " out-edges, _ep is set to " << _ep << std::endl; + return nodes; + } - if (load_tags) { - if (_enable_tags == false) - diskann::cout << "Enabling tags." << std::endl; - _enable_tags = true; - std::ifstream tag_file; - if (tag_filename == NULL) - tag_file = std::ifstream(std::string(filename) + std::string(".tags")); - else - tag_file = std::ifstream(std::string(tag_filename)); - if (!tag_file.is_open()) { - std::cerr << "Tag file not found." << std::endl; - return; - } - unsigned id = 0; - TagT tag; - while (tag_file >> tag) { - _location_to_tag[id] = tag; - _tag_to_location[tag] = id++; - } - tag_file.close(); - assert(id == _nd); + template + int Index::get_vector_by_tag(TagT &tag, T *vec) { + std::shared_lock lock(_tag_lock); + if (_tag_to_location.find(tag) == _tag_to_location.end()) { + diskann::cout << "Tag " << tag << " does not exist" << std::endl; + return -1; + } + unsigned location = _tag_to_location[tag]; + // memory should be allocated for vec before calling this function + memcpy((void *) vec, (void *) (_data + (size_t)(location * _aligned_dim)), + (size_t) _aligned_dim * sizeof(T)); + return 0; + } + + template + const T *Index::get_vector_by_tag(const TagT &tag) { + std::shared_lock lock(_tag_lock); + if (_tag_to_location.find(tag) == _tag_to_location.end()) { + diskann::cout << "Tag " << tag << " does not exist in the index." + << std::endl; + return nullptr; + } else { + unsigned location = _tag_to_location[tag]; + return _data + (size_t)(location * _aligned_dim); } } @@ -366,13 +561,13 @@ namespace diskann { for (size_t i = 0; i < _nd; i++) for (size_t j = 0; j < _aligned_dim; j++) - center[j] += _data[i * _aligned_dim + j]; + center[j] += (float) _data[i * _aligned_dim + j]; for (size_t j = 0; j < _aligned_dim; j++) - center[j] /= _nd; + center[j] /= (float) _nd; // compute all to one distance - float * distances = new float[_nd](); + float *distances = new float[_nd](); #pragma omp parallel for schedule(static, 65536) for (_s64 i = 0; i < (_s64) _nd; i++) { // extract point and distance reference @@ -381,7 +576,8 @@ namespace diskann { dist = 0; float diff = 0; for (size_t j = 0; j < _aligned_dim; j++) { - diff = (center[j] - cur_vec[j]) * (center[j] - cur_vec[j]); + diff = + (center[j] - (float) cur_vec[j]) * (center[j] - (float) cur_vec[j]); dist += diff; } } @@ -417,8 +613,11 @@ namespace diskann { const std::vector &init_ids, std::vector & expanded_nodes_info, tsl::robin_set & expanded_nodes_ids, - std::vector & best_L_nodes) { + std::vector &best_L_nodes, bool ret_frozen) { best_L_nodes.resize(Lsize + 1); + for (unsigned i = 0; i < Lsize + 1; i++) { + best_L_nodes[i].distance = std::numeric_limits::max(); + } expanded_nodes_info.reserve(10 * Lsize); expanded_nodes_ids.reserve(10 * Lsize); @@ -428,7 +627,7 @@ namespace diskann { inserted_into_pool.reserve(Lsize * 20); for (auto id : init_ids) { - assert(id < _max_points); + assert(id < _max_points + _num_frozen_pts); nn = Neighbor(id, _distance->compare(_data + _aligned_dim * (size_t) id, node_coords, (unsigned) _aligned_dim), @@ -453,16 +652,46 @@ namespace diskann { if (best_L_nodes[k].flag) { best_L_nodes[k].flag = false; auto n = best_L_nodes[k].id; - expanded_nodes_info.emplace_back(best_L_nodes[k]); - expanded_nodes_ids.insert(n); + if (!(best_L_nodes[k].id == _ep && _num_frozen_pts > 0 && + !ret_frozen)) { + expanded_nodes_info.emplace_back(best_L_nodes[k]); + expanded_nodes_ids.insert(n); + } + std::vector des; + if (_dynamic_index) { + LockGuard guard(_locks[n]); + for (unsigned m = 0; m < _final_graph[n].size(); m++) { + if (_final_graph[n][m] >= _max_points + _num_frozen_pts) { + diskann::cerr << "Wrong id found: " << _final_graph[n][m] + << std::endl; + throw diskann::ANNException( + std::string("Wrong id found") + + std::to_string(_final_graph[n][m]), + -1, __FUNCSIG__, __FILE__, __LINE__); + } + des.emplace_back(_final_graph[n][m]); + } + } else { + for (unsigned m = 0; m < _final_graph[n].size(); m++) { + if (_final_graph[n][m] >= _max_points + _num_frozen_pts) { + diskann::cerr << "Wrong id found: " << _final_graph[n][m] + << std::endl; + throw diskann::ANNException( + std::string("Wrong id found") + + std::to_string(_final_graph[n][m]), + -1, __FUNCSIG__, __FILE__, __LINE__); + } + des.emplace_back(_final_graph[n][m]); + } + } - for (unsigned m = 0; m < _final_graph[n].size(); ++m) { - unsigned id = _final_graph[n][m]; + for (unsigned m = 0; m < des.size(); ++m) { + unsigned id = des[m]; if (inserted_into_pool.find(id) == inserted_into_pool.end()) { inserted_into_pool.insert(id); - if ((m + 1) < _final_graph[n].size()) { - auto nextn = _final_graph[n][m + 1]; + if ((m + 1) < des.size()) { + auto nextn = des[m + 1]; diskann::prefetch_vector( (const char *) _data + _aligned_dim * (size_t) nextn, sizeof(T) * _aligned_dim); @@ -495,6 +724,25 @@ namespace diskann { return std::make_pair(hops, cmps); } + template + void Index::iterate_to_fixed_point( + const T *node_coords, const unsigned Lindex, + std::vector & expanded_nodes_info, + tsl::robin_map &coord_map, bool return_frozen_pt) { + std::vector init_ids; + init_ids.push_back(this->_ep); + std::vector best_L_nodes; + tsl::robin_set expanded_nodes_ids; + this->iterate_to_fixed_point(node_coords, Lindex, init_ids, + expanded_nodes_info, expanded_nodes_ids, + best_L_nodes, return_frozen_pt); + for (Neighbor &einf : expanded_nodes_info) { + T *coords = + this->_data + (uint64_t) einf.id * (uint64_t) this->_aligned_dim; + coord_map.insert(std::make_pair(einf.id, coords)); + } + } + template void Index::get_expanded_nodes( const size_t node_id, const unsigned Lindex, @@ -555,12 +803,12 @@ namespace diskann { } start++; } - cur_alpha *= 1.2; + cur_alpha *= 1.2f; } } template - void Index::prune_neighbors(const unsigned location, + void Index::prune_neighbors(const unsigned location, std::vector &pool, const Parameters & parameter, std::vector &pruned_list) { @@ -568,8 +816,10 @@ namespace diskann { unsigned maxc = parameter.Get("C"); float alpha = parameter.Get("alpha"); - if (pool.size() == 0) - return; + if (pool.size() == 0) { + throw diskann::ANNException("Pool passed to prune_neighbors is empty", + -1); + } _width = (std::max)(_width, range); @@ -618,7 +868,7 @@ namespace diskann { if (des == n) continue; /* des.id is the id of the neighbors of n */ - assert(des >= 0 && des < _max_points); + assert(des >= 0 && des < _max_points + _num_frozen_pts); if (des > _max_points) diskann::cout << "error. " << des << " exceeds max_pts" << std::endl; /* des_pool contains the neighbors of the neighbors of n */ @@ -640,19 +890,20 @@ namespace diskann { * the current node n. */ template - void Index::inter_insert(unsigned n, + void Index::inter_insert(unsigned n, std::vector &pruned_list, const Parameters & parameter, bool update_in_graph) { const auto range = parameter.Get("R"); - assert(n >= 0 && n < _nd); + assert(n >= 0 && n < _nd + _num_frozen_pts); + const auto &src_pool = pruned_list; assert(!src_pool.empty()); for (auto des : src_pool) { /* des.id is the id of the neighbors of n */ - assert(des >= 0 && des < _max_points); + assert(des >= 0 && des < _max_points + _num_frozen_pts); /* des_pool contains the neighbors of the neighbors of n */ auto & des_pool = _final_graph[des]; std::vector copy_of_neighbors; @@ -660,14 +911,11 @@ namespace diskann { { LockGuard guard(_locks[des]); if (std::find(des_pool.begin(), des_pool.end(), n) == des_pool.end()) { - if (des_pool.size() < SLACK_FACTOR * range) { + if (des_pool.size() < (_u64)(SLACK_FACTOR * range)) { des_pool.emplace_back(n); if (update_in_graph) { - // USE APPROPRIATE LOCKS FOR IN_GRAPH - if (std::find(_in_graph[n].begin(), _in_graph[n].end(), des) == - _in_graph[n].end()) { - _in_graph[n].emplace_back(des); - } + LockGuard guard(_locks_in[n]); + _in_graph[n].emplace_back(des); } prune_needed = false; } else { @@ -701,14 +949,30 @@ namespace diskann { prune_neighbors(des, dummy_pool, parameter, new_out_neighbors); { LockGuard guard(_locks[des]); - // DELETE IN-EDGES FROM IN_GRAPH USING APPROPRIATE LOCKS + // updating in_graph of out-neighbors of des + if (update_in_graph) { + for (auto out_nbr : _final_graph[des]) { + { + LockGuard guard(_locks_in[out_nbr]); + for (unsigned i = 0; i < _in_graph[out_nbr].size(); i++) { + if (_in_graph[out_nbr][i] == des) { + _in_graph[out_nbr].erase(_in_graph[out_nbr].begin() + i); + break; + } + } + } + } + } + _final_graph[des].clear(); for (auto new_nbr : new_out_neighbors) { _final_graph[des].emplace_back(new_nbr); if (update_in_graph) { + LockGuard guard(_locks_in[new_nbr]); _in_graph[new_nbr].emplace_back(des); } } + _final_graph[des].shrink_to_fit(); } } } @@ -750,14 +1014,16 @@ namespace diskann { parameters.Set("alpha", 1); /* visit_order is a vector that is initialized to the entire graph */ - std::vector visit_order; + std::vector visit_order; + std::vector pool, tmp; + tsl::robin_set visited; visit_order.reserve(_nd + _num_frozen_pts); for (unsigned i = 0; i < (unsigned) _nd; i++) { visit_order.emplace_back(i); } - for (unsigned i = 0; i < (unsigned) _num_frozen_pts; ++i) - visit_order.emplace_back((unsigned) (_max_points + i)); + if (_num_frozen_pts > 0) + visit_order.emplace_back((unsigned) _max_points); // if there are frozen points, the first such one is set to be the _ep if (_num_frozen_pts > 0) @@ -765,8 +1031,6 @@ namespace diskann { else _ep = calculate_entry_point(); - _final_graph.reserve(_max_points + _num_frozen_pts); - _final_graph.resize(_max_points + _num_frozen_pts); if (_support_eager_delete) { _in_graph.reserve(_max_points + _num_frozen_pts); _in_graph.resize(_max_points + _num_frozen_pts); @@ -925,20 +1189,24 @@ namespace diskann { progress_counter += 5; } } -// Gopal. Splittng diskann_dll into separate DLLs for search and build. +// Splittng diskann_dll into separate DLLs for search and build. // This code should only be available in the "build" DLL. -#ifdef DISKANN_BUILD +#ifdef USE_TCMALLOC MallocExtension::instance()->ReleaseFreeMemory(); #endif - diskann::cout << "Completed Pass " << rnd_no << " of data using L=" << L - << " and alpha=" << parameters.Get("alpha") - << ". Stats: "; - diskann::cout << "search+prune_time=" << total_sync_time - << "s, inter_time=" << total_inter_time - << "s, inter_count=" << total_inter_count << std::endl; + if (_nd > 0) { + diskann::cout << "Completed Pass " << rnd_no << " of data using L=" << L + << " and alpha=" << parameters.Get("alpha") + << ". Stats: "; + diskann::cout << "search+prune_time=" << total_sync_time + << "s, inter_time=" << total_inter_time + << "s, inter_count=" << total_inter_count << std::endl; + } } - diskann::cout << "Starting final cleanup.." << std::flush; + if (_nd > 0) { + diskann::cout << "Starting final cleanup.." << std::flush; + } #pragma omp parallel for schedule(dynamic, 65536) for (_s64 node_ctr = 0; node_ctr < (_s64)(visit_order.size()); node_ctr++) { auto node = visit_order[node_ctr]; @@ -965,30 +1233,141 @@ namespace diskann { _final_graph[node].emplace_back(id); } } - diskann::cout << "done. Link time: " - << ((double) link_timer.elapsed() / (double) 1000000) << "s" - << std::endl; + if (_nd > 0) { + diskann::cout << "done. Link time: " + << ((double) link_timer.elapsed() / (double) 1000000) << "s" + << std::endl; + } } template - void Index::build(Parameters ¶meters, - const std::vector &tags) { - if (_enable_tags) { - if (tags.size() != _nd) { - std::cerr << "#Tags should be equal to #points" << std::endl; - throw diskann::ANNException("#Tags must be equal to #points", -1, - __FUNCSIG__, __FILE__, __LINE__); - } - for (size_t i = 0; i < tags.size(); ++i) { - _tag_to_location[tags[i]] = (unsigned) i; - _location_to_tag[(unsigned) i] = tags[i]; - } - } - diskann::cout << "Starting index build..." << std::endl; - link(parameters); // Primary func for creating graph + void Index::prune_all_nbrs(const Parameters ¶meters) { + const unsigned range = parameters.Get("R"); - if (_support_eager_delete) { - update_in_graph(); // copying values to in_graph + diskann::Timer timer; +#pragma omp parallel for + for (_s64 node = 0; node < (_s64)(_max_points + _num_frozen_pts); node++) { + if ((size_t) node < _nd || (size_t) node == _max_points) { + if (_final_graph[node].size() > range) { + tsl::robin_set dummy_visited(0); + std::vector dummy_pool(0); + std::vector new_out_neighbors; + + for (auto cur_nbr : _final_graph[node]) { + if (dummy_visited.find(cur_nbr) == dummy_visited.end() && + cur_nbr != node) { + float dist = + _distance->compare(_data + _aligned_dim * (size_t) node, + _data + _aligned_dim * (size_t) cur_nbr, + (unsigned) _aligned_dim); + dummy_pool.emplace_back(Neighbor(cur_nbr, dist, true)); + dummy_visited.insert(cur_nbr); + } + } + prune_neighbors((_u32) node, dummy_pool, parameters, + new_out_neighbors); + + _final_graph[node].clear(); + for (auto id : new_out_neighbors) + _final_graph[node].emplace_back(id); + } + } + } + + diskann::cout << "Prune time : " << timer.elapsed() / 1000 << "ms" + << std::endl; + size_t max = 0, min = 1 << 30, total = 0, cnt = 0; + for (size_t i = 0; i < (_nd + _num_frozen_pts); i++) { + auto &pool = _final_graph[i]; + max = (std::max)(max, pool.size()); + min = (std::min)(min, pool.size()); + total += pool.size(); + if (pool.size() < 2) + cnt++; + } + if (min > max) + min = max; + if (_nd > 0) { + diskann::cout << "Index built with degree: max:" << max << " avg:" + << (float) total / (float) (_nd + _num_frozen_pts) + << " min:" << min << " count(deg<2):" << cnt << std::endl; + } + } + + template + void Index::build(const char * filename, + const size_t num_points_to_load, + Parameters & parameters, + const std::vector &tags) { + if (!file_exists(filename)) { + diskann::cerr << "Data file " << filename + << " does not exist!!! Exiting...." << std::endl; + std::stringstream stream; + stream << "Data file " << filename << " does not exist." << std::endl; + std::cerr << stream.str() << std::endl; + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } + + size_t file_num_points, file_dim; + if (filename == nullptr) { + diskann::cout << "Starting with an empty index." << std::endl; + _nd = 0; + } else { + diskann::get_bin_metadata(filename, file_num_points, file_dim); + if (file_num_points > _max_points || + num_points_to_load > file_num_points) { + std::stringstream stream; + stream << "ERROR: Driver requests loading " << num_points_to_load + << " points and file has " << file_num_points << " points, but " + << "index can support only " << _max_points + << " points as specified in constructor." << std::endl; + std::cerr << stream.str() << std::endl; + aligned_free(_data); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } + if (file_dim != _dim) { + std::stringstream stream; + stream << "ERROR: Driver requests loading " << _dim << " dimension," + << "but file has " << file_dim << " dimension." << std::endl; + std::cerr << stream.str() << std::endl; + aligned_free(_data); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } + + copy_aligned_data_from_file(std::string(filename), _data, + file_num_points, file_dim, _aligned_dim); + + diskann::cout << "Loading only first " << num_points_to_load + << " from file.. " << std::endl; + _nd = num_points_to_load; + + if (_enable_tags && tags.size() != num_points_to_load) { + std::stringstream stream; + stream << "ERROR: Driver requests loading " << num_points_to_load + << " points from file," + << "but tags vector is of size " << tags.size() << "." + << std::endl; + std::cerr << stream.str() << std::endl; + aligned_free(_data); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } + if (_enable_tags) { + for (size_t i = 0; i < tags.size(); ++i) { + _tag_to_location[tags[i]] = (unsigned) i; + _location_to_tag[(unsigned) i] = tags[i]; + } + } + } + + generate_frozen_point(); + link(parameters); // Primary func for creating nsg graph + + if (_support_eager_delete) { + update_in_graph(); // copying values to in_graph } size_t max = 0, min = 1 << 30, total = 0, cnt = 0; @@ -1000,38 +1379,204 @@ namespace diskann { if (pool.size() < 2) cnt++; } - diskann::cout << "Degree: max:" << max - << " avg:" << (float) total / (float) _nd << " min:" << min - << " count(deg<2):" << cnt << "\n" - << "Index built." << std::endl; + if (min > max) + min = max; + if (_nd > 0) { + diskann::cout << "Index built with degree: max:" << max << " avg:" + << (float) total / (float) (_nd + _num_frozen_pts) + << " min:" << min << " count(deg<2):" << cnt << std::endl; + } _width = (std::max)((unsigned) max, _width); _has_built = true; } template - std::pair Index::search(const T *query, + void Index::build(const char * filename, + const size_t num_points_to_load, + Parameters ¶meters, const char *tag_filename) { + if (!file_exists(filename)) { + std::cerr << "Data file provided " << filename << " does not exist." + << std::endl; + std::stringstream stream; + stream << "Data file provided " << filename << " does not exist." + << std::endl; + std::cerr << stream.str() << std::endl; + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } + + size_t file_num_points, file_dim; + if (filename == nullptr) { + diskann::cout << "Starting with an empty index." << std::endl; + _nd = 0; + } else { + diskann::get_bin_metadata(filename, file_num_points, file_dim); + if (file_num_points > _max_points || + num_points_to_load > file_num_points) { + std::stringstream stream; + stream << "ERROR: Driver requests loading " << num_points_to_load + << " points and file has " << file_num_points << " points, but " + << "index can support only " << _max_points + << " points as specified in constructor." << std::endl; + std::cerr << stream.str() << std::endl; + aligned_free(_data); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } + if (file_dim != _dim) { + std::stringstream stream; + stream << "ERROR: Driver requests loading " << _dim << " dimension," + << "but file has " << file_dim << " dimension." << std::endl; + std::cerr << stream.str() << std::endl; + aligned_free(_data); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } + + copy_aligned_data_from_file(std::string(filename), _data, + file_num_points, file_dim, _aligned_dim); + + diskann::cout << "Loading only first " << num_points_to_load + << " from file.. " << std::endl; + _nd = num_points_to_load; + if (_enable_tags) { + if (tag_filename == nullptr) { + for (unsigned i = 0; i < num_points_to_load; i++) { + _tag_to_location[i] = i; + _location_to_tag[i] = i; + } + } else { + if (file_exists(tag_filename)) { + diskann::cout << "Loading tags from " << tag_filename + << " for vamana index build" << std::endl; + TagT * tag_data = nullptr; + size_t npts, ndim; + diskann::load_bin(tag_filename, tag_data, npts, ndim); + if (npts != num_points_to_load) { + std::stringstream sstream; + sstream << "Loaded " << npts + << " tags instead of expected number: " + << num_points_to_load; + diskann::cerr << sstream.str() << std::endl; + throw diskann::ANNException(sstream.str(), -1, __FUNCSIG__, + __FILE__, __LINE__); + } + for (size_t i = 0; i < npts; i++) { + _tag_to_location[tag_data[i]] = (unsigned) i; + _location_to_tag[(unsigned) i] = tag_data[i]; + } + delete[] tag_data; + } else { + diskann::cerr << "Tag file " << tag_filename + << " does not exist. Exiting..." << std::endl; + throw diskann::ANNException( + std::string("Tag file") + tag_filename + " does not exist", -1, + __FUNCSIG__, __FILE__, __LINE__); + } + } + } + } + + generate_frozen_point(); + link(parameters); // Primary func for creating nsg graph + + if (_support_eager_delete) { + update_in_graph(); // copying values to in_graph + } + + size_t max = 0, min = 1 << 30, total = 0, cnt = 0; + for (size_t i = 0; i < _nd; i++) { + auto &pool = _final_graph[i]; + max = (std::max)(max, pool.size()); + min = (std::min)(min, pool.size()); + total += pool.size(); + if (pool.size() < 2) + cnt++; + } + if (min > max) + min = max; + if (_nd > 0) { + diskann::cout << "Index built with degree: max:" << max << " avg:" + << (float) total / (float) (_nd + _num_frozen_pts) + << " min:" << min << " count(deg<2):" << cnt << std::endl; + } + _width = (std::max)((unsigned) max, _width); + _has_built = true; + } + + template + std::pair Index::search( + const T *query, const size_t K, const unsigned L, + std::vector> &best_K_tags) { + std::shared_lock ulock(_update_lock); + assert(best_K_tags.size() == 0); + std::vector init_ids; + tsl::robin_set visited(10 * L); + std::vector best, expanded_nodes_info; + tsl::robin_set expanded_nodes_ids; + + if (init_ids.size() == 0) { + init_ids.emplace_back(_ep); + } + + T * aligned_query; + size_t allocSize = _aligned_dim * sizeof(T); + alloc_aligned(((void **) &aligned_query), allocSize, 8 * sizeof(T)); + memset(aligned_query, 0, _aligned_dim * sizeof(T)); + memcpy(aligned_query, query, _dim * sizeof(T)); + auto retval = + iterate_to_fixed_point(aligned_query, L, init_ids, expanded_nodes_info, + expanded_nodes_ids, best, false); + + std::shared_lock lock(_tag_lock); + for (auto iter : best) { + if (_location_to_tag.find(iter.id) != _location_to_tag.end()) + best_K_tags.emplace_back( + Neighbor_Tag(_location_to_tag[iter.id], iter.distance)); + if (best_K_tags.size() == K) + break; + } + aligned_free(aligned_query); + return retval; + } + + template + std::pair Index::search(const T * query, const size_t K, const unsigned L, - unsigned * indices) { + unsigned * indices, + float *distances) { std::vector init_ids; tsl::robin_set visited(10 * L); std::vector best_L_nodes, expanded_nodes_info; tsl::robin_set expanded_nodes_ids; + std::shared_lock lock(_update_lock); + if (init_ids.size() == 0) { init_ids.emplace_back(_ep); } + T * aligned_query; + size_t allocSize = _aligned_dim * sizeof(T); + alloc_aligned(((void **) &aligned_query), allocSize, 8 * sizeof(T)); + memset(aligned_query, 0, _aligned_dim * sizeof(T)); + memcpy(aligned_query, query, _dim * sizeof(T)); auto retval = - iterate_to_fixed_point(query, L, init_ids, expanded_nodes_info, + iterate_to_fixed_point(aligned_query, L, init_ids, expanded_nodes_info, expanded_nodes_ids, best_L_nodes); size_t pos = 0; for (auto it : best_L_nodes) { - indices[pos] = it.id; - pos++; + if (it.id < _max_points) { + indices[pos] = it.id; + if (distances != nullptr) + distances[pos] = it.distance; + pos++; + } if (pos == K) break; } + aligned_free(aligned_query); return retval; } @@ -1043,10 +1588,17 @@ namespace diskann { std::vector best_L_nodes, expanded_nodes_info; tsl::robin_set expanded_nodes_ids; + std::shared_lock lock(_update_lock); + if (init_ids.size() == 0) { init_ids.emplace_back(_ep); } - auto retval = iterate_to_fixed_point(query, (unsigned) L, init_ids, + T * aligned_query; + size_t allocSize = _aligned_dim * sizeof(T); + alloc_aligned(((void **) &aligned_query), allocSize, 8 * sizeof(T)); + memset(aligned_query, 0, _aligned_dim * sizeof(T)); + memcpy(aligned_query, query, _dim * sizeof(T)); + auto retval = iterate_to_fixed_point(aligned_query, (unsigned) L, init_ids, expanded_nodes_info, expanded_nodes_ids, best_L_nodes); @@ -1058,145 +1610,84 @@ namespace diskann { if (pos == K) break; } + aligned_free(aligned_query); return retval; } template - std::pair Index::search_with_tags( - const T *query, const size_t K, const unsigned L, TagT *tags, - unsigned frozen_pts, unsigned *indices_buffer) { - const bool alloc = indices_buffer == NULL; - auto indices = alloc ? new unsigned[K] : indices_buffer; - auto ret = search(query, K, L, indices); - for (int i = 0; i < (int) K; ++i) - tags[i] = _location_to_tag[indices[i]]; - if (alloc) - delete[] indices; - return ret; + size_t Index::search_with_tags(const T *query, const uint64_t K, + const unsigned L, TagT *tags, + float * distances, + std::vector &res_vectors) { + _u32 * indices = new unsigned[L]; + float *dist_interim = new float[L]; + search(query, L, L, indices, dist_interim); + + std::shared_lock ulock(_update_lock); + std::shared_lock lock(_tag_lock); + size_t pos = 0; + for (int i = 0; i < (int) L; ++i) + if (_location_to_tag.find(indices[i]) != _location_to_tag.end()) { + tags[pos] = _location_to_tag[indices[i]]; + res_vectors[i] = _data + indices[i] * _aligned_dim; + + if (distances != nullptr) + distances[pos] = dist_interim[i]; + pos++; + if (pos == K) + break; + } + delete[] indices; + delete[] dist_interim; + return pos; } template - void Index::optimize_graph() { // use after build or load - _data_len = (_aligned_dim + 1) * sizeof(float); - _neighbor_len = (_width + 1) * sizeof(unsigned); - _node_size = _data_len + _neighbor_len; - _opt_graph = (char *) malloc(_node_size * _nd); - DistanceFastL2 *dist_fast = (DistanceFastL2 *) _distance; - for (unsigned i = 0; i < _nd; i++) { - char *cur_node_offset = _opt_graph + i * _node_size; - float cur_norm = dist_fast->norm(_data + i * _aligned_dim, _aligned_dim); - std::memcpy(cur_node_offset, &cur_norm, sizeof(float)); - std::memcpy(cur_node_offset + sizeof(float), _data + i * _aligned_dim, - _data_len - sizeof(float)); - - cur_node_offset += _data_len; - unsigned k = _final_graph[i].size(); - std::memcpy(cur_node_offset, &k, sizeof(unsigned)); - std::memcpy(cur_node_offset + sizeof(unsigned), _final_graph[i].data(), - k * sizeof(unsigned)); - std::vector().swap(_final_graph[i]); - } - _final_graph.clear(); - _final_graph.shrink_to_fit(); + size_t Index::search_with_tags(const T *query, const size_t K, + const unsigned L, TagT *tags, + float *distances) { + _u32 * indices = new unsigned[L]; + float *dist_interim = new float[L]; + search(query, L, L, indices, dist_interim); + + std::shared_lock ulock(_update_lock); + std::shared_lock lock(_tag_lock); + size_t pos = 0; + for (int i = 0; i < (int) L; ++i) { + if (_location_to_tag.find(indices[i]) != _location_to_tag.end()) { + tags[pos] = _location_to_tag[indices[i]]; + if (distances != nullptr) + distances[pos] = dist_interim[i]; + pos++; + if (pos == K) + break; + } + } + delete[] indices; + delete[] dist_interim; + return pos; } template - void Index::search_with_opt_graph(const T *query, size_t K, size_t L, - unsigned *indices) { - DistanceFastL2 *dist_fast = (DistanceFastL2 *) _distance; - - std::vector retset(L + 1); - std::vector init_ids(L); - // std::mt19937 rng(rand()); - // GenRandom(rng, init_ids.data(), L, (unsigned) nd_); - - boost::dynamic_bitset<> flags{_nd, 0}; - unsigned tmp_l = 0; - unsigned * neighbors = - (unsigned *) (_opt_graph + _node_size * _ep + _data_len); - unsigned MaxM_ep = *neighbors; - neighbors++; - - for (; tmp_l < L && tmp_l < MaxM_ep; tmp_l++) { - init_ids[tmp_l] = neighbors[tmp_l]; - flags[init_ids[tmp_l]] = true; - } - - while (tmp_l < L) { - unsigned id = rand() % _nd; - if (flags[id]) - continue; - flags[id] = true; - init_ids[tmp_l] = id; - tmp_l++; - } - - for (unsigned i = 0; i < init_ids.size(); i++) { - unsigned id = init_ids[i]; - if (id >= _nd) - continue; - _mm_prefetch(_opt_graph + _node_size * id, _MM_HINT_T0); - } - L = 0; - for (unsigned i = 0; i < init_ids.size(); i++) { - unsigned id = init_ids[i]; - if (id >= _nd) - continue; - T * x = (T *) (_opt_graph + _node_size * id); - float norm_x = *x; - x++; - float dist = - dist_fast->compare(x, query, norm_x, (unsigned) _aligned_dim); - retset[i] = Neighbor(id, dist, true); - flags[id] = true; - L++; - } - // std::cout<compare(query, data, norm, (unsigned) _aligned_dim); - if (dist >= retset[L - 1].distance) - continue; - Neighbor nn(id, dist, true); - int r = InsertIntoPool(retset.data(), L, nn); + size_t Index::get_num_points() { + return _nd; + } - // if(L+1 < retset.size()) ++L; - if (r < nk) - nk = r; - } - } - if (nk <= k) - k = nk; - else - ++k; - } - for (size_t i = 0; i < K; i++) { - indices[i] = retset[i].id; - } + template + T *Index::get_data() { + if (_num_frozen_pts > 0) { + T * ret_data = nullptr; + size_t allocSize = _nd * _aligned_dim * sizeof(T); + alloc_aligned(((void **) &ret_data), allocSize, 8 * sizeof(T)); + memset(ret_data, 0, _nd * _aligned_dim * sizeof(T)); + memcpy(ret_data, _data, _nd * _aligned_dim * sizeof(T)); + return ret_data; + } + return _data; + } + template + size_t Index::return_max_points() { + return _max_points; } /************************************************* @@ -1206,149 +1697,177 @@ namespace diskann { // in case we add ''frozen'' auxiliary points to the dataset, these are not // visible to external world, we generate them here and update our dataset template - int Index::generate_random_frozen_points(const char *filename) { - if (_has_built) { - diskann::cout << "Index already built. Cannot add more points" - << std::endl; - return -1; - } + int Index::generate_frozen_point() { + if (_num_frozen_pts == 0) + return 0; - if (filename) { // user defined frozen points - T *frozen_pts; - load_aligned_bin(std::string(filename), frozen_pts, _num_frozen_pts, - _dim, _aligned_dim); - for (_u64 i = 0; i < _num_frozen_pts; i++) { - for (unsigned d = 0; d < _dim; d++) - _data[(i + _max_points) * _aligned_dim + d] = - frozen_pts[i * _dim + d]; - for (_u64 d = _dim; d < _aligned_dim; d++) - _data[(i + _max_points) * _aligned_dim + d] = 0; - } - } else { // random frozen points - - std::random_device device; - std::mt19937 generator(device()); - std::uniform_real_distribution dist(0, 1); - // Harsha: Should the distribution change with the distance metric? - - for (_u64 i = 0; i < _num_frozen_pts; ++i) { - for (_u64 d = 0; d < _dim; d++) - _data[(i + _max_points) * _aligned_dim + d] = (T) dist(generator); - for (_u64 d = _dim; d < _aligned_dim; d++) - _data[(i + _max_points) * _aligned_dim + d] = 0; - } + if (_nd == 0) { + memset(_data + (_max_points) *_aligned_dim, 0, _aligned_dim * sizeof(T)); + return 1; } - + size_t res = calculate_entry_point(); + memcpy(_data + _max_points * _aligned_dim, _data + res * _aligned_dim, + _aligned_dim * sizeof(T)); return 0; } template int Index::enable_delete() { - LockGuard guard(_change_lock); - assert(!_can_delete); assert(_enable_tags); - if (_can_delete) { - std::cerr << "Delete already enabled" << std::endl; - return -1; - } if (!_enable_tags) { std::cerr << "Tags must be instantiated for deletions" << std::endl; return -2; } - if (_consolidated_order && _compacted_order) { - assert(_empty_slots.size() == 0); - for (unsigned slot = (unsigned) _nd; slot < _max_points; ++slot) + if (_data_compacted) { + for (unsigned slot = (unsigned) _nd; slot < _max_points; ++slot) { _empty_slots.insert(slot); - _consolidated_order = false; - _compacted_order = false; + } } _lazy_done = false; _eager_done = false; - _can_delete = true; + if (_support_eager_delete) { + _in_graph.resize(_max_points + _num_frozen_pts); + _in_graph.reserve(_max_points + _num_frozen_pts); + update_in_graph(); + } return 0; } template - int Index::eager_delete(const TagT tag, - const Parameters ¶meters) { - if (_lazy_done && (!_consolidated_order)) { - diskann::cout << "Lazy delete reuests issued but data not consolidated, " + void Index::release_location() { + LockGuard guard(_change_lock); + _nd--; + } + + template + int Index::eager_delete(const TagT tag, const Parameters ¶meters, + int delete_mode) { + if (_lazy_done && (!_data_compacted)) { + diskann::cout << "Lazy delete requests issued but data not consolidated, " "cannot proceed with eager deletes." << std::endl; return -1; } - LockGuard guard(_change_lock); - if (_tag_to_location.find(tag) == _tag_to_location.end()) { - std::cerr << "Delete tag not found" << std::endl; - return -1; + + unsigned id = -1; + { + std::shared_lock lock(_tag_lock); + if (_tag_to_location.find(tag) == _tag_to_location.end()) { + std::cerr << "Delete tag not found" << std::endl; + return -1; + } + id = _tag_to_location[tag]; } - unsigned id = _tag_to_location[tag]; - _location_to_tag.erase(_tag_to_location[tag]); - _tag_to_location.erase(tag); - _delete_set.insert(id); - _empty_slots.insert(id); + { + std::unique_lock lock(_tag_lock); + _location_to_tag.erase(_tag_to_location[tag]); + _tag_to_location.erase(tag); + } + + { + // id will be valid because if not, it'll return in the {} above. + std::unique_lock lock(_delete_lock); + _delete_set.insert(id); + _empty_slots.insert(id); + } const unsigned range = parameters.Get("R"); const unsigned maxc = parameters.Get("C"); const float alpha = parameters.Get("alpha"); // delete point from out-neighbors' in-neighbor list - for (auto j : _final_graph[id]) - for (unsigned k = 0; k < _in_graph[j].size(); k++) - if (_in_graph[j][k] == id) { - _in_graph[j].erase(_in_graph[j].begin() + k); - break; + { + LockGuard guard(_locks[id]); + for (size_t i = 0; i < _final_graph[id].size(); i++) { + unsigned j = _final_graph[id][i]; + { + LockGuard guard(_locks_in[j]); + for (unsigned k = 0; k < _in_graph[j].size(); k++) { + if (_in_graph[j][k] == id) { + _in_graph[j].erase(_in_graph[j].begin() + k); + break; + } + } } + } + } tsl::robin_set in_nbr; - for (unsigned i = 0; i < _in_graph[id].size(); i++) - in_nbr.insert(_in_graph[id][i]); + { + LockGuard guard(_locks_in[id]); + for (unsigned i = 0; i < _in_graph[id].size(); i++) + in_nbr.insert(_in_graph[id][i]); + } assert(_in_graph[id].size() == in_nbr.size()); - tsl::robin_set candidate_set; - std::vector expanded_nghrs; - std::vector result; std::vector pool, tmp; tsl::robin_set visited; + std::vector intersection; + unsigned Lindex = parameters.Get("L"); + std::vector init_ids; + if (delete_mode == 2) { + // constructing list of in-neighbors to be processed + get_expanded_nodes(id, Lindex, init_ids, pool, visited); - unsigned Lindex = parameters.Get("L"); - std::vector init_ids; - - get_expanded_nodes(id, Lindex, init_ids, pool, visited); - - for (unsigned i = 0; i < pool.size(); i++) - if (pool[i].id == id) { - pool.erase(pool.begin() + i); - break; + for (auto node : visited) { + if (in_nbr.find(node) != in_nbr.end()) { + intersection.push_back(node); + } } + } + // deleting deleted point from all in-neighbors' out-neighbor list for (auto it : in_nbr) { + LockGuard guard(_locks[it]); _final_graph[it].erase( std::remove(_final_graph[it].begin(), _final_graph[it].end(), id), _final_graph[it].end()); } - for (auto it : visited) { - auto ngh = it; - if (in_nbr.find(ngh) != in_nbr.end()) { - candidate_set.clear(); - expanded_nghrs.clear(); - result.clear(); + tsl::robin_set candidate_set; + std::vector expanded_nghrs; + std::vector result; - for (auto j : _final_graph[id]) - if ((j != id) && (j != ngh) && - (_delete_set.find(j) == _delete_set.end())) - candidate_set.insert(j); + for (size_t i = 0; i < intersection.size(); i++) { + auto ngh = intersection[i]; - for (auto j : _final_graph[ngh]) - if ((j != id) && (j != ngh) && - (_delete_set.find(j) == _delete_set.end())) - candidate_set.insert(j); + candidate_set.clear(); + expanded_nghrs.clear(); + result.clear(); + + { + std::shared_lock lock(_delete_lock); + if (_delete_set.find(ngh) != _delete_set.end()) + continue; + } + + { + LockGuard guard(_locks[ngh]); + + // constructing candidate set from out-neighbors and out-neighbors of + // ngh and id + { // should a shared reader lock on delete_lock be held here at the + // beginning of the two for loops or should it be held and release + // for ech iteration of the for loops? Which is faster? + + std::shared_lock lock(_delete_lock); + for (auto j : _final_graph[id]) { + if ((j != id) && (j != ngh) && + (_delete_set.find(j) == _delete_set.end())) + candidate_set.insert(j); + } + + for (auto j : _final_graph[ngh]) { + if ((j != id) && (j != ngh) && + (_delete_set.find(j) == _delete_set.end())) + candidate_set.insert(j); + } + } for (auto j : candidate_set) expanded_nghrs.push_back( @@ -1360,78 +1879,79 @@ namespace diskann { std::sort(expanded_nghrs.begin(), expanded_nghrs.end()); occlude_list(expanded_nghrs, alpha, range, maxc, result); - for (auto iter : _final_graph[ngh]) - for (unsigned k = 0; k < _in_graph[iter].size(); k++) - if (_in_graph[iter][k] == ngh) { - _in_graph[iter].erase(_in_graph[iter].begin() + k); + // deleting ngh from its old out-neighbors' in-neighbor list + for (auto iter : _final_graph[ngh]) { + { + LockGuard guard(_locks_in[iter]); + for (unsigned k = 0; k < _in_graph[iter].size(); k++) { + if (_in_graph[iter][k] == ngh) { + _in_graph[iter].erase(_in_graph[iter].begin() + k); + break; + } } + } + } _final_graph[ngh].clear(); - for (auto j : result) { - if (_delete_set.find(j.id) == _delete_set.end()) - _final_graph[ngh].push_back(j.id); - if (std::find(_in_graph[j.id].begin(), _in_graph[j.id].end(), ngh) == - _in_graph[j.id].end()) - _in_graph[j.id].emplace_back(ngh); + // updating out-neighbors and in-neighbors of ngh + { + std::shared_lock lock(_delete_lock); + for (size_t i = 0; i < result.size(); i++) { + auto j = result[i]; + if (_delete_set.find(j.id) == _delete_set.end()) { + _final_graph[ngh].push_back(j.id); + { + LockGuard guard(_locks_in[j.id]); + if (std::find(_in_graph[j.id].begin(), _in_graph[j.id].end(), + ngh) == _in_graph[j.id].end()) { + _in_graph[j.id].emplace_back(ngh); + } + } + } + } } } } + _final_graph[id].clear(); - _nd--; + _in_graph[id].clear(); + + release_location(); _eager_done = true; + _data_compacted = false; return 0; } template void Index::update_in_graph() { - diskann::cout << "Updating in_graph....." << std::flush; + // diskann::cout << "Updating in_graph....."; for (unsigned i = 0; i < _in_graph.size(); i++) _in_graph[i].clear(); - for (unsigned i = 0; i < _final_graph.size(); + for (size_t i = 0; i < _final_graph.size(); i++) // copying to in-neighbor graph - - for (unsigned j = 0; j < _final_graph[i].size(); j++) { - if (std::find(_in_graph[_final_graph[i][j]].begin(), - _in_graph[_final_graph[i][j]].end(), - i) != _in_graph[_final_graph[i][j]].end()) - diskann::cout << "Duplicates found" << std::endl; - _in_graph[_final_graph[i][j]].emplace_back(i); - } - - size_t max_in, min_in, avg_in; - max_in = 0; - min_in = _max_points + 1; - avg_in = 0; - for (unsigned i = 0; i < _in_graph.size(); i++) { - avg_in += _in_graph[i].size(); - if (_in_graph[i].size() > max_in) - max_in = _in_graph[i].size(); - if ((_in_graph[i].size() < min_in) && (i != _ep)) - min_in = _in_graph[i].size(); - } - - diskann::cout << std::endl - << "Max in_degree = " << max_in - << "; Min in_degree = " << min_in << "; Average in_degree = " - << (float) (avg_in) / (float) (_nd + _num_frozen_pts) - << std::endl; + for (size_t j = 0; j < _final_graph[i].size(); j++) + _in_graph[_final_graph[i][j]].emplace_back((_u32) i); } // Do not call consolidate_deletes() if you have not locked _change_lock. // Returns number of live points left after consolidation + // proxy inserts all nghrs of deleted points + // original approach template size_t Index::consolidate_deletes(const Parameters ¶meters) { if (_eager_done) { - diskann::cout << "No consolidation required, eager deletes done" - << std::endl; + diskann::cout + << "In consolidate_deletes(), _eager_done is true. So exiting." + << std::endl; return 0; } - assert(!_consolidated_order); - assert(_can_delete); + diskann::cout << "Inside Index::consolidate_deletes()" << std::endl; + std::cout << "Empty slots size: " << _empty_slots.size() << " _nd: " << _nd + << " max_points: " << _max_points << std::endl; assert(_enable_tags); assert(_delete_set.size() <= _nd); assert(_empty_slots.size() + _nd == _max_points); @@ -1440,89 +1960,185 @@ namespace diskann { const unsigned maxc = parameters.Get("C"); const float alpha = parameters.Get("alpha"); - std::vector new_location; - new_location.resize(_max_points + _num_frozen_pts, - (unsigned) (_max_points + _num_frozen_pts)); - unsigned active = 0; - for (unsigned old = 0; old < _max_points + _num_frozen_pts; ++old) - if (_empty_slots.find(old) == _empty_slots.end() && - _delete_set.find(old) == _delete_set.end()) - new_location[old] = active++; - assert(active + _delete_set.size() == _max_points + _num_frozen_pts); + _u64 total_pts = _max_points + _num_frozen_pts; + unsigned block_size = 1 << 10; + _s64 total_blocks = DIV_ROUND_UP(total_pts, block_size); - tsl::robin_set candidate_set; - std::vector expanded_nghrs; - std::vector result; - - for (unsigned i = 0; i < _max_points + _num_frozen_pts; ++i) { - if (new_location[i] < _max_points + _num_frozen_pts) { - candidate_set.clear(); - expanded_nghrs.clear(); - result.clear(); - - bool modify = false; - for (auto ngh : _final_graph[i]) { - if (new_location[ngh] >= _max_points + _num_frozen_pts) { - modify = true; - - // Add outgoing links from - for (auto j : _final_graph[ngh]) - if (_delete_set.find(j) == _delete_set.end()) - candidate_set.insert(j); - } else { - candidate_set.insert(ngh); + auto start = std::chrono::high_resolution_clock::now(); +#pragma omp parallel for schedule(dynamic) + for (_s64 block = 0; block < total_blocks; ++block) { + tsl::robin_set candidate_set; + std::vector expanded_nghrs; + std::vector result; + + for (_s64 i = block * block_size; + i < (_s64)((block + 1) * block_size) && + i < (_s64)(_max_points + _num_frozen_pts); + i++) { + if ((_delete_set.find((_u32) i) == _delete_set.end()) && + (_empty_slots.find((_u32) i) == _empty_slots.end())) { + candidate_set.clear(); + expanded_nghrs.clear(); + result.clear(); + + bool modify = false; + for (auto ngh : _final_graph[(_u32) i]) { + if (_delete_set.find(ngh) != _delete_set.end()) { + modify = true; + + // Add outgoing links from + for (auto j : _final_graph[ngh]) + if (_delete_set.find(j) == _delete_set.end()) + candidate_set.insert(j); + } else { + candidate_set.insert(ngh); + } } - } + if (modify) { + for (auto j : candidate_set) { + expanded_nghrs.push_back( + Neighbor(j, + _distance->compare(_data + _aligned_dim * i, + _data + _aligned_dim * (size_t) j, + (unsigned) _aligned_dim), + true)); + } - if (modify) { - for (auto j : candidate_set) - expanded_nghrs.push_back( - Neighbor(j, - _distance->compare(_data + _aligned_dim * (size_t) i, - _data + _aligned_dim * (size_t) j, - (unsigned) _aligned_dim), - true)); - std::sort(expanded_nghrs.begin(), expanded_nghrs.end()); - occlude_list(expanded_nghrs, alpha, range, maxc, result); - - _final_graph[i].clear(); - for (auto j : result) { - if (j.id != i) - _final_graph[i].push_back(j.id); + std::sort(expanded_nghrs.begin(), expanded_nghrs.end()); + occlude_list(expanded_nghrs, alpha, range, maxc, result); + + _final_graph[(_u32) i].clear(); + for (auto j : result) { + if (j.id != (_u32) i && + (_delete_set.find(j.id) == _delete_set.end())) + _final_graph[(_u32) i].push_back(j.id); + } } } - } else - _final_graph[i].clear(); + } } if (_support_eager_delete) update_in_graph(); + for (auto iter : _delete_set) { + _empty_slots.insert(iter); + } _nd -= _delete_set.size(); - compact_data(new_location, active, _consolidated_order); + + _data_compacted = _delete_set.size() == 0; + + auto stop = std::chrono::high_resolution_clock::now(); + std::cout << "Time taken for consolidate_deletes() " + << std::chrono::duration_cast>( + stop - start) + .count() + << "s." << std::endl; + return _nd; } template - std::vector Index::get_new_location(unsigned &active) { - std::vector new_location; - new_location.resize(_max_points + _num_frozen_pts, - (unsigned) (_max_points + _num_frozen_pts)); - - for (unsigned old = 0; old < _max_points + _num_frozen_pts; ++old) - if (_empty_slots.find(old) == _empty_slots.end() && - _delete_set.find(old) == _delete_set.end()) - new_location[old] = active++; - assert(active + _delete_set.size() == _max_points + _num_frozen_pts); - - return new_location; + void Index::consolidate(Parameters ¶meters) { + consolidate_deletes(parameters); + compact_data(); } template - void Index::compact_data(std::vector new_location, - unsigned active, bool &mode) { + void Index::compact_frozen_point() { + if (_nd < _max_points) { + if (_num_frozen_pts > 0) { + // set new _ep to be frozen point + _ep = (_u32) _nd; + if (!_final_graph[_max_points].empty()) { + for (unsigned i = 0; i < _nd; i++) + for (unsigned j = 0; j < _final_graph[i].size(); j++) + if (_final_graph[i][j] == _max_points) + _final_graph[i][j] = (_u32) _nd; + + _final_graph[_nd].clear(); + for (unsigned k = 0; k < _final_graph[_max_points].size(); k++) + _final_graph[_nd].emplace_back(_final_graph[_max_points][k]); + + _final_graph[_max_points].clear(); + if (_support_eager_delete) + update_in_graph(); + + memcpy((void *) (_data + (size_t) _aligned_dim * _nd), + _data + (size_t) _aligned_dim * _max_points, sizeof(T) * _dim); + memset((_data + (size_t) _aligned_dim * _max_points), 0, + sizeof(T) * _aligned_dim); + } + } + } + } + + template + void Index::compact_data_for_search() { + compact_data(); + compact_frozen_point(); + } + + template + void Index::compact_data_for_insert() { + compact_data(); + + size_t max = 0, min = 1 << 30, total = 0, cnt = 0; + for (size_t i = 0; i < (_nd + _num_frozen_pts); i++) { + auto &pool = _final_graph[i]; + max = (std::max)(max, pool.size()); + min = (std::min)(min, pool.size()); + total += pool.size(); + if (pool.size() < 2) + cnt++; + } + if (min > max) + min = max; + if (_nd > 0) { + diskann::cout << "Index built with degree: max:" << max << " avg:" + << (float) total / (float) (_nd + _num_frozen_pts) + << " min:" << min << " count(deg<2):" << cnt << std::endl; + } + } + + template + void Index::compact_data() { + if (!_dynamic_index) + return; + + if (!_lazy_done && !_eager_done) + return; + + if (_data_compacted) { + diskann::cerr + << "Warning! Calling compact_data() when _data_compacted is true!" + << std::endl; + return; + } + + auto start = std::chrono::high_resolution_clock::now(); + auto fnstart = start; + + std::vector new_location = std::vector( + _max_points + _num_frozen_pts, (_u32) _max_points); + + _u32 new_counter = 0; + + for (_u32 old_counter = 0; old_counter < _max_points + _num_frozen_pts; + old_counter++) { + if (_location_to_tag.find(old_counter) != _location_to_tag.end()) { + new_location[old_counter] = new_counter; + new_counter++; + } + } + + auto stop = std::chrono::high_resolution_clock::now(); + std::cout << "Time taken for initial setup: " + << std::chrono::duration_cast>( + stop - start) + .count() + << "s." << std::endl; // If start node is removed, replace it. - assert(!mode); if (_delete_set.find(_ep) != _delete_set.end()) { std::cerr << "Replacing start node which has been deleted... " << std::flush; @@ -1541,20 +2157,41 @@ namespace diskann { __FUNCSIG__, __FILE__, __LINE__); } else { assert(_delete_set.find(_ep) == _delete_set.end()); - diskann::cout << "New start node is " << _ep << std::endl; } } - diskann::cout << "Re-numbering nodes and edges and consolidating data... " - << std::flush; - diskann::cout << "active = " << active << std::endl; - for (unsigned old = 0; old < _max_points + _num_frozen_pts; ++old) { - if (new_location[old] < - _max_points + _num_frozen_pts) { // If point continues to exist + start = std::chrono::high_resolution_clock::now(); + double copy_time = 0; + for (unsigned old = 0; old <= _max_points; ++old) { + if ((new_location[old] < _max_points) || + (old == _max_points)) { // If point continues to exist // Renumber nodes to compact the order for (size_t i = 0; i < _final_graph[old].size(); ++i) { - assert(new_location[_final_graph[old][i]] <= _final_graph[old][i]); + if (new_location[_final_graph[old][i]] > _final_graph[old][i]) { + std::stringstream sstream; + sstream << "Error in compact_data(). Found point: " << old + << " whose " << i << "th neighbor has new location " + << new_location[_final_graph[old][i]] + << " that is greater than its old location: " + << _final_graph[old][i]; + if (_delete_set.find(_final_graph[old][i]) != _delete_set.end()) { + sstream << std::endl + << " Point: " << old << " index: " << i + << " neighbor: " << _final_graph[old][i] + << " found in delete set of size: " << _delete_set.size() + << std::endl; + } else { + sstream << " Point: " << old + << " neighbor: " << _final_graph[old][i] + << " NOT found in delete set of size: " + << _delete_set.size() << std::endl; + } + + diskann::cerr << sstream.str() << std::endl; + throw diskann::ANNException(sstream.str(), -1, __FUNCSIG__, + __FILE__, __LINE__); + } _final_graph[old][i] = new_location[_final_graph[old][i]]; } @@ -1565,6 +2202,7 @@ namespace diskann { } // Move the data and adj list to the correct position + auto c_start = std::chrono::high_resolution_clock::now(); if (new_location[old] != old) { assert(new_location[old] < old); _final_graph[new_location[old]].swap(_final_graph[old]); @@ -1574,58 +2212,79 @@ namespace diskann { (void *) (_data + _aligned_dim * (size_t) old), _aligned_dim * sizeof(T)); } + auto c_stop = std::chrono::high_resolution_clock::now(); + copy_time += std::chrono::duration_cast>( + c_stop - c_start) + .count(); + + } else { + _final_graph[old].clear(); } } - diskann::cout << "done." << std::endl; + stop = std::chrono::high_resolution_clock::now(); + std::cout << "Time taken for moving data around: " + << std::chrono::duration_cast>( + stop - start) + .count() + << "s. Of which copy_time: " << copy_time << "s." << std::endl; - diskann::cout << "Updating mapping between tags and ids... " << std::flush; - // Update the location pointed to by tag + start = std::chrono::high_resolution_clock::now(); _tag_to_location.clear(); - for (auto iter : _location_to_tag) + for (auto iter : _location_to_tag) { _tag_to_location[iter.second] = new_location[iter.first]; + } _location_to_tag.clear(); - for (auto iter : _tag_to_location) + for (auto iter : _tag_to_location) { _location_to_tag[iter.second] = iter.first; - diskann::cout << "done." << std::endl; + } - for (unsigned old = active; old < _max_points + _num_frozen_pts; ++old) + for (_u64 old = _nd; old < _max_points; ++old) { _final_graph[old].clear(); + } _delete_set.clear(); _empty_slots.clear(); - mode = true; - diskann::cout << "Consolidated the index" << std::endl; - - /* for(unsigned i = 0; i < _nd + _num_frozen_pts; i++){ - int flag = 0; - for(unsigned j = 0; j < _final_graph[i].size(); j++) - if(_final_graph[i][j] == i){ - diskann::cout << "Self loop found just after compacting inside the - function" << std::endl; - flag = 1; - break; - } - if(flag == 1) - break; - } */ + for (_u32 i = _nd; i < _max_points; i++) { + _empty_slots.insert(i); + } + + _lazy_done = false; + _eager_done = false; + _data_compacted = true; + stop = std::chrono::high_resolution_clock::now(); + std::cout << "Time taken for tag<->index consolidation: " + << std::chrono::duration_cast>( + stop - start) + .count() + << "s." << std::endl; + std::cout << "Time taken for compact_data(): " + << std::chrono::duration_cast>( + stop - fnstart) + .count() + << "s." << std::endl; } // Do not call reserve_location() if you have not locked _change_lock. // It is not thread safe. template - unsigned Index::reserve_location() { - assert(_nd < _max_points); - + int Index::reserve_location() { + LockGuard guard(_change_lock); + if (_nd >= _max_points) { + return -1; + } unsigned location; - if (_consolidated_order || _compacted_order) + if (_data_compacted) { location = (unsigned) _nd; - else { + _empty_slots.erase(location); + } else { + // no need of delete_lock here, _change_lock will ensure no other thread + // executes this block of code assert(_empty_slots.size() != 0); assert(_empty_slots.size() + _nd == _max_points); auto iter = _empty_slots.begin(); location = *iter; _empty_slots.erase(iter); - _delete_set.erase(iter); + _delete_set.erase(location); } ++_nd; @@ -1633,75 +2292,179 @@ namespace diskann { } template - void Index::readjust_data(unsigned _num_frozen_pts) { - if (_num_frozen_pts > 0) { - if (_final_graph[_max_points].empty()) { - diskann::cout << "Readjusting data to correctly position frozen point" - << std::endl; - for (unsigned i = 0; i < _nd; i++) - for (unsigned j = 0; j < _final_graph[i].size(); j++) - if (_final_graph[i][j] >= _nd) - _final_graph[i][j] = - (unsigned) (_max_points + (_final_graph[i][j] - _nd)); - for (unsigned i = 0; i < _num_frozen_pts; i++) { - for (unsigned k = 0; k < _final_graph[_nd + i].size(); k++) - _final_graph[_max_points + i].emplace_back( - _final_graph[_nd + i][k]); - _final_graph[_nd + i].clear(); - } + void Index::reposition_point(unsigned old_location, + unsigned new_location) { + for (unsigned i = 0; i < _nd; i++) + for (unsigned j = 0; j < _final_graph[i].size(); j++) + if (_final_graph[i][j] == old_location) + _final_graph[i][j] = (unsigned) new_location; - if (_support_eager_delete) - update_in_graph(); + _final_graph[new_location].clear(); + for (unsigned k = 0; k < _final_graph[_nd].size(); k++) + _final_graph[new_location].emplace_back(_final_graph[old_location][k]); - diskann::cout << "Finished updating graph, updating data now" - << std::endl; - for (unsigned i = 0; i < _num_frozen_pts; i++) { - memcpy((void *) (_data + (size_t) _aligned_dim * (_max_points + i)), - _data + (size_t) _aligned_dim * (_nd + i), - sizeof(float) * _dim); - memset((_data + (size_t) _aligned_dim * (_nd + i)), 0, - sizeof(float) * _aligned_dim); - } - diskann::cout << "Readjustment done" << std::endl; - } - } else - diskann::cout << "No frozen points. No re-adjustment required" - << std::endl; + _final_graph[old_location].clear(); + + if (_support_eager_delete) { + update_in_graph(); + } + memcpy((void *) (_data + (size_t) _aligned_dim * new_location), + _data + (size_t) _aligned_dim * old_location, + sizeof(T) * _aligned_dim); + memset((_data + (size_t) _aligned_dim * old_location), 0, + sizeof(T) * _aligned_dim); + } + + template + void Index::reposition_frozen_point_to_end() { + if (_num_frozen_pts == 0) + return; + + if (_nd == _max_points) { + diskann::cout + << "Not repositioning frozen point as it is already at the end." + << std::endl; + return; + } + reposition_point(_nd, _max_points); + _ep = (_u32) _max_points; + } + + template + void Index::resize(uint32_t new_max_points) { + // TODO: Check if the _change_lock and _update_lock are both locked. + + auto start = std::chrono::high_resolution_clock::now(); + assert(_empty_slots.size() == + 0); // should not resize if there are empty slots. +#ifndef _WINDOWS + T *new_data; + /* alloc_aligned((void **) &new_data, + (new_max_points + 1) * _aligned_dim * sizeof(T), + 8 * sizeof(T)); + memcpy(new_data, _data, (_max_points + 1) * _aligned_dim * sizeof(T)); + aligned_free(_data); + _data = new_data; + */ + realloc_aligned((void **) &_data, (void **) &new_data, + (_max_points + 1) * _aligned_dim * sizeof(T), + (new_max_points + 1) * _aligned_dim * sizeof(T), + 8 * sizeof(T)); +#else + realloc_aligned((void **) &_data, + (new_max_points + 1) * _aligned_dim * sizeof(T), + 8 * sizeof(T)); +#endif + _final_graph.resize(new_max_points + 1); + _locks = std::vector(new_max_points + 1); + if (_support_eager_delete) { + _in_graph.resize(new_max_points + 1); + _locks_in = std::vector(new_max_points + 1); + } + + reposition_point(_max_points, new_max_points); + _max_points = new_max_points; + _ep = new_max_points; + + for (_u32 i = _nd; i < _max_points; i++) { + _empty_slots.insert(i); + } + + auto stop = std::chrono::high_resolution_clock::now(); + diskann::cout << "Resizing took: " + << std::chrono::duration(stop - start).count() << "s" + << std::endl; + } + + template + void Index::printTagToLocation() { + std::shared_lock lock(_tag_lock); + + std::cout << "Thread: " << std::this_thread::get_id() + << " _tag_to_location: " << std::endl; + for (auto tl : _tag_to_location) { + std::cout << "(" << tl.first << "," << tl.second << "),"; + } + std::cout << std::endl + << "Thread: " << std::this_thread::get_id() + << " _location_to_tag: " << std::endl; + for (auto lt : _location_to_tag) { + std::cout << "(" << lt.first << "," << lt.second << "),"; + } + std::cout << std::endl; } template int Index::insert_point(const T *point, const Parameters ¶meters, - std::vector & pool, - std::vector & tmp, - tsl::robin_set &visited, - vecNgh &cut_graph, const TagT tag) { + const TagT tag) { + std::shared_lock lock(_update_lock); unsigned range = parameters.Get("R"); - assert(_has_built); + // assert(_has_built); + std::vector pool; + std::vector tmp; + tsl::robin_set visited; - LockGuard guard(_change_lock); - if (_enable_tags && - (_tag_to_location.find(tag) != _tag_to_location.end())) { - std::cerr << "Entry with the tag " << tag << " exists already" - << std::endl; - return -1; - } - if (_nd == _max_points) { - std::cerr << "Can not insert, reached maximum(" << _max_points - << ") points." << std::endl; - return -2; + { + std::shared_lock lock(_tag_lock); + std::shared_lock tsl(_tag_lock); + if (_enable_tags && + (_tag_to_location.find(tag) != _tag_to_location.end())) { + // TODO! This is a repeat of lazy_delete, but we can't call + // that function because we are taking many locks here. Hence + // the repeated code. + tsl.unlock(); + std::unique_lock tul(_tag_lock); + std::unique_lock tdl(_delete_lock); + _lazy_done = true; + _delete_set.insert(_tag_to_location[tag]); + _location_to_tag.erase(_tag_to_location[tag]); + _tag_to_location.erase(tag); + } } auto location = reserve_location(); - _tag_to_location[tag] = location; - _location_to_tag[location] = tag; + if (location == -1) { + std::cout << "Thread: " << std::this_thread::get_id() + << " location == -1. Waiting for unique_lock. " << std::endl + << std::flush; + lock.unlock(); + std::unique_lock growth_lock(_update_lock); + + std::cout << "Thread: " << std::this_thread::get_id() + << " Obtained unique_lock. " << std::endl; + if (_nd >= _max_points) { + auto new_max_points = (size_t)(_max_points * INDEX_GROWTH_FACTOR); + diskann::cerr << "Thread: " << std::this_thread::get_id() + << ": Increasing _max_points from " << _max_points + << " to " << new_max_points << " _nd is: " << _nd + << std::endl; + resize(new_max_points); + } + growth_lock.unlock(); + lock.lock(); + location = reserve_location(); + // TODO: Consider making this a while/do_while loop so that we retry + // instead of terminating. + if (location == -1) { + throw diskann::ANNException( + "Cannot reserve location even after expanding graph. Terminating.", + -1, __FUNCSIG__, __FILE__, __LINE__); + } + } + + { + std::unique_lock lock(_tag_lock); + + _tag_to_location[tag] = location; + _location_to_tag[location] = tag; + } auto offset_data = _data + (size_t) _aligned_dim * location; - memset((void *) offset_data, 0, sizeof(float) * _aligned_dim); - memcpy((void *) offset_data, point, sizeof(float) * _dim); + memset((void *) offset_data, 0, sizeof(T) * _aligned_dim); + memcpy((void *) offset_data, point, sizeof(T) * _dim); pool.clear(); tmp.clear(); - cut_graph.clear(); visited.clear(); std::vector pruned_list; unsigned Lindex = parameters.Get("L"); @@ -1710,32 +2473,45 @@ namespace diskann { get_expanded_nodes(location, Lindex, init_ids, pool, visited); for (unsigned i = 0; i < pool.size(); i++) - if (pool[i].id == location) { + if (pool[i].id == (unsigned) location) { pool.erase(pool.begin() + i); - visited.erase(location); + visited.erase((unsigned) location); break; } prune_neighbors(location, pool, parameters, pruned_list); - assert(_final_graph.size() == _max_points + _num_frozen_pts); - for (unsigned i = 0; i < _final_graph[location].size(); i++) - _in_graph[_final_graph[location][i]].erase( - std::remove(_in_graph[_final_graph[location][i]].begin(), - _in_graph[_final_graph[location][i]].end(), location), - _in_graph[_final_graph[location][i]].end()); + if (_support_eager_delete) { + for (unsigned i = 0; i < _final_graph[location].size(); i++) { + { + LockGuard guard(_locks_in[_final_graph[location][i]]); + _in_graph[_final_graph[location][i]].erase( + std::remove(_in_graph[_final_graph[location][i]].begin(), + _in_graph[_final_graph[location][i]].end(), location), + _in_graph[_final_graph[location][i]].end()); + } + } + } _final_graph[location].clear(); - _final_graph[location].reserve(range); + _final_graph[location].shrink_to_fit(); + _final_graph[location].reserve((_u64)(range * SLACK_FACTOR * 1.05)); + + if (pruned_list.empty()) { + std::cout << "Thread: " << std::this_thread::get_id() << "Tag id: " << tag + << " pruned_list.size(): " << pruned_list.size() << std::endl; + } assert(!pruned_list.empty()); - for (auto link : pruned_list) { - _final_graph[location].emplace_back(link); - if (_support_eager_delete) - if (std::find(_in_graph[link].begin(), _in_graph[link].end(), - location) == _in_graph[link].end()) { + { + LockGuard guard(_locks[location]); + for (auto link : pruned_list) { + _final_graph[location].emplace_back(link); + if (_support_eager_delete) { + LockGuard guard(_locks_in[link]); _in_graph[link].emplace_back(location); } + } } assert(_final_graph[location].size() <= range); @@ -1743,77 +2519,207 @@ namespace diskann { inter_insert(location, pruned_list, parameters, 1); else inter_insert(location, pruned_list, parameters, 0); - return 0; } template - int Index::disable_delete(const Parameters ¶meters, - const bool consolidate) { - LockGuard guard(_change_lock); - if (!_can_delete) { - diskann::cerr << "Delete not currently enabled" << std::endl; - return -1; - } - if (!_enable_tags) { - diskann::cerr << "Point tag array not instantiated" << std::endl; - throw diskann::ANNException("Point tag array not instantiated", -1, - __FUNCSIG__, __FILE__, __LINE__); - } - if (_eager_done) { - diskann::cout << "#Points after eager_delete : " << _nd + _num_frozen_pts - << std::endl; - if (_tag_to_location.size() != _nd) { - diskann::cerr << "Tags to points array wrong sized" << std::endl; - return -2; - } - } else if (_tag_to_location.size() + _delete_set.size() != _nd) { - diskann::cerr << "Tags to points array wrong sized" << std::endl; + int Index::lazy_delete(const TagT &tag) { + if ((_eager_done) && (!_data_compacted)) { + std::cerr << "Eager delete requests were issued but data was not " + "compacted, cannot proceed with lazy_deletes" + << std::endl; return -2; } - if (_eager_done) { - if (_location_to_tag.size() != _nd) { - diskann::cerr << "Points to tags array wrong sized" << std::endl; - return -3; + std::shared_lock lock(_update_lock); + _lazy_done = true; + + { + std::shared_lock l(_tag_lock); + + if (_tag_to_location.find(tag) == _tag_to_location.end()) { + // diskann::cerr << "Delete tag not found" << std::endl; + return -1; } - } else if (_location_to_tag.size() + _delete_set.size() != _nd) { - diskann::cerr << "Points to tags array wrong sized" << std::endl; - return -3; + assert(_tag_to_location[tag] < _max_points); } - if (consolidate) { - auto nd = consolidate_deletes(parameters); - if (nd >= 0) - diskann::cout << "#Points after consolidation: " << nd + _num_frozen_pts - << std::endl; + { + std::unique_lock l(_delete_lock); + std::shared_lock tl(_tag_lock); + _delete_set.insert(_tag_to_location[tag]); + } + + { + std::unique_lock l(_tag_lock); + _location_to_tag.erase(_tag_to_location[tag]); + _tag_to_location.erase(tag); } - _can_delete = false; return 0; } + // TODO: Check if this function needs a shared_lock on _tag_lock. template - int Index::delete_point(const TagT tag) { - if ((_eager_done) && (!_compacted_order)) { + int Index::lazy_delete(const tsl::robin_set &tags, + std::vector & failed_tags) { + if (failed_tags.size() > 0) { + std::cerr << "failed_tags should be passed as an empty list" << std::endl; + return -3; + } + if ((_eager_done) && (!_data_compacted)) { diskann::cout << "Eager delete requests were issued but data was not " "compacted, cannot proceed with lazy_deletes" << std::endl; - return -1; + return -2; } - LockGuard guard(_change_lock); - if (_tag_to_location.find(tag) == _tag_to_location.end()) { - diskann::cerr << "Delete tag not found" << std::endl; + std::shared_lock lock(_update_lock); + _lazy_done = true; + + for (auto tag : tags) { + // assert(_tag_to_location[tag] < _max_points); + if (_tag_to_location.find(tag) == _tag_to_location.end()) { + failed_tags.push_back(tag); + } else { + _delete_set.insert(_tag_to_location[tag]); + _location_to_tag.erase(_tag_to_location[tag]); + _tag_to_location.erase(tag); + } + } + + return 0; + } + + template + int Index::extract_data( + T *ret_data, std::unordered_map &tag_to_location) { + if (!_data_compacted) { + std::cerr + << "Error! Data not compacted. Cannot give access to private data." + << std::endl; return -1; } - assert(_tag_to_location[tag] < _max_points); - _delete_set.insert(_tag_to_location[tag]); - _location_to_tag.erase(_tag_to_location[tag]); - _tag_to_location.erase(tag); + std::memset(ret_data, 0, (size_t) _aligned_dim * _nd * sizeof(T)); + std::memcpy(ret_data, _data, (size_t)(_aligned_dim) *_nd * sizeof(T)); + tag_to_location = _tag_to_location; return 0; } + template + void Index::get_location_to_tag( + std::unordered_map &ret_loc_to_tag) { + ret_loc_to_tag = _location_to_tag; + } + + template + bool Index::hasIndexBeenSaved() { + return _is_saved; + } + + template + void Index::get_active_tags(tsl::robin_set &active_tags) { + active_tags.clear(); + for (auto iter : _tag_to_location) { + active_tags.insert(iter.first); + } + } + + template + void Index::print_delete_set() const { + diskann::cout << "Delete set is of size: " << _delete_set.size() + << std::endl; + std::vector sorted_delete_set; + for (auto i : _delete_set) { + sorted_delete_set.push_back(i); + } + std::sort(sorted_delete_set.begin(), sorted_delete_set.end()); + diskann::cout << "Sorted Delete set is of size: " << _delete_set.size() + << std::endl; + + // TODO: Debugging ONLY + size_t counter = 60000; + std::vector<_u32> missing_ids; + for (auto i : sorted_delete_set) { + if (i != counter) { + missing_ids.push_back(i); + } + counter++; + } + diskann::cout << "Missing ids in delete set for the 60k-80k case:" + << std::endl; + for (auto i : missing_ids) { + diskann::cout << i << ","; + } + diskann::cout << std::endl; + } + template + void Index::are_deleted_points_in_graph() const { + std::vector> start_end_pairs; + for (size_t i = 0; i < _nd; i++) { + for (size_t j = 0; j < _final_graph[i].size(); j++) { + if (_delete_set.find(_final_graph[i][j]) != _delete_set.end()) { + start_end_pairs.push_back( + std::pair<_u32, _u32>(i, _final_graph[i][j])); + } + } + } + + if (start_end_pairs.size() > 0) { + diskann::cout << "Found " << start_end_pairs.size() + << " references to deleted vertices" << std::endl; + std::sort(start_end_pairs.begin(), start_end_pairs.end(), + [](const std::pair<_u32, _u32> &val1, + const std::pair<_u32, _u32> &val2) { + return val1.first < val2.first; + }); + diskann::cout << "Min source id: " << start_end_pairs[0].first + << " Max source id: " + << start_end_pairs[start_end_pairs.size() - 1].first + << std::endl; + std::sort(start_end_pairs.begin(), start_end_pairs.end(), + [](const std::pair<_u32, _u32> &val1, + const std::pair<_u32, _u32> &val2) { + return val1.second < val2.second; + }); + diskann::cout << "Min target id: " << start_end_pairs[0].second + << " Max target id: " + << start_end_pairs[start_end_pairs.size() - 1].second + << std::endl; + } + } + + template + void Index::print_status() const { + diskann::cout << "------------------- Index object: " << (uint64_t) this + << " -------------------" << std::endl; + diskann::cout << "Number of points: " << _nd << std::endl; + diskann::cout << "Graph size: " << _final_graph.size() << std::endl; + diskann::cout << "Location to tag size: " << _location_to_tag.size() + << std::endl; + diskann::cout << "Tag to location size: " << _tag_to_location.size() + << std::endl; + diskann::cout << "Number of empty slots: " << _empty_slots.size() + << std::endl; + diskann::cout << std::boolalpha + << "Data compacted: " << this->_data_compacted + << " Lazy done: " << this->_lazy_done + << " Eager done: " << this->_eager_done << std::endl; + diskann::cout << "---------------------------------------------------------" + "------------" + << std::endl; + } + + /* Internals of the library */ // EXPORTS - template DISKANN_DLLEXPORT class Index; - template DISKANN_DLLEXPORT class Index; - template DISKANN_DLLEXPORT class Index; + template DISKANN_DLLEXPORT class Index; + template DISKANN_DLLEXPORT class Index; + template DISKANN_DLLEXPORT class Index; + template DISKANN_DLLEXPORT class Index; + template DISKANN_DLLEXPORT class Index; + template DISKANN_DLLEXPORT class Index; + template DISKANN_DLLEXPORT class Index; + template DISKANN_DLLEXPORT class Index; + template DISKANN_DLLEXPORT class Index; + template DISKANN_DLLEXPORT class Index; + template DISKANN_DLLEXPORT class Index; + template DISKANN_DLLEXPORT class Index; } // namespace diskann diff --git a/src/linux_aligned_file_reader.cpp b/src/linux_aligned_file_reader.cpp index 69eae9b60e..c09c15dc6a 100644 --- a/src/linux_aligned_file_reader.cpp +++ b/src/linux_aligned_file_reader.cpp @@ -1,8 +1,7 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. #include "linux_aligned_file_reader.h" +#include #include #include #include @@ -19,7 +18,7 @@ namespace { #ifdef DEBUG for (auto &req : read_reqs) { assert(IS_ALIGNED(req.len, 512)); - // std::cout << "request:"<= req.len); @@ -56,8 +55,8 @@ namespace { if (ret != (int64_t) n_ops) { std::cerr << "io_submit() failed; returned " << ret << ", expected=" << n_ops << ", ernno=" << errno << "=" - << ::strerror(-ret) << ", try #" << n_tries + 1; - std::cout << "ctx: " << ctx << "\n"; + << ::strerror((int) -ret) << ", try #" << n_tries + 1; + diskann::cout << "ctx: " << ctx << "\n"; exit(-1); } else { // wait on io_getevents @@ -67,29 +66,16 @@ namespace { if (ret != (int64_t) n_ops) { std::cerr << "io_getevents() failed; returned " << ret << ", expected=" << n_ops << ", ernno=" << errno << "=" - << ::strerror(-ret) << ", try #" << n_tries + 1; + << ::strerror((int) -ret) << ", try #" << n_tries + 1; exit(-1); } else { break; } } } - // disabled since req.buf could be an offset into another buf - /* - for (auto &req : read_reqs) { - // corruption check - assert(malloc_usable_size(req.buf) >= req.len); - } - */ - } - - /* - for(unsigned i=0;i<64;i++){ - std::cout << *((unsigned*)read_reqs[0].buf + i) << " "; } - std::cout << std::endl;*/ } -} +} // namespace LinuxAlignedFileReader::LinuxAlignedFileReader() { this->file_desc = -1; @@ -128,8 +114,6 @@ void LinuxAlignedFileReader::register_thread() { auto my_id = std::this_thread::get_id(); std::unique_lock lk(ctx_mut); if (ctx_map.find(my_id) != ctx_map.end()) { - std::cerr << "multiple calls to register_thread from the same thread" - << std::endl; return; } io_context_t ctx = 0; @@ -141,8 +125,6 @@ void LinuxAlignedFileReader::register_thread() { std::cerr << "io_setup() failed; returned " << ret << ", errno=" << errno << ":" << ::strerror(errno) << std::endl; } else { - std::cerr << "allocating ctx: " << ctx << " to thread-id:" << my_id - << std::endl; ctx_map[my_id] = ctx; } lk.unlock(); @@ -159,16 +141,40 @@ void LinuxAlignedFileReader::deregister_thread() { // assert(ret == 0); lk.lock(); ctx_map.erase(my_id); - std::cerr << "returned ctx from thread-id:" << my_id << std::endl; + // std::cerr << "returned ctx from thread-id:" << my_id << std::endl; + lk.unlock(); +} + +void LinuxAlignedFileReader::deregister_all_threads() { + std::unique_lock lk(ctx_mut); + + for (auto &iter : ctx_map) { + io_context_t ctx = iter.second; + io_destroy(ctx); + // assert(ret == 0); + // std::cerr << "returned ctx from thread-id:" << iter.first << + // std::endl; + } + ctx_map.clear(); lk.unlock(); } -void LinuxAlignedFileReader::open(const std::string &fname) { - int flags = O_DIRECT | O_RDONLY | O_LARGEFILE; +void LinuxAlignedFileReader::open(const std::string &fname, + bool enable_writes = false, + bool enable_create = false) { + int flags = O_DIRECT | O_LARGEFILE; + if (!enable_writes) { + flags |= O_RDONLY; + } else { + flags |= O_RDWR; + } + if (enable_create) { + flags |= O_CREAT; + } this->file_desc = ::open(fname.c_str(), flags); // error checks assert(this->file_desc != -1); - std::cerr << "Opened file : " << fname << std::endl; + // std::cerr << "Opened file : " << fname << std::endl; } void LinuxAlignedFileReader::close() { @@ -185,9 +191,46 @@ void LinuxAlignedFileReader::close() { void LinuxAlignedFileReader::read(std::vector &read_reqs, io_context_t &ctx, bool async) { assert(this->file_desc != -1); - //#pragma omp critical - // std::cout << "thread: " << std::this_thread::get_id() << ", crtx: " << - // ctx - //<< "\n"; execute_io(ctx, this->file_desc, read_reqs); + if (async == true) { + std::cerr << "async only supported in Windows for now." << std::endl; + } +} + +void LinuxAlignedFileReader::sequential_write(AlignedRead &write_req, + IOContext & ctx) { + assert(this->file_desc != -1); + // check inputs + assert(IS_ALIGNED(write_req.offset, 4096)); + assert(IS_ALIGNED(write_req.buf, 4096)); + assert(IS_ALIGNED(write_req.len, 4096)); + + // create write request + io_event_t evt; + struct iocb cb; + iocb_t * cbs = &cb; + io_prep_pwrite(&cb, this->file_desc, write_req.buf, write_req.len, + write_req.offset); + + uint64_t n_tries = 0; + // issue reads + int64_t ret = io_submit(ctx, (int64_t) 1, &cbs); + // if requests didn't get accepted + if (ret != (int64_t) 1) { + std::cerr << "io_submit() failed; returned " << ret << ", expected=" << 1 + << ", ernno=" << errno << "=" << ::strerror((int) -ret) + << ", try #" << n_tries + 1; + diskann::cout << "ctx: " << ctx << "\n"; + exit(-1); + } else { + // wait on io_getevents + ret = io_getevents(ctx, (int64_t) 1, (int64_t) 1, &evt, nullptr); + // if requests didn't complete + if (ret != (int64_t) 1) { + std::cerr << "io_getevents() failed; returned " << ret + << ", expected=" << 1 << ", ernno=" << errno << "=" + << ::strerror((int) -ret) << ", try #" << n_tries + 1; + exit(-1); + } + } } diff --git a/src/math_utils.cpp b/src/math_utils.cpp index de63d93241..bdb3b51b5c 100644 --- a/src/math_utils.cpp +++ b/src/math_utils.cpp @@ -152,7 +152,10 @@ namespace math_utils { if (!is_norm_given_for_pts) pts_norms_squared = new float[num_points]; - size_t PAR_BLOCK_SIZE = num_points; + size_t PAR_BLOCK_SIZE = std::min((size_t) 1 << 23, num_points); + // (num_points > 1 << 20) ? 1 << 13 : (num_points / 16); + + // size_t PAR_BLOCK_SIZE = num_points; size_t N_BLOCKS = (num_points % PAR_BLOCK_SIZE) == 0 ? (num_points / PAR_BLOCK_SIZE) : (num_points / PAR_BLOCK_SIZE) + 1; @@ -204,9 +207,6 @@ namespace math_utils { void process_residuals(float* data_load, size_t num_points, size_t dim, float* cur_pivot_data, size_t num_centers, uint32_t* closest_centers, bool to_subtract) { - diskann::cout << "Processing residuals of " << num_points << " points in " - << dim << " dimensions using " << num_centers << " centers " - << std::endl; #pragma omp parallel for schedule(static, 8192) for (int64_t n_iter = 0; n_iter < (_s64) num_points; n_iter++) { for (size_t d_iter = 0; d_iter < dim; d_iter++) { @@ -334,13 +334,15 @@ namespace kmeans { residual = lloyds_iter(data, num_points, dim, centers, num_centers, docs_l2sq, closest_docs, closest_center); - diskann::cout << "Lloyd's iter " << i - << " dist_sq residual: " << residual << std::endl; + // diskann::cout << "Lloyd's iter " << i + // << " dist_sq residual: " << residual << std::endl; - if (((i != 0) && ((old_residual - residual) / residual) < 0.00001) || + if (((i != 0) && + ((double) ((old_residual - residual) / residual)) < 0.00001) || (residual < std::numeric_limits::epsilon())) { - diskann::cout << "Residuals unchanged: " << old_residual << " becomes " - << residual << ". Early termination." << std::endl; + // diskann::cout << "Residuals unchanged: " << old_residual << " becomes + // " + // << residual << ". Early termination." << std::endl; break; } } @@ -359,12 +361,9 @@ namespace kmeans { float* pivot_data, size_t num_centers) { // pivot_data = new float[num_centers * dim]; - std::vector picked; - diskann::cout << "Selecting " << num_centers << " pivots from " - << num_points << " points using "; - std::random_device rd; - auto x = rd(); - diskann::cout << "random seed " << x << std::endl; + std::vector picked; + std::random_device rd; + auto x = rd(); std::mt19937 generator(x); std::uniform_int_distribution distribution(0, num_points - 1); @@ -391,12 +390,9 @@ namespace kmeans { return; } - std::vector picked; - diskann::cout << "Selecting " << num_centers << " pivots from " - << num_points << " points using "; - std::random_device rd; - auto x = rd(); - diskann::cout << "random seed " << x << ": " << std::flush; + std::vector picked; + std::random_device rd; + auto x = rd(); std::mt19937 generator(x); std::uniform_real_distribution<> distribution(0, 1); std::uniform_int_distribution int_dist(0, num_points - 1); @@ -423,7 +419,7 @@ namespace kmeans { double sum = 0; for (size_t i = 0; i < num_points; i++) { - sum = sum + dist[i]; + sum = sum + (double) dist[i]; } if (sum == 0) sum_flag = true; @@ -433,11 +429,12 @@ namespace kmeans { double prefix_sum = 0; for (size_t i = 0; i < (num_points); i++) { tmp_pivot = i; - if (dart_val >= prefix_sum && dart_val < prefix_sum + dist[i]) { + if (dart_val >= prefix_sum && + dart_val < prefix_sum + (double) dist[i]) { break; } - prefix_sum += dist[i]; + prefix_sum += (double) dist[i]; } if (std::find(picked.begin(), picked.end(), tmp_pivot) != picked.end() && @@ -454,10 +451,10 @@ namespace kmeans { data + tmp_pivot * dim, dim)); } num_picked++; - if (num_picked % 32 == 0) - diskann::cout << "." << std::flush; + // if (num_picked % 32 == 0) + // diskann::cout << "." << std::flush; } - diskann::cout << "done." << std::endl; + // diskann::cout << "done." << std::endl; delete[] dist; } diff --git a/src/partition_and_pq.cpp b/src/partition_and_pq.cpp index 9da49b2203..1918996be0 100644 --- a/src/partition_and_pq.cpp +++ b/src/partition_and_pq.cpp @@ -15,11 +15,14 @@ #include #include #include + #include "logger.h" +#include "cached_io.h" #include "exceptions.h" #include "index.h" #include "parameters.h" #include "tsl/robin_set.h" +#include "tcmalloc/malloc_extension.h" #include "utils.h" #include @@ -35,13 +38,16 @@ #include #endif -#define BLOCK_SIZE 5000000 +// For fresh index, keep it at 1m instead of 5m. +#define MAX_BLOCK_SIZE 1000000 template void gen_random_slice(const std::string base_file, - const std::string output_prefix, double sampling_rate) { + const std::string output_prefix, double sampling_rate, + size_t offset) { _u64 read_blk_size = 64 * 1024 * 1024; - cached_ifstream base_reader(base_file.c_str(), read_blk_size); + cached_ifstream base_reader(base_file.c_str(), read_blk_size, + (uint32_t) offset); std::ofstream sample_writer(std::string(output_prefix + "_data.bin").c_str(), std::ios::binary); std::ofstream sample_id_writer( @@ -75,13 +81,19 @@ void gen_random_slice(const std::string base_file, for (size_t i = 0; i < npts; i++) { base_reader.read((char *) cur_row.get(), sizeof(T) * nd); float sample = distribution(generator); - if (sample < sampling_rate) { + if (sample < (float) sampling_rate) { sample_writer.write((char *) cur_row.get(), sizeof(T) * nd); uint32_t cur_i_u32 = (_u32) i; sample_id_writer.write((char *) &cur_i_u32, sizeof(uint32_t)); num_sampled_pts_u32++; } } + + if (num_sampled_pts_u32 == 0) { + // We have read something from file, so write it. + sample_writer.write((char *) cur_row.get(), sizeof(T) * nd); + num_sampled_pts_u32 = 1; + } sample_writer.seekp(0, std::ios::beg); sample_writer.write((char *) &num_sampled_pts_u32, sizeof(uint32_t)); sample_id_writer.seekp(0, std::ios::beg); @@ -101,6 +113,14 @@ void gen_random_slice(const std::string base_file, * Reimplement using gen_random_slice(const T* inputdata,...) ************************************/ +template +void gen_random_slice(const std::string data_file, double p_val, + std::unique_ptr &sampled_data, + size_t &slice_size, size_t &ndims) { + float *sampled_ptr = sampled_data.get(); + gen_random_slice(data_file, p_val, sampled_ptr, slice_size, ndims); + sampled_data.reset(sampled_ptr); +} template void gen_random_slice(const std::string data_file, double p_val, float *&sampled_data, size_t &slice_size, size_t &ndims) { @@ -130,7 +150,7 @@ void gen_random_slice(const std::string data_file, double p_val, for (size_t i = 0; i < npts; i++) { base_reader.read((char *) cur_vector_T.get(), ndims * sizeof(T)); float rnd_val = distribution(generator); - if (rnd_val < p_val) { + if (rnd_val < (float) p_val) { std::vector cur_vector_float; for (size_t d = 0; d < ndims; d++) cur_vector_float.push_back(cur_vector_T[d]); @@ -138,7 +158,14 @@ void gen_random_slice(const std::string data_file, double p_val, } } slice_size = sampled_vectors.size(); + if (slice_size == 0) { + slice_size = 1; + std::vector cur_vector_float(cur_vector_T.get(), + cur_vector_T.get() + ndims); + sampled_vectors.push_back(cur_vector_float); + } sampled_data = new float[slice_size * ndims]; + for (size_t i = 0; i < slice_size; i++) { for (size_t j = 0; j < ndims; j++) { sampled_data[i * ndims + j] = sampled_vectors[i][j]; @@ -166,15 +193,23 @@ void gen_random_slice(const T *inputdata, size_t npts, size_t ndims, for (size_t i = 0; i < npts; i++) { cur_vector_T = inputdata + ndims * i; float rnd_val = distribution(generator); - if (rnd_val < p_val) { + if (rnd_val < (float) p_val) { std::vector cur_vector_float; for (size_t d = 0; d < ndims; d++) cur_vector_float.push_back(cur_vector_T[d]); sampled_vectors.push_back(cur_vector_float); } } + // If we got zero sample vectors, simply add the first vector + // as the sample vector. slice_size = sampled_vectors.size(); + if (slice_size == 0) { + slice_size = 1; + std::vector cur_vector_float(inputdata, inputdata + ndims); + sampled_vectors.push_back(cur_vector_float); + } sampled_data = new float[slice_size * ndims]; + for (size_t i = 0; i < slice_size; i++) { for (size_t j = 0; j < ndims; j++) { sampled_data[i * ndims + j] = sampled_vectors[i][j]; @@ -187,12 +222,33 @@ void gen_random_slice(const T *inputdata, size_t npts, size_t ndims, // num_pq_chunks (if it divides dimension, else rounded) chunks, and runs // k-means in each chunk to compute the PQ pivots and stores in bin format in // file pq_pivots_path as a s num_centers*dim floating point binary file +template +int generate_pq_pivots(const std::unique_ptr &passed_train_data, + size_t num_train, unsigned dim, unsigned num_centers, + unsigned num_pq_chunks, unsigned max_k_means_reps, + std::string pq_pivots_path) { + std::unique_ptr train_float = + std::make_unique(num_train * (size_t)(dim)); + float *flt_ptr = train_float.get(); + T * T_ptr = passed_train_data.get(); + + for (_u64 i = 0; i < num_train; i++) { + for (_u64 j = 0; j < (_u64) dim; j++) { + flt_ptr[i * (_u64) dim + j] = (float) T_ptr[i * (_u64) dim + j]; + } + } + if (generate_pq_pivots(flt_ptr, num_train, dim, num_centers, num_pq_chunks, + max_k_means_reps, pq_pivots_path) != 0) + return -1; + return 0; +} + int generate_pq_pivots(const float *passed_train_data, size_t num_train, unsigned dim, unsigned num_centers, unsigned num_pq_chunks, unsigned max_k_means_reps, std::string pq_pivots_path) { if (num_pq_chunks > dim) { - diskann::cout << " Error: number of chunks more than dimension" + diskann::cerr << " Error: number of chunks more than dimension" << std::endl; return -1; } @@ -205,23 +261,12 @@ int generate_pq_pivots(const float *passed_train_data, size_t num_train, for (uint64_t i = 0; i < num_train; i++) { for (uint64_t j = 0; j < dim; j++) { if (passed_train_data[i * dim + j] != train_data[i * dim + j]) - diskann::cout << "error in copy" << std::endl; + diskann::cerr << "error in copy" << std::endl; } } std::unique_ptr full_pivot_data; - if (file_exists(pq_pivots_path)) { - size_t file_dim, file_num_centers; - diskann::load_bin(pq_pivots_path, full_pivot_data, file_num_centers, - file_dim); - if (file_dim == dim && file_num_centers == num_centers) { - diskann::cout << "PQ pivot file exists. Not generating again" - << std::endl; - return -1; - } - } - // Calculate centroid and center the training data std::unique_ptr centroid = std::make_unique(dim); for (uint64_t d = 0; d < dim; d++) { @@ -229,7 +274,7 @@ int generate_pq_pivots(const float *passed_train_data, size_t num_train, for (uint64_t p = 0; p < num_train; p++) { centroid[d] += train_data[p * dim + d]; } - centroid[d] /= num_train; + centroid[d] /= (float) num_train; } // std::memset(centroid, 0 , dim*sizeof(float)); @@ -251,7 +296,7 @@ int generate_pq_pivots(const float *passed_train_data, size_t num_train, std::vector> bin_to_dims(num_pq_chunks); tsl::robin_map dim_to_bin; - std::vector bin_loads(num_pq_chunks, 0); + std::vector bin_loads(num_pq_chunks, 0); // Process dimensions not inserted by previous loop for (uint32_t d = 0; d < dim; d++) { @@ -266,8 +311,6 @@ int generate_pq_pivots(const float *passed_train_data, size_t num_train, cur_best_load = bin_loads[b]; } } - diskann::cout << " Pushing " << d << " into bin #: " << cur_best - << std::endl; bin_to_dims[cur_best].push_back(d); if (bin_to_dims[cur_best].size() == high_val) { cur_num_high++; @@ -281,26 +324,20 @@ int generate_pq_pivots(const float *passed_train_data, size_t num_train, chunk_offsets.push_back(0); for (uint32_t b = 0; b < num_pq_chunks; b++) { - diskann::cout << "[ "; for (auto p : bin_to_dims[b]) { rearrangement.push_back(p); - diskann::cout << p << ","; } - diskann::cout << "] " << std::endl; if (b > 0) chunk_offsets.push_back(chunk_offsets[b - 1] + (unsigned) bin_to_dims[b - 1].size()); } chunk_offsets.push_back(dim); - diskann::cout << "\nCross-checking rearranged order of coordinates:" - << std::endl; - for (auto p : rearrangement) - diskann::cout << p << " "; - diskann::cout << std::endl; - full_pivot_data.reset(new float[num_centers * dim]); + // DEBUG ONLY + double kmeans_time = 0.0, lloyds_time = 0.0, copy_time = 0.0; + for (size_t i = 0; i < num_pq_chunks; i++) { size_t cur_chunk_size = chunk_offsets[i + 1] - chunk_offsets[i]; @@ -313,42 +350,79 @@ int generate_pq_pivots(const float *passed_train_data, size_t num_train, std::unique_ptr closest_center = std::make_unique(num_train); - diskann::cout << "Processing chunk " << i << " with dimensions [" - << chunk_offsets[i] << ", " << chunk_offsets[i + 1] << ")" - << std::endl; + memset((void *) cur_pivot_data.get(), 0, + num_centers * cur_chunk_size * sizeof(float)); + auto start = std::chrono::high_resolution_clock::now(); #pragma omp parallel for schedule(static, 65536) for (int64_t j = 0; j < (_s64) num_train; j++) { std::memcpy(cur_data.get() + j * cur_chunk_size, train_data.get() + j * dim + chunk_offsets[i], cur_chunk_size * sizeof(float)); } + auto end = std::chrono::high_resolution_clock::now(); + copy_time += std::chrono::duration(end - start).count(); - kmeans::kmeanspp_selecting_pivots(cur_data.get(), num_train, cur_chunk_size, - cur_pivot_data.get(), num_centers); + start = std::chrono::high_resolution_clock::now(); + // kmeans::kmeanspp_selecting_pivots(cur_data.get(), num_train, + // cur_chunk_size, + // cur_pivot_data.get(), num_centers); + kmeans::selecting_pivots(cur_data.get(), num_train, cur_chunk_size, + cur_pivot_data.get(), num_centers); + + unsigned k_means_reps = max_k_means_reps; kmeans::run_lloyds(cur_data.get(), num_train, cur_chunk_size, - cur_pivot_data.get(), num_centers, max_k_means_reps, - NULL, closest_center.get()); + cur_pivot_data.get(), num_centers, k_means_reps, nullptr, + closest_center.get()); + end = std::chrono::high_resolution_clock::now(); + kmeans_time += std::chrono::duration(end - start).count(); + + start = std::chrono::high_resolution_clock::now(); + if (num_train > 2 * num_centers) { + kmeans::run_lloyds(cur_data.get(), num_train, cur_chunk_size, + cur_pivot_data.get(), num_centers, max_k_means_reps, + NULL, closest_center.get()); + } + end = std::chrono::high_resolution_clock::now(); + lloyds_time += std::chrono::duration(end - start).count(); + start = std::chrono::high_resolution_clock::now(); for (uint64_t j = 0; j < num_centers; j++) { std::memcpy(full_pivot_data.get() + j * dim + chunk_offsets[i], cur_pivot_data.get() + j * cur_chunk_size, cur_chunk_size * sizeof(float)); } + end = std::chrono::high_resolution_clock::now(); + copy_time += std::chrono::duration(end - start).count(); } + diskann::cout << "Kmeans time: " << kmeans_time + << " Lloyds time: " << lloyds_time + << " Copy time: " << copy_time << std::endl; + + std::vector cumul_bytes(5, 0); + cumul_bytes[0] = METADATA_SIZE; + cumul_bytes[1] = + cumul_bytes[0] + + diskann::save_bin(pq_pivots_path.c_str(), full_pivot_data.get(), + (size_t) num_centers, dim, cumul_bytes[0]); + cumul_bytes[2] = cumul_bytes[1] + diskann::save_bin( + pq_pivots_path.c_str(), centroid.get(), + (size_t) dim, 1, cumul_bytes[1]); + cumul_bytes[3] = + cumul_bytes[2] + + diskann::save_bin(pq_pivots_path.c_str(), rearrangement.data(), + rearrangement.size(), 1, cumul_bytes[2]); + cumul_bytes[4] = + cumul_bytes[3] + + diskann::save_bin(pq_pivots_path.c_str(), chunk_offsets.data(), + chunk_offsets.size(), 1, cumul_bytes[3]); + diskann::save_bin<_u64>(pq_pivots_path.c_str(), cumul_bytes.data(), + cumul_bytes.size(), 1, 0); + + diskann::cout << "Saved pq pivot data to " << pq_pivots_path << " of size " + << cumul_bytes[cumul_bytes.size() - 1] << "B." << std::endl; - diskann::save_bin(pq_pivots_path.c_str(), full_pivot_data.get(), - (size_t) num_centers, dim); - std::string centroids_path = pq_pivots_path + "_centroid.bin"; - diskann::save_bin(centroids_path.c_str(), centroid.get(), (size_t) dim, - 1); - std::string rearrangement_path = pq_pivots_path + "_rearrangement_perm.bin"; - diskann::save_bin(rearrangement_path.c_str(), rearrangement.data(), - rearrangement.size(), 1); - std::string chunk_offsets_path = pq_pivots_path + "_chunk_offsets.bin"; - diskann::save_bin(chunk_offsets_path.c_str(), chunk_offsets.data(), - chunk_offsets.size(), 1); return 0; } @@ -361,9 +435,10 @@ template int generate_pq_data_from_pivots(const std::string data_file, unsigned num_centers, unsigned num_pq_chunks, std::string pq_pivots_path, - std::string pq_compressed_vectors_path) { + std::string pq_compressed_vectors_path, + size_t offset) { _u64 read_blk_size = 64 * 1024 * 1024; - cached_ifstream base_reader(data_file, read_blk_size); + cached_ifstream base_reader(data_file, read_blk_size, (uint32_t) offset); _u32 npts32; _u32 basedim32; base_reader.read((char *) &npts32, sizeof(uint32_t)); @@ -371,6 +446,12 @@ int generate_pq_data_from_pivots(const std::string data_file, size_t num_points = npts32; size_t dim = basedim32; +#ifdef SAVE_INFLATED_PQ + std::string inflated_pq_file = pq_compressed_vectors_path + "_full.bin"; +#endif + + size_t BLOCK_SIZE = (std::min)((size_t) MAX_BLOCK_SIZE, num_points); + std::unique_ptr full_pivot_data; std::unique_ptr centroid; std::unique_ptr rearrangement; @@ -380,67 +461,92 @@ int generate_pq_data_from_pivots(const std::string data_file, diskann::cout << "ERROR: PQ k-means pivot file not found" << std::endl; throw diskann::ANNException("PQ k-means pivot file not found", -1); } else { - uint64_t numr, numc; - - std::string centroids_path = pq_pivots_path + "_centroid.bin"; - diskann::load_bin(centroids_path.c_str(), centroid, numr, numc); - - if (numr != dim || numc != 1) { - diskann::cout << "Error reading centroid file." << std::endl; - throw diskann::ANNException("Error reading centroid file.", -1, - __FUNCSIG__, __FILE__, __LINE__); + _u64 nr, nc; + std::unique_ptr<_u64[]> file_offset_data; + + diskann::load_bin<_u64>(pq_pivots_path.c_str(), file_offset_data, nr, nc, + 0); + + if (nr != 5) { + diskann::cout << "Error reading pq_pivots file " << pq_pivots_path + << ". Offsets dont contain correct metadata, # offsets = " + << nr << ", but expecting 5."; + throw diskann::ANNException( + "Error reading pq_pivots file at offsets data.", -1, __FUNCSIG__, + __FILE__, __LINE__); } - std::string rearrangement_path = pq_pivots_path + "_rearrangement_perm.bin"; - diskann::load_bin(rearrangement_path.c_str(), rearrangement, numr, - numc); - if (numr != dim || numc != 1) { - diskann::cout << "Error reading rearrangement file." << std::endl; - throw diskann::ANNException("Error reading rearrangement file.", -1, - __FUNCSIG__, __FILE__, __LINE__); + + diskann::load_bin(pq_pivots_path.c_str(), full_pivot_data, nr, nc, + file_offset_data[0]); + + if ((nr != num_centers) || (nc != dim)) { + diskann::cout << "Error reading pq_pivots file " << pq_pivots_path + << ". file_num_centers = " << nr << ", file_dim = " << nc + << " but expecting " << num_centers << " centers in " << dim + << " dimensions."; + throw diskann::ANNException( + "Error reading pq_pivots file at pivots data.", -1, __FUNCSIG__, + __FILE__, __LINE__); } - std::string chunk_offsets_path = pq_pivots_path + "_chunk_offsets.bin"; - diskann::load_bin(chunk_offsets_path.c_str(), chunk_offsets, numr, - numc); - if (numr != (uint64_t) num_pq_chunks + 1 || numc != 1) { - diskann::cout << "Error reading chunk offsets file." << std::endl; - throw diskann::ANNException("Error reading chunk offsets file.", -1, - __FUNCSIG__, __FILE__, __LINE__); + + diskann::load_bin(pq_pivots_path.c_str(), centroid, nr, nc, + file_offset_data[1]); + + if ((nr != dim) || (nc != 1)) { + diskann::cout << "Error reading pq_pivots file " << pq_pivots_path + << ". file_dim = " << nr << ", file_cols = " << nc + << " but expecting " << dim << " entries in 1 dimension."; + throw diskann::ANNException( + "Error reading pq_pivots file at centroid data.", -1, __FUNCSIG__, + __FILE__, __LINE__); } - size_t file_num_centers; - size_t file_dim; - diskann::load_bin(pq_pivots_path, full_pivot_data, file_num_centers, - file_dim); - - if (file_num_centers != num_centers) { - std::stringstream stream; - stream << "ERROR: file number of PQ centers " << file_num_centers - << " does " - "not match input argument " - << num_centers << std::endl; - diskann::cout << stream.str() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); + diskann::load_bin(pq_pivots_path.c_str(), rearrangement, nr, nc, + file_offset_data[2]); + + if ((nr != dim) || (nc != 1)) { + diskann::cout << "Error reading pq_pivots file " << pq_pivots_path + << ". file_dim = " << nr << ", file_cols = " << nc + << " but expecting " << dim << " entries in 1 dimension."; + throw diskann::ANNException( + "Error reading pq_pivots file at re-arrangement data.", -1, + __FUNCSIG__, __FILE__, __LINE__); } - if (file_dim != dim) { - std::stringstream stream; - stream << "ERROR: PQ pivot dimension does " - "not match base file dimension" - << std::endl; - diskann::cout << stream.str() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); + + diskann::load_bin(pq_pivots_path.c_str(), chunk_offsets, nr, nc, + file_offset_data[3]); + + if (nr != (uint64_t) num_pq_chunks + 1 || nc != 1) { + diskann::cout + << "Error reading pq_pivots file at chunk offsets; file has nr=" << nr + << ",nc=" << nc << ", expecting nr=" << num_pq_chunks + 1 << ", nc=1." + << std::endl; + throw diskann::ANNException( + "Error reading pq_pivots file at chunk offsets.", -1, __FUNCSIG__, + __FILE__, __LINE__); } + diskann::cout << "Loaded PQ pivot information" << std::endl; } std::ofstream compressed_file_writer(pq_compressed_vectors_path, std::ios::binary); - _u32 num_pq_chunks_u32 = num_pq_chunks; + _u32 num_pq_chunks_u32 = num_pq_chunks; compressed_file_writer.write((char *) &num_points, sizeof(uint32_t)); compressed_file_writer.write((char *) &num_pq_chunks_u32, sizeof(uint32_t)); +#ifdef SAVE_INFLATED_PQ + std::ofstream inflated_file_writer(inflated_pq_file, std::ios::binary); + inflated_file_writer.write((char *) &npts32, sizeof(uint32_t)); + inflated_file_writer.write((char *) &basedim32, sizeof(uint32_t)); + + std::unique_ptr block_inflated_base = + std::make_unique(BLOCK_SIZE * (_u64) dim); + std::memset(block_inflated_base.get(), 0, + BLOCK_SIZE * (_u64) dim * sizeof(float)); +#endif + size_t block_size = num_points <= BLOCK_SIZE ? num_points : BLOCK_SIZE; std::unique_ptr<_u32[]> block_compressed_base = std::make_unique<_u32[]>(block_size * (_u64) num_pq_chunks); @@ -465,8 +571,11 @@ int generate_pq_data_from_pivots(const std::string data_file, diskann::convert_types(block_data_T.get(), block_data_tmp.get(), cur_blk_size, dim); - diskann::cout << "Processing points [" << start_id << ", " << end_id - << ").." << std::flush; + /* diskann::cout << "Processing points [" << start_id << ", " << end_id + << ")... Used memory: " + << getMemoryUsage() / (1024.0 * 1024.0) << "MB" << + std::flush; + */ for (uint64_t p = 0; p < cur_blk_size; p++) { for (uint64_t d = 0; d < dim; d++) { @@ -510,13 +619,23 @@ int generate_pq_data_from_pivots(const std::string data_file, math_utils::compute_closest_centers(cur_data.get(), cur_blk_size, cur_chunk_size, cur_pivot_data.get(), num_centers, 1, closest_center.get()); - #pragma omp parallel for schedule(static, 8192) for (int64_t j = 0; j < (_s64) cur_blk_size; j++) { block_compressed_base[j * num_pq_chunks + i] = closest_center[j]; +#ifdef SAVE_INFLATED_PQ + for (uint64_t k = 0; k < cur_chunk_size; k++) + block_inflated_base[j * dim + chunk_offsets[i] + k] = + cur_pivot_data[closest_center[j] * cur_chunk_size + k] + + centroid[chunk_offsets[i] + k]; +#endif } } +#ifdef SAVE_INFLATED_PQ + inflated_file_writer.write((char *) block_inflated_base.get(), + cur_blk_size * dim * sizeof(float)); +#endif + if (num_centers > 256) { compressed_file_writer.write( (char *) (block_compressed_base.get()), @@ -530,14 +649,17 @@ int generate_pq_data_from_pivots(const std::string data_file, (char *) (pVec.get()), cur_blk_size * num_pq_chunks * sizeof(uint8_t)); } - diskann::cout << ".done." << std::endl; + // diskann::cout << ".done." << std::endl; } -// Gopal. Splittng diskann_dll into separate DLLs for search and build. +// Splittng diskann_dll into separate DLLs for search and build. // This code should only be available in the "build" DLL. -#ifdef DISKANN_BUILD +#ifdef USE_TCMALLOC MallocExtension::instance()->ReleaseFreeMemory(); #endif compressed_file_writer.close(); +#ifdef SAVE_INFLATED_PQ + inflated_file_writer.close(); +#endif return 0; } @@ -567,6 +689,7 @@ int estimate_cluster_sizes(const std::string data_file, float *pivots, shard_counts[i] = 0; } + size_t BLOCK_SIZE = (std::min)((size_t) MAX_BLOCK_SIZE, num_test); size_t num_points = 0, num_dim = 0; diskann::get_bin_metadata(data_file, num_points, num_dim); size_t block_size = num_points <= BLOCK_SIZE ? num_points : BLOCK_SIZE; @@ -649,6 +772,7 @@ int shard_data_into_clusters(const std::string data_file, float *pivots, shard_counts[i] = 0; } + size_t BLOCK_SIZE = (std::min)((size_t) MAX_BLOCK_SIZE, num_points); size_t block_size = num_points <= BLOCK_SIZE ? num_points : BLOCK_SIZE; std::unique_ptr<_u32[]> block_closest_centers = std::make_unique<_u32[]>(block_size * k_base); @@ -840,19 +964,28 @@ int partition_with_ram_budget(const std::string data_file, } // Instantations of supported templates +template void DISKANN_DLLEXPORT gen_random_slice( + const std::string data_file, double p_val, + std::unique_ptr &sampled_data, size_t &slice_size, size_t &ndims); +template void DISKANN_DLLEXPORT gen_random_slice( + const std::string data_file, double p_val, + std::unique_ptr &sampled_data, size_t &slice_size, size_t &ndims); +template void DISKANN_DLLEXPORT gen_random_slice( + const std::string data_file, double p_val, + std::unique_ptr &sampled_data, size_t &slice_size, size_t &ndims); -template void DISKANN_DLLEXPORT -gen_random_slice(const std::string base_file, - const std::string output_prefix, double sampling_rate); +template void DISKANN_DLLEXPORT gen_random_slice( + const std::string base_file, const std::string output_prefix, + double sampling_rate, size_t offset); template void DISKANN_DLLEXPORT gen_random_slice( const std::string base_file, const std::string output_prefix, - double sampling_rate); -template void DISKANN_DLLEXPORT -gen_random_slice(const std::string base_file, - const std::string output_prefix, double sampling_rate); + double sampling_rate, size_t offset); +template void DISKANN_DLLEXPORT gen_random_slice( + const std::string base_file, const std::string output_prefix, + double sampling_rate, size_t offset); template void DISKANN_DLLEXPORT -gen_random_slice(const float *inputdata, size_t npts, size_t ndims, + gen_random_slice(const float *inputdata, size_t npts, size_t ndims, double p_val, float *&sampled_data, size_t &slice_size); template void DISKANN_DLLEXPORT gen_random_slice( const uint8_t *inputdata, size_t npts, size_t ndims, double p_val, @@ -891,12 +1024,28 @@ template DISKANN_DLLEXPORT int partition_with_ram_budget( const std::string data_file, const double sampling_rate, double ram_budget, size_t graph_degree, const std::string prefix_path, size_t k_base); +template DISKANN_DLLEXPORT int generate_pq_pivots( + const std::unique_ptr &passed_train_data, size_t num_train, + unsigned dim, unsigned num_centers, unsigned num_pq_chunks, + unsigned max_k_means_reps, std::string pq_pivots_path); +template DISKANN_DLLEXPORT int generate_pq_pivots( + const std::unique_ptr &passed_train_data, size_t num_train, + unsigned dim, unsigned num_centers, unsigned num_pq_chunks, + unsigned max_k_means_reps, std::string pq_pivots_path); +template DISKANN_DLLEXPORT int generate_pq_pivots( + const std::unique_ptr &passed_train_data, size_t num_train, + unsigned dim, unsigned num_centers, unsigned num_pq_chunks, + unsigned max_k_means_reps, std::string pq_pivots_path); + template DISKANN_DLLEXPORT int generate_pq_data_from_pivots( const std::string data_file, unsigned num_centers, unsigned num_pq_chunks, - std::string pq_pivots_path, std::string pq_compressed_vectors_path); + std::string pq_pivots_path, std::string pq_compressed_vectors_path, + size_t offset); template DISKANN_DLLEXPORT int generate_pq_data_from_pivots( const std::string data_file, unsigned num_centers, unsigned num_pq_chunks, - std::string pq_pivots_path, std::string pq_compressed_vectors_path); + std::string pq_pivots_path, std::string pq_compressed_vectors_path, + size_t offset); template DISKANN_DLLEXPORT int generate_pq_data_from_pivots( const std::string data_file, unsigned num_centers, unsigned num_pq_chunks, - std::string pq_pivots_path, std::string pq_compressed_vectors_path); + std::string pq_pivots_path, std::string pq_compressed_vectors_path, + size_t offset); diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index 21da3521d6..0f74239cb8 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -18,6 +18,10 @@ #include "timer.h" #include "utils.h" +#include +#include +#include "tcmalloc/malloc_extension.h" +#include "cosine_similarity.h" #include "tsl/robin_set.h" #ifdef _WINDOWS @@ -33,6 +37,7 @@ #define SECTOR_LEN 4096 #define READ_U64(stream, val) stream.read((char *) &val, sizeof(_u64)) +#define READ_U32(stream, val) stream.read((char *) &val, sizeof(_u32)) #define READ_UNSIGNED(stream, val) stream.read((char *) &val, sizeof(unsigned)) // sector # on disk where node_id is present @@ -84,127 +89,106 @@ namespace { } // namespace namespace diskann { - template<> - PQFlashIndex<_u8>::PQFlashIndex( - std::shared_ptr &fileReader) - : reader(fileReader) { - diskann::cout - << "dist_cmp function for _u8 uses slow implementation." - " Please contact gopalsr@microsoft.com if you need an AVX/AVX2" - " implementation." - << std::endl; - // TODO: No AVX2/AVX implementation available for uint8. - this->dist_cmp = new DistanceL2UInt8(); - if (Avx2SupportedCPU) { - diskann::cout << "Using AVX2 dist_cmp_float function." << std::endl; - this->dist_cmp_float = new DistanceL2(); - } else if (AvxSupportedCPU) { - diskann::cout << "Using AVX dist_cmp_float function" << std::endl; - this->dist_cmp_float = new AVXDistanceL2Float(); - } else { - diskann::cout << "No AVX/AVX2 support. Using Slow dist_cmp_float function" - << std::endl; - this->dist_cmp_float = new SlowDistanceL2Float(); - } + template + DiskNode::DiskNode(uint32_t id, T *coords, uint32_t *nhood) : id(id) { + this->coords = coords; + this->nnbrs = *nhood; + this->nbrs = nhood + 1; } - template<> - PQFlashIndex<_s8>::PQFlashIndex( - std::shared_ptr &fileReader) - : reader(fileReader) { - if (Avx2SupportedCPU) { - diskann::cout << "Using AVX2 function for dist_cmp and dist_cmp_float" - << std::endl; - this->dist_cmp = new DistanceL2Int8(); - this->dist_cmp_float = new DistanceL2(); - } else if (AvxSupportedCPU) { - diskann::cout << "No AVX2 support. Switching to AVX routines for " - "dist_cmp, dist_cmp_float." - << std::endl; - this->dist_cmp = new AVXDistanceL2Int8(); - this->dist_cmp_float = new AVXDistanceL2Float(); - } else { - diskann::cout << "No AVX/AVX2 support. Switching to slow routines for " - "dist_cmp, dist_cmp_float" - << std::endl; - this->dist_cmp = new SlowDistanceL2Int(); - this->dist_cmp_float = new SlowDistanceL2Float(); - } - } + // structs for DiskNode + template struct DiskNode; + template struct DiskNode; + template struct DiskNode; + + template + PQFlashIndex::PQFlashIndex( + diskann::Metric m, std::shared_ptr &fileReader, + bool single_file_index, bool tags) + : reader(fileReader), data_is_normalized(false), + single_index_file(single_file_index), enable_tags(tags) { + if (m == diskann::Metric::COSINE) { + if (std::is_floating_point::value) { + diskann::cout << "Cosine metric chosen for (normalized) float data." + "Changing distance to L2 to boost accuracy." + << std::endl; + m = diskann::Metric::L2; + data_is_normalized = true; - template<> - PQFlashIndex::PQFlashIndex( - std::shared_ptr &fileReader) - : reader(fileReader) { - if (Avx2SupportedCPU) { - diskann::cout << "Using AVX2 functions for dist_cmp and dist_cmp_float" - << std::endl; - this->dist_cmp = new DistanceL2(); - this->dist_cmp_float = new DistanceL2(); - } else if (AvxSupportedCPU) { - diskann::cout << "No AVX2 support. Switching to AVX functions for " - "dist_cmp and dist_cmp_float." - << std::endl; - this->dist_cmp = new AVXDistanceL2Float(); - this->dist_cmp_float = new AVXDistanceL2Float(); - } else { - diskann::cout << "No AVX/AVX2 support. Switching to slow implementations " - "for dist_cmp and dist_cmp_float" - << std::endl; - this->dist_cmp = new AVXDistanceL2Float(); - this->dist_cmp_float = new AVXDistanceL2Float(); + } else { + diskann::cerr << "WARNING: Cannot normalize integral data types." + << " This may result in erroneous results or poor recall." + << " Consider using L2 distance with integral data types." + << std::endl; + } } + + this->dist_cmp.reset(diskann::get_distance_function(m)); + this->dist_cmp_float.reset(diskann::get_distance_function(m)); } - template - PQFlashIndex::~PQFlashIndex() { + template + PQFlashIndex::~PQFlashIndex() { #ifndef EXEC_ENV_OLS if (data != nullptr) { delete[] data; } #endif + diskann::cout << "Thread Data size: " << this->thread_data.size() << "\n"; + assert(!this->thread_data.empty()); + +#ifndef EXEC_ENV_OLS if (centroid_data != nullptr) aligned_free(centroid_data); +#endif + // delete backing bufs for nhood and coord cache if (nhood_cache_buf != nullptr) { delete[] nhood_cache_buf; diskann::aligned_free(coord_cache_buf); } - delete this->dist_cmp; - delete this->dist_cmp_float; if (load_flag) { this->destroy_thread_data(); reader->close(); // delete reader; //not deleting reader because it is now passed by ref. } + + if (this->tags != nullptr) { + delete[] tags; + } + +#ifndef EXEC_ENV_OLS + if (medoids != nullptr) + delete[] medoids; +#endif +#ifdef USE_TCMALLOC + MallocExtension::instance()->ReleaseFreeMemory(); +#endif } - template - void PQFlashIndex::setup_thread_data(_u64 nthreads) { + template + void PQFlashIndex::setup_thread_data(_u64 nthreads) { diskann::cout << "Setting up thread-specific contexts for nthreads: " - << nthreads << std::endl; + << nthreads << "\n"; // omp parallel for to generate unique thread IDs #pragma omp parallel for num_threads((int) nthreads) for (_s64 thread = 0; thread < (_s64) nthreads; thread++) { #pragma omp critical { this->reader->register_thread(); - IOContext &ctx = this->reader->get_ctx(); - // diskann::cout << "ctx: " << ctx << "\n"; + IOContext & ctx = this->reader->get_ctx(); QueryScratch scratch; _u64 coord_alloc_size = ROUND_UP(MAX_N_CMPS * this->aligned_dim, 256); diskann::alloc_aligned((void **) &scratch.coord_scratch, coord_alloc_size, 256); - // scratch.coord_scratch = new T[MAX_N_CMPS * this->aligned_dim]; - // //Gopal. Commenting out the reallocation! diskann::alloc_aligned((void **) &scratch.sector_scratch, MAX_N_SECTOR_READS * SECTOR_LEN, SECTOR_LEN); diskann::alloc_aligned((void **) &scratch.aligned_scratch, 256 * sizeof(float), 256); diskann::alloc_aligned((void **) &scratch.aligned_pq_coord_scratch, - 25600 * sizeof(_u8), 256); + 32768 * 32 * sizeof(_u8), 256); diskann::alloc_aligned((void **) &scratch.aligned_pqtable_dist_scratch, 25600 * sizeof(float), 256); diskann::alloc_aligned((void **) &scratch.aligned_dist_scratch, @@ -215,8 +199,9 @@ namespace diskann { this->aligned_dim * sizeof(float), 8 * sizeof(float)); + memset(scratch.sector_scratch, 0, MAX_N_SECTOR_READS * SECTOR_LEN); memset(scratch.aligned_scratch, 0, 256 * sizeof(float)); - memset(scratch.coord_scratch, 0, MAX_N_CMPS * this->aligned_dim); + memset(scratch.coord_scratch, 0, coord_alloc_size); memset(scratch.aligned_query_T, 0, this->aligned_dim * sizeof(T)); memset(scratch.aligned_query_float, 0, this->aligned_dim * sizeof(float)); @@ -225,21 +210,15 @@ namespace diskann { data.ctx = ctx; data.scratch = scratch; this->thread_data.push(data); + this->thread_data_backing_buf.push_back(data); } } load_flag = true; } - template - void PQFlashIndex::destroy_thread_data() { - diskann::cout << "Clearing scratch" << std::endl; - assert(this->thread_data.size() == this->max_nthreads); - while (this->thread_data.size() > 0) { - ThreadData data = this->thread_data.pop(); - while (data.scratch.sector_scratch == nullptr) { - this->thread_data.wait_for_push_notify(); - data = this->thread_data.pop(); - } + template + void PQFlashIndex::destroy_thread_data() { + for (auto &data : this->thread_data_backing_buf) { auto &scratch = data.scratch; diskann::aligned_free((void *) scratch.coord_scratch); diskann::aligned_free((void *) scratch.sector_scratch); @@ -250,11 +229,13 @@ namespace diskann { diskann::aligned_free((void *) scratch.aligned_query_float); diskann::aligned_free((void *) scratch.aligned_query_T); } + this->reader->deregister_all_threads(); } - template - void PQFlashIndex::load_cache_list(std::vector &node_list) { - diskann::cout << "Loading the cache list into memory.." << std::flush; + template + void PQFlashIndex::load_cache_list( + std::vector &node_list) { + // diskann::cout << "Loading the cache list into memory.." << std::flush; _u64 num_cached_nodes = node_list.size(); // borrow thread data @@ -263,7 +244,6 @@ namespace diskann { this->thread_data.wait_for_push_notify(); this_thread_data = this->thread_data.pop(); } - IOContext &ctx = this_thread_data.ctx; nhood_cache_buf = new unsigned[num_cached_nodes * (max_degree + 1)]; @@ -280,7 +260,7 @@ namespace diskann { for (_u64 block = 0; block < num_blocks; block++) { _u64 start_idx = block * BLOCK_SIZE; _u64 end_idx = (std::min)(num_cached_nodes, (block + 1) * BLOCK_SIZE); - std::vector read_reqs; + std::vector read_reqs; std::vector> nhoods; for (_u64 node_idx = start_idx; node_idx < end_idx; node_idx++) { AlignedRead read; @@ -323,14 +303,14 @@ namespace diskann { } #ifdef EXEC_ENV_OLS - template - void PQFlashIndex::generate_cache_list_from_sample_queries( + template + void PQFlashIndex::generate_cache_list_from_sample_queries( MemoryMappedFiles &files, std::string sample_bin, _u64 l_search, _u64 beamwidth, _u64 num_nodes_to_cache, uint32_t nthreads, std::vector &node_list) { #else - template - void PQFlashIndex::generate_cache_list_from_sample_queries( + template + void PQFlashIndex::generate_cache_list_from_sample_queries( std::string sample_bin, _u64 l_search, _u64 beamwidth, _u64 num_nodes_to_cache, uint32_t nthreads, std::vector &node_list) { @@ -363,13 +343,13 @@ namespace diskann { return; } - std::vector tmp_result_ids_64(sample_num, 0); - std::vector tmp_result_dists(sample_num, 0); + std::vector tmp_result_tags(sample_num, 0); + std::vector tmp_result_dists(sample_num, 0); #pragma omp parallel for schedule(dynamic, 1) num_threads(nthreads) for (_s64 i = 0; i < (int64_t) sample_num; i++) { cached_beam_search(samples + (i * sample_aligned_dim), 1, l_search, - tmp_result_ids_64.data() + (i * 1), + tmp_result_tags.data() + (i * 1), tmp_result_dists.data() + (i * 1), beamwidth); } @@ -388,15 +368,26 @@ namespace diskann { diskann::aligned_free(samples); } - template - void PQFlashIndex::cache_bfs_levels(_u64 num_nodes_to_cache, - std::vector &node_list) { - // Gopal. random_shuffle() is deprecated. + template + void PQFlashIndex::cache_bfs_levels( + _u64 num_nodes_to_cache, std::vector &node_list) { + // random_shuffle() is deprecated. std::random_device rng; std::mt19937 urng(rng()); node_list.clear(); + // Do not cache more than 10% of the nodes in the index + _u64 tenp_nodes = (_u64)(std::round(this->num_points * 0.1)); + if (num_nodes_to_cache > tenp_nodes) { + diskann::cout << "Reducing nodes to cache from: " << num_nodes_to_cache + << " to: " << tenp_nodes + << "(10 percent of total nodes:" << this->num_points << ")" + << std::endl; + num_nodes_to_cache = tenp_nodes == 0 ? 1 : tenp_nodes; + } + diskann::cout << "Caching " << num_nodes_to_cache << "..." << std::endl; + // borrow thread data ThreadData this_thread_data = this->thread_data.pop(); while (this_thread_data.scratch.sector_scratch == nullptr) { @@ -434,7 +425,7 @@ namespace diskann { nodes_to_expand.push_back(id); } - // Gopal. random_shuffle() is deprecated. + // random_shuffle() is deprecated. std::shuffle(nodes_to_expand.begin(), nodes_to_expand.end(), urng); diskann::cout << "Level: " << lvl << std::flush; @@ -447,7 +438,7 @@ namespace diskann { size_t start = block * BLOCK_SIZE; size_t end = (std::min)((block + 1) * BLOCK_SIZE, nodes_to_expand.size()); - std::vector read_reqs; + std::vector read_reqs; std::vector> nhoods; for (size_t cur_pt = start; cur_pt < end; cur_pt++) { char *buf = nullptr; @@ -459,8 +450,10 @@ namespace diskann { read.offset = NODE_SECTOR_NO(nodes_to_expand[cur_pt]) * SECTOR_LEN; read_reqs.push_back(read); } + // issue read requests reader->read(read_reqs, ctx); + // process each nhood buf for (auto &nhood : nhoods) { // insert node coord into coord_cache @@ -492,7 +485,7 @@ namespace diskann { for (const unsigned &p : *cur_level) cur_level_node_list.push_back(p); - // Gopal. random_shuffle() is deprecated + // random_shuffle() is deprecated std::shuffle(cur_level_node_list.begin(), cur_level_node_list.end(), urng); size_t residual = num_nodes_to_cache - node_list.size(); @@ -503,13 +496,12 @@ namespace diskann { diskann::cout << "Level: " << lvl << std::flush; diskann::cout << ". #nodes: " << node_list.size() - prev_node_list_size << ", #nodes thus far: " << node_list.size() << std::endl; - // return thread data this->thread_data.push(this_thread_data); } - template - void PQFlashIndex::use_medoids_data_as_centroids() { + template + void PQFlashIndex::use_medoids_data_as_centroids() { if (centroid_data != nullptr) aligned_free(centroid_data); alloc_aligned(((void **) ¢roid_data), @@ -546,6 +538,7 @@ namespace diskann { for (uint32_t i = 0; i < data_dim; i++) centroid_data[cur_m * aligned_dim + i] = medoid_coords[i]; + delete[] medoid_coords; aligned_free(medoid_buf); } @@ -556,97 +549,166 @@ namespace diskann { } #ifdef EXEC_ENV_OLS - template - int PQFlashIndex::load(MemoryMappedFiles &files, uint32_t num_threads, - const char *pq_prefix, - const char *disk_index_file) { + template + int PQFlashIndex::load(const char *index_prefix, _u32 num_threads, + MemoryMappedFiles &files, + bool new_index_format) { #else - template - int PQFlashIndex::load(uint32_t num_threads, const char *pq_prefix, - const char *disk_index_file) { + template + int PQFlashIndex::load(const char *index_prefix, _u32 num_threads, + bool new_index_format) { #endif - std::string pq_table_bin = std::string(pq_prefix) + "_pivots.bin"; - std::string pq_compressed_vectors = - std::string(pq_prefix) + "_compressed.bin"; - std::string medoids_file = std::string(disk_index_file) + "_medoids.bin"; - std::string centroids_file = - std::string(disk_index_file) + "_centroids.bin"; - - size_t pq_file_dim, pq_file_num_centroids; + std::string pq_table_bin, pq_compressed_vectors, disk_index_file, + medoids_file, centroids_file; + + if (false == this->single_index_file) { + std::string iprefix = std::string(index_prefix); + pq_table_bin = iprefix + "_pq_pivots.bin"; + pq_compressed_vectors = iprefix + "_pq_compressed.bin"; + disk_index_file = iprefix + "_disk.index"; + this->_disk_index_file = disk_index_file; + medoids_file = disk_index_file + "_medoids.bin"; + centroids_file = disk_index_file + "_centroids.bin"; + } else { + // Since incremental index which uses single file index is never + // a result of merging multiple indices, we won't have medoids + // or centroids file (or medoids/centroid data stored in disk_index_file). + pq_table_bin = pq_compressed_vectors = disk_index_file = + std::string(index_prefix); + this->_disk_index_file = disk_index_file; + } + #ifdef EXEC_ENV_OLS - get_bin_metadata(files, pq_table_bin, pq_file_num_centroids, pq_file_dim); + // This is a bit tricky. We have to read the header from the + // disk_index_file. But this is now exclusively a preserve of the + // DiskPriorityIO class. So, we need to estimate how many bytes are needed + // to store the header and read in that many using our 'standard' aligned + // file reader approach. + reader->open(disk_index_file, false, false); + // this->setup_thread_data(num_threads); + this->max_nthreads = num_threads; + + char * bytes = getHeaderBytes(); + ContentBuf buf(bytes, HEADER_SIZE); + std::basic_istream index_metadata(&buf); + #else - get_bin_metadata(pq_table_bin, pq_file_num_centroids, pq_file_dim); + std::ifstream index_metadata(disk_index_file, std::ios::binary); #endif - this->disk_index_file = std::string(disk_index_file); + size_t tags_offset = 0; + size_t pq_pivots_offset = 0; + size_t pq_vectors_offset = 0; + _u64 disk_nnodes; + _u64 disk_ndims; + size_t medoid_id_on_file; + _u64 file_frozen_id; + + if (new_index_format) { + _u32 nr, nc; + + READ_U32(index_metadata, nr); + READ_U32(index_metadata, nc); + + READ_U64(index_metadata, disk_nnodes); + READ_U64(index_metadata, disk_ndims); + + READ_U64(index_metadata, medoid_id_on_file); + READ_U64(index_metadata, max_node_len); + READ_U64(index_metadata, nnodes_per_sector); + data_dim = disk_ndims; + max_degree = + ((max_node_len - data_dim * sizeof(T)) / sizeof(unsigned)) - 1; + + diskann::cout << "Disk-Index File Meta-data: " + << "# nodes per sector: " << nnodes_per_sector + << ", max node len (bytes): " << max_node_len + << ", max node degree: " << max_degree << ", npts: " << nr + << ", dim: " << nc << " disk_nnodes: " << disk_nnodes + << " disk_ndims: " << disk_ndims << std::endl; + + READ_U64(index_metadata, this->num_frozen_points); + READ_U64(index_metadata, file_frozen_id); + if (this->num_frozen_points == 1) { + this->frozen_location = file_frozen_id; + // if (this->num_frozen_points == 1) { + diskann::cout << " Detected frozen point in index at location " + << this->frozen_location + << ". Will not output it at search time." << std::endl; + } + READ_U64(index_metadata, tags_offset); + READ_U64(index_metadata, pq_pivots_offset); + READ_U64(index_metadata, pq_vectors_offset); + + diskann::cout << "Tags offset: " << tags_offset + << " PQ Pivots offset: " << pq_pivots_offset + << " PQ Vectors offset: " << pq_vectors_offset << std::endl; + } else { // old index file format + size_t actual_index_size = get_file_size(disk_index_file); + size_t expected_file_size; + READ_U64(index_metadata, expected_file_size); + if (actual_index_size != expected_file_size) { + diskann::cout << "File size mismatch for " << disk_index_file + << " (size: " << actual_index_size << ")" + << " with meta-data size: " << expected_file_size + << std::endl; + return -1; + } - if (pq_file_num_centroids != 256) { - diskann::cout << "Error. Number of PQ centroids is not 256. Exitting." - << std::endl; - return -1; + READ_U64(index_metadata, disk_nnodes); + READ_U64(index_metadata, medoid_id_on_file); + READ_U64(index_metadata, max_node_len); + READ_U64(index_metadata, nnodes_per_sector); + max_degree = + ((max_node_len - data_dim * sizeof(T)) / sizeof(unsigned)) - 1; + + diskann::cout << "Disk-Index File Meta-data: "; + diskann::cout << "# nodes per sector: " << nnodes_per_sector; + diskann::cout << ", max node len (bytes): " << max_node_len; + diskann::cout << ", max node degree: " << max_degree << std::endl; } - this->data_dim = pq_file_dim; - this->aligned_dim = ROUND_UP(pq_file_dim, 8); +#ifdef EXEC_ENV_OLS + delete[] bytes; +#else + index_metadata.close(); +#endif + + if (this->single_index_file) { + pq_table_bin = disk_index_file; + pq_compressed_vectors = disk_index_file; + } else { + pq_pivots_offset = 0; + pq_vectors_offset = 0; + } + diskann::cout << "After single file index check, Tags offset: " + << tags_offset << " PQ Pivots offset: " << pq_pivots_offset + << " PQ Vectors offset: " << pq_vectors_offset << std::endl; size_t npts_u64, nchunks_u64; #ifdef EXEC_ENV_OLS diskann::load_bin<_u8>(files, pq_compressed_vectors, this->data, npts_u64, - nchunks_u64); + nchunks_u64, pq_vectors_offset); #else diskann::load_bin<_u8>(pq_compressed_vectors, this->data, npts_u64, - nchunks_u64); + nchunks_u64, pq_vectors_offset); #endif - this->num_points = npts_u64; this->n_chunks = nchunks_u64; -#ifdef EXEC_ENV_OLS - pq_table.load_pq_centroid_bin(files, pq_table_bin.c_str(), nchunks_u64); -#else - pq_table.load_pq_centroid_bin(pq_table_bin.c_str(), nchunks_u64); -#endif - - diskann::cout - << "Loaded PQ centroids and in-memory compressed vectors. #points: " - << num_points << " #dim: " << data_dim - << " #aligned_dim: " << aligned_dim << " #chunks: " << n_chunks - << std::endl; + diskann::cout << "Load compressed vectors from file: " + << pq_compressed_vectors << " offset: " << pq_vectors_offset + << " num points: " << npts_u64 << " n_chunks: " << nchunks_u64 + << std::endl; -// read index metadata #ifdef EXEC_ENV_OLS - // This is a bit tricky. We have to read the header from the - // disk_index_file. But this is - // now exclusively a preserve of the DiskPriorityIO class. So, we need to - // estimate how many - // bytes are needed to store the header and read in that many using our - // 'standard' aligned - // file reader approach. - reader->open(disk_index_file); - this->setup_thread_data(num_threads); - this->max_nthreads = num_threads; - - char * bytes = getHeaderBytes(); - ContentBuf buf(bytes, HEADER_SIZE); - std::basic_istream index_metadata(&buf); + pq_table.load_pq_centroid_bin(files, pq_table_bin.c_str(), nchunks_u64, + pq_pivots_offset); #else - std::ifstream index_metadata(disk_index_file, std::ios::binary); + pq_table.load_pq_centroid_bin(pq_table_bin.c_str(), nchunks_u64, + pq_pivots_offset); #endif - size_t actual_index_size = get_file_size(disk_index_file); - size_t expected_file_size; - READ_U64(index_metadata, expected_file_size); - if (actual_index_size != expected_file_size) { - diskann::cout << "File size mismatch for " << disk_index_file - << " (size: " << actual_index_size << ")" - << " with meta-data size: " << expected_file_size - << std::endl; - return -1; - } - - _u64 disk_nnodes; - READ_U64(index_metadata, disk_nnodes); if (disk_nnodes != num_points) { diskann::cout << "Mismatch in #points for compressed data file and disk " "index file: " @@ -654,32 +716,23 @@ namespace diskann { return -1; } - size_t medoid_id_on_file; - READ_U64(index_metadata, medoid_id_on_file); - READ_U64(index_metadata, max_node_len); - READ_U64(index_metadata, nnodes_per_sector); - max_degree = ((max_node_len - data_dim * sizeof(T)) / sizeof(unsigned)) - 1; + this->data_dim = pq_table.get_dim(); + this->aligned_dim = ROUND_UP(this->data_dim, 8); - diskann::cout << "Disk-Index File Meta-data: "; - diskann::cout << "# nodes per sector: " << nnodes_per_sector; - diskann::cout << ", max node len (bytes): " << max_node_len; - diskann::cout << ", max node degree: " << max_degree << std::endl; - -#ifdef EXEC_ENV_OLS - delete[] bytes; -#else - index_metadata.close(); -#endif + diskann::cout + << "Loaded PQ centroids and in-memory compressed vectors. #points: " + << num_points << " #dim: " << data_dim + << " #aligned_dim: " << aligned_dim << " #chunks: " << n_chunks + << std::endl; +// read index metadata #ifndef EXEC_ENV_OLS // open AlignedFileReader handle to index_file std::string index_fname(disk_index_file); - reader->open(index_fname); + reader->open(index_fname, false, false); +#endif this->setup_thread_data(num_threads); this->max_nthreads = num_threads; - -#endif - #ifdef EXEC_ENV_OLS if (files.fileExists(medoids_file)) { size_t tmp_dim; @@ -725,8 +778,9 @@ namespace diskann { stream << "Error loading centroids data file. Expected bin format of " "m times data_dim vector of float, where m is number of " "medoids " - "in medoids file." - << std::endl; + "in medoids file."; + + diskann::cerr << stream.str() << std::endl; throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } @@ -738,10 +792,28 @@ namespace diskann { use_medoids_data_as_centroids(); } - diskann::cout << "done.." << std::endl; + // load tags + if (this->enable_tags) { + diskann::cout << "Loading tags..."; + if (!this->single_index_file) { + std::string tag_file = disk_index_file; + tag_file = tag_file + ".tags"; + diskann::cout << "from " << tag_file << std::endl; + this->load_tags(tag_file); + } else { + diskann::cout << "from " << disk_index_file << std::endl; + this->load_tags(disk_index_file, tags_offset); + } + } + diskann::cout << "PQFlashIndex loaded successfully." << std::endl; return 0; } + template + _u64 PQFlashIndex::return_nd() { + return this->num_points; + } + #ifdef USE_BING_INFRA bool getNextCompletedRequest(const IOContext &ctx, size_t size, int &completedIndex) { @@ -760,23 +832,138 @@ namespace diskann { } #endif - template - void PQFlashIndex::cached_beam_search(const T *query1, const _u64 k_search, - const _u64 l_search, _u64 *indices, - float * distances, - const _u64 beam_width, - QueryStats * stats, - Distance *output_dist_func) { - ThreadData data = this->thread_data.pop(); - while (data.scratch.sector_scratch == nullptr) { - this->thread_data.wait_for_push_notify(); + template + size_t PQFlashIndex::cached_beam_search_vectors( + const T *query, const _u64 k_search, const _u64 l_search, TagT *res_tags, + float *distances, std::vector &res_vectors, const _u64 beam_width, + QueryStats *stats) { + // iterate to fixed point + std::vector expanded_nodes_info; + expanded_nodes_info.reserve(2 * l_search); + tsl::robin_map coord_map; + + this->disk_iterate_to_fixed_point(query, (_u32) l_search, (_u32) beam_width, + expanded_nodes_info, &coord_map, stats); + + _u32 res_count = 0; + for (uint32_t i = 0; i < l_search && res_count < k_search; i++) { + if (this->num_frozen_points == 1 && + expanded_nodes_info[i].id == this->frozen_location) + continue; + + if (distances != nullptr) { + distances[res_count] = expanded_nodes_info[i].distance; + } + + memcpy(res_vectors[res_count], coord_map[expanded_nodes_info[i].id], + data_dim * sizeof(T)); + + if (res_tags != nullptr && this->tags != nullptr) { + res_tags[res_count] = this->tags[expanded_nodes_info[i].id]; + } else if (enable_tags) { + throw diskann::ANNException("Could not find tags", -1, __FUNCSIG__, + __FILE__, __LINE__); + } + res_count++; + } + return res_count; + } + + template + size_t PQFlashIndex::cached_beam_search( + const T *query, const _u64 k_search, const _u64 l_search, TagT *res_tags, + float *distances, const _u64 beam_width, QueryStats *stats) { + // iterate to fixed point + std::vector expanded_nodes_info; + expanded_nodes_info.reserve(2 * l_search); + tsl::robin_map coord_map; + + this->disk_iterate_to_fixed_point(query, (_u32) l_search, (_u32) beam_width, + expanded_nodes_info, &coord_map, stats); + // fill in `indices`, `distances` + _u64 res_count = 0; + for (uint32_t i = 0; i < l_search && res_count < k_search; i++) { + if (this->num_frozen_points == 1 && + expanded_nodes_info[i].id == this->frozen_location) + continue; + + if (distances != nullptr) { + distances[res_count] = expanded_nodes_info[i].distance; + } + if (res_tags != nullptr && this->tags != nullptr) { + res_tags[res_count] = this->tags[expanded_nodes_info[i].id]; + } + res_count++; + } + return res_count; + } + + template + size_t PQFlashIndex::cached_beam_search_ids( + const T *query, const _u64 k_search, const _u64 l_search, + uint64_t *res_ids, float *distances, const _u64 beam_width, + QueryStats *stats) { + // iterate to fixed point + std::vector expanded_nodes_info; + expanded_nodes_info.reserve(2 * l_search); + this->disk_iterate_to_fixed_point(query, (_u32) l_search, (_u32) beam_width, + expanded_nodes_info, nullptr, stats); + + _u64 res_count = 0; + for (uint32_t i = 0; i < l_search && res_count < k_search; i++) { + if (this->num_frozen_points == 1 && + expanded_nodes_info[i].id == this->frozen_location) { + continue; + } + + res_ids[res_count] = expanded_nodes_info[i].id; + + if (distances != nullptr) { + distances[res_count] = expanded_nodes_info[i].distance; + } + + res_count++; + } + + return res_count; + } + + template + void PQFlashIndex::disk_iterate_to_fixed_point( + const T *query1, const uint32_t l_search, const uint32_t beam_width, + std::vector & expanded_nodes_info, + tsl::robin_map *coord_map, QueryStats *stats, + ThreadData * passthrough_data, + tsl::robin_set *exclude_nodes) { + // only pull from sector scratch if ThreadData not passed as arg + + ThreadData data; + if (passthrough_data == nullptr) { data = this->thread_data.pop(); + while (data.scratch.sector_scratch == nullptr) { + this->thread_data.wait_for_push_notify(); + data = this->thread_data.pop(); + } + } else { + data = *passthrough_data; } - for (uint32_t i = 0; i < this->data_dim; i++) { - data.scratch.aligned_query_float[i] = query1[i]; + if (data_is_normalized) { + // Data has been normalized. Normalize search vector too. + float norm = diskann::compute_l2_norm(query1, this->data_dim); + for (uint32_t i = 0; i < this->data_dim; i++) { + data.scratch.aligned_query_float[i] = query1[i] / norm; + } + for (uint32_t i = 0; i < this->data_dim; i++) { + data.scratch.aligned_query_T[i] = + (T) data.scratch.aligned_query_float[i]; + } + } else { + for (uint32_t i = 0; i < this->data_dim; i++) { + data.scratch.aligned_query_float[i] = query1[i]; + } + memcpy(data.scratch.aligned_query_T, query1, this->data_dim * sizeof(T)); } - memcpy(data.scratch.aligned_query_T, query1, this->data_dim * sizeof(T)); const T * query = data.scratch.aligned_query_T; const float *query_float = data.scratch.aligned_query_float; @@ -808,21 +995,23 @@ namespace diskann { _u8 * pq_coord_scratch = query_scratch->aligned_pq_coord_scratch; // lambda to batch compute query<-> node distances in PQ space - auto compute_dists = [this, pq_coord_scratch, pq_dists]( - const unsigned *ids, const _u64 n_ids, float *dists_out) { + auto compute_dists = [this, pq_coord_scratch, pq_dists](const unsigned *ids, + const _u64 n_ids, + float *dists_out) { ::aggregate_coords(ids, n_ids, this->data, this->n_chunks, pq_coord_scratch); ::pq_dist_lookup(pq_coord_scratch, n_ids, this->n_chunks, pq_dists, dists_out); }; + Timer query_timer, io_timer, cpu_timer; - std::vector retset(l_search + 1); - tsl::robin_set<_u64> visited(4096); + std::vector retset; + retset.resize(l_search + 1); + tsl::robin_set<_u64> visited(4096); - std::vector full_retset; + // re-naming `expanded_nodes_info` to not change rest of the code + std::vector &full_retset = expanded_nodes_info; full_retset.reserve(4096); - tsl::robin_map<_u64, T *> fp_coords; - _u32 best_medoid = 0; float best_dist = (std::numeric_limits::max)(); std::vector medoid_dists; @@ -852,42 +1041,26 @@ namespace diskann { unsigned k = 0; // cleared every iteration - std::vector frontier; + std::vector frontier; std::vector> frontier_nhoods; - std::vector frontier_read_reqs; + std::vector frontier_read_reqs; std::vector>> cached_nhoods; while (k < cur_list_size) { auto nk = cur_list_size; - // clear iteration state frontier.clear(); frontier_nhoods.clear(); frontier_read_reqs.clear(); cached_nhoods.clear(); sector_scratch_idx = 0; - // find new beam // WAS: _u64 marker = k - 1; _u32 marker = k; _u32 num_seen = 0; - - /* - bool marker_set = false; - diskann::cout << "hop " << hops << ": "; - for (_u32 i = 0; i < cur_list_size; i++) { - diskann::cout << retset[i].id << "( " << retset[i].distance; - if (retset[i].flag && !marker_set) { - diskann::cout << ",*) "; - marker_set = true; - } else - diskann::cout << ") "; - } - diskann::cout << std::endl; - */ while (marker < cur_list_size && frontier.size() < beam_width && - num_seen < beam_width + 2) { + num_seen < beam_width) { if (retset[marker].flag) { num_seen++; auto iter = nhood_cache.find(retset[marker].id); @@ -915,7 +1088,7 @@ namespace diskann { if (stats != nullptr) stats->n_hops++; for (_u64 i = 0; i < frontier.size(); i++) { - auto id = frontier[i]; + auto id = frontier[i]; std::pair<_u32, char *> fnhood; fnhood.first = id; fnhood.second = sector_scratch + sector_scratch_idx * SECTOR_LEN; @@ -934,22 +1107,37 @@ namespace diskann { #ifdef USE_BING_INFRA reader->read(frontier_read_reqs, ctx, true); // async reader windows. #else - reader->read(frontier_read_reqs, ctx); // synchronous IO linux + reader->read(frontier_read_reqs, ctx, false); // synchronous IO linux #endif if (stats != nullptr) { - stats->io_us += io_timer.elapsed(); + stats->io_us += (double) io_timer.elapsed(); } } // process cached nhoods for (auto &cached_nhood : cached_nhoods) { - auto global_cache_iter = coord_cache.find(cached_nhood.first); - T * node_fp_coords_copy = global_cache_iter->second; + auto global_cache_iter = coord_cache.find(cached_nhood.first); + T * node_fp_coords = global_cache_iter->second; + T * node_fp_coords_copy = data_buf + (data_buf_idx * aligned_dim); + data_buf_idx++; + memcpy(node_fp_coords_copy, node_fp_coords, data_dim * sizeof(T)); float cur_expanded_dist = dist_cmp->compare(query, node_fp_coords_copy, (unsigned) aligned_dim); - full_retset.push_back( - Neighbor((unsigned) cached_nhood.first, cur_expanded_dist, true)); - + bool exclude_cur_node = false; + if (exclude_nodes != nullptr) { + exclude_cur_node = + (exclude_nodes->find(cached_nhood.first) != exclude_nodes->end()); + } + // only figure in final list if + if (!exclude_cur_node) { + // added for StreamingMerger calls + if (coord_map != nullptr) { + coord_map->insert( + std::make_pair(cached_nhood.first, node_fp_coords_copy)); + } + full_retset.push_back( + Neighbor((unsigned) cached_nhood.first, cur_expanded_dist, true)); + } _u64 nnbrs = cached_nhood.second.first; unsigned *node_nbrs = cached_nhood.second.second; @@ -957,8 +1145,8 @@ namespace diskann { cpu_timer.reset(); compute_dists(node_nbrs, nnbrs, dist_scratch); if (stats != nullptr) { - stats->n_cmps += nnbrs; - stats->cpu_us += cpu_timer.elapsed(); + stats->n_cmps += (double) nnbrs; + stats->cpu_us += (double) cpu_timer.elapsed(); } // process prefetched nhood @@ -1015,19 +1203,30 @@ namespace diskann { T *node_fp_coords_copy = data_buf + (data_buf_idx * aligned_dim); data_buf_idx++; memcpy(node_fp_coords_copy, node_fp_coords, data_dim * sizeof(T)); - float cur_expanded_dist = dist_cmp->compare(query, node_fp_coords_copy, (unsigned) aligned_dim); - full_retset.push_back( - Neighbor(frontier_nhood.first, cur_expanded_dist, true)); - + bool exclude_cur_node = false; + if (exclude_nodes != nullptr) { + exclude_cur_node = (exclude_nodes->find(frontier_nhood.first) != + exclude_nodes->end()); + } + // if node is to be excluded from final search results + if (!exclude_cur_node) { + // added for StreamingMerger calls + if (coord_map != nullptr) { + coord_map->insert( + std::make_pair(frontier_nhood.first, node_fp_coords_copy)); + } + full_retset.push_back( + Neighbor(frontier_nhood.first, cur_expanded_dist, true)); + } unsigned *node_nbrs = (node_buf + 1); // compute node_nbrs <-> query dist in PQ space cpu_timer.reset(); compute_dists(node_nbrs, nnbrs, dist_scratch); if (stats != nullptr) { - stats->n_cmps += nnbrs; - stats->cpu_us += cpu_timer.elapsed(); + stats->n_cmps += (double) nnbrs; + stats->cpu_us += (double) cpu_timer.elapsed(); } cpu_timer.reset(); @@ -1040,9 +1239,6 @@ namespace diskann { visited.insert(id); cmps++; float dist = dist_scratch[m]; - // diskann::cout << "cmp: " << id << ", dist: " << dist << - // std::endl; - // diskann::cout << "dist: " << dist << std::endl; if (stats != nullptr) { stats->n_cmps++; } @@ -1063,7 +1259,7 @@ namespace diskann { } if (stats != nullptr) { - stats->cpu_us += cpu_timer.elapsed(); + stats->cpu_us += (double) cpu_timer.elapsed(); } } @@ -1083,25 +1279,420 @@ namespace diskann { return left.distance < right.distance; }); - // copy k_search values - for (_u64 i = 0; i < k_search; i++) { - indices[i] = full_retset[i].id; - if (distances != nullptr) { - distances[i] = full_retset[i].distance; + // return data to ConcurrentQueue only if popped from it + if (passthrough_data == nullptr) { + this->thread_data.push(data); + this->thread_data.push_notify_all(); + } + + if (stats != nullptr) { + stats->total_us = (double) query_timer.elapsed(); + } + } + + template + void PQFlashIndex::compute_pq_dists(const T *query, const _u32 *ids, + float * fp_dists, + const _u32 count) { + // TODO (perf) :: more efficient impl without using populate_chunk_distances + ThreadData data = this->thread_data.pop(); + while (data.scratch.sector_scratch == nullptr) { + this->thread_data.wait_for_push_notify(); + data = this->thread_data.pop(); + } + + T *aligned_query = data.scratch.aligned_query_T; + memcpy(aligned_query, query, this->data_dim * sizeof(T)); + auto query_scratch = &(data.scratch); + + // reset query + query_scratch->reset(); + + // query <-> PQ chunk centers distances + float *pq_dists = query_scratch->aligned_pqtable_dist_scratch; + pq_table.populate_chunk_distances(aligned_query, pq_dists); + + _u8 *pq_coord_scratch = query_scratch->aligned_pq_coord_scratch; + + // lambda to batch compute query<-> node distances in PQ space + auto compute_dists = [this, pq_coord_scratch, pq_dists](const unsigned *ids, + const _u64 n_ids, + float *dists_out) { + ::aggregate_coords(ids, n_ids, this->data, this->n_chunks, + pq_coord_scratch); + ::pq_dist_lookup(pq_coord_scratch, n_ids, this->n_chunks, pq_dists, + dists_out); + }; + compute_dists(ids, count, fp_dists); + + // return scratch + this->thread_data.push(data); + this->thread_data.push_notify_all(); + } + + template + void PQFlashIndex::compute_pq_dists(const _u32 src, const _u32 *ids, + float * fp_dists, + const _u32 count, + uint8_t * aligned_scratch) { + const _u8 * src_ptr = this->data + (this->n_chunks * src); + ThreadData data; + bool popped = false; + if (aligned_scratch == nullptr) { + assert(false); + // get buffer to store aggregated coords + data = this->thread_data.pop(); + while (data.scratch.sector_scratch == nullptr) { + this->thread_data.wait_for_push_notify(); + data = this->thread_data.pop(); } + popped = true; + + auto query_scratch = &(data.scratch); + aligned_scratch = query_scratch->aligned_pq_coord_scratch; } + // aggregate PQ coords into scratch + ::aggregate_coords(ids, count, this->data, this->n_chunks, aligned_scratch); + + // compute distances + this->pq_table.compute_distances(src_ptr, aligned_scratch, fp_dists, count); + if (popped) { + // return scratch + this->thread_data.push(data); + this->thread_data.push_notify_all(); + } + } + + template + _u32 PQFlashIndex::merge_read(std::vector> &disk_nodes, + _u32 & start_id, + const _u32 sector_count, + char * scratch) { + assert(start_id % this->nnodes_per_sector == 0); + assert(IS_ALIGNED(scratch, SECTOR_LEN)); + disk_nodes.clear(); + assert(scratch != nullptr); + + ThreadData data = this->thread_data.pop(); + while (data.scratch.sector_scratch == nullptr) { + this->thread_data.wait_for_push_notify(); + data = this->thread_data.pop(); + } + + IOContext & ctx = data.ctx; + std::vector read_req(1); + _u64 start_off = NODE_SECTOR_NO(((size_t) start_id)) * SECTOR_LEN; + _u64 n_sectors = ROUND_UP(this->num_points - start_id, nnodes_per_sector) / + nnodes_per_sector; + n_sectors = std::min(n_sectors, (uint64_t) sector_count); + assert(n_sectors > 0); + read_req[0].buf = scratch; + read_req[0].len = n_sectors * SECTOR_LEN; + read_req[0].offset = start_off; + + // big sequential read + this->reader->read(read_req, ctx); + + // create disk nodes + _u32 cur_node_id = start_id; + for (_u32 i = 0; i < n_sectors; i++) { + char *sector_buf = scratch + (i * SECTOR_LEN); + for (_u32 j = 0; j < nnodes_per_sector && cur_node_id < this->num_points; + j++) { + char *node_buf = OFFSET_TO_NODE(sector_buf, cur_node_id); + disk_nodes.emplace_back(cur_node_id, OFFSET_TO_NODE_COORDS(node_buf), + OFFSET_TO_NODE_NHOOD(node_buf)); + cur_node_id++; + } + } + + // return scratch this->thread_data.push(data); this->thread_data.push_notify_all(); - if (stats != nullptr) { - stats->total_us = (double) query_timer.elapsed(); + // return cur_node_id as starting point for next iteration + return cur_node_id; + } + + template + void PQFlashIndex::scan_deleted_nodes( + const tsl::robin_set &delete_set, + std::vector> &deleted_nodes, char *buf, char *backing_buf, + const uint32_t sectors_per_scan) { + assert(buf != nullptr); + assert(IS_ALIGNED(buf, 4096)); + assert(IS_ALIGNED(backing_buf, 32)); + + uint64_t backing_buf_unit_size = ROUND_UP(this->max_node_len, 32); + uint64_t backing_buf_idx = 0; + + // TODO (perf) :: remove this memset + memset(buf, 0, sectors_per_scan * SECTOR_LEN); + + // get ctx + ThreadData data = this->thread_data.pop(); + while (data.scratch.sector_scratch == nullptr) { + this->thread_data.wait_for_push_notify(); + data = this->thread_data.pop(); + } + + IOContext &ctx = data.ctx; + + uint32_t n_scanned = 0; + uint32_t base_offset = (uint32_t)(NODE_SECTOR_NO(0) * SECTOR_LEN); + std::vector reads(1); + reads[0].buf = buf; + reads[0].len = sectors_per_scan * SECTOR_LEN; + reads[0].offset = base_offset; + while (n_scanned < this->num_points) { + memset(buf, 0, sectors_per_scan * SECTOR_LEN); + assert(this->reader); + + this->reader->read(reads, ctx); + reads[0].offset += reads[0].len; + + // scan each sector + for (uint32_t i = 0; i < sectors_per_scan && n_scanned < this->num_points; + i++) { + char *sector_buf = buf + i * SECTOR_LEN; + // scan each node + for (uint32_t j = 0; + j < nnodes_per_sector && n_scanned < this->num_points; j++) { + char *node_buf = OFFSET_TO_NODE(sector_buf, n_scanned); + // if in delete_set, add to deleted_nodes + if (delete_set.find(n_scanned) != delete_set.end()) { + char *buf_start = + backing_buf + (backing_buf_idx * backing_buf_unit_size); + backing_buf_idx++; + memcpy(buf_start, node_buf, max_node_len); + // create disk node object from backing buf instead of `buf` + DiskNode node(n_scanned, OFFSET_TO_NODE_COORDS(buf_start), + OFFSET_TO_NODE_NHOOD(buf_start)); + /* if ((!(node.nnbrs > 0)) || + ((n_scanned >= 325000) && (n_scanned < 325100)) || + ((n_scanned >= 300000) && (n_scanned < 300100))) { + std::cout << "#neighbors of " << n_scanned << " : " + << node.nnbrs << std::endl; + std::cout << NODE_SECTOR_NO(n_scanned) << std::endl; + std::cout << uint32_t(node_buf - sector_buf) << + std::endl; for (size_t i = 0; i < 128; i++) { std::cout << + *(OFFSET_TO_NODE_COORDS(buf_start) + i) + << " ; "; + } + std::cout << std::endl; + if (((n_scanned % 5) != 0) && (node.nnbrs == 0)) { + std::cout << "Previous vector : " << std::endl; + std::cout << NODE_SECTOR_NO(n_scanned - 1) << + std::endl; char *buf_prev = OFFSET_TO_NODE(sector_buf, n_scanned + - 1); std::cout << uint32_t(buf_prev - sector_buf) << std::endl; + for (size_t i = 0; i < 128; i++) { + std::cout << *(OFFSET_TO_NODE_COORDS(buf_prev) + + i) + << " ; "; + } + std::cout << std::endl; + } + } + */ + assert(node.nnbrs < 512); + assert(node.nnbrs > 0); + deleted_nodes.push_back(node); + } + n_scanned++; + } + } + } + + // return ctx + this->thread_data.push(data); + this->thread_data.push_notify_all(); + } + + template + std::vector<_u8> PQFlashIndex::deflate_vector(const T *vec) { + std::vector<_u8> pq_coords(this->n_chunks); + std::vector fp_vec(this->data_dim); + for (uint32_t i = 0; i < this->data_dim; i++) { + fp_vec[i] = (float) vec[i]; + } + this->pq_table.deflate_vec(fp_vec.data(), pq_coords.data()); + return pq_coords; + } + + template<> + std::vector<_u8> PQFlashIndex::deflate_vector(const float *vec) { + std::vector<_u8> pq_coords(this->n_chunks); + this->pq_table.deflate_vec(vec, pq_coords.data()); + return pq_coords; + } + + template + void PQFlashIndex::reload_index( + const std::string &disk_index_file, + const std::string &pq_compressed_vectors, const std::string &tags_file) { + // reload PQ coords + size_t npts_u64, nchunks_u64; + delete this->data; + + diskann::cout << "RELOAD: Loading compressed vectors from " + << pq_compressed_vectors << "\n"; + diskann::load_bin<_u8>(pq_compressed_vectors, data, npts_u64, nchunks_u64); + + this->num_points = npts_u64; + this->n_chunks = nchunks_u64; + + // close current FP + // this->reader->deregister_all_threads(); + this->reader->close(); + + diskann::cout << "RELOAD: Loading graph from " << disk_index_file << "\n"; + // read new graph from disk + std::ifstream diskann_meta(disk_index_file, std::ios::binary); + + // size_t actual_index_size = get_file_size(disk_index_file); + // size_t expected_file_size; + // READ_U64(diskann_meta, expected_file_size); + /* if (actual_index_size != expected_file_size) { + diskann::cout << "File size mismatch for " << disk_index_file + << " (size: " << actual_index_size << ")" + << " with meta-data size: " << expected_file_size + << std::endl; + exit(-1); + } */ + + _u32 nr, nc; + READ_U32(diskann_meta, nr); + READ_U32(diskann_meta, nc); + + _u64 disk_nnodes, ndims; + READ_U64(diskann_meta, disk_nnodes); + READ_U64(diskann_meta, ndims); + + size_t medoid_id_on_file; + READ_U64(diskann_meta, medoid_id_on_file); + diskann::cout << "Medoid-ID: " << medoid_id_on_file << "\n"; + this->medoids[0] = (_u32) medoid_id_on_file; + this->num_medoids = 1; + READ_U64(diskann_meta, max_node_len); + READ_U64(diskann_meta, nnodes_per_sector); + max_degree = ((max_node_len - data_dim * sizeof(T)) / sizeof(unsigned)) - 1; + READ_U64(diskann_meta, this->num_frozen_points); + _u64 file_frozen_id; + READ_U64(diskann_meta, file_frozen_id); + if (this->num_frozen_points == 1) + this->frozen_location = file_frozen_id; + if (this->num_frozen_points == 1) { + diskann::cout << " Detected frozen point in index at location " + << this->frozen_location + << ". Will not output it at search time." << std::endl; + } + size_t tags_offset; + READ_U64(diskann_meta, tags_offset); + + diskann_meta.close(); + + // point AlignedFileReader handle to diskann_file + std::string diskann_fname(disk_index_file); + reader->open(diskann_fname, true, false); +#ifdef _WINDOWS + setup_thread_data(1); +#endif + // skip setup_thread_data() and other load() calls here + + // re-load tags + this->load_tags(tags_file); + + // number of nodes to cache + uint32_t node_cache_count = std::min((_u32) this->nhood_cache.size(), + (_u32) this->coord_cache.size()); + + // clear cache + diskann::cout << "RELOAD: Clearing cache.\n"; + delete[] this->nhood_cache_buf; + this->nhood_cache.clear(); + // buffers allocated via alloc_aligned must be freed using + // aligned_free + // TODO: Revisit why we need alloc_aligned for coord_cache_buf. + aligned_free(this->coord_cache_buf); + // delete[] this->coord_cache_buf; + this->coord_cache.clear(); + + // refresh cache + diskann::cout << "RELOAD: Caching " << node_cache_count << " nodes.\n"; + std::vector node_list; + this->cache_bfs_levels(node_cache_count, node_list); + this->load_cache_list(node_list); + } + + template + void PQFlashIndex::load_tags(const std::string &tag_file_name, + size_t offset) { + if (file_exists(tag_file_name)) { + diskann::cout << "Loading tags from " << tag_file_name << std::endl; + size_t tag_num, tag_dim; + diskann::load_bin(tag_file_name, tags, tag_num, tag_dim, offset); + } else { + diskann::cout << "Tags file " << tag_file_name << " does not exist. " + << std::endl; + } + } + + template + void PQFlashIndex::get_active_tags( + tsl::robin_set &active_tags) { + active_tags.clear(); + if (this->enable_tags) { + for (_u64 i = 0; i < num_points; i++) { + if (this->num_frozen_points > 0) { + if (i != frozen_location) + active_tags.insert(tags[i]); + } else { + active_tags.insert(tags[i]); + } + } } } + template + int PQFlashIndex::get_vector_by_tag(const TagT &tag, + T * vector_coords) { + if (!enable_tags) { + diskann::cout << "Tags are disabled, cannot retrieve vector" << std::endl; + return -1; + } + size_t pos; + bool tag_found = false; + for (size_t i = 0; i < num_points; i++) { + if (tags[i] == tag) { + pos = i; + tag_found = true; + break; + } + } + if (!tag_found) { + diskann::cout << "Tag not found." << std::endl; + return -2; + } + + size_t num_sectors = NODE_SECTOR_NO(pos); + std::ifstream disk_reader(_disk_index_file.c_str(), std::ios::binary); + std::unique_ptr sector_buf = std::make_unique(SECTOR_LEN); + + disk_reader.seekg(SECTOR_LEN * num_sectors, std::ios::beg); + + disk_reader.read(sector_buf.get(), SECTOR_LEN); + + char *node_coords = (OFFSET_TO_NODE(sector_buf.get(), pos)); + memcpy((void *) vector_coords, (void *) node_coords, data_dim * sizeof(T)); + + return 0; + } + #ifdef EXEC_ENV_OLS - template - char *PQFlashIndex::getHeaderBytes() { + template + char *PQFlashIndex::getHeaderBytes() { + this->reader->register_thread(); IOContext & ctx = reader->get_ctx(); AlignedRead readReq; readReq.buf = new char[PQFlashIndex::HEADER_SIZE]; @@ -1118,8 +1709,16 @@ namespace diskann { #endif // instantiations - template class PQFlashIndex<_u8>; - template class PQFlashIndex<_s8>; - template class PQFlashIndex; - + template class PQFlashIndex; + template class PQFlashIndex<_s8, int32_t>; + template class PQFlashIndex<_u8, int32_t>; + template class PQFlashIndex; + template class PQFlashIndex<_s8, uint32_t>; + template class PQFlashIndex<_u8, uint32_t>; + template class PQFlashIndex; + template class PQFlashIndex<_s8, int64_t>; + template class PQFlashIndex<_u8, int64_t>; + template class PQFlashIndex; + template class PQFlashIndex<_s8, uint64_t>; + template class PQFlashIndex<_u8, uint64_t>; } // namespace diskann diff --git a/src/utils.cpp b/src/utils.cpp index c3eec1d918..2b1ce301a8 100644 --- a/src/utils.cpp +++ b/src/utils.cpp @@ -49,12 +49,157 @@ bool cpuHasAvx2Support() { } return false; } + #endif -#ifndef _WINDOWS -bool AvxSupportedCPU = false; -bool Avx2SupportedCPU = true; -#else +#ifdef _WINDOWS bool AvxSupportedCPU = cpuHasAvxSupport(); bool Avx2SupportedCPU = cpuHasAvx2Support(); +#else +bool Avx2SupportedCPU = true; +bool AvxSupportedCPU = false; #endif + +namespace diskann { + // Get the right distance function for the given metric. + template<> + diskann::Distance* get_distance_function(diskann::Metric m) { + if (m == diskann::Metric::L2) { + if (Avx2SupportedCPU) { + diskann::cout << "L2: Using AVX2 distance computation" << std::endl; + return new diskann::DistanceL2(); + } else if (AvxSupportedCPU) { + diskann::cout + << "L2: AVX2 not supported. Using AVX distance computation" + << std::endl; + return new diskann::AVXDistanceL2Float(); + } else { + diskann::cout << "L2: Older CPU. Using slow distance computation" + << std::endl; + return new diskann::SlowDistanceL2Float(); + } + } else if (m == diskann::Metric::COSINE) { + diskann::cout << "Cosine: Using either AVX or AVX2 implementation" + << std::endl; + return new diskann::DistanceCosineFloat(); + } else { + std::stringstream stream; + stream << "Only L2 and cosine metric supported as of now. Email " + "gopalsr@microsoft.com if you need support for any other metric" + << std::endl; + std::cerr << stream.str() << std::endl; + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } + } + + template<> + diskann::Distance* get_distance_function(diskann::Metric m) { + if (m == diskann::Metric::L2) { + if (Avx2SupportedCPU) { + diskann::cout << "Using AVX2 distance computation" << std::endl; + return new diskann::DistanceL2Int8(); + } else if (AvxSupportedCPU) { + diskann::cout << "AVX2 not supported. Using AVX distance computation" + << std::endl; + return new diskann::AVXDistanceL2Int8(); + } else { + diskann::cout << "Older CPU. Using slow distance computation" + << std::endl; + return new diskann::SlowDistanceL2Int(); + } + } else if (m == diskann::Metric::COSINE) { + diskann::cout << "Using either AVX or AVX2 for Cosine similarity" + << std::endl; + return new diskann::DistanceCosineInt8(); + } else { + std::stringstream stream; + stream << "Only L2 and cosine metric supported as of now. Email " + "gopalsr@microsoft.com if you need support for any other metric" + << std::endl; + std::cerr << stream.str() << std::endl; + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } + } + + template<> + diskann::Distance* get_distance_function(diskann::Metric m) { + if (m == diskann::Metric::L2) { + diskann::cout + << "AVX/AVX2 distance function not defined for Uint8. Using " + "slow version. " + "Contact gopalsr@microsoft.com if you need AVX/AVX2 support." + << std::endl; + return new diskann::DistanceL2UInt8(); + } else if (m == diskann::Metric::COSINE) { + diskann::cout + << "AVX/AVX2 distance function not defined for Uint8. Using " + "slow version. " + "Contact gopalsr@microsoft.com if you need AVX/AVX2 support." + << std::endl; + return new diskann::SlowDistanceCosineUInt8(); + } else { + std::stringstream stream; + stream << "Only L2 and Cosine metric supported as of now. Email " + "gopalsr@microsoft.com if you need any support for any other " + "metric" + << std::endl; + std::cerr << stream.str() << std::endl; + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } + } + + void block_convert(std::ofstream& writr, std::ifstream& readr, + float* read_buf, _u64 npts, _u64 ndims) { + readr.read((char*) read_buf, npts * ndims * sizeof(float)); + _u32 ndims_u32 = (_u32) ndims; +#pragma omp parallel for + for (_s64 i = 0; i < (_s64) npts; i++) { + float norm_pt = std::numeric_limits::epsilon(); + for (_u32 dim = 0; dim < ndims_u32; dim++) { + norm_pt += + *(read_buf + i * ndims + dim) * *(read_buf + i * ndims + dim); + } + norm_pt = std::sqrt(norm_pt); + for (_u32 dim = 0; dim < ndims_u32; dim++) { + *(read_buf + i * ndims + dim) = *(read_buf + i * ndims + dim) / norm_pt; + } + } + writr.write((char*) read_buf, npts * ndims * sizeof(float)); + } + + void normalize_data_file(const std::string& inFileName, + const std::string& outFileName) { + std::ifstream readr(inFileName, std::ios::binary); + std::ofstream writr(outFileName, std::ios::binary); + + int npts_s32, ndims_s32; + readr.read((char*) &npts_s32, sizeof(_s32)); + readr.read((char*) &ndims_s32, sizeof(_s32)); + + writr.write((char*) &npts_s32, sizeof(_s32)); + writr.write((char*) &ndims_s32, sizeof(_s32)); + + _u64 npts = (_u64) npts_s32, ndims = (_u64) ndims_s32; + diskann::cout << "Normalizing FLOAT vectors in file: " << inFileName + << std::endl; + diskann::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims + << std::endl; + + _u64 blk_size = 131072; + _u64 nblks = ROUND_UP(npts, blk_size) / blk_size; + diskann::cout << "# blks: " << nblks << std::endl; + + float* read_buf = new float[npts * ndims]; + for (_u64 i = 0; i < nblks; i++) { + _u64 cblk_size = std::min(npts - i * blk_size, blk_size); + block_convert(writr, readr, read_buf, cblk_size, ndims); + } + delete[] read_buf; + + diskann::cout << "Wrote normalized points to file: " << outFileName + << std::endl; + } +} // namespace diskann diff --git a/src/v2/delete_set.cpp b/src/v2/delete_set.cpp new file mode 100644 index 0000000000..a7a8247a1d --- /dev/null +++ b/src/v2/delete_set.cpp @@ -0,0 +1,103 @@ +#include "v2/delete_set.h" +#include "tsl/robin_map.h" +#include "tsl/robin_set.h" +#include +#include +#include +#include +#include + +#include "logger.h" + +namespace diskann { + DeleteSet::DeleteSet(uint32_t max_merges) : max_merges(max_merges) { + this->primary = new tsl::robin_map(); + this->secondary = new tsl::robin_map(); + } + + DeleteSet::~DeleteSet() { + std::lock_guard lk(this->lock); + // assert no deleted entries remaining to reconcile + assert(this->primary->empty()); + assert(this->secondary->empty()); + // free both primary & secondary delete lists + delete this->primary; + delete this->secondary; + lk.~lock_guard(); + } + + void DeleteSet::add_delete(uint32_t id) { + std::lock_guard l(this->lock); + // check if already in delete list + bool deleted = (this->primary->find(id) != this->primary->end()) || (this->secondary->find(id) != this->secondary->end()); + if (!deleted) { + this->primary->insert(std::make_pair(id, 0)); + } + } + + bool DeleteSet::is_dead(uint32_t id) { + std::lock_guard l(this->lock); + // fast short-circuit if empty + if(this->secondary->empty()) { + return this->primary->find(id) != this->primary->end(); + } else { + return (this->primary->find(id) != this->primary->end()) || (this->secondary->find(id) != this->secondary->end()); + } + } + + void DeleteSet::batch_is_dead(const uint32_t *ids, bool *dead, + const uint32_t count) { + std::lock_guard l(this->lock); + if (this->secondary->empty()) { + for (uint32_t i = 0; i < count; i++) { + // fast short-circuit if empty + dead[i] = this->primary->find(ids[i]) != this->primary->end(); + } + } else { + for (uint32_t i = 0; i < count; i++) { + // fast short-circuit if empty + dead[i] = (this->primary->find(ids[i]) != this->primary->end()) || + (this->secondary->find(ids[i]) != this->secondary->end()); + } + return; + } + } + + void DeleteSet::merge_start() { + std::lock_guard l(this->lock); + assert(this->secondary->empty()); + std::swap(this->primary, this->secondary); + } + + std::vector DeleteSet::track_merge() { + std::lock_guard l(this->lock); + // increment counts for secondary + for (auto &k_v : *this->secondary) { + this->secondary->operator[](k_v.first)++; + } + + // swap primary & secondary + std::swap(this->primary, this->secondary); + + // merge secondary into primary + for(auto &k_v : *this->secondary) { + this->primary->insert(k_v); + } + + // clear duplicates + this->secondary->clear(); + + // get reclaimed IDs + std::vector reclaimed; + for(auto &k_v : *this->primary) { + if (k_v.second >= this->max_merges) { + reclaimed.push_back(k_v.first); + } + } + for(auto id : reclaimed) { + this->primary->erase(id); + } + + return reclaimed; + } +} // namespace diskann diff --git a/src/v2/fs_allocator.cpp b/src/v2/fs_allocator.cpp new file mode 100644 index 0000000000..4c2001b654 --- /dev/null +++ b/src/v2/fs_allocator.cpp @@ -0,0 +1,65 @@ +#pragma once + +#include "v2/fs_allocator.h" +#include "tsl/robin_set.h" +#include +#include +#include +#include + +#include "utils.h" +#include "logger.h" + +namespace diskann { + + template + FixedSizeAlignedAllocator::FixedSizeAlignedAllocator(const uint32_t ndims, const uint32_t max_count) { + assert(IS_ALIGNED(ndims, 32)); + this->count = ROUND_UP(max_count, 32); + alloc_aligned( (void**)&this->buf, this->count * sizeof(T), 32); + std::vector ids(this->count); + std::iota(ids.begin(), ids.end(), 0); + this->free_set.insert(ids.begin(), ids.end()); + ids.clear(); + } + + template + FixedSizeAlignedAllocator::~FixedSizeAlignedAllocator() { + std::lock_guard lk(this->lock); + assert(this->free_set.size() == this->count); + aligned_free(this->buf); + } + + template + T* FixedSizeAlignedAllocator::allocate(){ + std::lock_guard lk(this->lock); + uint32_t id = std::numeric_limits::max(); + for(auto &v : this->free_set) { + id = v; + break; + } + + if(id == std::numeric_limits::max()) { + std::cerr << "UNABLE TO ALLOCATE MEMORY" << std::endl; + return nullptr; + } else{ + this->free_set.erase(id); + } + return this->buf + (id * ndims); + } + + template + void FixedSizeAlignedAllocator::deallocate(T* ptr) { + assert(IS_ALIGNED(ptr, 32)); + uint32_t id = (uint32_t) (ptr - this->buf) / ndims; + std::lock_guard lk(this->lock); + this->free_set.insert(id); + } + + // vectors + template class FixedSizeAlignedAllocator; + template class FixedSizeAlignedAllocator; + template class FixedSizeAlignedAllocator; + // nhoods + template class FixedSizeAlignedAllocator; +} // namespace diskann diff --git a/src/v2/graph_delta.cpp b/src/v2/graph_delta.cpp new file mode 100644 index 0000000000..07ca956f3a --- /dev/null +++ b/src/v2/graph_delta.cpp @@ -0,0 +1,105 @@ +#include "v2/graph_delta.h" +#include +#include +#include + +#include "utils.h" +#include "logger.h" + +namespace diskann { + GraphDelta::GraphDelta(const uint32_t offset, const uint32_t max_nodes) : offset(offset), max_nodes(max_nodes) { + diskann::cout << "GraphDelta: Allocating " << max_nodes << " deltas\n"; + this->graph.resize(max_nodes); + for (size_t i = 0; i < this->graph.size();i++) { + this->graph[i].shrink_to_fit(); + } + this->locks = std::make_unique(max_nodes); + } + + bool GraphDelta::is_relevant(const uint32_t id) { + return (id < offset + max_nodes && id >= offset); + } + + void GraphDelta::insert_vector(const uint32_t id, const uint32_t*nhood, const uint32_t nnbrs) { + if (!this->is_relevant(id)) { + return; + } + + assert(nhood != nullptr); + + uint32_t local_id = id - offset; + // check valid id + assert(local_id < this->graph.size()); + + // acquire lock + std::lock_guard lock(this->locks[local_id]); + + // copy nhood into graph[local_id], clear existing nbrs + this->graph[local_id].clear(); + this->graph[local_id].insert(this->graph[local_id].end(), nhood, nhood + nnbrs); + this->graph[local_id].shrink_to_fit(); + } + + void GraphDelta::inter_insert(const uint32_t dest, const uint32_t* srcs, const uint32_t src_count){ + assert(srcs != nullptr); + for(uint32_t i=0;i lock(this->locks[local_src_id]); + // add src->dest edge + this->graph[local_src_id].push_back(dest); + this->graph[local_src_id].shrink_to_fit(); + } + } + + const std::vector GraphDelta::get_nhood(const uint32_t id) { + if (!is_relevant(id)) { + return std::vector(); + } + uint32_t local_id = id - offset; + // acquire lock + std::lock_guard lock(this->locks[local_id]); + // add src->dest edge + return this->graph[local_id]; + } + + void GraphDelta::rename_edges(const tsl::robin_map& rename_map) { + std::atomic count; + count.store(0); +#pragma omp parallel for schedule(dynamic, 128) + for(int64_t i=0;i < (int64_t)this->graph.size(); i++) { + std::vector &delta = this->graph[ (uint32_t)i]; + count += delta.size(); + for(uint32_t j=0; j < delta.size(); j++) { + auto iter = rename_map.find(delta[j]); + if (iter != rename_map.end()) { + delta[j] = iter->second; + } + } + } + diskann::cout << "Renamed "<< count.load() << " edges.\n"; + } + + void GraphDelta::rename_edges(const std::function &rename_func) { + std::atomic count; + count.store(0); + #pragma omp parallel for schedule(dynamic, 128) + for(int64_t i=0;i < (int64_t)this->graph.size(); i++) { + std::vector &delta = this->graph[ (uint32_t)i]; + count += delta.size(); + for(uint32_t j=0; j < delta.size(); j++) { + uint32_t renamed_id = rename_func(delta[j]); + if (renamed_id != std::numeric_limits::max()) { + delta[j] = renamed_id; + } + } + } + diskann::cout << "Renamed "<< count.load() << " edges.\n"; + } +}; diff --git a/src/v2/index_merger.cpp b/src/v2/index_merger.cpp new file mode 100644 index 0000000000..ebe9f33de4 --- /dev/null +++ b/src/v2/index_merger.cpp @@ -0,0 +1,1647 @@ +#include "neighbor.h" +#include "timer.h" +#include "tsl/robin_map.h" +#include "tsl/robin_set.h" +#include "utils.h" +#include "v2/index_merger.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "tcmalloc/malloc_extension.h" +#include +#include +#include "logger.h" +#include "ann_exception.h" + +#define SECTORS_PER_MERGE (uint64_t) 65536 +// max number of points per mem index being merged -- 32M +#define MAX_PTS_PER_MEM_INDEX (uint64_t)(1 << 25) +#define INDEX_OFFSET (uint64_t)(MAX_PTS_PER_MEM_INDEX * 4) +#define MAX_INSERT_THREADS (uint64_t) 18 +#define MAX_N_THREADS (uint64_t) 18 +#define NUM_INDEX_LOAD_THREADS (uint64_t) 18 +#define PER_THREAD_BUF_SIZE (uint64_t)(65536 * 64 * 4) + +#define PQ_FLASH_INDEX_MAX_NODES_TO_CACHE 200000 + +namespace diskann { + template + StreamingMerger::StreamingMerger( + const uint32_t ndims, Distance *dist, diskann::Metric dist_metric, const uint32_t beam_width, + const uint32_t range, const uint32_t l_index, const float alpha, + const uint32_t maxc, bool single_file_index) { + // book keeping + this->ndims = ndims; + this->aligned_ndims = (_u32) ROUND_UP(this->ndims, 8); + this->range = range; + this->l_index = l_index; + this->beam_width = beam_width; + this->maxc = maxc; + this->alpha = alpha; + this->dist_metric = dist_metric; + this->dist_cmp = dist; + this->_single_file_index = single_file_index; + + std::cout << "StreamingMerger created with R=" << this->range + << " L=" << this->l_index << " BW=" << this->beam_width + << " MaxC=" << this->maxc << " alpha=" << this->alpha + << " ndims: " << this->ndims << std::endl; + } + + template + StreamingMerger::~StreamingMerger() { + // release scratch alloc memory + // delete this->fp_alloc; + // delete this->pq_alloc; + + if (this->disk_index != nullptr) + delete this->disk_index; + + delete this->disk_delta; + + for (auto &delta : this->mem_deltas) { + delete delta; + } + aligned_free((void *) this->thread_pq_scratch); + for (auto &data : this->mem_data) { + //delete[] data; + aligned_free((void *)data); + } + } + + template + void StreamingMerger::process_inserts_pq() { + Timer total_insert_timer; + this->insert_times.resize(MAX_N_THREADS, 0.0); + this->delta_times.resize(MAX_N_THREADS, 0.0); + // iterate through each vector in each mem index + for (uint32_t i = 0; i < this->mem_data.size(); i++) { + diskann::cout << "Processing pq of inserts from mem-DiskANN #" << i + 1 + << "\n"; + const tsl::robin_set &deleted_set = this->mem_deleted_ids[i]; + const T* coords = this->mem_data[i]; + const uint32_t offset = this->offset_ids[i]; + const uint32_t count = this->mem_npts[i]; +// TODO (perf) :: trivially parallelizes ?? +#pragma omp parallel for schedule(dynamic, 1) num_threads(MAX_N_THREADS) + // iteratively insert each point into full index + for (int32_t j = 0; j < (int32_t) count; j++) { + // filter out -- `j` is deleted + if (deleted_set.find((uint32_t) j) != deleted_set.end()) { + continue; + } + + // data for jth point + const T *j_coords = + coords + ((uint64_t)(this->aligned_ndims) * (uint64_t) j); + const uint32_t j_id = offset + (uint32_t) j; + + // get renamed ID + const uint32_t j_renamed = this->rename(j_id); + assert(j_renamed != std::numeric_limits::max()); + + // compute PQ coords + std::vector j_pq_coords = + this->disk_index->deflate_vector(j_coords); + // std::vector j_pq_coords(this->pq_nchunks,0); + + // directly copy into PQFlashIndex PQ data + const uint64_t j_pq_offset = + (uint64_t) j_renamed * (uint64_t) this->pq_nchunks; + memcpy(this->pq_data + j_pq_offset, j_pq_coords.data(), + this->pq_nchunks * sizeof(uint8_t)); + } + } + + diskann::cout << "Finished deflating all points\n"; + double e2e_time = ((double) total_insert_timer.elapsed()) / (1000000.0); + double insert_time = std::accumulate(this->insert_times.begin(), + this->insert_times.end(), 0.0); + double delta_time = std::accumulate(this->delta_times.begin(), + this->delta_times.end(), 0.0); + uint32_t n_inserts = + std::accumulate(this->mem_npts.begin(), this->mem_npts.end(), 0); + diskann::cout << "TIMER:: PQ time per point = " << insert_time / n_inserts + << ", Delta = " << delta_time / n_inserts << "\n"; + diskann::cout << " E2E pq time: " << e2e_time << " sec" << std::endl; + } + + template + void StreamingMerger::process_inserts() { + Timer total_insert_timer; + this->insert_times.resize(MAX_INSERT_THREADS, 0.0); + this->delta_times.resize(MAX_INSERT_THREADS, 0.0); + // iterate through each vector in each mem index + for (uint32_t i = 0; i < this->mem_data.size(); i++) { + diskann::cout << "Processing inserts from mem-DiskANN #" << i + 1 << "\n"; + const tsl::robin_set &deleted_set = this->mem_deleted_ids[i]; + const T* coords = this->mem_data[i]; + const uint32_t offset = this->offset_ids[i]; + const uint32_t count = this->mem_npts[i]; + + size_t cur_cache_size = 0; +#ifdef USE_TCMALLOC + MallocExtension::instance()->GetNumericProperty( + "tcmalloc.max_total_thread_cache_bytes", &cur_cache_size); + //diskann::cout << "Current cache size : " << (cur_cache_size >> 10) + // << " KiB\n" + // << std::endl; + MallocExtension::instance()->SetNumericProperty( + "tcmalloc.max_total_thread_cache_bytes", 128 * 1024 * 1024); + +#endif + diskann::Timer timer; +#pragma omp parallel for schedule(dynamic, 1) num_threads(MAX_INSERT_THREADS) + // iteratively insert each point into full index + for (int32_t j = 0; j < (int32_t) count; j++) { + // filter out -- `j` is deleted + if (deleted_set.find((uint32_t) j) != deleted_set.end()) { + continue; + } + + if (((j % 100000) == 0) && (j > 0)) { + diskann::cout << "Finished inserting " << j << " points" << std::endl; + std::cout << "When j = " << j + << " elapsed time: " << timer.elapsed() / 1000000 << "s" + << std::endl; + } + // data for jth point + const T *j_coords = + coords + ((uint64_t)(this->aligned_ndims) * (uint64_t) j); + const uint32_t j_id = offset + (uint32_t) j; + + // insert into index + this->insert_mem_vec(j_coords, j_id); + } + } + + diskann::cout << "Finished inserting all points\n"; + double e2e_time = ((double) total_insert_timer.elapsed()) / (1000000.0); + double insert_time = std::accumulate(this->insert_times.begin(), + this->insert_times.end(), 0.0); + double delta_time = std::accumulate(this->delta_times.begin(), + this->delta_times.end(), 0.0); + uint32_t n_inserts = + std::accumulate(this->mem_npts.begin(), this->mem_npts.end(), 0); + diskann::cout << "TIMER:: Insert time per point = " + << insert_time / n_inserts + << ", Delta = " << delta_time / n_inserts << "\n"; + diskann::cout << " E2E insert time: " << e2e_time << " sec" << std::endl; + } + + template + void StreamingMerger::insert_mem_vec(const T * mem_vec, + const uint32_t offset_id) { + Timer timer; + float insert_time, delta_time; + // START: mem_vec has no ID, no presence in system + std::vector pool; + std::vector tmp; + tsl::robin_map coord_map; + + // std::cout << "TID: " << std::this_thread::get_id() + // << " Before offset_iterate_to_Fixed_point()" << std::endl; + // search on combined graph + this->offset_iterate_to_fixed_point(mem_vec, this->l_index, pool, + coord_map); + insert_time = (float) timer.elapsed(); + + // prune neighbors using alpha + std::vector new_nhood; + prune_neighbors(coord_map, pool, new_nhood); + + if (new_nhood.size() > range) { + std::cout << "***ERROR*** After prune, for offset_id: " << offset_id << " found " + << new_nhood.size() << " neighbors instead of range: " << range + << std::endl; + + } + + //this->disk_delta->insert_vector(offset_id, new_nhood.data(), + // (_u32) new_nhood.size()); + this->disk_delta->inter_insert(offset_id, new_nhood.data(), + (_u32) new_nhood.size()); + + // insert into graph + for (auto &delta : this->mem_deltas) { + delta->insert_vector(offset_id, new_nhood.data(), + (_u32) new_nhood.size()); + //delta->inter_insert(offset_id, new_nhood.data(), (_u32) new_nhood.size()); + } + delta_time = (float) timer.elapsed(); + // END: mem_vec now connected with new ID + uint32_t thread_no = omp_get_thread_num(); + this->insert_times[thread_no] += insert_time; + this->delta_times[thread_no] += delta_time; + // std::cout << "TID: " << std::this_thread::get_id() + // << " Exiting insert_mem_vec() " << std::endl; + } + + template + void StreamingMerger::offset_iterate_to_fixed_point( + const T *vec, const uint32_t Lsize, + std::vector & expanded_nodes_info, + tsl::robin_map &coord_map) { + std::vector exp_node_info; + exp_node_info.reserve(2 * Lsize); + tsl::robin_map cmap; + // first hit PQ iterate to fixed point + // NOTE :: handling deletes for disk-index inside this call + // this->disk_iterate_to_fixed_point(vec, this->l_index, exp_node_info, + // exp_node_id, best_l_nodes, cmap); + uint32_t omp_thread_no = omp_get_thread_num(); + if (this->disk_thread_data.size() <= omp_thread_no) { + throw ANNException(std::string("Found ") + std::to_string(omp_thread_no) + + " thread when only " + + std::to_string(this->disk_thread_data.size()) + " were expected", + -1); + } + ThreadData &thread_data = this->disk_thread_data[omp_thread_no]; + // ThreadData * thread_data = nullptr; + + cmap.reserve(2 * Lsize); + this->disk_index->disk_iterate_to_fixed_point( + vec, Lsize, this->beam_width, exp_node_info, &cmap, nullptr, + &thread_data, &this->disk_deleted_ids); + + // reduce and pick top maxc expanded nodes only + std::sort(exp_node_info.begin(), exp_node_info.end()); + // expanded_nodes_info.clear(); + expanded_nodes_info.reserve(this->maxc); + expanded_nodes_info.insert(expanded_nodes_info.end(), exp_node_info.begin(), + exp_node_info.end()); + + // insert only relevant coords into coord_map + for (auto &nbr : expanded_nodes_info) { + uint32_t id = nbr.id; + auto iter = cmap.find(id); + assert(iter != cmap.end()); + coord_map.insert(std::make_pair(iter->first, iter->second)); + } + } + + template + void StreamingMerger::prune_neighbors( + const tsl::robin_map &coord_map, + std::vector &pool, std::vector &pruned_list) { + if (pool.size() == 0) + return; + + // sort the pool based on distance to query + std::sort(pool.begin(), pool.end()); + + std::vector result; + result.reserve(range); + std::vector occlude_factor(pool.size(), 0); + + occlude_list(pool, coord_map, result, occlude_factor); + + pruned_list.clear(); + assert(result.size() <= range); + for (auto iter : result) { + pruned_list.emplace_back(iter.id); + } + + if (alpha > 1) { + for (uint32_t i = 0; i < pool.size() && pruned_list.size() < range; i++) { + if (std::find(pruned_list.begin(), pruned_list.end(), pool[i].id) == + pruned_list.end()) + pruned_list.emplace_back(pool[i].id); + } + } + } + + template + void StreamingMerger::prune_neighbors_pq( + std::vector &pool, std::vector &pruned_list, + uint8_t *scratch) { + if (pool.size() == 0) + return; + + // sort the pool based on distance to query + std::sort(pool.begin(), pool.end()); + + std::vector result; + result.reserve(this->range); + std::vector occlude_factor(pool.size(), 0); + + occlude_list_pq(pool, result, occlude_factor, scratch); + + pruned_list.clear(); + assert(result.size() <= range); + for (auto iter : result) { + pruned_list.emplace_back(iter.id); + } + + if (alpha > 1) { + for (uint32_t i = 0; i < pool.size() && pruned_list.size() < range; i++) { + if (std::find(pruned_list.begin(), pruned_list.end(), pool[i].id) == + pruned_list.end()) + pruned_list.emplace_back(pool[i].id); + } + } + } + + template + void StreamingMerger::occlude_list( + std::vector & pool, + const tsl::robin_map &coord_map, + std::vector &result, std::vector &occlude_factor) { + if (pool.empty()) + return; + assert(std::is_sorted(pool.begin(), pool.end())); + assert(!pool.empty()); + + float cur_alpha = 1; + while (cur_alpha <= alpha && result.size() < range) { + uint32_t start = 0; + while (result.size() < range && (start) < pool.size() && start < maxc) { + auto &p = pool[start]; + if (occlude_factor[start] > cur_alpha) { + start++; + continue; + } + occlude_factor[start] = std::numeric_limits::max(); + result.push_back(p); + for (uint32_t t = start + 1; t < pool.size() && t < maxc; t++) { + if (occlude_factor[t] > alpha) + continue; + auto iter_right = coord_map.find(p.id); + auto iter_left = coord_map.find(pool[t].id); + // HAS to be in coord_map since it was expanded during + // iterate_to_fixed_point + assert(iter_right != coord_map.end()); + assert(iter_left != coord_map.end()); + // WARNING :: correct, but not fast -- NO SIMD version if using MSVC, + // g++ should auto vectorize + float djk = this->dist_cmp->compare(iter_left->second, + iter_right->second, this->ndims); + occlude_factor[t] = + (std::max)(occlude_factor[t], pool[t].distance / djk); + } + start++; + } + cur_alpha *= 1.2f; + } + } + + template + void StreamingMerger::occlude_list_pq( + std::vector &pool, std::vector &result, + std::vector &occlude_factor, uint8_t *scratch) { + if (pool.empty()) + return; + assert(std::is_sorted(pool.begin(), pool.end())); + assert(!pool.empty()); + + float cur_alpha = 1; + while (cur_alpha <= alpha && result.size() < range) { + uint32_t start = 0; + while (result.size() < range && (start) < pool.size() && start < maxc) { + auto &p = pool[start]; + if (occlude_factor[start] > cur_alpha) { + start++; + continue; + } + occlude_factor[start] = std::numeric_limits::max(); + result.push_back(p); + for (uint32_t t = start + 1; t < pool.size() && t < maxc; t++) { + if (occlude_factor[t] > alpha) + continue; + // djk = dist(p.id, pool[t.id]) + float djk; + this->disk_index->compute_pq_dists(p.id, &(pool[t].id), &djk, 1, + scratch); + occlude_factor[t] = + (std::max)(occlude_factor[t], pool[t].distance / djk); + } + start++; + } + cur_alpha *= 1.2f; + } + } + + template + void StreamingMerger::dump_to_disk(const uint32_t start_id, + const char * buf, + const uint32_t n_sectors, + std::ofstream & output_writer) { + assert(start_id % this->nnodes_per_sector == 0); + uint32_t start_sector = (start_id / this->nnodes_per_sector) + 1; + uint64_t start_off = start_sector * (uint64_t) SECTOR_LEN; + + // seek fp + output_writer.seekp(start_off, std::ios::beg); + + // dump + output_writer.write(buf, (uint64_t) n_sectors * (uint64_t) SECTOR_LEN); + + uint64_t nb_written = + (uint64_t) output_writer.tellp() - (uint64_t) start_off; + if (nb_written != (uint64_t) n_sectors * (uint64_t) SECTOR_LEN) { + std::stringstream sstream; + sstream << "ERROR!!! Wrote " << nb_written << " bytes to disk instead of " + << ((uint64_t) n_sectors) * SECTOR_LEN; + diskann::cerr << sstream.str() << std::endl; + throw diskann::ANNException(sstream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } + } + + template + void StreamingMerger::compute_deleted_ids() { + // process disk deleted tags + for (uint32_t i = 0; i < this->disk_npts; i++) { + TagT i_tag = this->disk_tags[i]; + if (this->deleted_tags.find(i_tag) != this->deleted_tags.end()) { + this->disk_deleted_ids.insert(i); + } + } + diskann::cout << "Found " << this->disk_deleted_ids.size() + << " tags to delete from SSD-DiskANN\n"; + +// this->mem_deleted_ids.resize(this->mem_data.size()); + for (uint32_t i = 0; i < this->mem_data.size(); i++) { + tsl::robin_set &deleted_ids = this->mem_deleted_ids[i]; + for (uint32_t id = 0; id < this->mem_npts[i]; id++) { + if(deleted_ids.find(id) != deleted_ids.end()) + continue; + const TagT tag = this->mem_tags[i][id]; + //if (this->deleted_tags.find(tag) != this->deleted_tags.end()) { + // deleted_ids.insert(id); + //} + if (this->latter_deleted_tags[i].find(tag) != + this->latter_deleted_tags[i].end()) { + deleted_ids.insert(id); + } + } + diskann::cout << "Found " << deleted_ids.size() + << " tags to delete from mem-DiskANN #" << i + 1 << "\n"; + } + } + + template + void StreamingMerger::process_deletes() { + // buf to hold data being read + char *buf = nullptr; + alloc_aligned((void **) &buf, SECTORS_PER_MERGE * SECTOR_LEN, SECTOR_LEN); + + // open output file for writing + // Was: this->disk_index_out_path + "_disk.index"; + + diskann::cout << "Writing delete consolidated graph to " + << this->temp_disk_index_path << std::endl; + std::ofstream output_writer(this->temp_disk_index_path, + std::ios::out | std::ios::binary); + assert(output_writer.is_open()); + // skip writing header for now + // this->output_writer.seekp(SECTOR_LEN, std::ios::beg); + std::unique_ptr sector_buf = std::make_unique(SECTOR_LEN); + output_writer.write(sector_buf.get(), SECTOR_LEN); + + Timer delete_timer; + // batch consolidate deletes + std::vector> disk_nodes; + std::vector>> id_nhoods; + uint32_t start_id = 0, new_start_id; + diskann::cout << "Consolidating deletes\n"; + + while (start_id < this->disk_npts) { + new_start_id = this->disk_index->merge_read(disk_nodes, start_id, + SECTORS_PER_MERGE, buf); +#pragma omp parallel for schedule(dynamic, 128) num_threads(MAX_N_THREADS) + for (int64_t i = 0; i < (int64_t) disk_nodes.size(); i++) { + // get thread-specific scratch + int omp_thread_no = omp_get_thread_num(); + uint8_t *pq_coord_scratch = this->thread_bufs[omp_thread_no]; + assert(pq_coord_scratch != nullptr); + DiskNode &disk_node = disk_nodes[i]; + this->consolidate_deletes(disk_node, pq_coord_scratch); + } + for (auto &disk_node : disk_nodes) { + if (this->is_deleted(disk_node)) { + this->free_ids.push_back(disk_node.id); + } + } + + uint64_t prev_pos = output_writer.tellp(); + this->dump_to_disk(start_id, buf, SECTORS_PER_MERGE, output_writer); + output_writer.flush(); + uint64_t cur_pos = output_writer.tellp(); + if (!(cur_pos - prev_pos == (SECTORS_PER_MERGE * SECTOR_LEN))) + diskann::cout << "ERROR! During dumping to disk" << std::endl; + + // advance to next block + disk_nodes.clear(); + id_nhoods.clear(); + diskann::cout << new_start_id << " / " << this->disk_npts + << " nodes processed.\n"; + start_id = new_start_id; + + } + double e2e_time = ((double) delete_timer.elapsed()) / (1000000.0); + diskann::cout << "Processed Deletes in " << e2e_time << " s." << std::endl; + diskann::cout << "Writing header.\n"; + + // write header + output_writer.seekp(0, std::ios::beg); + // HEADER --> [_u32 #metadata items][_u32 1][_u64 nnodes][_u64 ndims][_u64 + // medoid ID] + // [_u64 max_node_len][_u64 nnodes_per_sector][_u64 #frozen points in vamana + // index][_u64 frozen point location][_u64 file size] + uint64_t file_size = + SECTOR_LEN + (ROUND_UP(ROUND_UP(this->disk_npts, nnodes_per_sector) / + nnodes_per_sector, + SECTORS_PER_MERGE)) * + (uint64_t) SECTOR_LEN; + std::vector output_metadata; + output_metadata.push_back(this->disk_npts); + output_metadata.push_back((uint64_t) this->ndims); + // determine medoid + uint64_t medoid = this->init_ids[0]; + // TODO (correct?, misc) :: better way of selecting new medoid + while (this->disk_deleted_ids.find((_u32) medoid) != + this->disk_deleted_ids.end()) { + diskann::cout << "Medoid deleted. Choosing another start node.\n"; + auto iter = this->disk_deleted_nhoods.find((_u32) medoid); + assert(iter != this->disk_deleted_nhoods.end()); + medoid = iter->second[0]; + } + output_metadata.push_back((uint64_t) medoid); + uint64_t max_node_len = (this->ndims * sizeof(T)) + sizeof(uint32_t) + + (this->range * sizeof(uint32_t)); + uint64_t nnodes_per_sector = SECTOR_LEN / max_node_len; + output_metadata.push_back(max_node_len); + output_metadata.push_back(nnodes_per_sector); + output_metadata.push_back(this->disk_index_num_frozen); + output_metadata.push_back(this->disk_index_frozen_loc); + output_metadata.push_back(file_size); + + // close index + output_writer.close(); + diskann::save_bin<_u64>(this->temp_disk_index_path, output_metadata.data(), + output_metadata.size(), 1, 0); + // free buf + aligned_free((void *) buf); + + // free backing buf for deletes + aligned_free((void *) this->delete_backing_buf); + } + + template + void StreamingMerger::populate_deleted_nhoods() { + // buf for scratch + char *buf = nullptr; + alloc_aligned((void **) &buf, SECTORS_PER_MERGE * SECTOR_LEN, SECTOR_LEN); + + // scan deleted nodes and get + std::vector> deleted_nodes; + uint64_t backing_buf_size = (uint64_t) this->disk_deleted_ids.size() * + ROUND_UP(this->max_node_len, 32); + backing_buf_size = ROUND_UP(backing_buf_size, 256); + diskann::cout << "ALLOC: " << (backing_buf_size << 10) + << "KiB aligned buffer for deletes.\n"; + alloc_aligned((void **) &this->delete_backing_buf, backing_buf_size, 256); + memset(this->delete_backing_buf, 0, backing_buf_size); + this->disk_index->scan_deleted_nodes(this->disk_deleted_ids, deleted_nodes, + buf, this->delete_backing_buf, + SECTORS_PER_MERGE); + + // insert into deleted_nhoods + this->disk_deleted_nhoods.clear(); + this->disk_deleted_nhoods.reserve(deleted_nodes.size()); + for (auto &nhood : deleted_nodes) { + // WARNING :: ASSUMING DISK GRAPH DEGREE NEVER GOES OVER 512 + assert(nhood.nnbrs < 512); + std::vector non_deleted_nbrs; + for (uint32_t i = 0; i < nhood.nnbrs; i++) { + uint32_t id = nhood.nbrs[i]; + auto iter = this->disk_deleted_ids.find(id); + if (iter == this->disk_deleted_ids.end()) { + non_deleted_nbrs.push_back(id); + } + } + this->disk_deleted_nhoods.insert( + std::make_pair(nhood.id, non_deleted_nbrs)); + } + + // free buf + aligned_free((void *) buf); + assert(deleted_nodes.size() == this->disk_deleted_ids.size()); + assert(this->disk_deleted_nhoods.size() == this->disk_deleted_ids.size()); + } + + template + void StreamingMerger::consolidate_deletes(DiskNode &disk_node, + uint8_t * scratch) { + // if node is deleted + if (this->is_deleted(disk_node)) { + disk_node.nnbrs = 0; + *(disk_node.nbrs - 1) = 0; + return; + } + + const uint32_t id = disk_node.id; + + assert(disk_node.nnbrs < 512); + + + std::vector id_nhood(disk_node.nbrs, + disk_node.nbrs + disk_node.nnbrs); + + tsl::robin_set new_edges; + + bool change = false; + for (auto &nbr : id_nhood) { + auto iter = this->disk_deleted_nhoods.find(nbr); + if (iter != this->disk_deleted_nhoods.end()) { + change = true; + new_edges.insert(iter->second.begin(), iter->second.end()); + } else { + new_edges.insert(nbr); + } + } + // no refs to deleted nodes --> move to next node + if (!change) { + return; + } + + // refs to deleted nodes + id_nhood.clear(); + id_nhood.reserve(new_edges.size()); + for (auto &nbr : new_edges) { + // 2nd order deleted edge + auto iter = this->disk_deleted_ids.find(nbr); + if (iter != this->disk_deleted_ids.end()) { + continue; + } else { + id_nhood.push_back(nbr); + } + } + + // TODO (corner case) :: id_nhood might be empty in adversarial cases + if (id_nhood.empty()) { + diskann::cout << "Adversarial case -- all neighbors of node's neighbors " + "deleted -- ID : " + << id << "; exiting\n"; + exit(-1); + } + + // compute PQ dists and shrink + std::vector id_nhood_dists(id_nhood.size(), 0.0f); + assert(scratch != nullptr); + this->disk_index->compute_pq_dists(id, id_nhood.data(), + id_nhood_dists.data(), + (_u32) id_nhood.size(), scratch); + + // prune neighbor list using PQ distances + std::vector cand_nbrs(id_nhood.size()); + for (uint32_t i = 0; i < id_nhood.size(); i++) { + cand_nbrs[i].id = id_nhood[i]; + // auto iter = this->disk_deleted_ids.find(id_nhood[i]); + // assert(iter == this->disk_deleted_ids.end()); + cand_nbrs[i].distance = id_nhood_dists[i]; + } + // sort and keep only maxc neighbors + std::sort(cand_nbrs.begin(), cand_nbrs.end()); + if (cand_nbrs.size() > this->maxc) { + cand_nbrs.resize(this->maxc); + } + std::vector pruned_nbrs; + std::vector occlude_factor(cand_nbrs.size(), 0.0f); + pruned_nbrs.reserve(this->range); + this->occlude_list_pq(cand_nbrs, pruned_nbrs, occlude_factor, scratch); + + // copy back final nbrs + disk_node.nnbrs = (_u32) pruned_nbrs.size(); + *(disk_node.nbrs - 1) = disk_node.nnbrs; + for (uint32_t i = 0; i < (_u32) pruned_nbrs.size(); i++) { + disk_node.nbrs[i] = pruned_nbrs[i].id; + // auto iter = this->disk_deleted_ids.find(disk_node.nbrs[i]); + // assert(iter == this->disk_deleted_ids.end()); + } + } + + template + bool StreamingMerger::is_deleted(const DiskNode &disk_node) { + // short circuit when disk_node is a `hole` on disk + if (this->disk_tags[disk_node.id] == std::numeric_limits::max()) { + if (disk_node.nnbrs != 0) { + throw diskann::ANNException(std::string("Found node with id: ") + + std::to_string(disk_node.id) + + " that has non-zero degree.", + -1, __FUNCSIG__, __FILE__, __LINE__); + diskann::cerr << "Node with id " << disk_node.id + << " is a hole but has non-zero degree " + << disk_node.nnbrs << std::endl; + } else { + return true; + } + } + return (this->disk_deleted_ids.find(disk_node.id) != + this->disk_deleted_ids.end()); + } + + template + void StreamingMerger::compute_rename_map() { + uint32_t needed = 0; + for (auto &mem_npt : this->mem_npts) { + needed += mem_npt; + } + for (auto &del_set : this->mem_deleted_ids) { + needed -= (_u32) del_set.size(); + } + diskann::cout << "RENAME: Need " << needed + << ", free: " << this->free_ids.size() << "\n"; + + uint32_t last_id = this->disk_npts; + if (needed > this->free_ids.size()) { + this->free_ids.reserve(needed); + } + while (this->free_ids.size() < needed) { + this->free_ids.push_back(last_id); + last_id++; + } + + // assign free IDs to all new IDs + diskann::cout << "RENAME: Assigning IDs.\n"; + uint32_t next_free_index = 0; + this->rename_map.reserve(needed); + this->inverse_map.reserve(needed); + std::vector> rename_pairs(needed); + std::vector> inverse_pairs(needed); + for (uint32_t mem_id = 0; mem_id < this->mem_data.size(); mem_id++) { + diskann::cout << "Processing Mem-DiskANN #" << mem_id + 1 << "\n"; + uint32_t offset = this->offset_ids[mem_id]; + const tsl::robin_set &del_set = this->mem_deleted_ids[mem_id]; + std::vector deleted(this->mem_npts[mem_id], false); + for (auto &id : del_set) { + deleted[id] = true; + } + for (uint32_t j = 0; j < this->mem_npts[mem_id]; j++) { + // ignore any deleted points + if (deleted[j]) { + continue; + } + const uint32_t new_id = this->free_ids[next_free_index]; + assert(new_id < last_id); + rename_pairs[next_free_index].first = offset + j; + rename_pairs[next_free_index].second = new_id; + inverse_pairs[next_free_index].first = new_id; + inverse_pairs[next_free_index].second = offset + j; + next_free_index++; + } + } + diskann::cout << "RENAME: Storing mappings for " << next_free_index + << " points.\n"; + this->rename_list.clear(); + this->rename_list.reserve(next_free_index); + this->rename_list.insert(this->rename_list.end(), rename_pairs.begin(), + rename_pairs.end()); + this->inverse_list.clear(); + this->inverse_list.reserve(next_free_index); + this->inverse_list.insert(this->inverse_list.end(), inverse_pairs.begin(), + inverse_pairs.end()); + } + + template + uint32_t StreamingMerger::rename(uint32_t id) const { + auto iter = std::lower_bound( + this->rename_list.begin(), this->rename_list.end(), + std::make_pair(id, std::numeric_limits::max()), + [](const auto &left, const auto &right) { + return left.first < right.first; + }); + if (iter == this->rename_list.end()) { + return std::numeric_limits::max(); + } else { + uint32_t idx = (_u32) std::distance(this->rename_list.begin(), iter); + const std::pair &p = this->rename_list[idx]; + if (p.first == id) + return p.second; + else + return std::numeric_limits::max(); + } + } + + template + uint32_t StreamingMerger::rename_inverse(uint32_t renamed_id) const { + auto iter = std::lower_bound( + this->inverse_list.begin(), this->inverse_list.end(), + std::make_pair(renamed_id, std::numeric_limits::max()), + [](const auto &left, const auto &right) { + return left.first < right.first; + }); + if (iter == this->inverse_list.end()) { + return std::numeric_limits::max(); + } else { + uint32_t idx = (_u32) std::distance(this->inverse_list.begin(), iter); + const std::pair &p = this->inverse_list[idx]; + if (p.first == renamed_id) + return p.second; + else + return std::numeric_limits::max(); + } + } + + template + void StreamingMerger::rename(DiskNode &node) const { + uint32_t renamed_id = this->rename(node.id); + if (renamed_id != std::numeric_limits::max()) { + node.id = renamed_id; + } + uint32_t nnbrs = node.nnbrs; + for (uint32_t i = 0; i < nnbrs; i++) { + uint32_t renamed_nbr_i = this->rename(node.nbrs[i]); + if (renamed_nbr_i != std::numeric_limits::max()) { + node.nbrs[i] = renamed_nbr_i; + } + } + } + + template + void StreamingMerger::rename(std::vector &ids) const { + for (uint32_t i = 0; i < ids.size(); i++) { + uint32_t renamed_id = this->rename(ids[i]); + if (renamed_id != std::numeric_limits::max()) { + ids[i] = renamed_id; + } + } + } + + template + uint32_t StreamingMerger::get_index_id( + const uint32_t offset_id) const { + if (offset_id < this->offset_ids[0]) { + return std::numeric_limits::max(); + } + // should not happen unless some buffer is corrupted + if (offset_id > this->offset_ids.back() + INDEX_OFFSET) { + diskann::cout << "Seen: " << offset_id << ", min: " << offset_ids[0] + << ", max: " << offset_ids.back() << "\n"; + } + assert(offset_id < this->offset_ids.back() + INDEX_OFFSET); + uint32_t index_no = + (uint32_t)((offset_id - this->offset_ids[0]) / INDEX_OFFSET); + assert(index_no < this->offset_ids.size()); + return index_no; + } + + template + std::vector StreamingMerger::get_edge_list( + const uint32_t offset_id) { + const uint32_t index_no = this->get_index_id(offset_id); + if (index_no == std::numeric_limits::max()) { + assert(offset_id < this->offset_ids[0]); + return this->disk_delta->get_nhood(offset_id); + } + // uint32_t local_id = offset_id - this->offset_ids[index_no]; + // assert(local_id < this->mem_npts[index_no]); + std::vector ret = + this->mem_deltas[index_no]->get_nhood(offset_id); + // this->rename(ret); + return ret; + } + + template + const T *StreamingMerger::get_mem_data(const uint32_t offset_id) { + const uint32_t index_no = this->get_index_id(offset_id); + if (index_no == std::numeric_limits::max()) { + assert(offset_id < this->offset_ids[0]); + return nullptr; + } + uint32_t local_id = offset_id - this->offset_ids[index_no]; + assert(local_id < this->mem_npts[index_no]); + return this->mem_data[index_no] + + ((uint64_t) local_id * (uint64_t) this->aligned_ndims); + } + + template + void StreamingMerger::write_tag_file( + const std::string &tag_out_filename, const uint32_t npts) { + diskann::Timer timer; + diskann::cout << "Writing new tags to " << tag_out_filename << "\n"; + + TagT *cur_tags; + + size_t allocSize = npts * sizeof(TagT); + alloc_aligned(((void **) &cur_tags), allocSize, 8 * sizeof(TagT)); + + //TODO: We must detect holes in a better way. Currently, it is possible + //that one of the tags will be uint32_t::max() and will fail. + for (uint32_t i = 0; i < npts; i++) { + TagT cur_tag; + // check if `i` is in inverse map + const uint32_t offset_id = this->rename_inverse(i); + if (offset_id == std::numeric_limits::max()) { + cur_tag = this->disk_tags[i]; + if (this->deleted_tags.find(cur_tag) != this->deleted_tags.end()) { + *(cur_tags + i) = std::numeric_limits::max(); + } else + *(cur_tags + i) = cur_tag; + } else { + const uint32_t index_no = this->get_index_id(offset_id); + const uint32_t index_local_id = offset_id - this->offset_ids[index_no]; + cur_tag = this->mem_tags[index_no][index_local_id]; + if (this->latter_deleted_tags[index_no].find(cur_tag) != + this->latter_deleted_tags[index_no].end()) { + *(cur_tags + i) = std::numeric_limits::max(); + } else + *(cur_tags + i) = cur_tag; + } + } + diskann::save_bin(tag_out_filename, cur_tags, npts, 1); + + diskann::cout << "Tags written to " << tag_out_filename << " in " << timer.elapsed() + << " microsec" << std::endl; + + //Should not mix delete with alloc aligned + //TODO: This will work because we are dealing with uint64 at the moment. + //If we ever have string tags, this'll fail spectacularly. + //delete[] cur_tags; + aligned_free(cur_tags); + // release all tags -- automatically deleted since using `unique_ptr` + this->mem_tags.clear(); + } + + template + void StreamingMerger::process_merges() { + // buf to hold data being read + char *buf = nullptr; + alloc_aligned((void **) &buf, SECTORS_PER_MERGE * SECTOR_LEN, SECTOR_LEN); + + Timer merge_timer; + std::unique_ptr sector_buf = std::make_unique(SECTOR_LEN); + + std::ofstream output_writer(this->final_index_file, + std::ios::binary | std::ios::ate); + + output_writer.write(sector_buf.get(), SECTOR_LEN); + // start at sector=1 in output file for reading + writing + uint64_t cur_offset = SECTOR_LEN; + + // batch consolidate deletes + std::vector> disk_nodes; + std::vector>> id_nhoods; + uint32_t start_id = 0, new_start_id; + diskann::cout << "Merging inserts into SSD-DiskANN.\n"; + uint64_t delta_avg = 0, delta_max = 0, delta_count = 0; + std::atomic counts; + counts.store(0); + + // aux_writer << "buf_min = " << (uint64_t) buf << ", buf_max = " << + // (uint64_t) (buf + SECTORS_PER_MERGE * SECTOR_LEN) << "\n"; + while (start_id < this->disk_npts) { + // zero buf for better consistency + memset(buf, 0, SECTORS_PER_MERGE * SECTOR_LEN); + new_start_id = this->disk_index->merge_read(disk_nodes, start_id, + SECTORS_PER_MERGE, buf); +#pragma omp parallel for schedule(dynamic, 128) num_threads(MAX_N_THREADS) + for (int64_t idx = 0; idx < (int64_t) disk_nodes.size(); idx++) { + // get thread-specific scratch + int omp_thread_no = omp_get_thread_num(); + uint8_t *thread_scratch = this->thread_bufs[omp_thread_no]; + + DiskNode &disk_node = disk_nodes[idx]; + uint32_t id = disk_node.id; + + std::vector nhood; + std::vector deltas; + uint32_t offset_id = this->rename_inverse(id); + // replaced by new vector, copy coords and proceed as normal + if (offset_id != std::numeric_limits::max()) { + // copy coords + const T *vec = this->get_mem_data(offset_id); + assert(vec != nullptr); + memcpy(disk_node.coords, vec, this->ndims * sizeof(T)); + disk_node.nnbrs = 0; + *(disk_node.nbrs - 1) = 0; // also set on buffer + deltas = this->get_edge_list(offset_id); + // delta_count++; + // delta_avg += deltas.size(); + // delta_max = std::max(delta_max, (uint64_t) deltas.size()); + } else { + // not replaced + deltas = this->get_edge_list(id); + delta_count++; + delta_avg += deltas.size(); + delta_max = std::max(delta_max, (uint64_t) deltas.size()); + } + + // if no edges to add, continue + if (deltas.empty()) { + continue; + } + + uint32_t nnbrs = disk_node.nnbrs; + nhood.insert(nhood.end(), disk_node.nbrs, disk_node.nbrs + nnbrs); + nhood.insert(nhood.end(), deltas.begin(), deltas.end()); + // rename nbrs in nhood to use PQ dist comparisons + // this->rename(nhood); // skipping since get_edge_list() renames delta + // edges + + // prune neighbor list ONLY if exceeding graph max out-degree (`range`) + if (nhood.size() > this->range) { + std::vector dists(nhood.size(), 0.0f); + std::vector pool(nhood.size()); + this->disk_index->compute_pq_dists(id, nhood.data(), dists.data(), + (_u32) nhood.size(), + thread_scratch); + for (uint32_t k = 0; k < nhood.size(); k++) { + pool[k].id = nhood[k]; + pool[k].distance = dists[k]; + // diskann::cout << k << ": id=" << nhood[k] << ", dist=" << + // dists[k] << "\n"; + } + nhood.clear(); + // prune pool + std::sort(pool.begin(), pool.end()); + this->prune_neighbors_pq(pool, nhood, thread_scratch); + } + // copy edges from nhood to disk node + disk_node.nnbrs = (_u32) nhood.size(); + // *(disk_node.nbrs - 1) = nhood.size(); // write to buf + *(disk_node.nbrs - 1) = (_u32) nhood.size(); // write to buf + for (uint32_t i = 0; i < disk_node.nnbrs; i++) { + disk_node.nbrs[i] = nhood[i]; + } + memcpy(disk_node.nbrs, nhood.data(), + disk_node.nnbrs * sizeof(uint32_t)); + // aux_writer << "disk_node.nbrs = " << (uint64_t) disk_node.nbrs << + // ", id = "; + uint32_t lcounts = 0; + for (auto &nbr : nhood) { + if (nbr >= 980000) + lcounts++; + } + counts += lcounts; + } + + cur_offset += SECTORS_PER_MERGE * SECTOR_LEN; + output_writer.write(buf, SECTORS_PER_MERGE * SECTOR_LEN); + diskann::cout << new_start_id << " / " << this->disk_npts + << " nodes processed.\n"; + start_id = new_start_id; + } + diskann::cout << "Delta statistics:\nMax: " << delta_max + << ", Avg: " << (float) delta_avg / (float) delta_count + << "\n"; + diskann::cout << "Old -> new edges: " << counts.load() << "\n"; + + // close index + output_writer.close(); + + // [_u64 file size][_u64 nnodes][_u64 medoid ID][_u64 max_node_len][_u64 + // nnodes_per_sector] + uint64_t file_size = + SECTOR_LEN + (ROUND_UP(ROUND_UP(this->disk_npts, nnodes_per_sector) / + nnodes_per_sector, + SECTORS_PER_MERGE)) * + (uint64_t) SECTOR_LEN; + memset(buf, 0, SECTORS_PER_MERGE * SECTOR_LEN); + std::vector output_metadata; + output_metadata.push_back((uint64_t) this->disk_npts); + output_metadata.push_back((uint64_t) this->ndims); + // determine medoid + uint64_t medoid = this->init_ids[0]; + output_metadata.push_back((uint64_t) medoid); + uint64_t max_node_len = this->ndims * sizeof(T) + sizeof(uint32_t) + + this->range * sizeof(uint32_t); + uint64_t nnodes_per_sector = SECTOR_LEN / max_node_len; + output_metadata.push_back(max_node_len); + output_metadata.push_back(nnodes_per_sector); + output_metadata.push_back(this->disk_index_num_frozen); + output_metadata.push_back(this->disk_index_frozen_loc); + output_metadata.push_back(file_size); + + diskann::save_bin<_u64>(final_index_file, output_metadata.data(), + output_metadata.size(), 1, 0); + // free buf + aligned_free((void *) buf); + double e2e_time = ((double) merge_timer.elapsed()) / (1000000.0); + diskann::cout << "Time to merge the inserts to disk: " << e2e_time << "s." + << std::endl; + } + + + template + void StreamingMerger::merge(const char * disk_in, + const std::vector &mem_in, + const char * disk_out, + std::vector*> &deleted_tags_vectors, + std::string &working_folder) { + // load disk index + this->disk_index_out_path = disk_out; + this->disk_index_in_path = disk_in; + this->TMP_FOLDER = working_folder; + std::cout << "Working folder : " << working_folder << std::endl; + this->temp_disk_index_path = getTempFilePath(working_folder, "temp_disk_index"); + this->temp_pq_coords_path = getTempFilePath(working_folder, "temp_pq_compressed"); + this->temp_tags_path = getTempFilePath(working_folder, "temp_tags"); + std::cout << this->temp_disk_index_path << " , " << this->temp_pq_coords_path << " , " << this->temp_tags_path << std::endl; + this->final_index_file = + this->_single_file_index ? this->disk_index_out_path + : this->disk_index_out_path + "_disk.index"; + this->final_pq_coords_file = + this->_single_file_index + ? this->disk_index_out_path + : this->disk_index_out_path + "_pq_compressed.bin"; + this->final_tags_file = + this->_single_file_index + ? this->disk_index_out_path + : this->disk_index_out_path + "_disk.index.tags"; + + +#ifndef _WINDOWS + std::shared_ptr reader = + std::make_shared(); +#else + std::shared_ptr reader = + std::make_shared(); +#endif + + // std::shared_ptr reader = + // std::make_shared(); + this->disk_index = new PQFlashIndex( + this->dist_metric, reader, this->_single_file_index, true); + diskann::cout << "Created PQFlashIndex inside index_merger " << std::endl; + + diskann::cout << "Loading PQFlashIndex from file: " << disk_in + << " into object: " << std::hex << (_u64) & + (this->disk_index) << std::dec << std::endl; + this->disk_index->load(disk_in, NUM_INDEX_LOAD_THREADS); + + uint32_t node_cache_count = 1 + (uint32_t) round(this->disk_index->return_nd() * 0.01); + node_cache_count = node_cache_count > PQ_FLASH_INDEX_MAX_NODES_TO_CACHE + ? PQ_FLASH_INDEX_MAX_NODES_TO_CACHE + : node_cache_count; + std::vector cache_node_list; + this->disk_index->cache_bfs_levels(node_cache_count, + cache_node_list); + this->disk_index->load_cache_list(cache_node_list); + this->disk_tags = this->disk_index->get_tags(); + this->init_ids = this->disk_index->get_init_ids(); + this->disk_npts = (_u32) this->disk_index->return_nd(); + this->disk_thread_data = this->disk_index->get_thread_data(); + auto res_pq = this->disk_index->get_pq_config(); + this->pq_data = res_pq.first; + this->pq_nchunks = res_pq.second; + this->nnodes_per_sector = (_u32) this->disk_index->nnodes_per_sector; + this->max_node_len = (_u32) this->disk_index->max_node_len; + _u32 max_degree = + (max_node_len - (sizeof(T) * this->ndims)) / sizeof(uint32_t) - 1; + this->range = max_degree; + diskann::cout << "Setting range to: " << this->range << std::endl; + this->disk_index_num_frozen = this->disk_index->get_num_frozen_points(); + this->disk_index_frozen_loc = this->disk_index->get_frozen_loc(); + + // create deltas + this->disk_delta = new GraphDelta(0, this->disk_npts); + uint64_t base_offset = ROUND_UP(this->disk_npts, INDEX_OFFSET); + + // load mem-indices + for (auto &mem_index_path : mem_in) { + uint32_t npts; + if (!(this->_single_file_index)) { + std::string ind_path = mem_index_path; + std::string data_path = mem_index_path + ".data"; + std::ifstream bin_reader(data_path, std::ios::binary); + uint32_t bin_npts, bin_ndims; + bin_reader.read((char *) &bin_npts, sizeof(uint32_t)); + bin_reader.read((char *) &bin_ndims, sizeof(uint32_t)); + bin_reader.close(); + diskann::cout << "Index Path: " << ind_path << "\n"; + diskann::cout << "Data Path: " << data_path << "\n"; + diskann::cout << "Detected # pts = " << bin_npts + << ", # dims = " << bin_ndims << "\n"; + + auto mem_index = std::make_unique>( + this->dist_metric, bin_ndims, bin_npts + 100, true, + this->_single_file_index, true, false); + _u64 n1, n2, n3; + T * data_load; + diskann::load_aligned_bin(data_path, data_load, n1, n2, n3); + npts = (_u32)(n1 - 1); + assert(npts < MAX_PTS_PER_MEM_INDEX); + this->mem_npts.push_back(npts); + this->mem_data.push_back(data_load); + uint32_t index_offset = (_u32) base_offset; + base_offset += INDEX_OFFSET; + this->offset_ids.push_back(index_offset); + this->mem_deltas.push_back(new GraphDelta(index_offset, npts)); + tsl::robin_set temp_del_set; + if(file_exists(mem_index_path + ".del")) + { + mem_index->load_delete_set(mem_index_path + ".del"); + mem_index->get_delete_set(temp_del_set); + } + this->mem_deleted_ids.push_back(temp_del_set); + mem_index->load_tags(mem_index_path + ".tags"); + // manage tags + std::unique_ptr index_tags; + index_tags.reset(new TagT[npts]); + + const std::unordered_map &loc_tag_map = + *mem_index->get_tags(); + for (uint32_t k = 0; k < npts; k++) { + auto iter = loc_tag_map.find(k); + if (iter == loc_tag_map.end()) { + index_tags[k] = (TagT) 0; + } else { + index_tags[k] = iter->second; + } + } + this->mem_tags.push_back(std::move(index_tags)); + } else { + // read metadata from single index file for npts and ndims + _u64 nr, nc; + std::unique_ptr<_u64[]> file_offset_data; + diskann::load_bin<_u64>(mem_index_path, file_offset_data, nr, nc, 0); + + size_t data_dim, data_num_points; + diskann::get_bin_metadata(mem_index_path, data_num_points, data_dim, + file_offset_data[1]); + diskann::cout << "Detected # pts = " << data_num_points + << ", # dims = " << data_dim << "\n"; + diskann::cout << "Since vamana index is dynamic, it will have one " + "frozen point, hence #pts = " + << data_num_points - 1 << std::endl; + + // load mem_index_data with appropriate offset + _u64 n1, n2, n3; + T * data_load; + diskann::load_aligned_bin(mem_index_path, data_load, n1, n2, n3, + file_offset_data[1]); + npts = (_u32)(n1 - 1); + assert(npts < MAX_PTS_PER_MEM_INDEX); + this->mem_npts.push_back(npts); + this->mem_data.push_back(data_load); + + // call mem_index constructor with dynamic index and single index file + // set to true + auto mem_index = std::make_unique>( + this->dist_metric, data_dim, data_num_points + 100, true, + this->_single_file_index, true, false); + + // load tags with appropriate offset + uint32_t index_offset = (_u32) base_offset; + base_offset += INDEX_OFFSET; + this->offset_ids.push_back(index_offset); + this->mem_deltas.push_back(new GraphDelta(index_offset, npts)); + mem_index->load_tags(mem_index_path, file_offset_data[2]); + // manage tags + std::unique_ptr index_tags; + index_tags.reset(new TagT[npts]); + + const std::unordered_map &loc_tag_map = + *mem_index->get_tags(); + for (uint32_t k = 0; k < npts; k++) { + auto iter = loc_tag_map.find(k); + if (iter == loc_tag_map.end()) { + diskann::cout << "Index # " << this->mem_data.size() + << " : missing tag for node #" << k << "\n"; + exit(-1); + index_tags[k] = (TagT) k; + } else { + index_tags[k] = iter->second; + } + } + this->mem_tags.push_back(std::move(index_tags)); + } + } + +#ifdef USE_TCMALLOC + MallocExtension::instance()->ReleaseFreeMemory(); +#endif + + for (size_t j = 0; j < deleted_tags_vectors.size(); j++) { + this->latter_deleted_tags.push_back(tsl::robin_set()); + for (size_t i = j+1; i < deleted_tags_vectors.size(); i++) { + for (size_t k = 0; k < deleted_tags_vectors[i]->size(); k++) { + this->latter_deleted_tags[j].insert((*deleted_tags_vectors[i])[k]); + } + } + } + + //TODO: See if this can be included in the previous loop + for (auto &deleted_tags_vector : deleted_tags_vectors) { + for (size_t i = 0; i < deleted_tags_vector->size(); i++) { + this->deleted_tags.insert((*deleted_tags_vector)[i]); + } + } + + diskann::cout << "Allocating thread scratch space -- " + << PER_THREAD_BUF_SIZE / (1 << 20) << " MB / thread.\n"; + alloc_aligned((void **) &this->thread_pq_scratch, + MAX_N_THREADS * PER_THREAD_BUF_SIZE, SECTOR_LEN); + this->thread_bufs.resize(MAX_N_THREADS); + for (uint32_t i = 0; i < thread_bufs.size(); i++) { + this->thread_bufs[i] = this->thread_pq_scratch + i * PER_THREAD_BUF_SIZE; + } + + mergeImpl(); + } + + template + void StreamingMerger::mergeImpl() { + // populate deleted IDs + this->compute_deleted_ids(); + // BEGIN -- graph on disk has deleted references, maybe some holes + // populate deleted nodes + this->populate_deleted_nhoods(); + // process all deletes + this->process_deletes(); + // END -- graph on disk has NO deleted references, maybe some holes + + diskann::cout << "Computing rename-map.\n"; + // compute rename map + this->compute_rename_map(); + + // get max ID + 1 in rename-map as new max pts + uint32_t new_max_pts = this->disk_npts - 1; + // alternative using list + new_max_pts = std::max(this->inverse_list.back().first, new_max_pts); + new_max_pts = new_max_pts + 1; + + // TODO (correct) :: figure out naming scheme + //std::string new_disk_out(this->disk_index_out_path + "_disk.index"); + //diskann::cout << "RELOAD: Creating new disk graph at " << new_disk_out + // << "\n"; + //std::string new_pq_prefix(this->disk_index_out_path + "_pq"); + //std::string new_pq_coords(new_pq_prefix + "_compressed.bin"); + diskann::cout << "RELOAD: Creating new PQ coords file " + << this->temp_pq_coords_path << std::endl; + + + // TODO (correct) :: write to the right file + //std::string tmp_file = TMP_FOLDER + "/index_ravi"; + +#ifdef USE_TCMALLOC + MallocExtension::instance()->ReleaseFreeMemory(); +#endif + + //this->output_writer.open(tmp_file, std::ios::out | std::ios::binary); + //assert(this->output_writer.is_open()); + + // BEGIN -- PQ data on disk not consistent, not in right order + // write outdated PQ data into pq writer with intentionally wrong header - + // all these updates are made to a separate file, to be merged later into + // thw index file if instructed + std::fstream pq_writer(this->temp_pq_coords_path, + std::ios::out | std::ios::binary | std::ios::trunc); + assert(pq_writer.is_open()); + uint64_t pq_file_size = + ((uint64_t) new_max_pts * (uint64_t) this->pq_nchunks) + + (2 * sizeof(uint32_t)); + + // inflate file size to accommodate new points + uint64_t dummy = 0; + pq_writer.seekp(pq_file_size - sizeof(uint64_t), std::ios::beg); + pq_writer.write((char *) (&dummy), sizeof(uint64_t)); + + // write PQ compressed coords bin and close file + pq_writer.seekp(0, std::ios::beg); + uint32_t npts_u32 = new_max_pts, ndims_u32 = this->pq_nchunks; + pq_writer.write((char *) &npts_u32, sizeof(uint32_t)); + pq_writer.write((char *) &ndims_u32, sizeof(uint32_t)); + pq_writer.write((char *) this->pq_data, + (uint64_t) this->disk_npts * (uint64_t) ndims_u32); + pq_writer.close(); + + // write out tags + //const std::string tag_file = new_disk_out + ".tags"; + this->write_tag_file(this->temp_tags_path, new_max_pts); + + // switch index to read-only mode + this->disk_index->reload_index(this->temp_disk_index_path, + this->temp_pq_coords_path, + this->temp_tags_path); +#ifdef USE_TCMALLOC + MallocExtension::instance()->ReleaseFreeMemory(); +#endif + + // re-acquire pointers + auto res = this->disk_index->get_pq_config(); + this->pq_nchunks = res.second; + this->pq_data = res.first; + this->disk_npts = (_u32) this->disk_index->return_nd(); + this->init_ids.clear(); + this->init_ids = this->disk_index->get_init_ids(); + assert(this->disk_npts == new_max_pts); + + std::cout << "AFTER RELOAD: PQ_NChunks: " << res.second + << " Disk points: " << this->disk_npts + << " Frozen point id: " << this->init_ids[0] << std::endl; + + // call inserts + this->process_inserts(); +#ifdef USE_TCMALLOC + MallocExtension::instance()->ReleaseFreeMemory(); +#endif + + this->process_inserts_pq(); +#ifdef USE_TCMALLOC + MallocExtension::instance()->ReleaseFreeMemory(); +#endif + + diskann::cout << "Dumping full compressed PQ vectors from memory.\n"; + // re-open PQ writer + pq_writer.open(this->temp_pq_coords_path, + std::ios::in | std::ios::out | std::ios::binary); + pq_writer.seekp(2 * sizeof(uint32_t), std::ios::beg); + // write all (old + new) PQ data to disk; no need to modify header + pq_writer.write((char *) this->pq_data, + ((uint64_t) new_max_pts * (uint64_t) this->pq_nchunks)); + pq_writer.close(); + // END -- PQ data on disk consistent and in correct order + + // batch rename all inserted edges in each delta + diskann::cout << "Renaming edges for easier access during merge.\n"; + // const std::function rename_func = + // std::bind(&StreamingMerger::rename, this); + const std::function rename_func = [this](uint32_t id) { + return this->rename(id); + }; + this->disk_delta->rename_edges(rename_func); + for (auto &delta : this->mem_deltas) { + delta->rename_edges(rename_func); + } + + // start merging + // BEGIN -- graph on disk has NO deleted references, NO newly inserted + // points + + + this->process_merges(); + // END -- graph on disk has NO deleted references, has newly inserted points + + /* copy output from temp_file -> new_disk_out */ + // reset temp_file ptr + //this->output_writer.close(); + + auto copy_file = [](const std::string &src, const std::string &dest) { + diskann::cout << "COPY :: " << src << " --> " << dest << "\n"; + std::ofstream dest_writer(dest, std::ios::binary); + std::ifstream src_reader(src, std::ios::binary); + dest_writer << src_reader.rdbuf(); + dest_writer.close(); + src_reader.close(); + }; + // copy index + //copy_file(tmp_file, this->disk_index_out_path); + + // merge files if needed + if (this->_single_file_index) { + // update metadata with pq_pivots_file_size, pq_vector_file_size + size_t nr, nc; + std::vector output_metadata; + uint64_t * out_metadata; + + diskann::load_bin(this->final_index_file, out_metadata, nr, + nc); + for (size_t i = 0; i < nr; i++) + output_metadata.push_back(out_metadata[i]); + + delete[] out_metadata; + + // tags + TagT * tags; + uint64_t tag_num, tag_dim; + diskann::load_bin(this->temp_tags_path, tags, tag_num, tag_dim); + size_t tag_bytes_written = + diskann::save_bin(this->final_index_file, tags, tag_num, tag_dim, + output_metadata[output_metadata.size() - 1]); + delete[] tags; + + output_metadata.push_back(output_metadata[output_metadata.size() - 1] + + tag_bytes_written); + + size_t nr_in, nc_in; + uint64_t * in_metadata; + std::string disk_in = this->_single_file_index + ? this->disk_index_in_path + : this->disk_index_in_path + "_disk.index"; + diskann::load_bin(disk_in, in_metadata, nr_in, nc_in); + + uint64_t *pq_metadata_in; + size_t nr_pq_in, nc_pq_in; + diskann::load_bin(disk_in, pq_metadata_in, nr_pq_in, nc_pq_in, + in_metadata[8]); + diskann::save_bin(this->final_index_file, pq_metadata_in, + nr_pq_in, + nc_pq_in, + output_metadata[output_metadata.size() - 1]); + + size_t pq_pivots_total_bytes_written = 0; + // pq_pivots + float * pq_pivots_data; + uint64_t pq_pts, pq_dims; + diskann::load_bin(disk_in, pq_pivots_data, pq_pts, pq_dims, + in_metadata[8] + pq_metadata_in[0]); + size_t pq_pivots_bytes = diskann::save_bin( + this->final_index_file, pq_pivots_data, pq_pts, pq_dims, + output_metadata[output_metadata.size() - 1] + pq_metadata_in[0]); + delete[] pq_pivots_data; + diskann::cout << "Written pivots to single index file" << std::endl; + + // pq centroids + float * pq_centroid_data; + uint64_t centroid_num, centroid_dim; + diskann::load_bin(disk_in, pq_centroid_data, centroid_num, + centroid_dim, + in_metadata[8] + pq_metadata_in[1]); + size_t pq_centroid_bytes = diskann::save_bin( + this->final_index_file, pq_centroid_data, centroid_num, + centroid_dim, + output_metadata[output_metadata.size() - 1] + pq_metadata_in[1]); + delete[] pq_centroid_data; + diskann::cout << "Written centroids to single index file" << std::endl; + + // pq_rearrangment_perm + uint32_t *pq_rearrange_data; + uint64_t rearrange_num, rearrange_dim; + diskann::load_bin(disk_in, pq_rearrange_data, rearrange_num, + rearrange_dim, + in_metadata[8] + pq_metadata_in[2]); + size_t pq_rearrange_bytes = diskann::save_bin( + this->final_index_file, pq_rearrange_data, rearrange_num, + rearrange_dim, + output_metadata[output_metadata.size() - 1] + pq_metadata_in[2]); + delete[] pq_rearrange_data; + diskann::cout << "Written rearrangement data to single index file" + << std::endl; + + // pq_chunk_offsets + uint32_t *pq_offset_data; + uint64_t chunk_offset_num, chunk_offset_dim; + diskann::load_bin(disk_in, pq_offset_data, chunk_offset_num, + chunk_offset_dim, + in_metadata[8] + pq_metadata_in[3]); + size_t pq_offset_bytes = diskann::save_bin( + this->final_index_file, pq_offset_data, chunk_offset_num, + chunk_offset_dim, + output_metadata[output_metadata.size() - 1] + pq_metadata_in[3]); + delete[] pq_offset_data; + diskann::cout << "Written offsets to single index file" << std::endl; + + pq_pivots_total_bytes_written = pq_pivots_bytes + pq_centroid_bytes + + pq_rearrange_bytes + pq_offset_bytes; + output_metadata.push_back(output_metadata[output_metadata.size() - 1] + + pq_pivots_total_bytes_written + + pq_metadata_in[0]); + + // pq vectors + size_t pq_vector_bytes = diskann::save_bin( + this->final_index_file, this->pq_data, (uint64_t) this->disk_npts, + (uint64_t) ndims_u32, output_metadata[output_metadata.size() - 1]); + + output_metadata.push_back(output_metadata[output_metadata.size() - 1] + + pq_vector_bytes); + + delete[] pq_metadata_in; + + diskann::save_bin(this->final_index_file, + output_metadata.data(), + output_metadata.size(), 1); + } else { + // update pq table related data into new files + /* copy PQ tables */ + std::string prefix_pq_in = this->disk_index_in_path + "_pq"; + std::string prefix_pq_out = this->disk_index_out_path + "_pq"; + // PQ pivots + + copy_file(prefix_pq_in + "_pivots.bin", prefix_pq_out + "_pivots.bin"); + diskann::save_bin( + this->final_pq_coords_file, this->pq_data, (uint64_t) this->disk_npts, + (uint64_t) ndims_u32); + + copy_file(this->temp_tags_path, this->final_tags_file); + + } + + // destruct PQFlashIndex + delete this->disk_index; + this->disk_index = nullptr; + diskann::cout << "Destroyed PQ Flash Index\n"; + + } + + // template class instantiations + template class StreamingMerger; + template class StreamingMerger; + template class StreamingMerger; + template class StreamingMerger; + template class StreamingMerger; + template class StreamingMerger; + template class StreamingMerger; + template class StreamingMerger; + template class StreamingMerger; +} // namespace diskann diff --git a/src/v2/merge_insert.cpp b/src/v2/merge_insert.cpp new file mode 100644 index 0000000000..f16fa5acb6 --- /dev/null +++ b/src/v2/merge_insert.cpp @@ -0,0 +1,588 @@ +#include "neighbor.h" +#include "timer.h" +#include "tsl/robin_map.h" +#include "tsl/robin_set.h" +#include "utils.h" +#include "v2/merge_insert.h" +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "tcmalloc/malloc_extension.h" + +#include +#include +#include + +#include "aux_utils.h" +#include "exceptions.h" +#include "index.h" +#include "pq_flash_index.h" +#include "math_utils.h" +#include "memory_mapper.h" +#include "parameters.h" +#include "partition_and_pq.h" + +#include "logger.h" + +#include "Neighbor_Tag.h" +#ifdef _WINDOWS +#ifdef USE_BING_INFRA +#include "bing_aligned_file_reader.h" +#else +#include "windows_aligned_file_reader.h" +#endif +#else +#include "linux_aligned_file_reader.h" +#endif + +namespace diskann { + template + MergeInsert::MergeInsert( + Parameters& parameters, size_t dim, const std::string mem_prefix, + const std::string disk_prefix_in, const std::string disk_prefix_out, + Distance* dist, diskann::Metric dist_metric, bool single_file_index, + std::string working_folder) + : _dim(dim), _dist_metric(dist_metric), _active_0(true), _active_1(false), _active_del_0(true), + _active_del_1(false), _clearing_index_0(false), + _clearing_index_1(false), _switching_disk_prefixes(false), + _check_switch_index(false), _check_switch_delete(false) { + _merge_th = MERGE_TH; + _single_file_index = single_file_index; + this->_dist_metric = dist_metric; + _mem_index_0 = std::make_shared>( + this->_dist_metric, dim, _merge_th * 2, 1, _single_file_index, 1); + _mem_index_1 = std::make_shared>( + this->_dist_metric, dim, _merge_th * 2, 1, _single_file_index, 1); + + _paras_mem.Set("L", parameters.Get("L_mem")); + _paras_mem.Set("R", parameters.Get("R_mem")); + _paras_mem.Set("C", parameters.Get("C")); + _paras_mem.Set("alpha", parameters.Get("alpha_mem")); + _paras_mem.Set("num_rnds", 2); + _paras_mem.Set("saturate_graph", 0); + + _paras_disk.Set("L", parameters.Get("L_disk")); + _paras_disk.Set("R", parameters.Get("R_disk")); + _paras_disk.Set("C", parameters.Get("C")); + _paras_disk.Set("alpha", parameters.Get("alpha_disk")); + _paras_disk.Set("num_rnds", 2); + _paras_disk.Set("saturate_graph", 0); + + _num_search_threads = parameters.Get<_u32>("num_search_threads"); + _beamwidth = parameters.Get("beamwidth"); + _num_nodes_to_cache = parameters.Get<_u32>("nodes_to_cache"); + + _search_tpool = new ThreadPool(_num_search_threads); + + _mem_index_prefix = mem_prefix; + _deleted_tags_file = mem_prefix + "_deleted.tags"; + _disk_index_prefix_in = disk_prefix_in; + _disk_index_prefix_out = disk_prefix_out; + _dist_comp = dist; +#ifdef _WINDOWS +#ifndef USE_BING_INFRA + reader.reset(new WindowsAlignedFileReader()); +#else + reader.reset(new diskann::BingAlignedFileReader()); +#endif +#else + reader.reset(new LinuxAlignedFileReader()); +#endif + + _disk_index = new diskann::PQFlashIndex( + this->_dist_metric, reader, _single_file_index, true); + + std::string pq_prefix = _disk_index_prefix_in + "_pq"; + std::string disk_index_file = _disk_index_prefix_in + "_disk.index"; + int res = + _disk_index->load(_disk_index_prefix_in.c_str(), _num_search_threads); + if (res != 0) { + diskann::cout << "Failed to load disk index in MergeInsert constructor" + << std::endl; + exit(-1); + } + + TMP_FOLDER = working_folder; + std::cout << "TMP_FOLDER inside MergeInsert : " << TMP_FOLDER << std::endl; + } + + template + MergeInsert::~MergeInsert() { + // put in destructor code + } + + template + void MergeInsert::construct_index_merger() + { + uint32_t range = _paras_disk.Get("R"); + uint32_t l_index = _paras_disk.Get("L"); + uint32_t maxc = _paras_disk.Get("C"); + float alpha = _paras_disk.Get("alpha"); + _merger = new diskann::StreamingMerger((uint32_t) _dim, _dist_comp, _dist_metric, (uint32_t) _beamwidth, range, l_index, alpha, maxc, _single_file_index); + diskann::cout << "Created index merger object" << std::endl; + + } + + template + void MergeInsert::destruct_index_merger() + { + delete(_merger); + _merger = nullptr; + } + + template + int MergeInsert::insert(const T* point, const TagT& tag) + { + while(_check_switch_index.load()) + { + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + } + std::shared_lock lock(_index_lock); + if((_active_index == 0) && (_active_0.load() == false)) + { + diskann::cout << "Active index indicated as _mem_index_0 but it cannot accept insertions" << std::endl; + return -1; + } + if((_active_index == 1) && (_active_1.load() == false)) + { + diskann::cout << "Active index indicated as _mem_index_1 but it cannot accept insertions" << std::endl; + return -1; + } + + if(_active_index == 0) + { + if(_mem_index_0->get_num_points() < _mem_index_0->return_max_points()) + { + if(_mem_index_0->insert_point(point, _paras_mem, tag) != 0) + { + diskann::cout << "Could not insert point with tag " << tag << std::endl; + return -3; + } + { + std::unique_lock lock(_change_lock); + _mem_points++; + } + return 0; + } + else + { + diskann::cout << "Capacity exceeded" << std::endl; + } + } + else + { + if(_mem_index_1->get_num_points() < _mem_index_1->return_max_points()) + { + if(_mem_index_1->insert_point(point, _paras_mem, tag) != 0) + { + diskann::cout << "Could not insert point with tag " << tag << std::endl; + return -3; + } + { + std::unique_lock lock(_change_lock); + _mem_points++; + } + return 0; + } + else + { + diskann::cout << "Capacity exceeded in mem_index 1" << std::endl; + } + } + + return -2; + } + + + template + void MergeInsert::lazy_delete(const TagT& tag) + { + std::unique_lock lock(_delete_lock); + if((_active_delete_set == 0) && (_active_del_0.load() == false)) + { + diskann::cout << "Active deletion set indicated as _deletion_set_0 but it cannot accept deletions" << std::endl; + } + if((_active_delete_set == 1) && (_active_del_1.load() == false)) + { + diskann::cout << "Active deletion set indicated as _deletion_set_1 but it cannot accept deletions" << std::endl; + } + + if(_active_delete_set == 0) + {_deletion_set_0.insert(tag); + _mem_index_0->lazy_delete(tag); + } + else{ + _deletion_set_1.insert(tag); + _mem_index_1->lazy_delete(tag); + } + } + + template + void MergeInsert::search_sync(const T* query, const uint64_t K, const uint64_t search_L, + TagT* tags, float * distances, QueryStats * stats) + { + std::set> best; + //search disk index and get top K tags + { + std::shared_lock lock(_disk_lock); + assert(_switching_disk_prefixes == false); + std::vector disk_result_ids_64(search_L); + std::vector disk_result_dists(search_L); + std::vector disk_result_tags(search_L); + _disk_index->cached_beam_search( + query, search_L, search_L, disk_result_tags.data(), disk_result_dists.data(), _beamwidth, + stats); + for(unsigned i = 0; i < disk_result_tags.size(); i++) + { + Neighbor_Tag n; + n = Neighbor_Tag(disk_result_tags[i], disk_result_dists[i]); +// best.insert(Neighbor_Tag(disk_result_tags[i], disk_result_dists[i])); + best.insert(n); + } + } + //check each memory index - if non empty and not being currently cleared - search and get top K active tags + { + if(_clearing_index_0.load() == false) + { + std::shared_lock lock(_clear_lock_0); + if(_mem_index_0->get_num_points() > 0) + { + std::vector> best_mem_index_0; + _mem_index_0->search(query, (uint32_t)search_L, (uint32_t)search_L, best_mem_index_0); + for(auto iter : best_mem_index_0) + best.insert(iter); + } + } + + if(_clearing_index_1.load() == false) + { + std::shared_lock lock(_clear_lock_1); + if(_mem_index_1->get_num_points() > 0) + { + std::vector> best_mem_index_1; + _mem_index_1->search(query, (uint32_t)search_L, (uint32_t)search_L, best_mem_index_1); + for(auto iter : best_mem_index_1) + best.insert(iter); + } + } + } + std::vector> best_vec; + for(auto iter : best) + best_vec.emplace_back(iter); +// std::sort(best_vec.begin(), best_vec.end()); + if (best_vec.size() > K) +// best_vec.erase(best_vec.begin() + K, best_vec.end()); + //aggregate results, sort and pick top K candidates + { + std::shared_lock lock(_delete_lock); + size_t pos = 0; + for (auto iter : best_vec) { + if((_deletion_set_0.find(iter.tag) == _deletion_set_0.end()) && (_deletion_set_1.find(iter.tag) == _deletion_set_1.end())) { + tags[pos] = iter.tag; + distances[pos] = iter.dist; + pos++; + } + if (pos == K) + break; + } + } + } + + template + int MergeInsert::trigger_merge() + { + if(_mem_points >= _merge_th) + { + save_del_set(); + switch_index(); + return 1; + } + return 0; + } + + template + void MergeInsert::final_merge() + { + diskann::cout << "Inside final_merge()." << std::endl; + diskann::cout << _mem_index_0->get_num_points() << " " << _mem_index_1->get_num_points() << std::endl; + if(_mem_points > 0) + { + save_del_set(); + switch_index(); + } + diskann::cout << _mem_index_0->get_num_points() << " " << _mem_index_1->get_num_points() << std::endl; + } + + template + void MergeInsert::merge() + { + std::vector mem_in; + if(_active_index == 0) + mem_in.push_back(_mem_index_prefix + "_1"); + else + mem_in.push_back(_mem_index_prefix + "_0"); + + _merger->merge(_disk_index_prefix_in.c_str(), mem_in, _disk_index_prefix_out.c_str(), _deleted_tags_vector, TMP_FOLDER); + + diskann::cout << "Merge done" << std::endl; + { + std::unique_lock lock(_disk_lock); + bool expected_value = false; + if (_switching_disk_prefixes.compare_exchange_strong(expected_value, true)) { + diskann::cout << "Switching to latest merged disk index " << std::endl; + } else { + diskann::cout << "Failed to switch" << std::endl; +// return -1; + } + + std::string temp = _disk_index_prefix_out; + _disk_index_prefix_out = _disk_index_prefix_in; + _disk_index_prefix_in = temp; + delete(_disk_index); + _disk_index = new diskann::PQFlashIndex(this->_dist_metric, reader, _single_file_index, true); + + std::string pq_prefix = _disk_index_prefix_in + "_pq"; + std::string disk_index_file = _disk_index_prefix_in + "_disk.index"; + int res = _disk_index->load(_disk_index_prefix_in.c_str(), _num_search_threads); + if(res != 0) + { + diskann::cout << "Failed to load new disk index after merge" << std::endl; + exit(-1); + } + expected_value = true; + _switching_disk_prefixes.compare_exchange_strong(expected_value,false); + } + } + + template + void MergeInsert::switch_index() + { + //unique lock throughout the function to ensure another thread does not flip the value of _active_index after it has been saved by one thread, + //and multiple threads do not save the same index + //unique lock is acquired when no other thread holds any shared lock over it, so this function will wait till any on-going insertions are completed + //and then change the value of all related flags + { + bool expected_value = false; + _check_switch_index.compare_exchange_strong(expected_value, true); + std::unique_lock lock(_index_lock); + //make new index active + if(_active_index == 0) + { + _mem_index_1 = std::make_shared>(this->_dist_metric, _dim, _merge_th * 2 , 1, _single_file_index, 1); + bool expected_active = false; + if (_active_1.compare_exchange_strong(expected_active, true)) { + diskann::cout << "Initialised new index for _mem_index_1 " << std::endl; + } else { + diskann::cout << "Failed to initialise new _mem_index_1" << std::endl; +// return -1; + } + + } + else + { + _mem_index_0 = std::make_shared>(this->_dist_metric, _dim, _merge_th * 2, 1, _single_file_index, 1); + bool expected_active = false; + if (_active_0.compare_exchange_strong(expected_active, true)) { + diskann::cout << "Initialised new index for _mem_index_0 " << std::endl; + } else { + diskann::cout << "Failed to initialise new _mem_index_0" << std::endl; + // return -1; + } + } + _active_index = 1 - _active_index; + _mem_points = 0; + expected_value = true; + _check_switch_index.compare_exchange_strong(expected_value, false); + + } + + save(); + //start timer + diskann::Timer timer; + construct_index_merger(); + merge(); + destruct_index_merger(); + diskann::cout << "Merge time : " << timer.elapsed()/1000 << " ms" << std::endl; + //end timer + + { + std::shared_lock lock(_index_lock); + //make older index inactive after merge is complete or before ? + if(_active_index == 0) + { + bool expected_clearing = false; + bool expected_active = true; + _clearing_index_1.compare_exchange_strong(expected_clearing, true); + { + std::unique_lock lock(_clear_lock_1); + _mem_index_1.reset(); + _mem_index_1 = nullptr; + _mem_index_1 = std::make_shared>( + _dist_metric, _dim, _merge_th * 2, 1, _single_file_index, 1); + } + expected_clearing = true; + assert(expected_clearing == true); + _clearing_index_1.compare_exchange_strong(expected_clearing, false); + assert(expected_active == true); + _active_1.compare_exchange_strong(expected_active, false); + } + else + { + bool expected_clearing = false; + bool expected_active = true; + _clearing_index_0.compare_exchange_strong(expected_clearing, true); + std::unique_lock lock(_clear_lock_0); + { + _mem_index_0.reset(); + _mem_index_0 = nullptr; + _mem_index_0 = std::make_shared>( + _dist_metric, _dim, _merge_th * 2, 1, _single_file_index, 1); + } + expected_clearing = true; + assert(expected_clearing == true); + _clearing_index_0.compare_exchange_strong(expected_clearing, false); + assert(expected_active == true); + _active_0.compare_exchange_strong(expected_active, false); + } + //if merge() has returned, clear older active index + } + + } + + template + int MergeInsert::save() + { + //only switch_index will call this function + bool expected_active = true; + if(_active_index == 1) + { + if (_active_0.compare_exchange_strong(expected_active, false)) { + diskann::cout << "Saving mem index 0 to merge it into disk index" << std::endl; + std::string save_path = _mem_index_prefix + "_0"; + _mem_index_0->save(save_path.c_str()); + } else { + diskann::cout << "Index 0 is already inactive" << std::endl; + return -1; + } + } + else + { + if (_active_1.compare_exchange_strong(expected_active, false)) { + diskann::cout << "Saving mem index 1 to merge it into disk index" << std::endl; + std::string save_path = _mem_index_prefix + "_1"; + _mem_index_1->save(save_path.c_str()); + } else { + diskann::cout << "Index 1 is already inactive" << std::endl; + return -1; + } + } + diskann::cout << "Saved mem index" << std::endl; + return 0; + } + + template + void MergeInsert::save_del_set() + { + { + bool expected_value = false; + _check_switch_delete.compare_exchange_strong(expected_value, true); + std::unique_lock lock(_delete_lock); + if(_active_delete_set == 0) + { + _deletion_set_1.clear(); + bool expected_active = false; + if(_active_del_1.compare_exchange_strong(expected_active, true)) { + diskann::cout << "Cleared _deletion_set_1 - ready to accept new points" << std::endl; + } + else + { + diskann::cout << "Failed to clear _deletion_set_1" << std::endl; + } + } + else + { + _deletion_set_0.clear(); + bool expected_active = false; + if(_active_del_0.compare_exchange_strong(expected_active, true)) { + diskann::cout << "Cleared _deletion_set_0 - ready to accept new points" << std::endl; + } + else + { + diskann::cout << "Failed to clear _deletion_set_0" << std::endl; + } + } + _active_delete_set = 1 - _active_delete_set; + bool expected_active = true; + if(_active_delete_set == 0) + _active_del_1.compare_exchange_strong(expected_active, false); + else + _active_del_0.compare_exchange_strong(expected_active, false); + expected_value = true; + _check_switch_delete.compare_exchange_strong(expected_value, false); + } + + if(_active_delete_set == 0) + { + std::vector * del_vec = new std::vector(_deletion_set_1.size()); + + size_t i = 0; + for(auto iter : _deletion_set_1) + { + (*del_vec)[i] = iter; + i++; + } + _deleted_tags_vector.clear(); + _deleted_tags_vector.push_back(del_vec); + } + else + { + std::vector * del_vec = new std::vector(_deletion_set_0.size()); + + size_t i = 0; + for(auto iter : _deletion_set_0) + { + (*del_vec)[i] = iter; + i++; + } + _deleted_tags_vector.clear(); + _deleted_tags_vector.push_back(del_vec); + } + + } + + template + std::string MergeInsert::ret_merge_prefix() + { + return _disk_index_prefix_in; + } + // template class instantiations + template class MergeInsert; + template class MergeInsert; + template class MergeInsert; + template class MergeInsert; + template class MergeInsert; + template class MergeInsert; + template class MergeInsert; + template class MergeInsert; + template class MergeInsert; +} diff --git a/src/windows_aligned_file_reader.cpp b/src/windows_aligned_file_reader.cpp index e8d4553602..ca4db06209 100644 --- a/src/windows_aligned_file_reader.cpp +++ b/src/windows_aligned_file_reader.cpp @@ -7,9 +7,8 @@ #include #include "utils.h" -#define SECTOR_LEN 4096 - -void WindowsAlignedFileReader::open(const std::string& fname) { +void WindowsAlignedFileReader::open(const std::string& fname, + bool enable_writes, bool enable_create) { m_filename = std::wstring(fname.begin(), fname.end()); this->register_thread(); } @@ -67,7 +66,7 @@ IOContext& WindowsAlignedFileReader::get_ctx() { } void WindowsAlignedFileReader::read(std::vector& read_reqs, - IOContext& ctx, bool async) { + IOContext& ctx, bool async) { using namespace std::chrono_literals; // execute each request sequentially _u64 n_reqs = read_reqs.size(); diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index fa6e8ea045..1f4300f547 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -3,6 +3,15 @@ set(CMAKE_CXX_STANDARD 14) +add_executable(lazy_reinsert_batch lazy_reinsert_batch.cpp ) +if(MSVC) + target_link_options(lazy_reinsert_batch PRIVATE /MACHINE:x64 /DEBUG:FULL "/INCLUDE:_tcmalloc") + target_link_libraries(lazy_reinsert_batch debug ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}/diskann_dll.lib ${PROJECT_SOURCE_DIR}/dependencies/windows/tcmalloc/libtcmalloc_minimal.lib) + target_link_libraries(lazy_reinsert_batch optimized ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}/diskann_dll.lib ${PROJECT_SOURCE_DIR}/dependencies/windows/tcmalloc/libtcmalloc_minimal.lib) +else() + target_link_libraries(lazy_reinsert_batch ${PROJECT_NAME} -ltcmalloc) +endif() + add_executable(build_memory_index build_memory_index.cpp ) if(MSVC) target_link_options(build_memory_index PRIVATE /MACHINE:x64 /DEBUG:FULL "/INCLUDE:_tcmalloc") @@ -40,4 +49,41 @@ else() target_link_libraries(search_disk_index ${PROJECT_NAME} aio -ltcmalloc) endif() +set(CMAKE_CXX_STANDARD 14) + +add_executable(test_concurr_merge_insert test_concurr_merge_insert.cpp) +if(MSVC) + + target_link_options(test_concurr_merge_insert PRIVATE /MACHINE:x64) + target_link_libraries(test_concurr_merge_insert debug ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}/nsg_dll.lib) + target_link_libraries(test_concurr_merge_insert optimized ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}/nsg_dll.lib) +else() + target_link_libraries(test_concurr_merge_insert ${PROJECT_NAME} aio) +endif() + +add_executable(test_incremental_index test_incremental_index.cpp) +if(MSVC) + + target_link_options(test_incremental_index PRIVATE /MACHINE:x64) + target_link_libraries(test_incremental_index debug ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}/nsg_dll.lib) + target_link_libraries(test_incremental_index optimized ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}/nsg_dll.lib) +else() + target_link_libraries(test_incremental_index ${PROJECT_NAME} -ltcmalloc aio) +endif() + +add_executable(build_stream_merger_disk_index build_stream_merger_disk_index.cpp) +if(MSVC) + target_link_options(build_stream_merger_disk_index PRIVATE /MACHINE:x64 /DEBUG:FULL "/INCLUDE:_tcmalloc") + target_link_libraries(build_stream_merger_disk_index debug ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}/nsg_dll.lib ${PROJECT_SOURCE_DIR}/dependencies/windows/tcmalloc/libtcmalloc_minimal.lib) + target_link_libraries(build_stream_merger_disk_index optimized ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}/nsg_dll.lib ${PROJECT_SOURCE_DIR}/dependencies/windows/tcmalloc/libtcmalloc_minimal.lib) +else() + target_link_libraries(build_stream_merger_disk_index ${PROJECT_NAME} -ltcmalloc aio) +endif() +if (MSVC) +add_executable(test_incr_intf test_incr_intf.cpp) + target_link_options(test_incr_intf PRIVATE /MACHINE:x64 /DEBUG:FULL) + target_link_libraries(test_incr_intf debug ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}/nsg_dll.lib) + target_link_libraries(test_incr_intf optimized ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}/nsg_dll.lib) +endif() +# formatter diff --git a/tests/align_query_file.cpp b/tests/align_query_file.cpp new file mode 100644 index 0000000000..e3ac804187 --- /dev/null +++ b/tests/align_query_file.cpp @@ -0,0 +1,44 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "aux_utils.h" +#include "utils.h" + +#ifndef _WINDOWS +#include +#include +#include +#endif + +#include "memory_mapper.h" + +int main(int argc, char** argv) { + if (argc < 3) { + diskann::cout << "Correct usage : " << argv[0] + << " " << std::endl; + exit(-1); + } + + std::string input(argv[1]); + std::string output(argv[2]); + + float* data = nullptr; + size_t in_num, in_dim, in_aligned_dim; + diskann::load_aligned_bin(input.c_str(), data, in_num, in_dim, + in_aligned_dim); + diskann::save_bin(output + ".bin", data, in_num, in_aligned_dim); + + return 0; +} diff --git a/tests/build_disk_index.cpp b/tests/build_disk_index.cpp index b29c3f0d1e..c28f0cb570 100644 --- a/tests/build_disk_index.cpp +++ b/tests/build_disk_index.cpp @@ -11,30 +11,44 @@ template bool build_index(const char* dataFilePath, const char* indexFilePath, - const char* indexBuildParameters) { - return diskann::build_disk_index( - dataFilePath, indexFilePath, indexBuildParameters, diskann::Metric::L2); + const char* indexBuildParameters, diskann::Metric m, + bool singleFile) { + return diskann::build_disk_index(dataFilePath, indexFilePath, + indexBuildParameters, m, singleFile); } int main(int argc, char** argv) { - if (argc != 9) { - std::cout << "Usage: " << argv[0] - << " [data_type] [data_file.bin] " - "[index_prefix_path] " - "[R] [L] [B] [M] [T]. See README for more information on " - "parameters." - << std::endl; + if (argc != 11) { + diskann::cout << "Usage: " << argv[0] + << " " + " " + " ." + " " + " See README for more information on parameters." + << std::endl; } else { std::string params = std::string(argv[4]) + " " + std::string(argv[5]) + " " + std::string(argv[6]) + " " + std::string(argv[7]) + " " + std::string(argv[8]); + std::string dist_metric(argv[9]); + bool single_file_index = std::atoi(argv[10]) != 0; + + diskann::Metric m = + dist_metric == "cosine" ? diskann::Metric::COSINE : diskann::Metric::L2; + if (dist_metric != "l2" && m == diskann::Metric::L2) { + diskann::cout << "Metric " << dist_metric << " is not supported. Using L2" + << std::endl; + } if (std::string(argv[1]) == std::string("float")) - build_index(argv[2], argv[3], params.c_str()); + build_index(argv[2], argv[3], params.c_str(), m, + single_file_index); else if (std::string(argv[1]) == std::string("int8")) - build_index(argv[2], argv[3], params.c_str()); + build_index(argv[2], argv[3], params.c_str(), m, + single_file_index); else if (std::string(argv[1]) == std::string("uint8")) - build_index(argv[2], argv[3], params.c_str()); + build_index(argv[2], argv[3], params.c_str(), m, + single_file_index); else - std::cout << "Error. wrong file type" << std::endl; + diskann::cout << "Error. wrong file type" << std::endl; } } diff --git a/tests/build_memory_index.cpp b/tests/build_memory_index.cpp index a5a13595e3..50f5d79e4d 100644 --- a/tests/build_memory_index.cpp +++ b/tests/build_memory_index.cpp @@ -4,6 +4,7 @@ #include #include #include +#include #include "utils.h" #ifndef _WINDOWS @@ -16,10 +17,12 @@ #include "memory_mapper.h" template -int build_in_memory_index(const std::string& data_path, const unsigned R, +int build_in_memory_index(const std::string& data_path, + const std::string& tags_file, const unsigned R, const unsigned L, const float alpha, const std::string& save_path, - const unsigned num_threads) { + const unsigned num_threads, bool dynamic_index, + bool single_file_index, diskann::Metric distMetric) { diskann::Parameters paras; paras.Set("R", R); paras.Set("L", L); @@ -29,46 +32,94 @@ int build_in_memory_index(const std::string& data_path, const unsigned R, paras.Set("saturate_graph", 0); paras.Set("num_threads", num_threads); - diskann::Index index(diskann::L2, data_path.c_str()); - auto s = std::chrono::high_resolution_clock::now(); - index.build(paras); - std::chrono::duration diff = - std::chrono::high_resolution_clock::now() - s; + _u64 data_num, data_dim; + diskann::get_bin_metadata(data_path, data_num, data_dim); + diskann::cout << "Building in-memory index with parameters: data_file: " + << data_path << "tags file: " << tags_file << " R: " << R + << " L: " << L << " alpha: " << alpha + << " index_path: " << save_path << " #threads: " << num_threads + << ", using distance metric: " + << (distMetric == diskann::Metric::COSINE ? "cosine " : "l2 "); - std::cout << "Indexing time: " << diff.count() << "\n"; + typedef int TagT; + + // TODO: Should not hard code enable tags. + diskann::Index index(distMetric, data_dim, data_num, dynamic_index, + single_file_index, + true); // enable_tags forced to true! + if (dynamic_index) { + std::vector tags(data_num); + std::iota(tags.begin(), tags.end(), 0); + + auto s = std::chrono::high_resolution_clock::now(); + index.build(data_path.c_str(), data_num, paras, tags); + std::chrono::duration diff = + std::chrono::high_resolution_clock::now() - s; + + diskann::cout << "Indexing time: " << diff.count() << "\n"; + } else { + std::vector tags(data_num); + std::iota(tags.begin(), tags.end(), 0); + + auto s = std::chrono::high_resolution_clock::now(); + index.build(data_path.c_str(), data_num, paras, tags); + std::chrono::duration diff = + std::chrono::high_resolution_clock::now() - s; + + diskann::cout << "Indexing time: " << diff.count() << "\n"; + } index.save(save_path.c_str()); return 0; } int main(int argc, char** argv) { - if (argc != 8) { - std::cout << "Usage: " << argv[0] - << " [data_type] [data_file.bin] " - "[output_index_file] " - << "[R] [L] [alpha]" - << " [num_threads_to_use]. See README for more information on " - "parameters." - << std::endl; + if (argc != 12) { + diskann::cout + << "Usage: " << argv[0] + << " " + " (use null if no tags file is used) " + " " + << " " + << " ." + << " See README for more information on parameters." << std::endl; exit(-1); } - const std::string data_path(argv[2]); - const std::string save_path(argv[3]); - const unsigned R = (unsigned) atoi(argv[4]); - const unsigned L = (unsigned) atoi(argv[5]); - const float alpha = (float) atof(argv[6]); - const unsigned num_threads = (unsigned) atoi(argv[7]); + int arg_no = 2; + + const std::string data_path(argv[arg_no++]); + const std::string tags_file(argv[arg_no++]); + const std::string save_path(argv[arg_no++]); + bool dynamic_index = (bool) atoi(argv[arg_no++]); + bool single_file_index = (bool) atoi(argv[arg_no++]); + const unsigned R = (unsigned) atoi(argv[arg_no++]); + const unsigned L = (unsigned) atoi(argv[arg_no++]); + const float alpha = (float) atof(argv[arg_no++]); + const unsigned num_threads = (unsigned) atoi(argv[arg_no++]); + const std::string dist_metric_str = argv[arg_no++]; + enum diskann::Metric distMetric = + dist_metric_str == "cosine" + ? diskann::Metric::COSINE + : diskann::Metric::L2; // set to l2 even if something else is chosen + + if (dist_metric_str != "l2" && distMetric == diskann::Metric::L2) { + std::cerr << "Unknown distance metric " << argv[argc - 1] + << ". Setting metric to L2" << std::endl; + } if (std::string(argv[1]) == std::string("int8")) - build_in_memory_index(data_path, R, L, alpha, save_path, - num_threads); + build_in_memory_index(data_path, tags_file, R, L, alpha, save_path, + num_threads, dynamic_index, single_file_index, + distMetric); else if (std::string(argv[1]) == std::string("uint8")) - build_in_memory_index(data_path, R, L, alpha, save_path, - num_threads); + build_in_memory_index(data_path, tags_file, R, L, alpha, save_path, + num_threads, dynamic_index, + single_file_index, distMetric); else if (std::string(argv[1]) == std::string("float")) - build_in_memory_index(data_path, R, L, alpha, save_path, - num_threads); + build_in_memory_index(data_path, tags_file, R, L, alpha, save_path, + num_threads, dynamic_index, single_file_index, + distMetric); else - std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; + diskann::cout << "Unsupported type. Use float/int8/uint8" << std::endl; } diff --git a/tests/build_stream_merger_disk_index.cpp b/tests/build_stream_merger_disk_index.cpp new file mode 100644 index 0000000000..65183907c4 --- /dev/null +++ b/tests/build_stream_merger_disk_index.cpp @@ -0,0 +1,60 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#include "omp.h" + +#include "aux_utils.h" +#include "index.h" +#include "math_utils.h" +#include "partition_and_pq.h" +#include "utils.h" + +template +bool build_index(const char* dataFilePath, const char* indexFilePath, + const char* indexBuildParameters, const char* tag_file, + int single_file_index) { + std::string tag_filename = std::string(tag_file); + + bool save_index_as_one_file; + if (single_file_index == 0) + save_index_as_one_file = false; + else + save_index_as_one_file = true; + + if (tag_filename == "null") + return diskann::build_disk_index( + dataFilePath, indexFilePath, indexBuildParameters, diskann::Metric::L2, + save_index_as_one_file, nullptr); + return diskann::build_disk_index(dataFilePath, indexFilePath, + indexBuildParameters, diskann::Metric::L2, + save_index_as_one_file, tag_file); +} + +int main(int argc, char** argv) { + if (argc != 11) { + diskann::cout << "Usage: " << argv[0] + << " [data_type] [data_file.bin] " + "[index_prefix_path] " + "[R] [L] [B] [M] [T] [] . See README for more " + "information on " + "parameters." + << std::endl; + } else { + std::string params = std::string(argv[4]) + " " + std::string(argv[5]) + + " " + std::string(argv[6]) + " " + + std::string(argv[7]) + " " + std::string(argv[8]); + int single_file_index = atoi(argv[10]); + if (std::string(argv[1]) == std::string("float")) + build_index(argv[2], argv[3], params.c_str(), argv[9], + single_file_index); + else if (std::string(argv[1]) == std::string("int8")) + build_index(argv[2], argv[3], params.c_str(), argv[9], + single_file_index); + else if (std::string(argv[1]) == std::string("uint8")) + build_index(argv[2], argv[3], params.c_str(), argv[9], + single_file_index); + else + diskann::cout << "Error. wrong file type" << std::endl; + } +} diff --git a/tests/lazy_reinsert_batch.cpp b/tests/lazy_reinsert_batch.cpp new file mode 100644 index 0000000000..88dcf71e8e --- /dev/null +++ b/tests/lazy_reinsert_batch.cpp @@ -0,0 +1,260 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "aux_utils.h" +#include "utils.h" +#include "tsl/robin_set.h" + +#ifndef _WINDOWS +#include +#include +#include +#endif + +#include "memory_mapper.h" + +template +void search_kernel( + T* query, size_t query_num, size_t query_aligned_dim, const int recall_at, + std::vector<_u64> Lvec, diskann::Index& index, + const std::string& truthset_file, + tsl::robin_set active_tags = tsl::robin_set()) { + unsigned* gt_ids = NULL; + unsigned* gt_tags = NULL; + float* gt_dists = NULL; + size_t gt_num, gt_dim; + diskann::load_truthset(truthset_file, gt_ids, gt_dists, gt_num, gt_dim, + >_tags); + + float* query_result_dists = new float[recall_at * query_num]; + unsigned* query_result_ids = new unsigned[recall_at * query_num]; + TagT* query_result_tags = new TagT[recall_at * query_num]; + memset(query_result_dists, 0, sizeof(float) * recall_at * query_num); + memset(query_result_tags, 0, sizeof(TagT) * recall_at * query_num); + memset(query_result_ids, 0, sizeof(unsigned) * recall_at * query_num); + + std::string recall_string = "Recall@" + std::to_string(recall_at); + diskann::cout << std::setw(4) << "Ls" << std::setw(12) << "QPS " + << std::setw(18) << "Mean Latency (ms)" << std::setw(15) + << "99.9 Latency" << std::setw(12) << recall_string + << std::endl; + + diskann::cout + << "===============================================================" + "===============" + << std::endl; + + for (uint32_t test_id = 0; test_id < Lvec.size(); test_id++) { + std::vector latency_stats(query_num, 0); + memset(query_result_dists, 0, sizeof(float) * recall_at * query_num); + memset(query_result_tags, 0, sizeof(TagT) * recall_at * query_num); + memset(query_result_ids, 0, sizeof(unsigned) * recall_at * query_num); + _u64 L = Lvec[test_id]; + auto s = std::chrono::high_resolution_clock::now(); +#pragma omp parallel for + for (int64_t i = 0; i < (int64_t) query_num; i++) { + auto qs = std::chrono::high_resolution_clock::now(); + index.search_with_tags(query + i * query_aligned_dim, recall_at, (_u32) L, + query_result_tags + i * recall_at, + query_result_dists + i * recall_at); + auto qe = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = qe - qs; + latency_stats[i] = diff.count() * 1000; + // std::this_thread::sleep_for(std::chrono::milliseconds(2)); + } + auto e = std::chrono::high_resolution_clock::now(); + + std::chrono::duration diff = e - s; + float qps = (float) (query_num / diff.count()); + + float recall; + if (active_tags.size() > 0) { + recall = (float) diskann::calculate_recall( + (_u32) query_num, gt_ids, gt_dists, (_u32) gt_dim, query_result_tags, + (_u32) recall_at, (_u32) recall_at, active_tags); + } else { + recall = (float) diskann::calculate_recall( + (_u32) query_num, gt_ids, gt_dists, (_u32) gt_dim, query_result_tags, + (_u32) recall_at, (_u32) recall_at); + } + + std::sort(latency_stats.begin(), latency_stats.end()); + diskann::cout << std::setw(4) << L << std::setw(12) << qps << std::setw(18) + << std::accumulate(latency_stats.begin(), latency_stats.end(), + 0) / + (float) query_num + << std::setw(15) + << (float) latency_stats[(_u64)(0.999 * query_num)] + << std::setw(12) << recall << std::endl; + } + delete[] query_result_dists; + delete[] query_result_ids; + delete[] query_result_tags; +} + +template +int build_incremental_index(const std::string& data_path, + const std::string& memory_index_file, + const unsigned L, const unsigned R, + const unsigned C, const unsigned num_rnds, + const float alpha, const std::string& save_path, + const unsigned num_cycles, int fraction, + const std::string& query_file, + const std::string& truthset_file, + const int recall_at, std::vector<_u64> Lvec) { + diskann::Parameters paras; + paras.Set("L", L); + paras.Set("R", R); + paras.Set("C", C); + paras.Set("alpha", alpha); + paras.Set("num_rnds", num_rnds); + + T* data_load = NULL; + size_t num_points, dim, aligned_dim; + + diskann::load_aligned_bin(data_path.c_str(), data_load, num_points, dim, + aligned_dim); + + diskann::Index index(diskann::Metric::L2, dim, num_points + 100, 1, + false, true, 0); + + auto tag_path = memory_index_file + ".tags"; + index.load(memory_index_file.c_str()); + diskann::cout << "Loaded index and tags and data" << std::endl; + T* query = NULL; + size_t query_num, query_dim, query_aligned_dim; + diskann::load_aligned_bin(query_file, query, query_num, query_dim, + query_aligned_dim); + diskann::cout << "Search on static index" << std::endl; + search_kernel(query, query_num, query_aligned_dim, recall_at, Lvec, index, + truthset_file); + unsigned i = 0; + while (i < num_cycles) { + size_t delete_size = (num_points / 100) * fraction; + tsl::robin_set delete_set; + while (delete_set.size() < delete_size) + delete_set.insert(rand() % num_points); + std::vector delete_vector(delete_set.begin(), delete_set.end()); + tsl::robin_set active_tags; + for (size_t j = 0; j < num_points; j++) { + if (delete_set.find(j) == delete_set.end()) { + active_tags.insert(j); + } + } + diskann::cout << "\nDeleting " << delete_vector.size() << " elements... "; + + { + index.enable_delete(); + std::vector failed_tags; + if (index.lazy_delete(delete_set, failed_tags) < 0) { + std::cerr << "Error in delete_points" << std::endl; + } + if (failed_tags.size() > 0) { + std::cerr << "Failed to delete " << failed_tags.size() << " tags" + << std::endl; + } + /* + std::string save_del_path = + save_path + "_" + std::to_string(i) + ".delete"; + index.save(save_del_path.c_str()); + index.load(save_del_path.c_str()); + */ + diskann::Timer del_timer; + diskann::cout + << "Starting consolidation of deletes and compacting data....."; + index.consolidate(paras); + diskann::cout << "completed in " << del_timer.elapsed() / 1000000.0 + << "sec." << std::endl; + + diskann::cout << "Search post deletion....." << std::endl; + search_kernel(query, query_num, query_aligned_dim, recall_at, Lvec, index, + truthset_file, active_tags); + } + + { + diskann::Timer timer; +#pragma omp parallel for + for (size_t i = 0; i < delete_vector.size(); i++) { + unsigned p = delete_vector[i]; + index.insert_point(data_load + (size_t) p * (size_t) aligned_dim, paras, + p); + } + diskann::cout << "Re-incremental time: " << timer.elapsed() / 1000 + << "ms\n"; + search_kernel(query, query_num, query_aligned_dim, recall_at, Lvec, index, + truthset_file); + } + i++; + } + + delete[] data_load; + + return 0; +} + +int main(int argc, char** argv) { + if (argc < 16) { + diskann::cout + << "Correct usage: " << argv[0] + << " [int8/uint8/float] " + " " + " " + << " <#batches> " + "<#batch_del_size> " + " ...." + << std::endl; + exit(-1); + } + + int arg_no = 4; + unsigned L = (unsigned) atoi(argv[arg_no++]); + unsigned R = (unsigned) atoi(argv[arg_no++]); + unsigned C = (unsigned) atoi(argv[arg_no++]); + float alpha = (float) std::atof(argv[arg_no++]); + unsigned num_rnds = (unsigned) std::atoi(argv[arg_no++]); + std::string save_path(argv[arg_no++]); + unsigned num_cycles = (unsigned) atoi(argv[arg_no++]); + int fraction = (int) atoi(argv[arg_no++]); + std::string query_file(argv[arg_no++]); + std::string truthset(argv[arg_no++]); + int recall_at = (int) std::atoi(argv[arg_no++]); + std::vector<_u64> Lvec; + + for (int ctr = 15; ctr < argc; ctr++) { + _u64 curL = std::atoi(argv[ctr]); + if (curL >= recall_at) + Lvec.push_back(curL); + } + + if (Lvec.size() == 0) { + diskann::cout + << "No valid Lsearch found. Lsearch must be at least recall_at." + << std::endl; + return -1; + } + + if (std::string(argv[1]) == std::string("int8")) + build_incremental_index( + argv[2], argv[3], L, R, C, num_rnds, alpha, save_path, num_cycles, + fraction, query_file, truthset, recall_at, Lvec); + else if (std::string(argv[1]) == std::string("uint8")) + build_incremental_index( + argv[2], argv[3], L, R, C, num_rnds, alpha, save_path, num_cycles, + fraction, query_file, truthset, recall_at, Lvec); + else if (std::string(argv[1]) == std::string("float")) + build_incremental_index( + argv[2], argv[3], L, R, C, num_rnds, alpha, save_path, num_cycles, + fraction, query_file, truthset, recall_at, Lvec); + else + diskann::cout << "Unsupported type. Use float/int8/uint8" << std::endl; +} diff --git a/tests/search_disk_index.cpp b/tests/search_disk_index.cpp index 4d5c79ca89..7d0b25460d 100644 --- a/tests/search_disk_index.cpp +++ b/tests/search_disk_index.cpp @@ -53,24 +53,36 @@ int search_disk_index(int argc, char** argv) { T* query = nullptr; unsigned* gt_ids = nullptr; float* gt_dists = nullptr; + uint32_t* tags = nullptr; size_t query_num, query_dim, query_aligned_dim, gt_num, gt_dim; std::vector<_u64> Lvec; - std::string index_prefix_path(argv[2]); - std::string pq_prefix = index_prefix_path + "_pq"; - std::string disk_index_file = index_prefix_path + "_disk.index"; + int index = 2; + std::string index_prefix_path(argv[index++]); std::string warmup_query_file = index_prefix_path + "_sample_data.bin"; - _u64 num_nodes_to_cache = std::atoi(argv[3]); - _u32 num_threads = std::atoi(argv[4]); - _u32 beamwidth = std::atoi(argv[5]); - std::string query_bin(argv[6]); - std::string truthset_bin(argv[7]); - _u64 recall_at = std::atoi(argv[8]); - std::string result_output_prefix(argv[9]); + bool single_file_index = std::atoi(argv[index++]) != 0; + bool tags_flag = std::atoi(argv[index++]) != 0; + _u64 num_nodes_to_cache = std::atoi(argv[index++]); + _u32 num_threads = std::atoi(argv[index++]); + _u32 beamwidth = std::atoi(argv[index++]); + std::string query_bin(argv[index++]); + std::string truthset_bin(argv[index++]); + _u64 recall_at = std::atoi(argv[index++]); + std::string result_output_prefix(argv[index++]); + std::string dist_metric(argv[index++]); + + diskann::Metric m = + dist_metric == "cosine" ? diskann::Metric::COSINE : diskann::Metric::L2; + if (dist_metric != "l2" && m == diskann::Metric::L2) { + diskann::cout << "Unknown distance metric: " << dist_metric + << ". Using default(L2) instead." << std::endl; + } + + std::string disk_index_tag_file = index_prefix_path + "_disk.index.tags"; bool calc_recall_flag = false; - for (int ctr = 10; ctr < argc; ctr++) { + for (int ctr = index; ctr < argc; ctr++) { _u64 curL = std::atoi(argv[ctr]); if (curL >= recall_at) Lvec.push_back(curL); @@ -93,7 +105,8 @@ int search_disk_index(int argc, char** argv) { query_aligned_dim); if (file_exists(truthset_bin)) { - diskann::load_truthset(truthset_bin, gt_ids, gt_dists, gt_num, gt_dim); + diskann::load_truthset(truthset_bin, gt_ids, gt_dists, gt_num, gt_dim, + &tags); if (gt_num != query_num) { diskann::cout << "Error. Mismatch in number of queries and ground truth data" @@ -111,34 +124,45 @@ int search_disk_index(int argc, char** argv) { #endif #else reader.reset(new LinuxAlignedFileReader()); +// reader.reset(new diskann::MemAlignedFileReader()); #endif std::unique_ptr> _pFlashIndex( - new diskann::PQFlashIndex(reader)); - - int res = _pFlashIndex->load(num_threads, pq_prefix.c_str(), - disk_index_file.c_str()); + new diskann::PQFlashIndex(m, reader, single_file_index, + tags_flag)); // no tags support yet. + int res = _pFlashIndex->load(index_prefix_path.c_str(), num_threads); if (res != 0) { return res; } + + if (tags_flag) { + tsl::robin_set<_u32> active_tags; + _pFlashIndex->get_active_tags(active_tags); + + diskann::cout << "Loaded " << active_tags.size() + << " tags from index for recall measurement." << std::endl; + } else { + diskann::cout << "Not loading tags since they are disabled." << std::endl; + } + // cache bfs levels std::vector node_list; diskann::cout << "Caching " << num_nodes_to_cache << " BFS nodes around medoid(s)" << std::endl; - // _pFlashIndex->cache_bfs_levels(num_nodes_to_cache, node_list); - _pFlashIndex->generate_cache_list_from_sample_queries( - warmup_query_file, 15, 6, num_nodes_to_cache, num_threads, node_list); + _pFlashIndex->cache_bfs_levels(num_nodes_to_cache, node_list); + //_pFlashIndex->generate_cache_list_from_sample_queries( + // warmup_query_file, 15, 6, num_nodes_to_cache, num_threads, node_list); _pFlashIndex->load_cache_list(node_list); node_list.clear(); node_list.shrink_to_fit(); omp_set_num_threads(num_threads); - uint64_t warmup_L = 20; + uint64_t warmup_L; + warmup_L = 20; uint64_t warmup_num = 0, warmup_dim = 0, warmup_aligned_dim = 0; T* warmup = nullptr; - if (WARMUP) { if (file_exists(warmup_query_file)) { diskann::load_aligned_bin(warmup_query_file, warmup, warmup_num, @@ -161,15 +185,15 @@ int search_disk_index(int argc, char** argv) { } } diskann::cout << "Warming up index... " << std::flush; - std::vector warmup_result_ids_64(warmup_num, 0); + std::vector warmup_result_tags_64(warmup_num, 0); std::vector warmup_result_dists(warmup_num, 0); #pragma omp parallel for schedule(dynamic, 1) for (_s64 i = 0; i < (int64_t) warmup_num; i++) { - _pFlashIndex->cached_beam_search(warmup + (i * warmup_aligned_dim), 1, - warmup_L, - warmup_result_ids_64.data() + (i * 1), - warmup_result_dists.data() + (i * 1), 4); + _pFlashIndex->cached_beam_search_ids( + warmup + (i * warmup_aligned_dim), 1, warmup_L, + warmup_result_tags_64.data() + (i * 1), + warmup_result_dists.data() + (i * 1), (uint64_t) 4); } diskann::cout << "..done" << std::endl; } @@ -192,12 +216,11 @@ int search_disk_index(int argc, char** argv) { << std::endl; std::vector> query_result_ids(Lvec.size()); + std::vector> query_result_tags(Lvec.size()); std::vector> query_result_dists(Lvec.size()); uint32_t optimized_beamwidth = 2; - // query_num = 1; - for (uint32_t test_id = 0; test_id < Lvec.size(); test_id++) { _u64 L = Lvec[test_id]; @@ -205,54 +228,67 @@ int search_disk_index(int argc, char** argv) { // diskann::cout<<"Tuning beamwidth.." << std::endl; optimized_beamwidth = optimize_beamwidth(_pFlashIndex, warmup, warmup_num, - warmup_aligned_dim, L, optimized_beamwidth); + warmup_aligned_dim, (_u32) L, optimized_beamwidth); } else optimized_beamwidth = beamwidth; query_result_ids[test_id].resize(recall_at * query_num); query_result_dists[test_id].resize(recall_at * query_num); + query_result_tags[test_id].resize(recall_at * query_num); diskann::QueryStats* stats = new diskann::QueryStats[query_num]; - std::vector query_result_ids_64(recall_at * query_num); + std::vector query_result_tags_64(recall_at * query_num); + std::vector query_result_tags_32(recall_at * query_num); auto s = std::chrono::high_resolution_clock::now(); -#pragma omp parallel for schedule(dynamic, 1) +#pragma omp parallel for schedule(dynamic, 1) for (_s64 i = 0; i < (int64_t) query_num; i++) { _pFlashIndex->cached_beam_search( - query + (i * query_aligned_dim), recall_at, L, - query_result_ids_64.data() + (i * recall_at), + query + (i * query_aligned_dim), (uint64_t) recall_at, (uint64_t) L, + query_result_tags_32.data() + (i * recall_at), query_result_dists[test_id].data() + (i * recall_at), - optimized_beamwidth, stats + i); + (uint64_t) optimized_beamwidth, stats + i); } auto e = std::chrono::high_resolution_clock::now(); std::chrono::duration diff = e - s; - float qps = (1.0 * query_num) / (1.0 * diff.count()); + float qps = + (float) ((1.0 * (double) query_num) / (1.0 * (double) diff.count())); - diskann::convert_types(query_result_ids_64.data(), - query_result_ids[test_id].data(), - query_num, recall_at); + diskann::convert_types( + query_result_tags_32.data(), query_result_tags[test_id].data(), + (size_t) query_num, (size_t) recall_at); - float mean_latency = diskann::get_mean_stats( + float mean_latency = (float) diskann::get_mean_stats( stats, query_num, [](const diskann::QueryStats& stats) { return stats.total_us; }); - float latency_999 = diskann::get_percentile_stats( - stats, query_num, 0.999, + /* float latency_90 = (float) diskann::get_percentile_stats( + stats, query_num, 0.900, + [](const diskann::QueryStats& stats) { return stats.total_us; }); + + float latency_95 = (float) diskann::get_percentile_stats( + stats, query_num, 0.950, + [](const diskann::QueryStats& stats) { return stats.total_us; }); + */ + float latency_999 = (float) diskann::get_percentile_stats( + stats, query_num, 0.999f, [](const diskann::QueryStats& stats) { return stats.total_us; }); - float mean_ios = diskann::get_mean_stats( + float mean_ios = (float) diskann::get_mean_stats( stats, query_num, [](const diskann::QueryStats& stats) { return stats.n_ios; }); - float mean_cpuus = diskann::get_mean_stats( + float mean_cpuus = (float) diskann::get_mean_stats( stats, query_num, [](const diskann::QueryStats& stats) { return stats.cpu_us; }); + delete[] stats; float recall = 0; if (calc_recall_flag) { - recall = diskann::calculate_recall(query_num, gt_ids, gt_dists, gt_dim, - query_result_ids[test_id].data(), - recall_at, recall_at); + recall = (float) diskann::calculate_recall( + (_u32) query_num, tags, gt_dists, (_u32) gt_dim, + query_result_tags[test_id].data(), (_u32) recall_at, + (_u32) recall_at); } diskann::cout << std::setw(6) << L << std::setw(12) << optimized_beamwidth @@ -261,9 +297,9 @@ int search_disk_index(int argc, char** argv) { << std::setw(16) << mean_cpuus; if (calc_recall_flag) { diskann::cout << std::setw(16) << recall << std::endl; - } else - diskann::cout << std::endl; + } } + std::this_thread::sleep_for(std::chrono::seconds(10)); diskann::cout << "Done searching. Now saving results " << std::endl; _u64 test_id = 0; @@ -273,31 +309,45 @@ int search_disk_index(int argc, char** argv) { diskann::save_bin<_u32>(cur_result_path, query_result_ids[test_id].data(), query_num, recall_at); + cur_result_path = + result_output_prefix + "_" + std::to_string(L) + "_tags_uint32.bin"; + diskann::save_bin<_u32>(cur_result_path, query_result_tags[test_id].data(), + query_num, recall_at); cur_result_path = result_output_prefix + "_" + std::to_string(L) + "_dists_float.bin"; diskann::save_bin(cur_result_path, query_result_dists[test_id++].data(), query_num, recall_at); } + diskann::aligned_free(query); if (warmup != nullptr) diskann::aligned_free(warmup); + delete[] gt_ids; + delete[] gt_dists; return 0; } int main(int argc, char** argv) { - if (argc < 11) { + if (argc < 14) { diskann::cout << "Usage: " << argv[0] - << " [index_type] [index_prefix_path] " - " [num_nodes_to_cache] [num_threads] [beamwidth (use 0 to " - "optimize internally)] " - " [query_file.bin] [truthset.bin (use \"null\" for none)] " - " [K] [result_output_prefix] " - " [L1] [L2] etc. See README for more information on parameters." + << " " + " " + " " + " " + " [L2] etc. See README for more information on parameters." << std::endl; exit(-1); } + + diskann::cout << "Attach debugger and press a key" << std::endl; + /* char x; + std::cin >> x; + */ + if (std::string(argv[1]) == std::string("float")) search_disk_index(argc, argv); else if (std::string(argv[1]) == std::string("int8")) diff --git a/tests/search_memory_index.cpp b/tests/search_memory_index.cpp index 32592710e5..3b26188247 100644 --- a/tests/search_memory_index.cpp +++ b/tests/search_memory_index.cpp @@ -23,127 +23,156 @@ template int search_memory_index(int argc, char** argv) { T* query = nullptr; unsigned* gt_ids = nullptr; + unsigned* gt_tags = nullptr; float* gt_dists = nullptr; size_t query_num, query_dim, query_aligned_dim, gt_num, gt_dim; std::vector<_u64> Lvec; - std::string data_file(argv[2]); - std::string memory_index_file(argv[3]); - _u64 num_threads = std::atoi(argv[4]); - std::string query_bin(argv[5]); - std::string truthset_bin(argv[6]); - _u64 recall_at = std::atoi(argv[7]); - std::string result_output_prefix(argv[8]); - bool use_optimized_search = std::atoi(argv[9]); - - if ((std::string(argv[1]) != std::string("float")) && - (use_optimized_search == true)) { - std::cout << "Error. Optimized search currently only supported for " - "floating point datatypes. Using un-optimized search." - << std::endl; - use_optimized_search = false; + int arg_no = 2; + + _u64 max_points = (_u64) std::atoi(argv[arg_no++]); + std::string memory_index_file(argv[arg_no++]); + bool dynamic_index = (bool) std::atoi(argv[arg_no++]); + bool single_index_file = (bool) std::atoi(argv[arg_no++]); + // std::string data_file(argv[arg_no++]); + std::string query_bin(argv[arg_no++]); + std::string truthset_bin(argv[arg_no++]); + _u64 recall_at = std::atoi(argv[arg_no++]); + std::string result_output_prefix(argv[arg_no++]); + _u32 num_threads = (_u32) std::atoi(argv[arg_no++]); + std::string distance_metric(argv[arg_no++]); + + diskann::Metric m = (distance_metric == "cosine" ? diskann::Metric::COSINE + : diskann::Metric::L2); + + if (distance_metric != "l2" && m == diskann::Metric::L2) { + std::cout << "Not processing metric: '" << distance_metric + << "'. Setting to default (L2)" << std::endl; } bool calc_recall_flag = false; - for (int ctr = 10; ctr < argc; ctr++) { + for (int ctr = arg_no; ctr < argc; ctr++) { _u64 curL = std::atoi(argv[ctr]); if (curL >= recall_at) Lvec.push_back(curL); } if (Lvec.size() == 0) { - std::cout << "No valid Lsearch found. Lsearch must be at least recall_at." - << std::endl; + diskann::cout + << "No valid Lsearch found. Lsearch must be at least recall_at." + << std::endl; return -1; } + std::cout << "Running search on type " << typeid(T).name() + << " max_points: " << max_points + << " index file:" << memory_index_file + << " is dynamic: " << dynamic_index + << " is single file: " << single_index_file + << " query file: " << query_bin + << " truthset file: " << truthset_bin << " K: " << recall_at + << " num threads: " << num_threads + << " save prefix: " << result_output_prefix + << " similarity metric: " + << (m == diskann::Metric::COSINE ? "cosine" : "l2") + << " first L: " << Lvec[0] << std::endl; + diskann::load_aligned_bin(query_bin, query, query_num, query_dim, query_aligned_dim); if (file_exists(truthset_bin)) { - diskann::load_truthset(truthset_bin, gt_ids, gt_dists, gt_num, gt_dim); + diskann::load_truthset(truthset_bin, gt_ids, gt_dists, gt_num, gt_dim, + >_tags); if (gt_num != query_num) { - std::cout << "Error. Mismatch in number of queries and ground truth data" - << std::endl; + diskann::cout + << "Error. Mismatch in number of queries and ground truth data" + << std::endl; } calc_recall_flag = true; } - std::cout.setf(std::ios_base::fixed, std::ios_base::floatfield); - std::cout.precision(2); + diskann::cout.setf(std::ios_base::fixed, std::ios_base::floatfield); + diskann::cout.precision(2); - auto metric = diskann::L2; - if (use_optimized_search) - metric = diskann::FAST_L2; - diskann::Index index(metric, data_file.c_str()); - index.load(memory_index_file.c_str()); // to load NSG - std::cout << "Index loaded" << std::endl; + diskann::Index index(m, query_dim, max_points, dynamic_index, + single_index_file, dynamic_index); + index.load(memory_index_file.c_str()); - if (use_optimized_search) - index.optimize_graph(); + tsl::robin_set active_tags; + index.get_active_tags(active_tags); + + diskann::cout << "Index loaded" << std::endl; diskann::Parameters paras; std::string recall_string = "Recall@" + std::to_string(recall_at); - std::cout << std::setw(4) << "Ls" << std::setw(12) << "QPS " << std::setw(18) - << "Mean Latency (mus)" << std::setw(15) << "99.9 Latency" - << std::setw(12) << recall_string << std::endl; - std::cout << "===============================================================" - "===============" - << std::endl; + diskann::cout << std::setw(4) << "Ls" << std::setw(12) << "QPS " + << std::setw(18) << "Mean Latency (ms)" << std::setw(15) + << "99.9 Latency" << std::setw(12) << recall_string + << std::endl; + diskann::cout + << "===============================================================" + "===============" + << std::endl; std::vector> query_result_ids(Lvec.size()); std::vector> query_result_dists(Lvec.size()); + std::vector> query_result_tags(Lvec.size()); std::vector latency_stats(query_num, 0); for (uint32_t test_id = 0; test_id < Lvec.size(); test_id++) { _u64 L = Lvec[test_id]; query_result_ids[test_id].resize(recall_at * query_num); + query_result_tags[test_id].resize(recall_at * query_num); + query_result_dists[test_id].resize(recall_at * query_num); auto s = std::chrono::high_resolution_clock::now(); - omp_set_num_threads(num_threads); -#pragma omp parallel for schedule(dynamic, 1) + +#pragma omp parallel for schedule(dynamic, 1) num_threads(num_threads) for (int64_t i = 0; i < (int64_t) query_num; i++) { auto qs = std::chrono::high_resolution_clock::now(); - if (use_optimized_search) { - index.search_with_opt_graph( - query + i * query_aligned_dim, recall_at, L, - query_result_ids[test_id].data() + i * recall_at); - } else { - index.search(query + i * query_aligned_dim, recall_at, L, + if (dynamic_index) + index.search_with_tags( + query + i * query_aligned_dim, (uint64_t) recall_at, (_u32) L, + query_result_ids[test_id].data() + i * recall_at, + query_result_dists[test_id].data() + i * recall_at); + else + index.search(query + i * query_aligned_dim, recall_at, (_u32) L, query_result_ids[test_id].data() + i * recall_at); - } auto qe = std::chrono::high_resolution_clock::now(); std::chrono::duration diff = qe - qs; - latency_stats[i] = diff.count() * 1000000; + latency_stats[i] = diff.count() * 1000; } auto e = std::chrono::high_resolution_clock::now(); std::chrono::duration diff = e - s; - float qps = (query_num / diff.count()); + float qps = (float) ((double) query_num / (double) diff.count()); float recall = 0; if (calc_recall_flag) - recall = diskann::calculate_recall(query_num, gt_ids, gt_dists, gt_dim, - query_result_ids[test_id].data(), - recall_at, recall_at); + + recall = (float) diskann::calculate_recall( + (_u32) query_num, gt_tags, gt_dists, (_u32) gt_dim, + query_result_ids[test_id].data(), (_u32) recall_at, (_u32) recall_at, + active_tags); std::sort(latency_stats.begin(), latency_stats.end()); double mean_latency = 0; for (uint64_t q = 0; q < query_num; q++) { mean_latency += latency_stats[q]; } - mean_latency /= query_num; + mean_latency /= (double) query_num; - std::cout << std::setw(4) << L << std::setw(12) << qps << std::setw(18) - << (float) mean_latency << std::setw(15) - << (float) latency_stats[(_u64)(0.999 * query_num)] - << std::setw(12) << recall << std::endl; + diskann::cout << std::setw(4) << L << std::setw(12) << qps << std::setw(18) + << (float) mean_latency << std::setw(15) + << (float) latency_stats[(_u64)(0.999 * (double) query_num)] + << std::setw(12) << recall << std::endl; } - std::cout << "Done searching. Now saving results " << std::endl; + diskann::cout << "Done searching. Now saving results " << std::endl; _u64 test_id = 0; + for (auto L : Lvec) { std::string cur_result_path = result_output_prefix + "_" + std::to_string(L) + "_idx_uint32.bin"; @@ -157,18 +186,21 @@ int search_memory_index(int argc, char** argv) { } int main(int argc, char** argv) { - if (argc < 11) { - std::cout + if (argc < 13) { + diskann::cout << "Usage: " << argv[0] - << " [index_type] [data_file.bin] " - "[memory_index_path] [num_threads] " - "[query_file.bin] [truthset.bin (use \"null\" for none)] " - " [K] [result_output_prefix] [use_optimized_search (for small ~1M " - "data)] " - " [L1] [L2] etc. See README for more information on parameters. " + << " " + " " + " " + " " + " " + " [L2] etc. See README for more information on parameters. " << std::endl; exit(-1); } + if (std::string(argv[1]) == std::string("int8")) search_memory_index(argc, argv); else if (std::string(argv[1]) == std::string("uint8")) @@ -176,5 +208,5 @@ int main(int argc, char** argv) { else if (std::string(argv[1]) == std::string("float")) search_memory_index(argc, argv); else - std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; + diskann::cout << "Unsupported type. Use float/int8/uint8" << std::endl; } diff --git a/tests/test_concurr_insert_search.cpp b/tests/test_concurr_insert_search.cpp new file mode 100644 index 0000000000..571eec22c5 --- /dev/null +++ b/tests/test_concurr_insert_search.cpp @@ -0,0 +1,172 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#include +#include +#include +#include +#include +#include +#include + +#include "utils.h" + +#ifndef _WINDOWS +#include +#include +#include +#endif + +#include "memory_mapper.h" + +template +void search_kernel(T* query, size_t query_num, size_t query_aligned_dim, + const int recall_at, std::vector<_u64> Lvec, + diskann::Index& index) { + uint32_t* results = new uint32_t[recall_at * query_num]; + for (uint32_t test_id = 0; test_id < Lvec.size(); test_id++) { +#pragma omp parallel for num_threads(8) + for (int64_t i = 0; i < (int64_t) query_num; i++) { + _u64 L = Lvec[test_id]; + index.search(query + i * query_aligned_dim, recall_at, L, + results + i * recall_at); + } + } + delete[] results; +} + +template +void insert_kernel(T* data_load, diskann::Index& index, + diskann::Parameters& parameters, size_t num_points, + size_t aligned_dim, unsigned start) { + diskann::cout << "Insertion thread" << std::endl; + diskann::Timer timer; +// do not parallelize without making _iterate_to_fixed_point() thread safe +#pragma omp parallel for num_threads(16) + for (unsigned insertions = start; insertions < num_points; ++insertions) { + index.insert_point(data_load + insertions * aligned_dim, parameters, + insertions); + } + + unsigned time_secs = timer.elapsed() / 1000000; + diskann::cout << "Insert time : " << time_secs << " s\n" + << "Inserts/sec : " << num_points / time_secs << std::endl; +} +template +int build_incremental_index(const std::string& data_path, const unsigned L, + const unsigned R, const unsigned C, + const unsigned num_rnds, const float alpha, + const std::string& save_path, + const std::string& query_file, const int recall_at, + std::vector<_u64> Lvec) { + diskann::Parameters paras; + paras.Set("L", L); + paras.Set("R", R); + paras.Set("C", C); + paras.Set("alpha", alpha); + paras.Set("num_rnds", num_rnds); + paras.Set("saturate_graph", 0); + + T* data_load = NULL; + size_t num_points, dim, aligned_dim; + + diskann::load_aligned_bin(data_path.c_str(), data_load, num_points, dim, + aligned_dim); + + unsigned num_prebuilt = 1; + + typedef int TagT; + diskann::Index index(diskann::L2, dim, num_points, 1, true, false); + + { + std::vector tags(1); + std::iota(tags.begin(), tags.end(), 0); + + diskann::Timer timer; + index.build(data_path.c_str(), num_points, paras, tags); + diskann::cout << "Index build time: " << timer.elapsed() / 1000 << "ms\n"; + } + + diskann::cout << "Saved batch build" << std::endl; + T* query = NULL; + size_t query_num, query_dim, query_aligned_dim; + + diskann::load_aligned_bin(query_file, query, query_num, query_dim, + query_aligned_dim); + { + std::future future = std::async( + std::launch::async, insert_kernel, data_load, std::ref(index), + std::ref(paras), num_points, aligned_dim, num_prebuilt); + + std::future_status status; + + unsigned total_queries = 0; + do { + status = future.wait_for(std::chrono::milliseconds(1)); + if (status == std::future_status::deferred) { + diskann::cout << "deferred\n"; + } else if (status == std::future_status::timeout) { + search_kernel(query, query_num, query_aligned_dim, recall_at, Lvec, + index); + total_queries += query_num; + diskann::cout << "total queries: " << total_queries << std::endl; + } else if (status == std::future_status::ready) { + diskann::cout << "Insertions complete!\n"; + } + } while (status != std::future_status::ready); + } + auto save_path_inc = save_path + ".inc"; + index.save(save_path_inc.c_str()); + diskann::cout << "Saved file post insertions" << std::endl; + + delete[] data_load; + + return 0; +} + +int main(int argc, char** argv) { + if (argc < 12) { + diskann::cout << "Correct usage: " << argv[0] + << " type[int8/uint8/float] data_file L R C alpha " + "num_rounds " + << "save_graph_file query_file recall@ L1 L2....." + << std::endl; + exit(-1); + } + + unsigned L = (unsigned) atoi(argv[3]); + unsigned R = (unsigned) atoi(argv[4]); + unsigned C = (unsigned) atoi(argv[5]); + float alpha = (float) std::atof(argv[6]); + unsigned num_rnds = (unsigned) std::atoi(argv[7]); + std::string save_path(argv[8]); + std::string query_file(argv[9]); + int recall_at = (int) atoi(argv[10]); + + std::vector<_u64> Lvec; + + for (int ctr = 11; ctr < argc; ctr++) { + _u64 curL = std::atoi(argv[ctr]); + if (curL >= (_u64) recall_at) + Lvec.push_back(curL); + } + + if (Lvec.size() == 0) { + diskann::cout + << "No valid Lsearch found. Lsearch must be at least recall_at." + << std::endl; + return -1; + } + + if (std::string(argv[1]) == std::string("int8")) + build_incremental_index(argv[2], L, R, C, num_rnds, alpha, + save_path, query_file, recall_at, Lvec); + else if (std::string(argv[1]) == std::string("uint8")) + build_incremental_index(argv[2], L, R, C, num_rnds, alpha, + save_path, query_file, recall_at, Lvec); + else if (std::string(argv[1]) == std::string("float")) + build_incremental_index(argv[2], L, R, C, num_rnds, alpha, save_path, + query_file, recall_at, Lvec); + else + diskann::cout << "Unsupported type. Use float/int8/uint8" << std::endl; +} diff --git a/tests/test_concurr_merge_insert.cpp b/tests/test_concurr_merge_insert.cpp new file mode 100644 index 0000000000..040831197a --- /dev/null +++ b/tests/test_concurr_merge_insert.cpp @@ -0,0 +1,749 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#include "v2/index_merger.h" +#include "v2/merge_insert.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "aux_utils.h" +#include "utils.h" +#include "math_utils.h" +#include "partition_and_pq.h" + +#ifndef _WINDOWS +#include +#include +#include +#endif + +#include +#include + +#define NUM_INSERT_THREADS 2 +#define NUM_DELETE_THREADS 1 +#define NUM_SEARCH_THREADS 6 +// random number generator +std::random_device dev; +std::mt19937 rng(dev()); + +tsl::robin_map params; +float mem_alpha, merge_alpha; +uint32_t medoid_id = std::numeric_limits::max(); +std::atomic_bool _insertions_done(true); +std::atomic_bool _del_done(true); +std::vector Lvec; +std::future delete_future; +std::future insert_future; +std::future merge_future; +diskann::Timer global_timer; +std::string all_points_file; +bool save_index_as_one_file; +std::string TMP_FOLDER; +std::string query_file = ""; +std::string truthset_file = ""; + +template +void seed_iter(tsl::robin_set &active_set, + tsl::robin_set &inactive_set, + const std::string & inserted_points_file, + const std::string & inserted_tags_file, + tsl::robin_set & deleted_tags) { + const uint32_t insert_count = params[std::string("insert_count")]; + const uint32_t delete_count = params[std::string("delete_count")]; + const uint32_t ndims = params[std::string("ndims")]; + std::cout << "ITER: start = " << active_set.size() << ", " + << inactive_set.size() << "\n"; + + // pick `delete_count` tags + std::vector active_vec(active_set.begin(), active_set.end()); + std::shuffle(active_vec.begin(), active_vec.end(), rng); + std::vector delete_vec; + if (active_vec.size() < delete_count) + delete_vec.insert(delete_vec.end(), active_vec.begin(), active_vec.end()); + else + delete_vec.insert(delete_vec.end(), active_vec.begin(), + active_vec.begin() + delete_count); + for (auto iter : delete_vec) + deleted_tags.insert(iter); + active_set.clear(); + active_set.insert(active_vec.begin() + delete_vec.size(), active_vec.end()); + std::cout << "ITER: DELETE - " << delete_vec.size() << " IDs\n"; + // pick `insert_count` tags + std::vector inactive_vec(inactive_set.begin(), inactive_set.end()); + std::shuffle(inactive_vec.begin(), inactive_vec.end(), rng); + std::vector insert_vec; + if (inactive_vec.size() < insert_count) + insert_vec.insert(insert_vec.end(), inactive_vec.begin(), + inactive_vec.end()); + else + insert_vec.insert(insert_vec.end(), inactive_vec.begin(), + inactive_vec.begin() + insert_count); + inactive_set.clear(); + + std::cout << "ITER: INSERT - " << insert_vec.size() << " IDs in " + << inserted_tags_file << "\n"; + inactive_set.insert(inactive_vec.begin() + insert_vec.size(), + inactive_vec.end()); + std::sort(insert_vec.begin(), insert_vec.end()); + TagT *tag_data = new TagT[insert_vec.size()]; + for (size_t i = 0; i < insert_vec.size(); i++) + tag_data[i] = insert_vec[i]; + diskann::save_bin(inserted_tags_file, tag_data, insert_vec.size(), 1); + delete[] tag_data; + + // use ifstream reader to load node coordinates + std::ifstream base_reader; + base_reader.open(::all_points_file, std::ios::binary | std::ios::ate); + + base_reader.seekg(2 * sizeof(uint32_t), std::ios::beg); + + std::ofstream inserted_points_writer(inserted_points_file, std::ios::binary); + T *new_pts = new T[(uint32_t) insert_vec.size() * (uint32_t) ndims]; + for (uint64_t idx = 0; idx < insert_vec.size(); idx++) { + uint32_t actual_idx = insert_vec[idx]; + T * point = new T[ndims]; + base_reader.seekg( + (2 * sizeof(uint32_t) + actual_idx * (uint64_t) ndims * sizeof(T)), + std::ios::beg); + base_reader.read((char *) point, ((uint64_t) ndims) * sizeof(T)); + T *dest_ptr = new_pts + idx * (uint64_t) ndims; + std::memcpy(dest_ptr, point, ndims * sizeof(T)); + delete[] point; + } + base_reader.close(); + + uint32_t npts_u32 = (uint32_t) insert_vec.size(); + uint32_t ndims_u32 = ndims; + inserted_points_writer.write((char *) &npts_u32, sizeof(uint32_t)); + inserted_points_writer.write((char *) &ndims_u32, sizeof(uint32_t)); + inserted_points_writer.write( + (char *) new_pts, + (uint64_t) insert_vec.size() * (uint64_t) ndims * sizeof(T)); + inserted_points_writer.close(); + delete[] new_pts; + + // balance tags + inactive_set.insert(delete_vec.begin(), delete_vec.end()); + active_set.insert(insert_vec.begin(), insert_vec.end()); + + diskann::cout << "ITER: end = " << active_set.size() << ", " + << inactive_set.size() << "\n"; +#ifndef _WINDOWS + std::cout << "ITER: end = " << active_set.size() << ", " + << inactive_set.size() << "\n"; + malloc_stats(); +#endif +} + +float compute_active_recall(const uint32_t *result_tags, + const uint32_t result_count, + const uint32_t *gs_tags, const uint64_t gs_count, + const tsl::robin_set &inactive_set) { + tsl::robin_set active_gs; + for (uint32_t i = 0; i < gs_count && active_gs.size() < result_count; i++) { + auto iter = inactive_set.find(gs_tags[i]); + if (iter == inactive_set.end()) { + active_gs.insert(gs_tags[i]); + } + } + uint32_t match = 0; + for (uint32_t i = 0; i < result_count; i++) { + match += (active_gs.find(result_tags[i]) != active_gs.end()); + } + return ((float) match / (float) result_count) * 100; +} + +template +void search_disk_index(const std::string & index_prefix_path, + const tsl::robin_set &inactive_tags, + const std::string & query_path, + const std::string & gs_path) { + std::string pq_prefix = index_prefix_path + "_pq"; + std::string disk_index_file = index_prefix_path + "_disk.index"; + std::string warmup_query_file = index_prefix_path + "_sample_data.bin"; + uint32_t beamwidth = params[std::string("beam_width")]; + uint32_t num_threads = 60; + std::string query_bin = query_path; + std::string truthset_bin = gs_path; + uint64_t recall_at = params[std::string("recall_k")]; + uint64_t search_L = ::Lvec[0]; + // hold data + T * query = nullptr; + unsigned *gt_ids = nullptr; + uint32_t *gt_tags = nullptr; + float * gt_dists = nullptr; + size_t query_num, query_dim, query_aligned_dim, gt_num, gt_dim; + + // load query + truthset + diskann::load_aligned_bin(query_bin, query, query_num, query_dim, + query_aligned_dim); + diskann::load_truthset(truthset_bin, gt_ids, gt_dists, gt_num, gt_dim); + if (gt_num != query_num) { + std::cout << "Error. Mismatch in number of queries and ground truth data" + << std::endl; + } + + // load PQ Flash Index + std::shared_ptr reader(new LinuxAlignedFileReader()); + std::unique_ptr> _pFlashIndex( + new diskann::PQFlashIndex(diskann::Metric::L2, reader, + ::save_index_as_one_file, true)); + int res = _pFlashIndex->load(num_threads, pq_prefix.c_str(), + disk_index_file.c_str()); + if (res != 0) { + std::cerr << "Failed to load index.\n"; + exit(-1); + } + + // prep for search + std::vector query_result_ids; + std::vector query_result_tags; + std::vector query_result_dists; + query_result_ids.resize(recall_at * query_num); + query_result_dists.resize(recall_at * query_num); + query_result_tags.resize(recall_at * query_num); + diskann::QueryStats * stats = new diskann::QueryStats[query_num]; + std::vector query_result_ids_64(recall_at * query_num); +#pragma omp parallel for schedule(dynamic, 1) // num_threads(1) + for (_s64 i = 0; i < (int64_t) query_num; i++) { + _pFlashIndex->cached_beam_search( + query + (i * query_aligned_dim), recall_at, search_L, + query_result_ids_64.data() + (i * recall_at), + query_result_dists.data() + (i * recall_at), beamwidth, stats + i, + query_result_tags.data() + (i * recall_at)); + } + + // compute mean recall, IOs + float mean_recall = 0.0f; + for (uint32_t i = 0; i < query_num; i++) { + auto *result_tags = query_result_tags.data() + (i * recall_at); + auto *gs_tags = gt_tags + (i * gt_dim); + float query_recall = compute_active_recall( + result_tags, (uint32_t) recall_at, gs_tags, gt_dim, inactive_tags); + mean_recall += query_recall; + } + mean_recall /= query_num; + + float mean_ios = (float) diskann::get_mean_stats( + stats, query_num, + [](const diskann::QueryStats &stats) { return stats.n_ios; }); + std::cout << "PQFlashIndex :: recall-" << recall_at << "@" << recall_at + << ": " << mean_recall << ", mean IOs: " << mean_ios << "\n"; + diskann::aligned_free(query); + delete[] stats; + delete[] gt_ids; + delete[] gt_dists; + delete[] gt_tags; +} + +template +void search_kernel(diskann::MergeInsert & merge_insert, + const tsl::robin_set &active_tags, + bool print_stats = false) { + uint64_t recall_at = params[std::string("recall_k")]; + + // hold data + T * query = nullptr; + unsigned *gt_ids = nullptr; + uint32_t *gt_tags = nullptr; + float * gt_dists = nullptr; + size_t query_num, query_dim, query_aligned_dim, gt_num, gt_dim; + + const std::string temp = "/mnt/t-adisin/sift_query.bin"; + std::cout << "Loading query : " << temp << std::endl; + // load query + truthset + diskann::load_aligned_bin(temp, query, query_num, query_dim, + query_aligned_dim); + std::cout << "Loaded query : " << temp << std::endl; + diskann::load_truthset(::truthset_file, gt_ids, gt_dists, gt_num, gt_dim, + >_tags); + std::cout << "Loaded gt" << std::endl; + if (gt_num != query_num) { + std::cout << "Error. Mismatch in number of queries and ground truth data" + << std::endl; + } + + if (print_stats) { + std::string recall_string = "SS-Recall@" + std::to_string(recall_at); + std::cout << std::setw(4) << "Ls" << std::setw(12) << "QPS " + << std::setw(18) << "Mean Latency (ms)" << std::setw(12) + << "90 Latency" << std::setw(12) << "95 Latency" << std::setw(12) + << "99 Latency" << std::setw(12) << "99.9 Latency" + << std::setw(12) << recall_string << std::setw(12) + << "Mean disk IOs" << std::endl; + + std::cout + + << "===============================================================" + "===============" + << std::endl; + } else { + std::string recall_string = "Recall@" + std::to_string(recall_at); + std::cout << std::setw(4) << "Ls" << std::setw(12) << "QPS " + << std::setw(18) << "Mean Latency (ms)" << std::setw(12) + << "90 Latency" << std::setw(12) << "95 Latency" << std::setw(12) + << "99 Latency" << std::setw(12) << "99.9 Latency" + << std::setw(12) << recall_string << std::setw(12) + << "Mean disk IOs" << std::endl; + std::cout + << "===============================================================" + "===============" + << std::endl; + } + + // prep for search + std::vector query_result_ids; + std::vector query_result_tags; + std::vector query_result_dists; + query_result_ids.resize(recall_at * query_num); + query_result_dists.resize(recall_at * query_num); + query_result_tags.resize(recall_at * query_num); + std::vector query_result_ids_32(recall_at * query_num); + + for (size_t test_id = 0; test_id < ::Lvec.size(); test_id++) { + diskann::QueryStats *stats = new diskann::QueryStats[query_num]; + uint32_t L = Lvec[test_id]; + std::vector latency_stats(query_num, 0); + auto s = std::chrono::high_resolution_clock::now(); +#pragma omp parallel for num_threads(NUM_SEARCH_THREADS) + for (_s64 i = 0; i < (int64_t) query_num; i++) { + auto qs = std::chrono::high_resolution_clock::now(); + merge_insert.search_sync(query + (i * query_aligned_dim), recall_at, L, + (query_result_tags.data() + (i * recall_at)), + query_result_dists.data() + (i * recall_at), + stats + i); + auto qe = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = qe - qs; + latency_stats[i] = diff.count() * 1000; + // std::this_thread::sleep_for(std::chrono::milliseconds(2)); + } + auto e = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = e - s; + float qps = (float) (((double) query_num) / diff.count()); + // compute mean recall, IOs + float mean_recall = 0.0f; + mean_recall = diskann::calculate_recall( + (unsigned) query_num, gt_ids, gt_dists, (unsigned) gt_dim, + query_result_tags.data(), (unsigned) recall_at, (unsigned) recall_at, + active_tags); + // mean_recall /= (float) query_num; + float mean_ios = (float) diskann::get_mean_stats( + stats, query_num, + [](const diskann::QueryStats &stats) { return stats.n_ios; }); + std::sort(latency_stats.begin(), latency_stats.end()); + std::cout << std::setw(4) << L << std::setw(12) << qps << std::setw(18) + << ((float) std::accumulate(latency_stats.begin(), + latency_stats.end(), 0)) / + (float) query_num + << std::setw(12) + << (float) latency_stats[(_u64)(0.90 * ((double) query_num))] + << std::setw(12) + << (float) latency_stats[(_u64)(0.95 * ((double) query_num))] + << std::setw(12) + << (float) latency_stats[(_u64)(0.99 * ((double) query_num))] + << std::setw(12) + << (float) latency_stats[(_u64)(0.999 * ((double) query_num))] + << std::setw(12) << mean_recall << std::setw(12) << mean_ios + << std::endl; + delete[] stats; + } + diskann::aligned_free(query); + delete[] gt_ids; + delete[] gt_dists; + delete[] gt_tags; +} + +template +void insertion_kernel(diskann::MergeInsert &merge_insert, + std::string mem_pts_file, std::string mem_tags_file) { + if (::_insertions_done.load()) { + std::cout << "Insertions_done is true at the beginning of insertion kernel" + << std::endl; + exit(-1); + } + T * data_insert = nullptr; + size_t npts, ndim, aligned_dim; + diskann::load_aligned_bin(mem_pts_file, data_insert, npts, ndim, + aligned_dim); + size_t tag_num, tag_dim; + TagT * tag_data; + diskann::load_bin(mem_tags_file, tag_data, tag_num, tag_dim); + if (tag_num != npts) { + std::cout << "In insertion_kernel(), number of tags loaded is not equal to " + "number of points loaded. Exiting....." + << std::endl; + exit(-1); + } + _s64 i; + std::vector insert_latencies(npts, 0); + diskann::Timer timer; +#pragma omp parallel for num_threads(NUM_INSERT_THREADS) + for (i = 0; i < (_s64) npts; i++) { + diskann::Timer insert_timer; + if (merge_insert.insert(data_insert + i * aligned_dim, tag_data[i]) == 0) { + insert_latencies[i] = ((double) insert_timer.elapsed()); + } else { + std::cout << "Point " << i << "could not be inserted." << std::endl; + } + if ((i % 1000000 == 0) && (i > 0)) + std::cout << "Inserted another 1M points" << std::endl; + } + std::cout << "Mem index insertion time : " << timer.elapsed() / 1000 << " ms" + << std::endl + << "10th percentile insertion time : " + << insert_latencies[(size_t)(0.10 * ((double) npts))] << " microsec" + << std::endl + << "50th percentile insertion time : " + << insert_latencies[(size_t)(0.5 * ((double) npts))] << " microsec" + << "90th percentile insertion time : " + << insert_latencies[(size_t)(0.90 * ((double) npts))] << " microsec" + << std::endl; + ::_insertions_done.store(true); + delete[] data_insert; + delete[] tag_data; +} +template +void deletion_kernel(diskann::MergeInsert &merge_insert, + tsl::robin_set del_tags) { + if (::_del_done.load()) { + std::cout << "_del_done is already true" << std::endl; + exit(-1); + } + std::this_thread::sleep_for(std::chrono::milliseconds(1000)); + diskann::Timer timer; + for (auto iter : del_tags) { + merge_insert.lazy_delete(iter); + } + std::cout << "Deletion time : " << timer.elapsed() / 1000 << " ms" + << std::endl; + ::_del_done.store(true); +} + +template +void merge_kernel(diskann::MergeInsert &merge_insert) { + merge_insert.final_merge(); +} + +template +void run_iter(diskann::MergeInsert & merge_insert, + const std::string & mem_prefix, + tsl::robin_set &active_set, + tsl::robin_set &inactive_set) { + // files for mem-DiskANN + std::string mem_pts_file = mem_prefix + ".data_orig"; + std::string mem_tags_file = mem_prefix + ".tags_orig"; + std::this_thread::sleep_for(std::chrono::seconds(10)); + + ::merge_future = + std::async(std::launch::async, merge_kernel, std::ref(merge_insert)); + + while (!(::_insertions_done.load() && ::_del_done.load())) { + std::cout << "Search at " << ::global_timer.elapsed() / 1000000 + << " seconds " << std::endl; + search_kernel(merge_insert, active_set); + std::this_thread::sleep_for(std::chrono::milliseconds(5000)); + } + + if (::_insertions_done.load() && ::_del_done.load()) { + ::_insertions_done.store(false); + ::_del_done.store(false); + + std::cout << "Searching all indices" << std::endl; + std::cout << "Search at " << ::global_timer.elapsed() / 1000000 + << " seconds " << std::endl; + search_kernel(merge_insert, active_set, true); + + std::cout << "ITER: Seeding iteration" + << "\n"; + // seed the iteration + tsl::robin_set deleted_tags; + seed_iter(active_set, inactive_set, mem_pts_file, mem_tags_file, + deleted_tags); + ::delete_future = std::async(std::launch::async, deletion_kernel, + std::ref(merge_insert), deleted_tags); + ::insert_future = + std::async(std::launch::async, insertion_kernel, + std::ref(merge_insert), mem_pts_file, mem_tags_file); + } + + std::future_status merge_status; + do { + merge_status = ::merge_future.wait_for(std::chrono::milliseconds(1)); + std::cout << "Search at " << ::global_timer.elapsed() / 1000000 + << " seconds " << std::endl; + search_kernel(merge_insert, active_set); + + std::this_thread::sleep_for(std::chrono::milliseconds(1000)); + } while ((merge_status != std::future_status::ready)); +} + +template +void run_single_iter(diskann::MergeInsert & merge_insert, + const std::string & base_prefix, + const std::string & merge_prefix, + const std::string & mem_prefix, + tsl::robin_set &active_set, + tsl::robin_set &inactive_set, + diskann::Distance * dist_cmp) { + // files for mem-DiskANN + std::string mem_pts_file = mem_prefix + ".data_orig"; + std::string mem_tags_file = mem_prefix + ".tags_orig"; + if (::_insertions_done.load() && ::_del_done.load()) { + ::_insertions_done.store(false); + ::_del_done.store(false); + + /* std::cout << "Searching all indices" << std::endl; + std::cout << "Search at " << ::global_timer.elapsed() / 1000000 + << " seconds " << std::endl; + search_kernel(merge_insert, active_set, true); + */ + std::cout << "ITER: Seeding iteration" + << "\n"; + // seed the iteration + tsl::robin_set deleted_tags; + seed_iter(active_set, inactive_set, mem_pts_file, mem_tags_file, + deleted_tags); + ::delete_future = std::async(std::launch::async, deletion_kernel, + std::ref(merge_insert), deleted_tags); + ::insert_future = + std::async(std::launch::async, insertion_kernel, + std::ref(merge_insert), mem_pts_file, mem_tags_file); + } + std::future_status insert_status, delete_status; + do { + insert_status = ::insert_future.wait_for(std::chrono::milliseconds(1)); + delete_status = ::delete_future.wait_for(std::chrono::milliseconds(1)); + std::this_thread::sleep_for(std::chrono::seconds(60)); + } while ((insert_status != std::future_status::ready) || + (delete_status != std::future_status::ready)); + + ::merge_future = + std::async(std::launch::async, merge_kernel, std::ref(merge_insert)); + + std::future_status merge_status; + do { + merge_status = ::merge_future.wait_for(std::chrono::milliseconds(1)); + /* std::cout << "Search at " << ::global_timer.elapsed() / 1000000 + << " seconds " << std::endl; + search_kernel(merge_insert, active_set); + */ + // std::this_thread::sleep_for(std::chrono::milliseconds(10000)); + } while ((merge_status != std::future_status::ready)); +} + +template +void run_all_iters(std::string base_prefix, std::string merge_prefix, + const std::string mem_prefix, const std::string data_file, + const std::string active_tags_file, + diskann::Distance *dist_cmp) { + // load all data points + uint64_t npts = 0, ndims = 0; + diskann::get_bin_metadata(data_file, npts, ndims); + std::cout << "Loaded base bin" << std::endl; + params[std::string("ndims")] = (uint32_t) ndims; + + uint32_t n_iters = params["n_iters"]; + // load active tags + tsl::robin_set active_tags; + TagT * tag_data; + size_t tag_num, tag_dim; + if (::save_index_as_one_file) { + uint64_t *metadata; + size_t nr, nc; + diskann::load_bin(active_tags_file, metadata, nr, nc); + diskann::load_bin(active_tags_file, tag_data, tag_num, tag_dim, + metadata[7]); + } else { + diskann::load_bin(active_tags_file, tag_data, tag_num, tag_dim); + } + + size_t tags_loaded = 0; + size_t del_tags_found = 0; + active_tags.reserve(tag_num); + for (size_t i = 0; i < tag_num; i++) { + if (tag_data[i] != std::numeric_limits::max()) { + active_tags.insert(tag_data[i]); + tags_loaded++; + } else { + if (del_tags_found < 5) + std::cout << "Driver file found invalid tag in active tag file : " + << tag_data[i] << std::endl; + del_tags_found++; + } + } + std::cout << "Loaded " << tags_loaded << " tags" << std::endl; + delete[] tag_data; + std::cout << del_tags_found + << " deleted/invalid tags found in active tags file" << std::endl; + // read medoid ID from base_prefix + std::ifstream disk_reader(base_prefix + "_disk.index", std::ios::binary); + disk_reader.seekg(2 * sizeof(uint32_t), std::ios::beg); + disk_reader.seekg(2 * sizeof(uint64_t), std::ios::cur); + uint64_t medoid = std::numeric_limits::max(); + disk_reader.read((char *) &medoid, sizeof(uint64_t)); + std::cout << "Detected medoid = " << medoid + << " ==> excluding from insert/deletes.\n"; + ::medoid_id = (uint32_t) medoid; + + // generate inactive tags + tsl::robin_set inactive_tags; + inactive_tags.reserve(npts - tag_num); + for (uint32_t i = 0; i < npts; i++) { + auto iter = active_tags.find(i); + if (iter == active_tags.end()) { + inactive_tags.insert(i); + } + } + std::cout << "Inactive tags : " << inactive_tags.size() << std::endl; + // remove medoid from active_set + active_tags.erase(::medoid_id); + + diskann::Parameters paras; + paras.Set("L_mem", params[std::string("mem_l_index")]); + paras.Set("R_mem", params[std::string("range")]); + paras.Set("alpha_mem", ::mem_alpha); + paras.Set("L_disk", params[std::string("merge_l_index")]); + paras.Set("R_disk", params[std::string("range")]); + paras.Set("alpha_disk", ::merge_alpha); + paras.Set("C", params[std::string("merge_maxc")]); + paras.Set("beamwidth", params[std::string("beam_width")]); + paras.Set("nodes_to_cache", + params[std::string("disk_search_node_cache_count")]); + paras.Set("num_search_threads", + params[std::string("disk_search_nthreads")]); + + const std::string working_folder = ::TMP_FOLDER; + diskann::Metric metric = diskann::Metric::L2; + diskann::MergeInsert merge_insert( + paras, ndims, mem_prefix, base_prefix, merge_prefix, dist_cmp, metric, + ::save_index_as_one_file, working_folder); + // search_kernel(merge_insert, active_tags, true); + for (size_t i = 0; i < n_iters; i++) { + std::cout << "ITER : " << i << std::endl; + run_iter(merge_insert, mem_prefix, active_tags, inactive_tags); + } + /* + std::cout << "Done running all iterations, now merging any leftover points." + << std::endl; + std::future_status merge_status, insert_status, delete_status; + do { + merge_status = ::merge_future.wait_for(std::chrono::milliseconds(1)); + insert_status = ::insert_future.wait_for(std::chrono::milliseconds(1)); + delete_status = ::delete_future.wait_for(std::chrono::milliseconds(1)); + + // search_kernel(merge_insert, active_tags, + // false); + } while ((merge_status != std::future_status::ready) || + (insert_status != std::future_status::ready) || + (delete_status != std::future_status::ready)); + merge_kernel(merge_insert); + */ + // search_kernel(merge_insert, active_tags, + // true); +} + +int main(int argc, char **argv) { + std::cout << "Entering main()" << std::endl; + if (argc < 20) { + std::cout << "Correct usage: " << argv[0] + << " " + " " + " " + << " " + << " " + " " + " ...." + << "\n WARNING: Other parameters set inside CPP source." + << std::endl; + exit(-1); + } else { + std::cout << "This driver file only works with uint32 type tags" + << std::endl; + } + std::cout.setf(std::ios::unitbuf); + + int arg_no = 1; + std::string index_type = argv[arg_no++]; + TMP_FOLDER = argv[arg_no++]; + std::string base_prefix(argv[arg_no++]); + std::string merge_prefix(argv[arg_no++]); + std::string mem_prefix(argv[arg_no++]); + unsigned L_mem = (unsigned) atoi(argv[arg_no++]); + float alpha_mem = (float) atof(argv[arg_no++]); + unsigned L_disk = (unsigned) atoi(argv[arg_no++]); + float alpha_disk = (float) atof(argv[arg_no++]); + std::string data_bin(argv[arg_no++]); + int single_file = atoi(argv[arg_no++]); + std::string query_path(argv[arg_no++]); + std::string gt_file(argv[arg_no++]); + int n_iters = atoi(argv[arg_no++]); + uint32_t insert_count = (uint32_t) atoi(argv[arg_no++]); + uint32_t delete_count = (uint32_t) atoi(argv[arg_no++]); + uint32_t range = (uint32_t) atoi(argv[arg_no++]); + uint32_t recall_k = (uint32_t) atoi(argv[arg_no++]); + + for (int ctr = arg_no; ctr < argc; ctr++) { + _u32 curL = std::atoi(argv[ctr]); + if (curL >= recall_k) + ::Lvec.push_back(curL); + } + + std::cout << "Assigning parameters" << std::endl; + params[std::string("n_iters")] = n_iters; + params[std::string("insert_count")] = insert_count; + params[std::string("delete_count")] = delete_count; + params[std::string("range")] = range; + params[std::string("recall_k")] = recall_k; + + // hard-coded params + params[std::string("disk_search_node_cache_count")] = 100; + params[std::string("disk_search_nthreads")] = 16; + params[std::string("beam_width")] = 4; + params[std::string("mem_l_index")] = L_mem; + mem_alpha = alpha_mem; + merge_alpha = alpha_disk; + params[std::string("mem_nthreads")] = 32; + params[std::string("merge_maxc")] = (uint32_t)(range * 2.5); + params[std::string("merge_l_index")] = L_disk; + + ::query_file = ::query_file + query_path; + ::truthset_file = gt_file; + if (single_file == 1) + ::save_index_as_one_file = true; + else + ::save_index_as_one_file = false; + + std::string active_tags_filename; + if (single_file) + active_tags_filename = base_prefix + "_disk.index"; + else + active_tags_filename = base_prefix + "_disk.index.tags"; + + std::cout << "Calling run_all_iters()" << std::endl; + if (index_type == std::string("float")) { + diskann::DistanceL2 dist_cmp; + run_all_iters(base_prefix, merge_prefix, mem_prefix, data_bin, + active_tags_filename, &dist_cmp); + } else if (index_type == std::string("uint8")) { + diskann::DistanceL2UInt8 dist_cmp; + run_all_iters(base_prefix, merge_prefix, mem_prefix, data_bin, + active_tags_filename, &dist_cmp); + } else if (index_type == std::string("int8")) { + diskann::DistanceL2Int8 dist_cmp; + run_all_iters(base_prefix, merge_prefix, mem_prefix, data_bin, + active_tags_filename, &dist_cmp); + } else { + std::cout << "Unsupported type : " << index_type << "\n"; + } + std::cout << "Exiting\n"; + return 0; +} diff --git a/tests/test_incr_intf.cpp b/tests/test_incr_intf.cpp new file mode 100644 index 0000000000..e880a6879d --- /dev/null +++ b/tests/test_incr_intf.cpp @@ -0,0 +1,102 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +// Testing insert +#include "index.h" +#include "UTILS.H" +#include "diskann_incr_index.h" + +const std::string IndexParams = + "IndexParams=64 75 1 20 20"; // R, L, SearchMem, IndexMem, ThreadCount +const std::string SearchParams = "SearchParams=14 10 6"; // L, Threads, BW + +template +using IncrIndex = diskann::DiskANNIncrementalIndex; + +template +class InMemoryIndexParallelTest { + public: + InMemoryIndexParallelTest(_u32 maxPoints) : _maxPoints(maxPoints) { + std::random_device device; + std::mt19937 generator(device()); + _pDistribution = + std::make_shared>(0, maxPoints); + } + + private: + _u32 _maxPoints; + tsl::robin_set<_u64> _insertedIds; + std::shared_ptr> _pDistribution; + + void insertThread(T* data, _u32 ndims, IncrIndex& incrIntf) { + for (_u32 i = 0; i < 2 * maxPoints; i++) + } +}; + +template +void insertData(T* data, _u32 startIndex, _u32 endIndex, _u32 ndims, + IncrIndex& incrIntf) { + for (_u32 i = startIndex; i < endIndex; i++) { + incrIntf.AddData((_u64) startIndex * ndims, data + startIndex * ndims); + } +} + +template +std::shared_ptr> initialCreateInMemoryIndex( + T* data, _u32 npts, _u32 ndims, _u32 pointsPerIndex, + const std::string& workingDir) { + std::string parameters = IndexParams + ";" + SearchParams + + ";Placement=INMEMORY;WorkingDir=" + workingDir; + std::shared_ptr> incrIndx = std::make_shared>( + dims, ANNIndex::DT_L2, maxMemIndexPoints, parameters); + +#pragma omp parallel for schedule(dynamic, 1) num_threads(nthreads) + for (int i = 0; i <) +} + +template +void runIncrementalIntfTest(const std::string& basefile, + const std::string& workingDir, + uint32_t maxMemIndexPoints) { + T* data; + size_t npts, ndim; + diskann::load_bin + // The aim of this code is to simply exercise the incremental interface. It + // does not + // attempt to do perf tests, e.t.c beyond a point. The goal is to mimic the + // store + if (argc != 4) { + std::cout << "Usage: " + " " + << std::endl; + } + int argi = 1; + std::string datatype = argv[argi++]; + std::string basefile = argv[argi++]; + std::string workingDir = argv[argi++]; + uint32_t maxMemIndexPoints = (_u32) std::atoi(argv[argi++]); + + if (datatype == "float") { + runIncrementalIntfTest(basefile, workingDir, maxMemIndexPoints); + } else if (datatype == "uint8") { + runIncrementalIntfTest(basefile, workingDir, maxMemIndexPoints); + } else if (datatype == "int8") { + runIncrementalIntfTest(basefile, workingDir, maxMemIndexPoints); + } else { + std::cerr << std::string("Unknown data type: ") << datatype + << ".Only float,uint8,and int8 are supported." << std::endl; + } + return 0; +} diff --git a/tests/test_incr_ramp_up.cpp b/tests/test_incr_ramp_up.cpp new file mode 100644 index 0000000000..7c9d1a3490 --- /dev/null +++ b/tests/test_incr_ramp_up.cpp @@ -0,0 +1,196 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#include +#include +#include +#include +#include +#include + +#include "utils.h" + +#ifndef _WINDOWS +#include +#include +#include +#endif + +#include "memory_mapper.h" + +template +int build_incremental_index(const std::string& data_path, const unsigned L, + const unsigned R, const unsigned C, + const unsigned num_rnds, const float alpha, + const std::string& save_path, + const unsigned num_incr, const unsigned num_frozen, + const int del_mode) { + diskann::Parameters paras; + paras.Set("L", L); + paras.Set("R", R); + paras.Set("C", C); + paras.Set("alpha", alpha); + paras.Set("num_rnds", num_rnds); + + T* data_load = NULL; + size_t num_points, dim, aligned_dim; + + diskann::load_aligned_bin(data_path.c_str(), data_load, num_points, dim, + aligned_dim); + + typedef int TagT; + + diskann::Index index(diskann::L2, dim, num_points, num_frozen, true, + true); + diskann::cout << "num_points = " << num_points; + { + std::vector tags(num_points - num_incr); + std::iota(tags.begin(), tags.end(), 0); + + diskann::Timer timer; + index.build(data_path.c_str(), num_points - num_incr, paras, tags); + diskann::cout << "Index build time: " << timer.elapsed() / 1000 << "ms\n"; + } + + unsigned insert_size = (num_points / 100) * 2; + unsigned rounds = 100 / 2; + + tsl::robin_set insert_list; + tsl::robin_set delete_list; + tsl::robin_set used_tags; + tsl::robin_set inserted_tags; + tsl::robin_set deleted_tags; + unsigned num_curr_pts = num_points - num_incr + num_frozen; + + for (unsigned i = 0; i < rounds; i++) { + diskann::cout << i << std::endl << std::endl; + /*___________________________Insertion____________________*/ + + while (insert_list.size() < insert_size) { + insert_list.insert(rand() % num_points); + } + unsigned tag_p; + size_t res = used_tags.size(); + while (used_tags.size() < res + insert_size) { + tag_p = rand() % num_points + 1; + size_t temp = used_tags.size(); + used_tags.insert(tag_p); + if (used_tags.size() > temp) + inserted_tags.insert(tag_p); + } + diskann::cout << "Inserting " << insert_size << " points" << std::endl; + tsl::robin_set::iterator it = inserted_tags.begin(); + diskann::Timer insert_timer; + for (auto p : insert_list) { + index.insert_point(data_load + p * aligned_dim, paras, *it); + it++; + } + diskann::cout << "Insertion time " << insert_timer.elapsed() / 1000 + << "ms\n"; + + auto save_path_inc = + save_path + ".inc" + std::to_string(i); // 10 -denotes that 10% of base + // points are being inserted at + // a time + index.save(save_path_inc.c_str()); + + num_curr_pts += insert_size; + diskann::cout << "Number of points in the index post insertion " + << num_curr_pts << std::endl; + insert_list.clear(); + + /*_________________________Deleting points___________________________*/ + + unsigned delete_size = (unsigned) (num_curr_pts / 10); + while (deleted_tags.size() < delete_size) { + auto r = rand() % used_tags.size(); + tsl::robin_set::iterator iter = used_tags.begin(); + for (unsigned j = 0; j < r; j++) + iter++; + size_t res = delete_list.size(); + delete_list.insert(*iter); + if (delete_list.size() > res) + deleted_tags.insert(*iter); + } + diskann::cout << "Deleting " << delete_size << " points from the index" + << std::endl; + + index.enable_delete(); + for (auto p : deleted_tags) { + if (del_mode == 0) { + if (index.eager_delete(p, paras) != 0) { + std::cerr << "Delete tag " << p << " not found" << std::endl; + } else + used_tags.erase(p); + } else { + if (index.lazy_delete(p) != 0) + std::cerr << "Delete tag" << p << "not found" << std::endl; + else + used_tags.erase(p); + } + } + if (index.disable_delete(paras, true) != 0) { + std::cerr << "Disable delete failed" << std::endl; + return -1; + } + diskann::cout << "Delete time in this phase " << del_timer.elapsed() / 1000 + << "ms\n"; + deleted_tags.clear(); + inserted_tags.clear(); + + auto save_path_del = save_path + ".del" + std::to_string(i); + index.save(save_path_del.c_str()); + + num_curr_pts -= delete_size; + diskann::cout << "Number of points in the graph currently = " + << num_curr_pts << std::endl; + } + + { + diskann::Timer timer; + for (size_t i = num_points - num_incr; i < num_points; ++i) { + index.insert_point(data_load + i * aligned_dim, paras, i); + } + diskann::cout << "Incremental time: " << timer.elapsed() / 1000 << "ms\n"; + auto save_path_inc = save_path + ".inc"; + index.save(save_path_inc.c_str()); + } + + delete[] data_load; + + return 0; +} + +int main(int argc, char** argv) { + if (argc != 12) { + diskann::cout << "Correct usage: " << argv[0] + << " type[int8/uint8/float] data_file L R C alpha " + "num_rounds " + << "save_graph_file #incr_points #frozen_points " + "delete_mode[0-eager/1-lazy]" + << std::endl; + exit(-1); + } + + unsigned L = (unsigned) atoi(argv[3]); + unsigned R = (unsigned) atoi(argv[4]); + unsigned C = (unsigned) atoi(argv[5]); + float alpha = (float) std::atof(argv[6]); + unsigned num_rnds = (unsigned) std::atoi(argv[7]); + std::string save_path(argv[8]); + unsigned num_incr = (unsigned) atoi(argv[9]); + unsigned num_frozen = (unsigned) atoi(argv[10]); + int del_mode = (int) atoi(argv[11]); + + if (std::string(argv[1]) == std::string("int8")) + build_incremental_index(argv[2], L, R, C, num_rnds, alpha, + save_path, num_incr, num_frozen, del_mode); + else if (std::string(argv[1]) == std::string("uint8")) + build_incremental_index(argv[2], L, R, C, num_rnds, alpha, + save_path, num_incr, num_frozen, del_mode); + else if (std::string(argv[1]) == std::string("float")) + build_incremental_index(argv[2], L, R, C, num_rnds, alpha, save_path, + num_incr, num_frozen, del_mode); + else + diskann::cout << "Unsupported type. Use float/int8/uint8" << std::endl; +} diff --git a/tests/test_incremental_index.cpp b/tests/test_incremental_index.cpp index 8b98233b28..f92dd175d1 100644 --- a/tests/test_incremental_index.cpp +++ b/tests/test_incremental_index.cpp @@ -18,104 +18,92 @@ #include "memory_mapper.h" -int main(int argc, char** argv) { - if (argc != 10) { - std::cout << "Correct usage: " << argv[0] - << " data_file L R C alpha num_rounds " - << "save_graph_file #incr_points #frozen_points" << std::endl; - exit(-1); - } - - float* data_load = NULL; - size_t num_points, dim, aligned_dim; - - diskann::load_aligned_bin(argv[1], data_load, num_points, dim, - aligned_dim); - - unsigned L = (unsigned) atoi(argv[2]); - unsigned R = (unsigned) atoi(argv[3]); - unsigned C = (unsigned) atoi(argv[4]); - float alpha = (float) std::atof(argv[5]); - unsigned num_rnds = (unsigned) std::atoi(argv[6]); - std::string save_path(argv[7]); - unsigned num_incr = (unsigned) atoi(argv[8]); - unsigned num_frozen = (unsigned) atoi(argv[9]); - +template +int build_incremental_index(const std::string& data_path, const unsigned L, + const unsigned R, const float alpha, + const std::string& save_path, + const unsigned num_incr) { diskann::Parameters paras; paras.Set("L", L); paras.Set("R", R); - paras.Set("C", C); + paras.Set("C", 750); paras.Set("alpha", alpha); paras.Set("saturate_graph", false); - paras.Set("num_rnds", num_rnds); + paras.Set("num_rnds", 2); + + T* data_load = NULL; + size_t num_points, dim, aligned_dim; + + diskann::load_aligned_bin(data_path.c_str(), data_load, num_points, dim, + aligned_dim); - typedef int TagT; - diskann::Index index(diskann::L2, argv[1], num_points, - num_points - num_incr, num_frozen, true, - true, true); + typedef uint32_t TagT; + + diskann::Index index(diskann::L2, dim, num_points, true, + false, // single file index + true); { std::vector tags(num_points - num_incr); std::iota(tags.begin(), tags.end(), 0); - if (argc > 10) { - std::string frozen_points_file(argv[10]); - index.generate_random_frozen_points(frozen_points_file.c_str()); - } else - index.generate_random_frozen_points(); - diskann::Timer timer; - index.build(paras, tags); - std::cout << "Index build time: " << timer.elapsed() / 1000 << "ms\n"; + index.build(data_path.c_str(), num_points - num_incr, paras, tags); + diskann::cout << "Index build time: " << timer.elapsed() / 1000 << "ms\n"; + index.save(save_path.c_str()); } - std::vector pool, tmp; - tsl::robin_set visited; - std::vector cut_graph; - index.readjust_data(num_frozen); - { diskann::Timer timer; - for (size_t i = num_points - num_incr; i < num_points; ++i) { - index.insert_point(data_load + i * aligned_dim, paras, pool, tmp, visited, - cut_graph, i); +#pragma omp parallel for + for (_s64 i = num_points - num_incr; i < (_s64) num_points; ++i) { + index.insert_point(data_load + i * aligned_dim, paras, (TagT) i); } - std::cout << "Incremental time: " << timer.elapsed() / 1000 << "ms\n"; + diskann::cout << "Incremental time: " << timer.elapsed() / 1000 << "ms\n"; auto save_path_inc = save_path + ".inc"; index.save(save_path_inc.c_str()); } tsl::robin_set delete_list; while (delete_list.size() < num_incr) - delete_list.insert(rand() % num_points); - std::cout << "Deleting " << delete_list.size() << " elements" << std::endl; + delete_list.insert((uint32_t)(rand() % num_points)); + diskann::cout << "Deleting " << delete_list.size() << " elements" + << std::endl; + std::vector delete_vector; + for (auto p : delete_list) { + delete_vector.emplace_back(p); + } + diskann::cout << "Size of delete_vector : " << delete_vector.size() + << std::endl; { - diskann::Timer timer; index.enable_delete(); - for (auto p : delete_list) - - if (index.eager_delete(p, paras) != 0) - // if (index.delete_point(p) != 0) + for (size_t i = 0; i < delete_vector.size(); i++) { + unsigned p = delete_vector[i]; + if (index.lazy_delete(p) != 0) std::cerr << "Delete tag " << p << " not found" << std::endl; - - if (index.disable_delete(paras, true) != 0) { - std::cerr << "Disable delete failed" << std::endl; - return -1; } - std::cout << "Delete time: " << timer.elapsed() / 1000 << "ms\n"; } - auto save_path_del = save_path + ".del"; + auto save_path_del = save_path + ".delete"; index.save(save_path_del.c_str()); + index.load(save_path_del.c_str()); + index.consolidate(paras); - index.readjust_data(num_frozen); { + index.reposition_frozen_point_to_end(); + std::vector<_u32> reinsert_vec; + for (auto p : delete_list) + reinsert_vec.emplace_back(p); + diskann::Timer timer; - for (auto p : delete_list) { - index.insert_point(data_load + (size_t) p * (size_t) aligned_dim, paras, - pool, tmp, visited, cut_graph, p); +#pragma omp parallel for + for (_s64 p = 0; p < (_s64) reinsert_vec.size(); p++) { + index.insert_point( + data_load + (size_t)(reinsert_vec[p]) * (size_t) aligned_dim, paras, + (reinsert_vec[p])); } - std::cout << "Re-incremental time: " << timer.elapsed() / 1000 << "ms\n"; + diskann::cout << "Re-incremental time: " << timer.elapsed() / 1000 + << "ms\n"; } auto save_path_reinc = save_path + ".reinc"; @@ -125,3 +113,27 @@ int main(int argc, char** argv) { return 0; } + +int main(int argc, char** argv) { + if (argc != 8) { + diskann::cout << "Correct usage: " << argv[0] + << " type[int8/uint8/float] data_file L R alpha " + << "save_graph_file #incr_points " << std::endl; + exit(-1); + } + + unsigned L = (unsigned) atoi(argv[3]); + unsigned R = (unsigned) atoi(argv[4]); + float alpha = (float) std::atof(argv[5]); + std::string save_path(argv[6]); + unsigned num_incr = (unsigned) atoi(argv[7]); + + if (std::string(argv[1]) == std::string("int8")) + build_incremental_index(argv[2], L, R, alpha, save_path, num_incr); + else if (std::string(argv[1]) == std::string("uint8")) + build_incremental_index(argv[2], L, R, alpha, save_path, num_incr); + else if (std::string(argv[1]) == std::string("float")) + build_incremental_index(argv[2], L, R, alpha, save_path, num_incr); + else + diskann::cout << "Unsupported type. Use float/int8/uint8" << std::endl; +} diff --git a/tests/utils/CMakeLists.txt b/tests/utils/CMakeLists.txt index e69722dcfc..388580f773 100644 --- a/tests/utils/CMakeLists.txt +++ b/tests/utils/CMakeLists.txt @@ -3,25 +3,40 @@ set(CMAKE_CXX_STANDARD 14) +add_executable(uint8_to_float uint8_to_float.cpp) +if(MSVC) + target_link_options(uint8_to_float PRIVATE /MACHINE:x64) + target_link_libraries(uint8_to_float debug ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}/diskann_dll.lib) + target_link_libraries(uint8_to_float optimized ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}/diskann_dll.lib) +else() + target_link_libraries(uint8_to_float ${PROJECT_NAME}) +endif() + add_executable(fvecs_to_bin fvecs_to_bin.cpp) if(MSVC) target_link_options(fvecs_to_bin PRIVATE /MACHINE:x64) target_link_libraries(fvecs_to_bin debug ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}/diskann_dll.lib) target_link_libraries(fvecs_to_bin optimized ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}/diskann_dll.lib) +else() + target_link_libraries(fvecs_to_bin ${PROJECT_NAME}) endif() -add_executable(float_bin_to_int8 float_bin_to_int8.cpp) +add_executable(bin_to_fvecs bin_to_fvecs.cpp) if(MSVC) - target_link_options(float_bin_to_int8 PRIVATE /MACHINE:x64) - target_link_libraries(float_bin_to_int8 debug ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}/diskann_dll.lib) - target_link_libraries(float_bin_to_int8 optimized ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}/diskann_dll.lib) + target_link_options(bin_to_fvecs PRIVATE /MACHINE:x64) + target_link_libraries(bin_to_fvecs debug ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}/diskann_dll.lib) + target_link_libraries(bin_to_fvecs optimized ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}/diskann_dll.lib) +else() + target_link_libraries(bin_to_fvecs ${PROJECT_NAME}) endif() -add_executable(ivecs_to_bin ivecs_to_bin.cpp) +add_executable(normalize_data normalize_data.cpp) if(MSVC) - target_link_options(ivecs_to_bin PRIVATE /MACHINE:x64) - target_link_libraries(ivecs_to_bin debug ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}/diskann_dll.lib) - target_link_libraries(ivecs_to_bin optimized ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}/diskann_dll.lib) + target_link_options(normalize_data PRIVATE /MACHINE:x64) + target_link_libraries(normalize_data debug ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}/diskann_dll.lib) + target_link_libraries(normalize_data optimized ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}/diskann_dll.lib) +else() + target_link_libraries(normalize_data ${PROJECT_NAME}) endif() add_executable(tsv_to_bin tsv_to_bin.cpp) @@ -29,13 +44,17 @@ if(MSVC) target_link_options(tsv_to_bin PRIVATE /MACHINE:x64) target_link_libraries(tsv_to_bin debug ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}/diskann_dll.lib) target_link_libraries(tsv_to_bin optimized ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}/diskann_dll.lib) +else() + target_link_libraries(tsv_to_bin ${PROJECT_NAME}) endif() -add_executable(bin_to_tsv bin_to_tsv.cpp) +add_executable(ivecs_to_bin ivecs_to_bin.cpp) if(MSVC) - target_link_options(bin_to_tsv PRIVATE /MACHINE:x64) - target_link_libraries(bin_to_tsv debug ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}/diskann_dll.lib) - target_link_libraries(bin_to_tsv optimized ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}/diskann_dll.lib) + target_link_options(ivecs_to_bin PRIVATE /MACHINE:x64) + target_link_libraries(ivecs_to_bin debug ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}/diskann_dll.lib) + target_link_libraries(ivecs_to_bin optimized ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}/diskann_dll.lib) +else() + target_link_libraries(ivecs_to_bin ${PROJECT_NAME}) endif() add_executable(int8_to_float int8_to_float.cpp) @@ -56,14 +75,14 @@ else() target_link_libraries(uint32_to_uint8 ${PROJECT_NAME}) endif() -add_executable(gen_random_slice gen_random_slice.cpp) -if(MSVC) - target_link_options(gen_random_slice PRIVATE /MACHINE:x64) - target_link_libraries(gen_random_slice debug ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}/diskann_dll.lib) - target_link_libraries(gen_random_slice optimized ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}/diskann_dll.lib) -else() - target_link_libraries(gen_random_slice ${PROJECT_NAME} -ltcmalloc) -endif() +#add_executable(gen_random_slice gen_random_slice.cpp) +#if(MSVC) +# target_link_options(gen_random_slice PRIVATE /MACHINE:x64) +# target_link_libraries(gen_random_slice debug ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}/diskann_dll.lib) +# target_link_libraries(gen_random_slice optimized ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}/diskann_dll.lib) +#else() +# target_link_libraries(gen_random_slice ${PROJECT_NAME} -ltcmalloc) +#endif() add_executable(calculate_recall calculate_recall.cpp) if(MSVC) @@ -85,6 +104,15 @@ else() endif() +add_executable(seed_index_merger seed_index_merger.cpp) +if(MSVC) + target_link_options(seed_index_merger PRIVATE /MACHINE:x64 /DEBUG:FULL) + target_link_libraries(seed_index_merger debug ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}/diskann_dll.lib) + target_link_libraries(seed_index_merger optimized ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}/diskann_dll.lib) +else() + target_link_libraries(seed_index_merger ${PROJECT_NAME} -ltcmalloc) +endif() + add_executable(generate_pq generate_pq.cpp) if(MSVC) target_link_options(generate_pq PRIVATE /MACHINE:x64 /DEBUG:FULL) @@ -103,7 +131,6 @@ else() target_link_libraries(partition_data ${PROJECT_NAME} -ltcmalloc) endif() - add_executable(partition_with_ram_budget partition_with_ram_budget.cpp) if(MSVC) target_link_options(partition_with_ram_budget PRIVATE /MACHINE:x64) @@ -134,8 +161,32 @@ else() endif() +if(MSVC) + add_executable(test_l2_avx test_l2_avx.cpp) + target_link_options(test_l2_avx PRIVATE /MACHINE:x64 /DEBUG:FULL) + target_link_libraries(test_l2_avx debug ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}/diskann_dll.lib) + target_link_libraries(test_l2_avx optimized ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}/diskann_dll.lib) + + add_executable(extract_rows extract_rows.cpp) + target_link_options(extract_rows PRIVATE /MACHINE:x64 /DEBUG:FULL) + target_link_libraries(extract_rows debug ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}/diskann_dll.lib) + target_link_libraries(extract_rows optimized ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}/diskann_dll.lib) + +endif() + +#add_executable(convert_axon_to_bin convert_axon_to_bin.cpp ${PROJECT_SOURCE_DIR}/src/ann_exception.cpp) +#if(MSVC) +# target_include_directories(convert_axon_to_bin PRIVATE "${BOOST_ROOT}") +# target_link_options(convert_axon_to_bin PRIVATE /MACHINE:x64) +#else() +# target_link_libraries(convert_axon_to_bin ${PROJECT_NAME}) +#endif() + + + + # formatter -if (LINUX) - add_custom_command(TARGET gen_random_slice PRE_BUILD COMMAND clang-format-4.0 -i ../../../include/*.h ../../../include/dll/*.h ../../../src/*.cpp ../../../tests/*.cpp ../../../src/dll/*.cpp ../../../tests/utils/*.cpp) +if (NOT MSVC) + add_custom_command(TARGET generate_pq PRE_BUILD COMMAND clang-format -i ../../../include/*.h ../../../src/*.cpp ../../../tests/*.cpp ../../../src/dll/*.cpp ../../../tests/utils/*.cpp) endif() diff --git a/tests/utils/bin_to_fvecs.cpp b/tests/utils/bin_to_fvecs.cpp index 307437a416..050f37f3b6 100644 --- a/tests/utils/bin_to_fvecs.cpp +++ b/tests/utils/bin_to_fvecs.cpp @@ -2,60 +2,79 @@ // Licensed under the MIT license. #include -#include "util.h" +#include +#include +#include "utils.h" -void block_convert(std::ifstream& writr, std::ofstream& readr, float* read_buf, - float* write_buf, _u64 npts, _u64 ndims) { - writr.write((char*) read_buf, - npts * (ndims * sizeof(float) + sizeof(unsigned))); +void block_convert(std::ofstream& writr, std::ifstream& readr, float* read_buf, + float* write_buf, _u64 npts, _u64 ndims, + bool normalize = false) { + readr.read((char*) read_buf, npts * ndims * sizeof(float)); + _u32 ndims_u32 = (_u32) ndims; #pragma omp parallel for - for (_u64 i = 0; i < npts; i++) { - memcpy(write_buf + i * ndims, (read_buf + i * (ndims + 1)) + 1, + for (_s64 i = 0; i < (_s64) npts; i++) { + if (normalize) { + float norm_pt = std::numeric_limits::epsilon(); + for (_u32 dim = 0; dim < ndims_u32; dim++) { + norm_pt += + *(read_buf + i * ndims + dim) * *(read_buf + i * ndims + dim); + } + norm_pt = std::sqrt(norm_pt); + for (_u32 dim = 0; dim < ndims_u32; dim++) { + *(read_buf + i * ndims + dim) = *(read_buf + i * ndims + dim) / norm_pt; + } + } + + memcpy(write_buf + i * (ndims + 1), &ndims_u32, sizeof(float)); + memcpy(write_buf + i * (ndims + 1) + 1, (read_buf + i * ndims), ndims * sizeof(float)); } - readr.read((char*) write_buf, npts * ndims * sizeof(float)); + writr.write((char*) write_buf, + npts * (ndims * sizeof(float) + sizeof(unsigned))); } int main(int argc, char** argv) { - if (argc != 3) { - std::cout << argv[0] << " input_bin output_fvecs" << std::endl; + if (argc != 4) { + diskann::cout << argv[0] + << ": [input_bin] [output_fvecs] [normalize? (1 for yes]" + << std::endl; exit(-1); } + bool normalize = (bool) std::atoi(argv[3]); std::ifstream readr(argv[1], std::ios::binary); int npts_s32; int ndims_s32; readr.read((char*) &npts_s32, sizeof(_s32)); readr.read((char*) &ndims_s32, sizeof(_s32)); - size_t npts = npts_s32; - size_t ndims = ndims_s32; - _u32 ndims_u32 = (_u32) ndims_s32; - // _u64 fsize = writr.tellg(); - readr.seekg(0, std::ios::beg); + // size_t npt = npts_s32; + // size_t ndim = ndims_s32; + _u32 ndims_u32 = (_u32) ndims_s32; + _u32 npts_u32 = (_u32) npts_s32; + // readr.seekg(0, std::ios::end); + //_u64 fsize = readr.tellg(); - unsigned ndims_u32; - writr.write((char*) &ndims_u32, sizeof(unsigned)); - writr.seekg(0, std::ios::beg); + std::ofstream writr(argv[2], std::ios::binary); + // writr.write((char*) &ndims_u32, sizeof(unsigned)); + // writr.seekg(0, std::ios::beg); _u64 ndims = (_u64) ndims_u32; - _u64 npts = fsize / ((ndims + 1) * sizeof(float)); - std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims - << std::endl; + _u64 npts = (_u64) npts_u32; + diskann::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims + << std::endl; _u64 blk_size = 131072; _u64 nblks = ROUND_UP(npts, blk_size) / blk_size; - std::cout << "# blks: " << nblks << std::endl; + diskann::cout << "# blks: " << nblks << std::endl; - std::ofstream writr(argv[2], std::ios::binary); - float* read_buf = new float[npts * (ndims + 1)]; - float* write_buf = new float[npts * ndims]; + float* read_buf = new float[npts * ndims]; + float* write_buf = new float[npts * (ndims + 1)]; for (_u64 i = 0; i < nblks; i++) { _u64 cblk_size = std::min(npts - i * blk_size, blk_size); - block_convert(writr, readr, read_buf, write_buf, cblk_size, ndims); - std::cout << "Block #" << i << " written" << std::endl; + block_convert(writr, readr, read_buf, write_buf, cblk_size, ndims, + normalize); + diskann::cout << "Block #" << i << " written" << std::endl; } - delete[] read_buf; delete[] write_buf; - writr.close(); readr.close(); } diff --git a/tests/utils/bin_to_tsv.cpp b/tests/utils/bin_to_tsv.cpp index 37874e2437..9a7e180a42 100644 --- a/tests/utils/bin_to_tsv.cpp +++ b/tests/utils/bin_to_tsv.cpp @@ -22,7 +22,8 @@ void block_convert(std::ofstream& writer, std::ifstream& reader, T* read_buf, int main(int argc, char** argv) { if (argc != 4) { - std::cout << argv[0] << " input_bin output_tsv" << std::endl; + std::cout << argv[0] << " input_bin output_tsv" + << std::endl; exit(-1); } std::string type_string(argv[1]); @@ -50,9 +51,10 @@ int main(int argc, char** argv) { for (_u64 i = 0; i < nblks; i++) { _u64 cblk_size = std::min(npts - i * blk_size, blk_size); if (type_string == std::string("float")) - block_convert(writer, reader, (float*)read_buf, cblk_size, ndims); + block_convert(writer, reader, (float*) read_buf, cblk_size, ndims); else if (type_string == std::string("int8")) - block_convert(writer, reader, (int8_t*) read_buf, cblk_size, ndims); + block_convert(writer, reader, (int8_t*) read_buf, cblk_size, + ndims); else if (type_string == std::string("uint8")) block_convert(writer, reader, (uint8_t*) read_buf, cblk_size, ndims); diff --git a/tests/utils/calculate_recall.cpp b/tests/utils/calculate_recall.cpp index 7fe03798ff..cb94039d72 100644 --- a/tests/utils/calculate_recall.cpp +++ b/tests/utils/calculate_recall.cpp @@ -14,8 +14,8 @@ int main(int argc, char** argv) { if (argc != 4) { - std::cout << argv[0] << " " - << std::endl; + diskann::cout << argv[0] << " " + << std::endl; return -1; } unsigned* gold_std = NULL; @@ -29,7 +29,7 @@ int main(int argc, char** argv) { diskann::load_truthset(argv[2], our_results, or_dist, points_num_or, dim_or); if (points_num_gs != points_num_or) { - std::cout + diskann::cout << "Error. Number of queries mismatch in ground truth and our results" << std::endl; return -1; @@ -39,15 +39,16 @@ int main(int argc, char** argv) { uint32_t recall_at = std::atoi(argv[3]); if ((dim_or < recall_at) || (recall_at > dim_gs)) { - std::cout << "ground truth has size " << dim_gs << "; our set has " - << dim_or << " points. Asking for recall " << recall_at - << std::endl; + diskann::cout << "ground truth has size " << dim_gs << "; our set has " + << dim_or << " points. Asking for recall " << recall_at + << std::endl; return -1; } - std::cout << "Calculating recall@" << recall_at << std::endl; - float recall_val = diskann::calculate_recall( - points_num, gold_std, gs_dist, dim_gs, our_results, dim_or, recall_at); + diskann::cout << "Calculating recall@" << recall_at << std::endl; + float recall_val = (float) diskann::calculate_recall( + (_u32) points_num, gold_std, gs_dist, (_u32) dim_gs, our_results, + (_u32) dim_or, recall_at); // double avg_recall = (recall*1.0)/(points_num*1.0); - std::cout << "Avg. recall@" << recall_at << " is " << recall_val << "\n"; + diskann::cout << "Avg. recall@" << recall_at << " is " << recall_val << "\n"; } diff --git a/tests/utils/compute_groundtruth.cpp b/tests/utils/compute_groundtruth.cpp index 8fef8c929c..1c3a10267f 100644 --- a/tests/utils/compute_groundtruth.cpp +++ b/tests/utils/compute_groundtruth.cpp @@ -31,10 +31,10 @@ #define ALIGNMENT 512 void command_line_help() { - std::cerr - << " " - << std::endl; + std::cerr << " " + " optional:" + << std::endl; } template @@ -131,8 +131,8 @@ void exact_knn(const size_t dim, const size_t k, distsq_to_points(dim, dist_matrix, npoints, points, points_l2sq, q_e - q_b, queries + (ptrdiff_t) q_b * (ptrdiff_t) dim, queries_l2sq + q_b); - std::cout << "Computed distances for queries: [" << q_b << "," << q_e << ")" - << std::endl; + diskann::cout << "Computed distances for queries: [" << q_b << "," << q_e + << ")" << std::endl; #pragma omp parallel for schedule(dynamic, 16) for (long long q = q_b; q < q_e; q++) { @@ -163,8 +163,8 @@ void exact_knn(const size_t dim, const size_t k, dist_closest_points + (ptrdiff_t) q * (ptrdiff_t) k, dist_closest_points + (ptrdiff_t)(q + 1) * (ptrdiff_t) k)); } - std::cout << "Computed exact k-NN for queries: [" << q_b << "," << q_e - << ")" << std::endl; + diskann::cout << "Computed exact k-NN for queries: [" << q_b << "," << q_e + << ")" << std::endl; } delete[] dist_matrix; @@ -176,16 +176,18 @@ void exact_knn(const size_t dim, const size_t k, template inline int get_num_parts(const char *filename) { std::ifstream reader(filename, std::ios::binary); - std::cout << "Reading bin file " << filename << " ...\n"; + diskann::cout << "Reading bin file " << filename << " ...\n"; int npts_i32, ndims_i32; reader.read((char *) &npts_i32, sizeof(int)); reader.read((char *) &ndims_i32, sizeof(int)); - std::cout << "#pts = " << npts_i32 << ", #dims = " << ndims_i32 << std::endl; + diskann::cout << "#pts = " << npts_i32 << ", #dims = " << ndims_i32 + << std::endl; reader.close(); - int num_parts = (npts_i32 % PARTSIZE) == 0 - ? npts_i32 / PARTSIZE - : std::floor(npts_i32 / PARTSIZE) + 1; - std::cout << "Number of parts: " << num_parts << std::endl; + uint32_t num_parts = + (npts_i32 % PARTSIZE) == 0 + ? (_u32)(npts_i32 / PARTSIZE) + : (_u32) std::floor((double) npts_i32 / (double) PARTSIZE) + 1; + diskann::cout << "Number of parts: " << num_parts << std::endl; return num_parts; } @@ -193,7 +195,7 @@ template inline void load_bin_as_float(const char *filename, float *&data, size_t &npts, size_t &ndims, int part_num) { std::ifstream reader(filename, std::ios::binary); - std::cout << "Reading bin file " << filename << " ...\n"; + diskann::cout << "Reading bin file " << filename << " ...\n"; int npts_i32, ndims_i32; reader.read((char *) &npts_i32, sizeof(int)); reader.read((char *) &ndims_i32, sizeof(int)); @@ -203,16 +205,18 @@ inline void load_bin_as_float(const char *filename, float *&data, size_t &npts, ndims = (unsigned) ndims_i32; uint64_t nptsuint64_t = (uint64_t) npts; uint64_t ndimsuint64_t = (uint64_t) ndims; - std::cout << "#pts in part = " << npts << ", #dims = " << ndims - << ", size = " << nptsuint64_t * ndimsuint64_t * sizeof(T) << "B" - << std::endl; + diskann::cout << "#pts in part = " << npts << ", #dims = " << ndims + << ", size = " << nptsuint64_t * ndimsuint64_t * sizeof(T) + << "B" << std::endl; reader.seekg(start_id * ndims * sizeof(T) + 2 * sizeof(uint32_t), std::ios::beg); + // data = new T[nptsuint64_t * ndimsuint64_t]; T *data_T = new T[nptsuint64_t * ndimsuint64_t]; reader.read((char *) data_T, sizeof(T) * nptsuint64_t * ndimsuint64_t); - std::cout << "Finished reading part of the bin file." << std::endl; + diskann::cout << "Finished reading part of the bin file." << std::endl; reader.close(); + // data = (nptsuint64_t*ndimsuint64_t, ALIGNMENT); data = aligned_malloc(nptsuint64_t * ndimsuint64_t, ALIGNMENT); #pragma omp parallel for schedule(dynamic, 32768) for (int64_t i = 0; i < (int64_t) nptsuint64_t; i++) { @@ -223,55 +227,62 @@ inline void load_bin_as_float(const char *filename, float *&data, size_t &npts, } } delete[] data_T; - std::cout << "Finished converting part data to float." << std::endl; + diskann::cout << "Finished converting part data to float." << std::endl; } template inline void save_bin(const std::string filename, T *data, size_t npts, size_t ndims) { std::ofstream writer(filename, std::ios::binary | std::ios::out); - std::cout << "Writing bin: " << filename << "\n"; + diskann::cout << "Writing bin: " << filename << "\n"; int npts_i32 = (int) npts, ndims_i32 = (int) ndims; writer.write((char *) &npts_i32, sizeof(int)); writer.write((char *) &ndims_i32, sizeof(int)); - std::cout << "bin: #pts = " << npts << ", #dims = " << ndims - << ", size = " << npts * ndims * sizeof(T) + 2 * sizeof(int) << "B" - << std::endl; + diskann::cout << "bin: #pts = " << npts << ", #dims = " << ndims + << ", size = " << npts * ndims * sizeof(T) + 2 * sizeof(int) + << "B" << std::endl; // data = new T[npts_u64 * ndims_u64]; writer.write((char *) data, npts * ndims * sizeof(T)); writer.close(); - std::cout << "Finished writing bin" << std::endl; + diskann::cout << "Finished writing bin" << std::endl; } inline void save_groundtruth_as_one_file(const std::string filename, int32_t *data, float *distances, - size_t npts, size_t ndims) { + size_t npts, size_t ndims, + uint32_t *tags = nullptr) { std::ofstream writer(filename, std::ios::binary | std::ios::out); int npts_i32 = (int) npts, ndims_i32 = (int) ndims; writer.write((char *) &npts_i32, sizeof(int)); writer.write((char *) &ndims_i32, sizeof(int)); - std::cout << "Saving truthset in one file (npts, dim, npts*dim id-matrix, " - "npts*dim dist-matrix) with npts = " - << npts << ", dim = " << ndims << ", size = " - << 2 * npts * ndims * sizeof(unsigned) + 2 * sizeof(int) << "B" - << std::endl; + diskann::cout + << "Saving truthset in one file (npts, dim, npts*dim id-matrix, " + "npts*dim dist-matrix) with npts = " + << npts << ", dim = " << ndims + << ", size = " << 2 * npts * ndims * sizeof(unsigned) + 2 * sizeof(int) + << "B" << std::endl; // data = new T[npts_u64 * ndims_u64]; writer.write((char *) data, npts * ndims * sizeof(uint32_t)); writer.write((char *) distances, npts * ndims * sizeof(float)); + if (tags != nullptr) { + writer.write((char *) tags, npts * ndims * sizeof(uint32_t)); + } else { + writer.write((char *) data, npts * ndims * sizeof(uint32_t)); + } + writer.close(); - std::cout << "Finished writing truthset" << std::endl; + diskann::cout << "Finished writing truthset" << std::endl; } template -int aux_main(int argv, char **argc) { - +int aux_main(int argc, char **argv) { size_t npoints, nqueries, dim; - std::string base_file(argc[2]); - std::string query_file(argc[3]); - size_t k = atoi(argc[4]); - std::string gt_file(argc[5]); + std::string base_file(argv[2]); + std::string query_file(argv[3]); + size_t k = atoi(argv[4]); + std::string gt_file(argv[5]); float *base_data; float *query_data; @@ -314,21 +325,38 @@ int aux_main(int argv, char **argc) { dist_closest_points[i * k + j] = cur_res[j].second; } } + uint32_t *tags = nullptr; + if (argc == 7) { + std::cout << "Loading tags from " << argv[6] << "\n"; + tags = new uint32_t[nqueries * k]; + uint32_t * all_tags; + std::string tag_file = std::string(argv[6]); + size_t tag_pts, tag_dim; + diskann::load_bin(tag_file, all_tags, tag_pts, tag_dim); + + diskann::cout << "Loaded tags for " << tag_pts << " points.\n"; + for (uint64_t i = 0; i < nqueries * k; i++) { + tags[i] = all_tags[closest_points[i]]; + } + } save_groundtruth_as_one_file(gt_file, closest_points, dist_closest_points, - nqueries, k); + nqueries, k, tags); diskann::aligned_free(query_data); delete[] closest_points; delete[] dist_closest_points; + if (tags != nullptr) { + delete[] tags; + } + return 0; } int main(int argc, char **argv) { - if (argc != 6) { + if (argc != 6 && argc != 7) { command_line_help(); return -1; } - if (std::string(argv[1]) == std::string("float")) aux_main(argc, argv); if (std::string(argv[1]) == std::string("int8")) diff --git a/tests/utils/convert_axon_to_bin.cpp b/tests/utils/convert_axon_to_bin.cpp new file mode 100644 index 0000000000..cfeca80c3e --- /dev/null +++ b/tests/utils/convert_axon_to_bin.cpp @@ -0,0 +1,314 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#include +#include +#include +#include +#include +#include +#include "cached_io.h" +#include "ann_exception.h" + +float (*atofLambda)(const std::string&) = [](const std::string& str) { + return (float) ::atof(str.c_str()); +}; +int8_t (*atoiLambda)(const std::string&) = [](const std::string& str) { + return (int8_t)::atoi(str.c_str()); +}; +uint8_t (*atouLambda)(const std::string&) = [](const std::string& str) { + return (uint8_t)::atoi(str.c_str()); +}; + +const int MIN_POINTS_FOR_VALUE_VERIFICATION = 10000; + +std::vector splitLine(const std::string& str, const char delimiter, + int suggestedCapacity) { + std::vector records; + records.reserve(suggestedCapacity); + + boost::split(records, str, [delimiter](char c) { return c == delimiter; }); + + return records; +} + +template +std::vector getVector(const std::string& str, + T convertFn(const std::string& str), + int suggestedSize) { + auto idsAndVector = splitLine(str, '\t', suggestedSize); + + std::vector typedVec; + typedVec.reserve(suggestedSize); + std::vector strVec; + + if (idsAndVector.size() >= 3) { + strVec = splitLine(idsAndVector[2], ',', suggestedSize); + } else if (idsAndVector.size() == 2) { + strVec = splitLine(idsAndVector[1], ',', suggestedSize); + } else if (idsAndVector.size() == 1) { + strVec = splitLine(idsAndVector[0], ',', suggestedSize); + } else { + std::cerr << "Found line " << str + << " that has non-standard number of delimiters: " + << idsAndVector.size() << std::endl; + } + std::transform( + strVec.begin(), strVec.end(), std::back_inserter(typedVec), + [&convertFn](const std::string& str) { return convertFn(str); }); + + return typedVec; +} + +template +uint32_t getDimension(const char* inFile, T convertFn(const std::string& val)) { + std::ifstream fin(inFile); + if (fin.is_open()) { + std::string line; + std::getline(fin, line); + + auto vector = getVector(line, convertFn, 1); + return (uint32_t) vector.size(); + } else { + char message[2048]; + std::cerr << "Could not open input file: " << inFile << ". " + << strerror_s(message, 2048, errno) << std::endl; + exit(1); + } +} + +uint32_t getRowCount(const char* infile) { + std::ifstream fin(infile); + std::string line; + + uint32_t count = 0; + while (std::getline(fin, line)) { + if (line.length() != 0) { + count++; + if (count % 10000 == 0) { + diskann::cout << "Counted " << count << " lines." << std::endl; + } + } + } + + return count; +} + +template +void writeBinaryVec(const char* outfile, uint32_t numDimensions, + const std::vector>& allVecs) { + uint32_t numPoints = (uint32_t) allVecs.size(); + diskann::cout << "Writing " << numPoints + << " records of dimension: " << numDimensions + << " to file: " << outfile << "..."; + + std::ofstream fout(outfile, std::ios::binary); + if (fout.is_open()) { + fout.write((const char*) &numPoints, sizeof(uint32_t)); + fout.write((const char*) &numDimensions, sizeof(uint32_t)); + + for (auto vec : allVecs) { + fout.write((const char*) vec.data(), numDimensions * sizeof(T)); + } + diskann::cout << "done." << std::endl; + } else { + char message[2048]; + std::cerr << "Could not open output file: " << outfile << ". " + << strerror_s(message, 2048, errno) << std::endl; + } +} + +// template +// void writeBinary(const char* outfile, uint32_t numDimensions, +// uint32_t numPoints, T* data) { +// diskann::cout << "Writing " << numPoints +// << " records of dimension: " << numDimensions +// << " to file: " << outfile << "..."; +// +// std::ofstream fout(outfile, std::ios::binary); +// if (fout.is_open()) { +// fout.write((const char*) &numPoints, sizeof(uint32_t)); +// fout.write((const char*) &numDimensions, sizeof(uint32_t)); +// +// fout.write((const char*) data, +// ((uint64_t) numDimensions) * numPoints * sizeof(T)); +// diskann::cout << "done." << std::endl; +// } else { +// char message[2048]; +// std::cerr << "Could not open output file: " << outfile << ". " +// << strerror_s(message, 2048, errno) << std::endl; +// } +//} + +uint32_t numPointsToCheck(uint32_t numPoints) { + uint32_t point01percent = (uint32_t)(0.01 * numPoints / 100); + uint32_t point1percent = (uint32_t)(0.1 * numPoints / 100); + uint32_t onepercent = (uint32_t)(1 * numPoints / 100); + + return point01percent > MIN_POINTS_FOR_VALUE_VERIFICATION + ? point01percent + : point1percent > MIN_POINTS_FOR_VALUE_VERIFICATION + ? point1percent + : onepercent > MIN_POINTS_FOR_VALUE_VERIFICATION ? onepercent + : numPoints; +} + +template +void sanityCheck(const char* origfile, const char* binfile, + uint32_t numDimensions, uint32_t numPoints, + T convertFn(const std::string& str)) { + uint32_t nd, np; + std::ifstream fin(binfile, std::ios::binary); + std::ifstream orig(origfile); + + fin.read((char*) &np, sizeof(uint32_t)); + fin.read((char*) &nd, sizeof(uint32_t)); + + assert(np == numPoints); + if (np != numPoints) { + throw diskann::ANNException( + "Verifying binary write failed! np != numPoints", -1, __FUNCSIG__, + __FILE__, __LINE__); + } + assert(nd == numDimensions); + if (nd != numDimensions) { + throw diskann::ANNException( + "Verifying binary write failed! nd != numDimensions", -1, __FUNCSIG__, + __FILE__, __LINE__); + } + + uint32_t mismatchCount = 0; + uint32_t vectorsToCheck = numPointsToCheck(numPoints); + T* data = new T[numDimensions * vectorsToCheck]; + fin.read((char*) data, sizeof(T) * numDimensions * vectorsToCheck); + + diskann::cout << "Checking first " << vectorsToCheck + << " records for byte-byte match." << std::endl; + + for (uint32_t i = 0; i < vectorsToCheck; i++) { + std::string line; + std::getline(orig, line); + + auto origvec = getVector(line, convertFn, nd); + assert(origvec.size() == nd); + + for (uint32_t j = 0; j < nd; j++) { + if (origvec[j] != data[i * nd + j]) { + std::cerr << "Found mismatch at line " << i << " orig: " << origvec[j] + << " bin val: " << data[i * nd + j] << std::endl; + mismatchCount++; + } + } + } + + delete[] data; + + if (mismatchCount == 0) { + diskann::cout << "Successfully verified byte-byte conversion." << std::endl; + } else { + diskann::cout << "Conversion failed!!! Found " << mismatchCount + << " mismatches in vector elements" << std::endl; + } + diskann::cout << "Successfully verified binary conversion." << std::endl; +} + +template +std::vector getVector(const std::string& str, + T convertFn(const std::string& val)) { + auto idsAndVector = splitLine(str, '\t', suggestedSize); + + std::vector typedVec; + typedVec.reserve(suggestedSize); + + if (idsAndVector.size() >= 3) { + auto strVec = splitLine(idsAndVector[2], ',', suggestedSize); + std::transform( + strVec.begin(), strVec.end(), std::back_inserter(typedVec), + [&convertFn](const std::string& str) { return convertFn(str); }); + } else { + std::cerr << "Found line " << str + << " that didn't have standard delimiters." << std::endl; + } + return typedVec; +} + +template +void convert(const char* infile, const char* outfile, + T convertFn(const std::string& val), uint32_t suggestedSize) { + uint32_t numDimensions = getDimension(infile, convertFn); + uint32_t numPoints = 0; + + std::vector> allVecs; + allVecs.reserve(suggestedSize > 0 ? suggestedSize : 1); + + // T* data = new T[((uint64_t) numDimensions) * numPoints]; + // int index = 0; + + std::ifstream fin(infile); + std::string line; + int lc = 0; + while (std::getline(fin, line)) { + lc++; + auto vector = getVector(line, convertFn, numDimensions); + if (vector.size() != 0) { + if (vector.size() != numDimensions) { + std::stringstream stream; + stream << "Found vector with dimension " << vector.size() + << " instead of " << numDimensions << " at line: " << lc; + throw diskann::ANNException(stream.str(), -1); + } + allVecs.push_back(vector); + if (allVecs.size() % 10000 == 0) { + diskann::cout << "Processed " << allVecs.size() << " records." + << std::endl; + } + // std::for_each(vector.begin(), vector.end(), + // [data, &index](const T& v) { + // data[index] = v; + // index++; + // }); + } + + // if ((index / numDimensions) % 5000 == 0) { + // diskann::cout << "Processed " << index / numDimensions << " + // records." + // << std::endl; + //} + } + + // writeBinary(outfile, numDimensions, numPoints, data); + writeBinaryVec(outfile, numDimensions, allVecs); + // delete[] data; + + sanityCheck(infile, outfile, numDimensions, (uint32_t) allVecs.size(), + convertFn); +} + +int main(int argc, char** argv) { + if (argc != 5) { + std::cerr << "Usage: " << argv[0] + << " " + "[num_of_points] (use zero to not specify)" + << std::endl; + return -1; + } + + diskann::cout << "Converting (" << argv[1] << ") axon file: " << argv[2] + << " with " << argv[4] << " points to diskann bin format file " + << argv[3] << std::endl; + + std::string datatype(argv[1]); + uint32_t numPoints = (uint32_t)::atoi(argv[4]); + + if (datatype == "uint8") { + convert(argv[2], argv[3], atouLambda, numPoints); + } else if (datatype == "int8") { + convert(argv[2], argv[3], atoiLambda, numPoints); + } else if (datatype == "float") { + convert(argv[2], argv[3], atofLambda, numPoints); + } else { + std::cerr << "Unknown data type " << argv[1] + << ". Supported data types are int8, uint8, and float." + << std::endl; + } +} \ No newline at end of file diff --git a/tests/utils/create_disk_layout.cpp b/tests/utils/create_disk_layout.cpp index ac378272d8..bb4ea2db6d 100644 --- a/tests/utils/create_disk_layout.cpp +++ b/tests/utils/create_disk_layout.cpp @@ -13,30 +13,40 @@ #include "utils.h" template -int create_disk_layout(int argc, char **argv) { - if (argc != 5) { - std::cout << argv[0] << " data_type data_bin " - "vamana_index_file output_diskann_index_file" - << std::endl; - exit(-1); - } - std::string base_file(argv[2]); - std::string vamana_file(argv[3]); - std::string output_file(argv[4]); - diskann::create_disk_layout(base_file, vamana_file, output_file); +int create_disk_layout(char **argv) { + std::string vamana_file(argv[2]); + std::string base_file(argv[3]); + std::string tags_file(argv[4]); + std::string pq_pivots_file(argv[5]); + std::string pq_vectors_file(argv[6]); + std::string output_file(argv[7]); + bool single_index_flag = false; + if (base_file == "null") + single_index_flag = true; + diskann::create_disk_layout(vamana_file, base_file, "", + pq_pivots_file, pq_vectors_file, + single_index_flag, output_file); return 0; } int main(int argc, char **argv) { + if (argc != 8) { + diskann::cout << argv[0] + << " data_type vamana_index_file " + " data_file tags_bin pq_pivots_file pq_vectors_file " + "output_diskann_file" + << std::endl; + exit(-1); + } int ret_val = -1; if (std::string(argv[1]) == std::string("float")) - ret_val = create_disk_layout(argc, argv); + ret_val = create_disk_layout(argv); else if (std::string(argv[1]) == std::string("int8")) - ret_val = create_disk_layout(argc, argv); + ret_val = create_disk_layout(argv); else if (std::string(argv[1]) == std::string("uint8")) - ret_val = create_disk_layout(argc, argv); + ret_val = create_disk_layout(argv); else { - std::cout << "unsupported type. use int8/uint8/float " << std::endl; + diskann::cout << "unsupported type. use int8/uint8/float " << std::endl; ret_val = -2; } return ret_val; diff --git a/tests/utils/extract_rows.cpp b/tests/utils/extract_rows.cpp new file mode 100644 index 0000000000..ff05e72343 --- /dev/null +++ b/tests/utils/extract_rows.cpp @@ -0,0 +1,143 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#include +#include +#include "utils.h" +#include "extract_rows.h" + +using namespace std; + +template +void extract_rows(const std::string& input_file, _u64 nrows, _u64 ndims, + _u64 offset, _u64 nrows_to_extract, + const std::string& output_file, bool has_id, + bool replace_content) { + std::cout << "Trying to extract: " << nrows_to_extract << " rows of " + << typeid(T).name() << " data in " << ndims + << " dimensions from offset: " << offset + << " and saving to file: " << output_file << "..."; + std::ifstream fin; + fin.open(input_file, std::ios::binary); + + _u32 rows32 = (_u32) nrows_to_extract, dims32 = (_u32) ndims; + std::ofstream fout(output_file, std::ios::binary); + fout.write((const char*) &rows32, sizeof(uint32_t)); + fout.write((const char*) &dims32, sizeof(uint32_t)); + + T* data = new T[nrows_to_extract * ndims]; + if (replace_content) { + std::cout << std::endl << "Replacing file content with random values..."; + srand((unsigned int) time((time_t*) nullptr)); + for (int i = 0; i < nrows_to_extract * ndims; i++) { + // since we won't allow replace_content if datatype != float, we can + // assume that T is float always. + data[i] = (T)(rand() / RAND_MAX); + } + std::cout << nrows_to_extract * ndims << " values replaced. " << std::endl; + } else { + fin.seekg(2 * sizeof(_u32) + offset * ndims * sizeof(T)); + fin.read((char*) data, nrows_to_extract * ndims * sizeof(T)); + } + + fout.write((const char*) data, nrows_to_extract * ndims * sizeof(T)); + + if (has_id) { + _u64* tags = new _u64[nrows_to_extract]; + _u64 tag_start = + 2 * sizeof(_u32) + nrows * ndims * sizeof(T) + offset * sizeof(_u64); + fin.seekg(tag_start); + fin.read((char*) tags, nrows_to_extract * sizeof(_u64)); + fout.write((const char*) tags, nrows_to_extract * sizeof(_u64)); + delete[] tags; + } + fout.close(); + fin.close(); + delete[] data; + + std::cout << " done." << std::endl; +} + +int main(int argc, char* argv[]) { + if (argc != 6 && argc != 7 && argc != 8) { + std::cout + << "Usage: (uint8|int8|float) " + "<#_rows_to_extract> " + "(zero if beginning) [has_tags(true/false) default is " + "false, file doesn't have ids.] [replace_random (true/false) " + "default is false, i.e.don't replace file content with random data]" + << std::endl; + } + + int count = 1; + std::string input_file = argv[count++]; + std::string datatype = argv[count++]; + uint64_t nrows_to_extract = atoi(argv[count++]); + uint64_t offset = atoi(argv[count++]); + std::string output_file = argv[count++]; + bool has_id = false; + bool replace_content = false; + + if (argc >= 7) { + if (std::string(argv[count++]) == "true") { + has_id = true; + } + } + if (argc >= 8) { + if (std::string(argv[count++]) == "true") { + if (!has_id || datatype != "float") { + std::cout << "Can replace content ONLY if has_id is true and datatype " + "== float. " + << std::endl; + return -6; + } else { + replace_content = true; + } + } + } + + if (!file_exists(input_file)) { + std::cerr << "Input file: " << input_file << " does not exist. Terminating." + << std::endl; + return -1; + } + uint64_t nrows, ncols; + diskann::get_bin_metadata(input_file, nrows, ncols, 0); + + if (nrows_to_extract > nrows) { + std::cerr << "Number of rows in file: " << nrows + << " is less than num rows to extract: " << nrows_to_extract + << std::endl; + return -2; + } + if (offset > nrows) { + std::cerr << "Start offset " << offset + << " is greater than or equal to rows in file: " << nrows + << std::endl; + return -3; + } + if (offset + nrows_to_extract > nrows) { + std::cerr << "Sum of start offset: " << offset + << " and # rows to extract: " << nrows_to_extract + << " is greater than or equal to rows in file: " << nrows + << std::endl; + return -4; + } + + if (datatype == "uint8") { + extract_rows(input_file, nrows, ncols, offset, nrows_to_extract, + output_file, has_id, replace_content); + } else if (datatype == "int8") { + extract_rows(input_file, nrows, ncols, offset, nrows_to_extract, + output_file, has_id, replace_content); + } else if (datatype == "float") { + extract_rows(input_file, nrows, ncols, offset, nrows_to_extract, + output_file, has_id, replace_content); + + } else { + std::cerr << "Unknown datatype: " << datatype; + return -5; + } + + return 0; +} \ No newline at end of file diff --git a/tests/utils/float_bin_to_int8.cpp b/tests/utils/float_bin_to_int8.cpp index 4f422a2336..0620730a51 100644 --- a/tests/utils/float_bin_to_int8.cpp +++ b/tests/utils/float_bin_to_int8.cpp @@ -4,7 +4,6 @@ #include #include "utils.h" - void block_convert(std::ofstream& writer, int8_t* write_buf, std::ifstream& reader, float* read_buf, _u64 npts, _u64 ndims, float bias, float scale) { diff --git a/tests/utils/fvecs_to_bin.cpp b/tests/utils/fvecs_to_bin.cpp index d5c7fe4bab..c26cf1d462 100644 --- a/tests/utils/fvecs_to_bin.cpp +++ b/tests/utils/fvecs_to_bin.cpp @@ -3,6 +3,7 @@ #include #include "utils.h" +#include "logger.h" void block_convert(std::ifstream& reader, std::ofstream& writer, float* read_buf, float* write_buf, _u64 npts, _u64 ndims) { @@ -17,7 +18,7 @@ void block_convert(std::ifstream& reader, std::ofstream& writer, int main(int argc, char** argv) { if (argc != 3) { - std::cout << argv[0] << " input_fvecs output_bin" << std::endl; + diskann::cout << argv[0] << " input_fvecs output_bin" << std::endl; exit(-1); } std::ifstream reader(argv[1], std::ios::binary | std::ios::ate); @@ -29,12 +30,12 @@ int main(int argc, char** argv) { reader.seekg(0, std::ios::beg); _u64 ndims = (_u64) ndims_u32; _u64 npts = fsize / ((ndims + 1) * sizeof(float)); - std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims - << std::endl; + diskann::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims + << std::endl; _u64 blk_size = 131072; _u64 nblks = ROUND_UP(npts, blk_size) / blk_size; - std::cout << "# blks: " << nblks << std::endl; + diskann::cout << "# blks: " << nblks << std::endl; std::ofstream writer(argv[2], std::ios::binary); int npts_s32 = (_s32) npts; int ndims_s32 = (_s32) ndims; @@ -45,7 +46,7 @@ int main(int argc, char** argv) { for (_u64 i = 0; i < nblks; i++) { _u64 cblk_size = std::min(npts - i * blk_size, blk_size); block_convert(reader, writer, read_buf, write_buf, cblk_size, ndims); - std::cout << "Block #" << i << " written" << std::endl; + diskann::cout << "Block #" << i << " written" << std::endl; } delete[] read_buf; diff --git a/tests/utils/gen_random_slice.cpp b/tests/utils/gen_random_slice.cpp index 0417c12c0d..27f7fec9dc 100644 --- a/tests/utils/gen_random_slice.cpp +++ b/tests/utils/gen_random_slice.cpp @@ -21,14 +21,7 @@ #include template -int aux_main(int argc, char** argv) { - if (argc != 5) { - std::cout << argv[0] << " data_type [fliat/int8/uint8] base_bin_file " - "sample_output_prefix sampling_probability" - << std::endl; - exit(-1); - } - +int aux_main(char** argv) { std::string base_file(argv[2]); std::string output_prefix(argv[3]); float sampling_rate = (float) (std::atof(argv[4])); @@ -37,13 +30,21 @@ int aux_main(int argc, char** argv) { } int main(int argc, char** argv) { + if (argc != 5) { + diskann::cout << argv[0] + << " data_type [float/int8/uint8] base_bin_file " + "sample_output_prefix sampling_probability" + << std::endl; + exit(-1); + } + if (std::string(argv[1]) == std::string("float")) { - aux_main(argc, argv); + aux_main(argv); } else if (std::string(argv[1]) == std::string("int8")) { - aux_main(argc, argv); + aux_main(argv); } else if (std::string(argv[1]) == std::string("uint8")) { - aux_main(argc, argv); + aux_main(argv); } else - std::cout << "Unsupported type. Use float/int8/uint8." << std::endl; + diskann::cout << "Unsupported type. Use float/int8/uint8." << std::endl; return 0; } diff --git a/tests/utils/generate_pq.cpp b/tests/utils/generate_pq.cpp index 9e7a37c597..4279df05d8 100644 --- a/tests/utils/generate_pq.cpp +++ b/tests/utils/generate_pq.cpp @@ -13,20 +13,22 @@ bool generate_pq(const std::string& data_path, const float sampling_rate) { std::string pq_pivots_path = index_prefix_path + "_pq_pivots.bin"; std::string pq_compressed_vectors_path = - index_prefix_path + "_compressed.bin"; + index_prefix_path + "_pq_compressed.bin"; // generates random sample and sets it to train_data and updates train_size size_t train_size, train_dim; float* train_data; gen_random_slice(data_path, sampling_rate, train_data, train_size, train_dim); - std::cout << "For computing pivots, loaded sample data of size " << train_size - << std::endl; + diskann::cout << "For computing pivots, loaded sample data of size " + << train_size << std::endl; - generate_pq_pivots(train_data, train_size, train_dim, num_pq_centers, - num_pq_chunks, KMEANS_ITERS_FOR_PQ, pq_pivots_path); - generate_pq_data_from_pivots(data_path, num_pq_centers, num_pq_chunks, - pq_pivots_path, pq_compressed_vectors_path); + generate_pq_pivots(train_data, train_size, (_u32) train_dim, + (_u32) num_pq_centers, (_u32) num_pq_chunks, + KMEANS_ITERS_FOR_PQ, pq_pivots_path); + generate_pq_data_from_pivots(data_path, (_u32) num_pq_centers, + (_u32) num_pq_chunks, pq_pivots_path, + pq_compressed_vectors_path); delete[] train_data; @@ -35,7 +37,7 @@ bool generate_pq(const std::string& data_path, int main(int argc, char** argv) { if (argc != 6) { - std::cout + diskann::cout << "Usage: \n" << argv[0] << " " @@ -46,7 +48,7 @@ int main(int argc, char** argv) { const std::string index_prefix_path(argv[3]); const size_t num_pq_centers = 256; const size_t num_pq_chunks = (size_t) atoi(argv[4]); - const float sampling_rate = atof(argv[5]); + const float sampling_rate = (float) atof(argv[5]); if (std::string(argv[1]) == std::string("float")) generate_pq(data_path, index_prefix_path, num_pq_centers, @@ -58,6 +60,6 @@ int main(int argc, char** argv) { generate_pq(data_path, index_prefix_path, num_pq_centers, num_pq_chunks, sampling_rate); else - std::cout << "Error. wrong file type" << std::endl; + diskann::cout << "Error. wrong file type" << std::endl; } } diff --git a/tests/utils/int8_to_float.cpp b/tests/utils/int8_to_float.cpp index 3a62ea5749..e654b1874a 100644 --- a/tests/utils/int8_to_float.cpp +++ b/tests/utils/int8_to_float.cpp @@ -6,7 +6,7 @@ int main(int argc, char** argv) { if (argc != 3) { - std::cout << argv[0] << " input_int8_bin output_float_bin" << std::endl; + diskann::cout << argv[0] << " input_int8_bin output_float_bin" << std::endl; exit(-1); } diff --git a/tests/utils/ivecs_to_bin.cpp b/tests/utils/ivecs_to_bin.cpp index e42eeae7ee..28bcedd2d2 100644 --- a/tests/utils/ivecs_to_bin.cpp +++ b/tests/utils/ivecs_to_bin.cpp @@ -17,7 +17,7 @@ void block_convert(std::ifstream& reader, std::ofstream& writer, _u32* read_buf, int main(int argc, char** argv) { if (argc != 3) { - std::cout << argv[0] << " input_ivecs output_bin" << std::endl; + diskann::cout << argv[0] << " input_ivecs output_bin" << std::endl; exit(-1); } std::ifstream reader(argv[1], std::ios::binary | std::ios::ate); @@ -29,12 +29,12 @@ int main(int argc, char** argv) { reader.seekg(0, std::ios::beg); _u64 ndims = (_u64) ndims_u32; _u64 npts = fsize / ((ndims + 1) * sizeof(_u32)); - std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims - << std::endl; + diskann::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims + << std::endl; _u64 blk_size = 131072; _u64 nblks = ROUND_UP(npts, blk_size) / blk_size; - std::cout << "# blks: " << nblks << std::endl; + diskann::cout << "# blks: " << nblks << std::endl; std::ofstream writer(argv[2], std::ios::binary); int npts_s32 = (_s32) npts; int ndims_s32 = (_s32) ndims; @@ -45,7 +45,7 @@ int main(int argc, char** argv) { for (_u64 i = 0; i < nblks; i++) { _u64 cblk_size = std::min(npts - i * blk_size, blk_size); block_convert(reader, writer, read_buf, write_buf, cblk_size, ndims); - std::cout << "Block #" << i << " written" << std::endl; + diskann::cout << "Block #" << i << " written" << std::endl; } delete[] read_buf; diff --git a/tests/utils/merge_shards.cpp b/tests/utils/merge_shards.cpp index 2f9eadec07..f6ce017649 100644 --- a/tests/utils/merge_shards.cpp +++ b/tests/utils/merge_shards.cpp @@ -16,7 +16,7 @@ int main(int argc, char **argv) { if (argc != 9) { - std::cout + diskann::cout << argv[0] << " vamana_index_prefix[1] vamana_index_suffix[2] idmaps_prefix[3] " "idmaps_suffix[4] n_shards[5] max_degree[6] output_vamana_path[7] " diff --git a/tests/utils/normalize_data.cpp b/tests/utils/normalize_data.cpp new file mode 100644 index 0000000000..ca74265912 --- /dev/null +++ b/tests/utils/normalize_data.cpp @@ -0,0 +1,70 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#include +#include +#include +#include "utils.h" + +void block_convert(std::ofstream& writr, std::ifstream& readr, float* read_buf, + _u64 npts, _u64 ndims) { + readr.read((char*) read_buf, npts * ndims * sizeof(float)); + _u32 ndims_u32 = (_u32) ndims; +#pragma omp parallel for + for (_s64 i = 0; i < (_s64) npts; i++) { + float norm_pt = std::numeric_limits::epsilon(); + for (_u32 dim = 0; dim < ndims_u32; dim++) { + norm_pt += *(read_buf + i * ndims + dim) * *(read_buf + i * ndims + dim); + } + norm_pt = std::sqrt(norm_pt); + for (_u32 dim = 0; dim < ndims_u32; dim++) { + *(read_buf + i * ndims + dim) = *(read_buf + i * ndims + dim) / norm_pt; + } + } + writr.write((char*) read_buf, npts * ndims * sizeof(float)); +} + +int main(int argc, char** argv) { + if (argc != 3) { + diskann::cout << argv[0] + << ": [input_bin (float data)] [output_bin (float data)] " + << std::endl; + exit(-1); + } + std::ifstream readr(argv[1], std::ios::binary); + int npts_s32; + int ndims_s32; + readr.read((char*) &npts_s32, sizeof(_s32)); + readr.read((char*) &ndims_s32, sizeof(_s32)); + // size_t npt = npts_s32; + // size_t ndim = ndims_s32; + _u32 ndims_u32 = (_u32) ndims_s32; + _u32 npts_u32 = (_u32) npts_s32; + // readr.seekg(0, std::ios::end); + //_u64 fsize = readr.tellg(); + + std::ofstream writr(argv[2], std::ios::binary); + writr.write((char*) &npts_s32, sizeof(_s32)); + writr.write((char*) &ndims_s32, sizeof(_s32)); + + // writr.write((char*) &ndims_u32, sizeof(unsigned)); + // writr.seekg(0, std::ios::beg); + _u64 ndims = (_u64) ndims_u32; + _u64 npts = (_u64) npts_u32; + diskann::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims + << std::endl; + + _u64 blk_size = 131072; + _u64 nblks = ROUND_UP(npts, blk_size) / blk_size; + diskann::cout << "# blks: " << nblks << std::endl; + + float* read_buf = new float[npts * ndims]; + for (_u64 i = 0; i < nblks; i++) { + _u64 cblk_size = std::min(npts - i * blk_size, blk_size); + block_convert(writr, readr, read_buf, cblk_size, ndims); + diskann::cout << "Block #" << i << " written" << std::endl; + } + delete[] read_buf; + writr.close(); + readr.close(); +} diff --git a/tests/utils/partition_data.cpp b/tests/utils/partition_data.cpp index 63476b15e2..0cc2ad46d8 100644 --- a/tests/utils/partition_data.cpp +++ b/tests/utils/partition_data.cpp @@ -10,17 +10,18 @@ int main(int argc, char** argv) { if (argc != 7) { - std::cout << "Usage:\n" - << argv[0] << " datatype " - " " - " " - << std::endl; + diskann::cout << "Usage:\n" + << argv[0] + << " datatype " + " " + " " + << std::endl; exit(-1); } const std::string data_path(argv[2]); const std::string prefix_path(argv[3]); - const float sampling_rate = atof(argv[4]); + const float sampling_rate = (float) atof(argv[4]); const size_t num_partitions = (size_t) std::atoi(argv[5]); const size_t max_reps = 15; const size_t k_index = (size_t) std::atoi(argv[6]); @@ -35,5 +36,6 @@ int main(int argc, char** argv) { partition(data_path, sampling_rate, num_partitions, max_reps, prefix_path, k_index); else - std::cout << "unsupported data format. use float/int8/uint8" << std::endl; + diskann::cout << "unsupported data format. use float/int8/uint8" + << std::endl; } diff --git a/tests/utils/partition_with_ram_budget.cpp b/tests/utils/partition_with_ram_budget.cpp index fee3c84e82..296e70153d 100644 --- a/tests/utils/partition_with_ram_budget.cpp +++ b/tests/utils/partition_with_ram_budget.cpp @@ -10,17 +10,18 @@ int main(int argc, char** argv) { if (argc != 8) { - std::cout << "Usage:\n" - << argv[0] << " datatype " - " " - " " - << std::endl; + diskann::cout << "Usage:\n" + << argv[0] + << " datatype " + " " + " " + << std::endl; exit(-1); } const std::string data_path(argv[2]); const std::string prefix_path(argv[3]); - const float sampling_rate = atof(argv[4]); + const float sampling_rate = (float) atof(argv[4]); const double ram_budget = (double) std::atof(argv[5]); const size_t graph_degree = (size_t) std::atoi(argv[6]); const size_t k_index = (size_t) std::atoi(argv[7]); @@ -35,5 +36,6 @@ int main(int argc, char** argv) { partition_with_ram_budget(data_path, sampling_rate, ram_budget, graph_degree, prefix_path, k_index); else - std::cout << "unsupported data format. use float/int8/uint8" << std::endl; + diskann::cout << "unsupported data format. use float/int8/uint8" + << std::endl; } diff --git a/tests/utils/prep_index_merger.sh b/tests/utils/prep_index_merger.sh new file mode 100644 index 0000000000..76a9a35ec0 --- /dev/null +++ b/tests/utils/prep_index_merger.sh @@ -0,0 +1,21 @@ +BASE_PREFIX="/dev/shm/test/sample_base" +MEM_PREFIX="/dev/shm/test/sample_mem" +DELETE_LIST="/dev/shm/sample_deleted.tags" +ONESHOT_PREFIX="/dev/shm/test/sample_oneshot" +MERGED_PREFIX="/dev/shm/test/sample_merged" +NUM_MEM_INDICES=5 +# copy tags from base -> base_index +cp ${BASE_PREFIX}.tags ${BASE_PREFIX}_index_disk.index.tags +cp ${ONESHOT_PREFIX}.tags ${ONESHOT_PREFIX}_index_disk.index.tags + +# copy tags file for mem indices +for i in $(seq 1 $NUM_MEM_INDICES) +do + cp ${MEM_PREFIX}_${i}.tags ${MEM_PREFIX}_${i}_index.tags +done + +# copy PQ stuff for merged from base +cp ${BASE_PREFIX}_index_pq_pivots.bin ${MERGED_PREFIX}_index_pq_pivots.bin +cp ${BASE_PREFIX}_index_pq_pivots.bin_centroid.bin ${MERGED_PREFIX}_index_pq_pivots.bin_centroid.bin +cp ${BASE_PREFIX}_index_pq_pivots.bin_chunk_offsets.bin ${MERGED_PREFIX}_index_pq_pivots.bin_chunk_offsets.bin +cp ${BASE_PREFIX}_index_pq_pivots.bin_rearrangement_perm.bin ${MERGED_PREFIX}_index_pq_pivots.bin_rearrangement_perm.bin diff --git a/tests/utils/script.sh b/tests/utils/script.sh new file mode 100644 index 0000000000..5cdca72927 --- /dev/null +++ b/tests/utils/script.sh @@ -0,0 +1,56 @@ +#BASE_FILE=/home/t-adisin/sift_base.bin +#QUERY_FILE=/home/t-adisin/sift_query.bin +#GT_FILE=~/gs500_sift1m +#NUM_BASE=900000 +#NUM_DELETE=50000 +#NUM_INCR=50000 +#NUM_CYCLES=30 +#SAVE_PATH=/mnt/rakri/save/SIFT1M +#LOGFILE=~/sift1m_${NUM_BASE}_${NUM_DELETE}_${NUM_INCR}_${NUM_CYCLES} + +BASE_FILE=/mnt/aditi/datasets/sift_rnd100m_data.bin +QUERY_FILE=/home/t-adisin/datasets/bigann_query_float.bin +GT_FILE=~/gs500_sift100m +NUM_BASE=90000000 +NUM_DELETE=4500000 +NUM_INCR=4500000 +NUM_CYCLES=50 +SAVE_PATH=/mnt/rakri/save/SIFT100M +LOGFILE=~/sift100m_${NUM_BASE}_${NUM_DELETE}_${NUM_INCR}_${NUM_CYCLES} + +WORKING_PATH=/mnt/rakri/test3/ +#WORKING_PATH=/dev/shm/test4 + +INDEX_FILE=merger_index_${NUM_BASE}_${NUM_DELETE}_${NUM_INCR}_${NUM_CYCLES} + +#rm -rf /mnt/rakri/*pq* +./tests/utils/seed_index_merger float $NUM_BASE 1 $NUM_DELETE $NUM_INCR $NUM_CYCLES $BASE_FILE ${SAVE_PATH}/${INDEX_FILE} ${SAVE_PATH}/${INDEX_FILE}_deleted.tags 2>&1 0${LOGFILE}_build.log + + +# train indices +numactl --interleave=all ./tests/build_disk_index float ${SAVE_PATH}/${INDEX_FILE}_base.data ${WORKING_PATH}/${INDEX_FILE}_cycle_0 64 75 100 500 64 2>&1 0>${LOGFILE}_build.log +cp ${SAVE_PATH}/${INDEX_FILE}_base.tags ${WORKING_PATH}/${INDEX_FILE}_cycle_0_disk.index.tags + +mkdir $SAVE_PATH/base_copy/ +cp ${WORKING_PATH}/${INDEX_FILE}_cycle_0* $SAVE_PATH/base_copy/ + +#./tests/utils/compute_groundtruth float $BASE_FILE $QUERY_FILE 500 $GT_FILE + +# search on base index +numactl --interleave=all ./tests/search_disk_index float $WORKING_PATH/${INDEX_FILE}_cycle_0 100000 32 4 $QUERY_FILE $GT_FILE 10 /tmp/abc 100 2>&1 0${LOGFILE}_search.log + +for (( c=0; c<$NUM_CYCLES; c++ )) +do + ((nextc = c + 1)) + cp ${SAVE_PATH}/${INDEX_FILE}_cycle_${nextc}_mem_1.data $WORKING_PATH/ + cp ${SAVE_PATH}/${INDEX_FILE}_cycle_${nextc}_mem_1.tags $WORKING_PATH/ + cp ${SAVE_PATH}/${INDEX_FILE}_deleted.tags_cycle_${nextc} $WORKING_PATH/ + numactl --interleave=all ./tests/build_memory_index float $WORKING_PATH/${INDEX_FILE}_cycle_${nextc}_mem_1.data $WORKING_PATH/${INDEX_FILE}_cycle_${nextc}_mem_1.index 64 75 1.2 64 + cp $WORKING_PATH/${INDEX_FILE}_cycle_${nextc}_mem_1.tags $WORKING_PATH/${INDEX_FILE}_cycle_${nextc}_mem_1.index.tags 2>&1 0>${LOGFILE}_build.log + /usr/bin/time numactl --interleave=all ./tests/test_index_merger float $WORKING_PATH/${INDEX_FILE}_cycle_${c} $WORKING_PATH/${INDEX_FILE}_cycle_${nextc} $WORKING_PATH/${INDEX_FILE}_deleted.tags_cycle_${nextc} 128 4 64 70 1.2 1000 $WORKING_PATH $WORKING_PATH/${INDEX_FILE}_cycle_${nextc}_mem_1.index 2>&1 0>${LOGFILE}_build.log + numactl --interleave=all ./tests/search_disk_index float $WORKING_PATH/${INDEX_FILE}_cycle_${nextc} 100000 32 4 $QUERY_FILE $GT_FILE 10 /tmp/def 100 2>&1 0>${LOGFILE}_search.log + sudo rm -rf $WORKING_PATH/${INDEX_FILE}_cycle_${c}_* + rm -rf $SAVE_PATH/${INDEX_FILE}_cycle_${c}_* + cp $WORKING_PATH/${INDEX_FILE}_cycle_${nextc}* $SAVE_PATH/ +done + diff --git a/tests/utils/seed_index_merger.cpp b/tests/utils/seed_index_merger.cpp new file mode 100644 index 0000000000..358477c23a --- /dev/null +++ b/tests/utils/seed_index_merger.cpp @@ -0,0 +1,219 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#include "utils.h" + +#include +#include +#include +#include +#include +#include +#include "tsl/robin_set.h" + +#define ENTRY_POINT 52292725 +//#define ENTRY_POINT 123742 + +template +void dump_to_disk(const T *all_pts, const uint64_t ndims, + const std::string & filename, + const std::vector &tags) { + T * new_data = new T[ndims * tags.size()]; + TagT *new_tags = new TagT[tags.size()]; + + std::string tag_filename = filename + ".tags"; + std::string data_filename = filename + ".data"; + diskann::cout << "# points : " << tags.size() << "\n"; + diskann::cout << "Tag file : " << tag_filename << "\n"; + diskann::cout << "Data file : " << data_filename << "\n"; + + std::ofstream tag_writer(tag_filename); + for (uint64_t i = 0; i < tags.size(); i++) { + // tag_writer << tags[i] << std::endl; + *(new_tags + i) = tags[i]; + memcpy(new_data + (i * ndims), all_pts + (tags[i] * ndims), + ndims * sizeof(float)); + } + // tag_writer.close(); + diskann::save_bin(tag_filename, new_tags, tags.size(), 1); + diskann::save_bin(data_filename, new_data, tags.size(), ndims); + delete new_data; + delete new_tags; +} + +template +void run(const uint32_t base_count, const uint32_t num_mem_indices, + const uint32_t delete_count, const uint32_t incr_count, + const uint32_t num_cycles, const std::string &in_file, + const std::string &out_prefix, const std::string &deleted_tags_file) { + // random number generator + std::random_device dev; + std::mt19937 rng(dev()); + + T * all_points = nullptr; + uint64_t npts, ndims; + diskann::load_bin(in_file, all_points, npts, ndims); + diskann::cout << "Loaded " << npts << " pts x " << ndims << " dims\n"; + // assert(npts >= base_count + num_mem_indices); + + std::vector tags(npts); + std::iota(tags.begin(), tags.end(), 0); + + diskann::cout << "Base Index : choosing " << base_count << " points\n"; + std::vector base_tags(tags.begin(), tags.begin() + base_count); + + tsl::robin_set active_tags; + tsl::robin_set inactive_tags; + + tsl::robin_set new_active_tags; + tsl::robin_set new_inactive_tags; + + for (uint32_t i = 0; i < base_count; i++) { + active_tags.insert(i); + } + + for (uint32_t i = base_count; i < npts; i++) { + inactive_tags.insert(i); + } + + diskann::cout << "Dumping base set \n"; + // write base + dump_to_disk(all_points, ndims, out_prefix + "_base", base_tags); + + std::vector delete_vec; + std::vector insert_vec; + uint32_t count = 0; + while (count++ < num_cycles) { + std::random_device rd; + std::mt19937 gen(rd()); + std::uniform_real_distribution dis(0, 1); + + new_active_tags.clear(); + new_inactive_tags.clear(); + + delete_vec.clear(); + insert_vec.clear(); + + float active_tags_sampling_rate = (float) ((std::min)( + (1.0 * delete_count) / (1.0 * ((double) active_tags.size())), 1.0)); + + for (auto iter = active_tags.begin(); iter != active_tags.end(); iter++) { + if (dis(gen) < active_tags_sampling_rate && *iter != ENTRY_POINT) { + delete_vec.emplace_back(*iter); + new_inactive_tags.insert(*iter); + } else + new_active_tags.insert(*iter); + } + + float inactive_tags_sampling_rate = (float) ((std::min)( + (1.0 * incr_count) / (1.0 * ((double) inactive_tags.size())), 1.0)); + + for (auto iter = inactive_tags.begin(); iter != inactive_tags.end(); + iter++) { + if (dis(gen) < inactive_tags_sampling_rate) { + insert_vec.emplace_back(*iter); + new_active_tags.insert(*iter); + } else + new_inactive_tags.insert(*iter); + } + + diskann::cout << "Merge program will insert " << insert_vec.size() + << " points and delete " << delete_vec.size() + << " points in round " << count << std::endl; + + active_tags.swap(new_active_tags); + inactive_tags.swap(new_inactive_tags); + + // TODO (correct) :: enable shuffling tags for better randomness + // std::shuffle(tags.begin(), tags.end(), rng); + + // split tags + + const uint64_t mem_count = + ROUND_UP(insert_vec.size(), num_mem_indices) / num_mem_indices; + + std::vector> mem_tags(num_mem_indices); + uint64_t cur_start = 0; + for (uint64_t i = 0; i < num_mem_indices; i++) { + std::vector &ith_tags = mem_tags[i]; + uint64_t new_start = std::min(cur_start + mem_count, insert_vec.size()); + diskann::cout << "Index #" << i + 1 << " : choosing " + << new_start - cur_start << " points\n"; + ith_tags.insert(ith_tags.end(), insert_vec.begin() + cur_start, + insert_vec.begin() + new_start); + cur_start = new_start; + } + + // write mem + for (uint64_t i = 0; i < num_mem_indices; i++) { + diskann::cout << "Dumping mem set #" << i + 1 << "\n"; + dump_to_disk(all_points, ndims, + out_prefix + "_cycle_" + std::to_string(count) + "_mem_" + + std::to_string(i + 1), + mem_tags[i]); + } + + // re-shuffle tags to get delete list + // std::shuffle(tags.begin(), tags.end(), rng); + std::ofstream deleted_tags_writer(deleted_tags_file + "_cycle_" + + std::to_string(count)); + for (uint64_t i = 0; i < delete_vec.size(); i++) { + deleted_tags_writer << delete_vec[i] << std::endl; + } + deleted_tags_writer.close(); + + // add remaining tags to final list + // std::vector rem_tags(tags.begin() + delete_count, tags.end()); + // std::sort(rem_tags.begin(), rem_tags.end()); + // write mem + // diskann::cout << "Dumping {all} \\ {deleted} set\n"; + // dump_to_disk(all_points, ndims, out_prefix + "_oneshot", rem_tags); + } + // free all points + delete all_points; +} + +int main(int argc, char **argv) { + if (argc != 10) { + diskann::cout << "Correct usage: " << argv[0] + << " " + << " " + " " + << std::endl; + exit(-1); + } + diskann::cout.setf(std::ios::unitbuf); + + int arg_no = 1; + std::string index_type = argv[arg_no++]; + const uint32_t base_count = std::atoi(argv[arg_no++]); + diskann::cout << "# base points : " << base_count << "\n"; + const uint32_t num_mem_indices = std::atoi(argv[arg_no++]); + diskann::cout << "# mem indices : " << num_mem_indices << "\n"; + const uint32_t delete_count = std::atoi(argv[arg_no++]); + diskann::cout << "# deleted tags per cycle : " << delete_count << "\n"; + const uint32_t incr_count = std::atoi(argv[arg_no++]); + diskann::cout << "# inserted tags per cycle : " << incr_count << "\n"; + const uint32_t num_cycles = std::atoi(argv[arg_no++]); + diskann::cout << "# cycles : " << num_cycles << "\n"; + const std::string in_file = argv[arg_no++]; + diskann::cout << "In file : " << in_file << "\n"; + const std::string out_prefix = argv[arg_no++]; + diskann::cout << "Out prefix : " << out_prefix << "\n"; + const std::string deleted_tags_file = argv[arg_no++]; + diskann::cout << "Deleted tags file prefix : " << deleted_tags_file << "\n"; + + if (index_type == std::string("float")) { + run(base_count, num_mem_indices, delete_count, incr_count, + num_cycles, in_file, out_prefix, deleted_tags_file); + } else if (index_type == std::string("uint8")) { + run(base_count, num_mem_indices, delete_count, incr_count, + num_cycles, in_file, out_prefix, deleted_tags_file); + } else if (index_type == std::string("int8")) { + run(base_count, num_mem_indices, delete_count, incr_count, + num_cycles, in_file, out_prefix, deleted_tags_file); + } else { + diskann::cout << "Unsupported type : " << index_type << "\n"; + } + diskann::cout << "Exiting\n"; +} diff --git a/tests/utils/test_l2_avx.cpp b/tests/utils/test_l2_avx.cpp new file mode 100644 index 0000000000..e315d99ca2 --- /dev/null +++ b/tests/utils/test_l2_avx.cpp @@ -0,0 +1,482 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include "utils.h" +#include "distance.h" +#include "logger.h" + +// 7-- OpenMP thread pool test + +void testOpenMPThreadPool() { + std::vector vec(8192); + std::iota(vec.begin(), vec.end(), 1); + std::vector squares(8192); + std::vector resource(8); + +#pragma omp parallel for schedule(dynamic, 64) num_threads(8) + for (int i = 0; i < vec.size(); i++) { + squares[i] = vec[i] * vec[i]; + resource[omp_get_thread_num()] = omp_get_thread_num(); + } + + // Anything to prevent the above code from being optimized away. + for (int i = 0; i < 8192; i++) { + if (i % 1000 == 0) { + std::cout << i << squares[i] << std::endl; + } + } + std::cout << "Omp is running: " << omp_get_num_threads() << std::endl; + +#pragma omp parallel for schedule(dynamic, 64) num_threads(8) + for (int i = 0; i < vec.size(); i++) { + squares[i] = vec[i] * vec[i]; + resource[omp_get_thread_num()] = omp_get_thread_num(); + } + + for (auto id : resource) { + std::cout << id << " "; + } + std::cout << std::endl; + std::cout << "After second omp, " << omp_get_num_threads() << std::endl; +} + +// For flexibility in switching between AVX/AVX2. +bool Avx2SupportedCPU = true; + +// Utility functions. +template +void printArray(T* arr, int size) { + for (int i = 0; i < size; i++) { + std::cout << (int) arr[i] << " "; + } + std::cout << std::endl; +} + +template<> +void printArray(float* arr, int size) { + for (int i = 0; i < size; i++) { + std::cout << arr[i] << " "; + } + std::cout << std::endl; +} + +// Lame but correct distance computations +float distanceL2_I(const int8_t* a, const int8_t* b, size_t size) { + float distance = 0; + for (int i = 0; i < size; i++) { + int16_t diff = ((int16_t) a[i] - (int16_t) b[i]); + distance += diff * diff; + } + return distance; +} + +float distanceL2_F(const float* a, const float* b, size_t size) { + float distance = 0; + for (int i = 0; i < size; i++) { + distance += (a[i] - b[i]) * (a[i] - b[i]); + } + return distance; +} + +float distanceCosine_I(const int8_t* a, const int8_t* b, size_t size) { + float aMag = 0.0f, bMag = 0.0f, scalarProduct = 0.0f; + for (int i = 0; i < size; i++) { + aMag += ((int32_t) a[i]) * ((int32_t) a[i]); + bMag += ((int32_t) b[i]) * ((int32_t) b[i]); + scalarProduct += ((int32_t) a[i]) * ((int32_t) b[i]); + } + std::cout << "aMag: " << aMag << " bMag: " << bMag + << " product: " << scalarProduct << std::endl; + return 1.0f - (float) (scalarProduct / (sqrt(aMag) * sqrt(bMag))); +} +float distanceCosine_F(const float* a, const float* b, size_t size) { + float aMag = 0.0f, bMag = 0.0f, scalarProduct = 0.0f; + for (int i = 0; i < size; i++) { + aMag += a[i] * a[i]; + bMag += b[i] * b[i]; + scalarProduct += a[i] * b[i]; + } + std::cout << "aMag: " << aMag << " bMag: " << bMag + << " product: " << scalarProduct << std::endl; + + return 1.0f - (float) (scalarProduct / (sqrt(aMag) * sqrt(bMag))); +} + +int8_t* createVector(int size) { + auto p = new int8_t[size]; + for (int i = 0; i < size; i++) { + p[i] = 0; + } + return p; +} + +uint8_t* getData(const char* infileName, uint8_t dataTypeSizeInBytes, + int32_t& count, int32_t& dimension) { + std::ifstream infile(infileName, std::ios::binary); + if (!infile.is_open()) { + std::cerr << "Could not open input file: " << infileName << std::endl; + } + + infile.read(reinterpret_cast(&count), sizeof(count)); + infile.read(reinterpret_cast(&dimension), sizeof(dimension)); + + diskann::cout << infileName << ": count: " << count + << " dimensions: " << dimension << std::endl; + + uint64_t sizeToRead = (uint64_t) count * dataTypeSizeInBytes * dimension; + uint8_t* bytes = new uint8_t[sizeToRead]; + infile.read((char*) bytes, sizeToRead); + + return bytes; +} + +void concat_files(int argc, char** argv) { + if (argc < 5) { + diskann::cout + << "Mode 1 requires 2 input files and 1 output file as argument." + << std::endl + << " This will concatenate the contents of the files specified " + "(in the same order) and save them to outfile." + << std::endl; + return; + } + + int dimension1, dimension2, count1, count2; + uint8_t* data1 = getData(argv[2], 1, count1, dimension1); + uint8_t* data2 = getData(argv[3], 1, count2, dimension2); + + if (dimension1 != dimension2) { + diskann::cout << "Error! Cannot combine vectors of differing dimensions (" + << dimension1 << "," << dimension2 << ")" << std::endl; + return; + } + + std::ofstream outFile(argv[4], std::ios::binary); + if (!outFile.is_open()) { + std::cerr << "Could not open output file: " << argv[4] << std::endl; + return; + } + outFile << count1 + count2; + outFile << dimension1; + outFile.write((char*) data1, count1); + outFile.write((char*) data2, count2); + outFile.close(); + + delete[] data1; + delete[] data2; + + diskann::cout << "Concatenated " << (count1 + count2) << " " << dimension1 + << "-dimension vectors from files: " << argv[2] << " and " + << argv[3] << " and saved to output file: " << argv[4] + << std::endl; +} + +void assignPtr(std::unique_ptr& data) { + float* ptr = new float[30]; + for (int i = 0; i < 30; i++) { + ptr[i] = i * 1.0f; + } + + data.reset(ptr); + // std::make_unique(30); +} + +void uniquePtrAssignment() { + std::unique_ptr data; + assignPtr(data); + + for (int i = 0; i < 30; i++) { + diskann::cout << data[i] << " "; + if (data[i] - i > 0.001) { + diskann::cout << data[i] << "," << i << " screwed " << std::endl; + break; + } + } + diskann::cout << std::endl; + diskann::cout << "safe." << std::endl; +} + +float epsilon = 0.01; +const int A_SIZE = 100; +const int A_COUNT = 20; +void compareDistanceComputationsInt() { + // int8_t vec1[] = {127, 127, 127, 127, 127, 127, 127, 127, 127, 127, 127, + // 127, 127, 127, 127, 127, 127, 127, 127, 127, 127, 127, + // 127, 127, 127, 127, 127, 127, 127, 127, 127, 127}; + // int8_t vec2[] = {-128, -128, -128, -128, -128, -128, -128, -128, + // -128, -128, -128, -128, -128, -128, -128, -128, + // -128, -128, -128, -128, -128, -128, -128, -128, + // -128, -128, -128, -128, -128, -128, -128, -128}; + + int8_t vec1[A_SIZE], vec2[A_SIZE]; + srand(time(0)); + bool neg = false; + for (int i = 0; i < A_SIZE; i++) { + auto a = rand() % 10; + auto b = rand() % 10; + vec1[i] = -a; // < 0 ? -a : a; + vec2[i] = -b; // < 0 ? -b : b; + + neg = neg || vec1[i] < 0 || vec2[i] < 0; + } + if (!neg) { + diskann::cout << "No negative numbers in test. " << std::endl; + } + + std::cout << "Starting distance computation test for float." << std::endl; + printArray(vec1, 16); + printArray(vec2, 16); + + { + // diskann::DistanceL2Int8 dist; + diskann::AVXDistanceL2Int8 dist; + float dist1 = dist.compare(vec1, vec2, 16); + float dist2 = distanceL2_I(vec1, vec2, 16); + + if (abs(dist1 - dist2) > epsilon) { + std::cout + << "compareDistanceComputationsInt(): L2 Test failed. AVX dist: " + << dist1 << " normal dist: " << dist2 << " difference > " << epsilon + << std::endl; + } else { + std::cout << "Two scores are the same. " << std::endl; + } + } + { + diskann::DistanceCosineInt8 dist; + float dist1 = dist.compare(vec1, vec2, 16); + float dist2 = distanceCosine_I(vec1, vec2, 16); + if (abs(dist1 - dist2) > epsilon) { + std::cout + << "compareDistanceComputationsInt(): Cosine Test failed. AVX dist: " + << dist1 << " normal dist: " << dist2 << " difference > " << epsilon + << std::endl; + } else { + std::cout << "Two scores are the same. " << std::endl; + } + } +} + +void compareDistanceComputationsFloat() { + srand(time(0)); + float a[A_SIZE], b[A_SIZE]; + for (int i = 0; i < A_SIZE; i++) { + a[i] = (float) (rand() / 10E5); + b[i] = (float) (rand() / 10E5); + } + + printArray(a, 16); + printArray(b, 16); + + std::cout << "Starting distance computation test for float." << std::endl; + + { + diskann::AVXDistanceL2Float dist; + float dist1 = dist.compare(a, b, 16); + float dist2 = distanceL2_F(a, b, 16); + if (abs(dist1 - dist2) > epsilon) { + std::cout + << "compareDistanceComputationsFloat(): L2 Test failed. AVX dist: " + << dist1 << " normal dist: " << dist2 << " difference > " << epsilon + << std::endl; + } else { + std::cout << "Two scores are the same." << std::endl; + } + } + { + diskann::DistanceCosineFloat dist; + float dist1 = dist.compare(a, b, 16); + float dist2 = distanceCosine_F(a, b, 16); + if (abs(dist1 - dist2) > epsilon) { + std::cout << "compareDistanceComputationsFloat(): Cosine Test failed. " + "AVX dist: " + << dist1 << " normal dist: " << dist2 << " difference > " + << epsilon << std::endl; + } else { + std::cout << "Two scores are the same." << std::endl; + } + } +} + +void testStreamBufImpl() { + std::vector v(100); + std::iota(v.begin(), v.end(), 1); + + diskann::cout << "Printing with diskann::cout" << std::endl; +#pragma omp parallel for schedule(dynamic, 64) + for (int i = 0; i < v.size(); i++) { + diskann::cout << std::to_string(i) + ","; + if (i != 0 && i % 10 == 0) { + diskann::cout << std::endl; + } + } + + // diskann::cout << "Printing with diskann::cout" << std::endl; + //#pragma omp parallel for schedule(dynamic, 64) + // for (int i = 0; i < v.size(); i++) { + // diskann::cout << std::to_string(i) + ","; + // if (i != 0 && i % 10 == 0) { + // diskann::cout << std::endl; + // } + // } +} + +//--5. Testing MemBuf Impl +class ContentBuf : public std::basic_streambuf { + public: + ContentBuf(char* ptr, size_t size) { + setg(ptr, ptr, ptr + size); + } +}; +void testMemBufImpl(int argc, char** argv) { + // Create a simple binary file. + std::string fileName = argv[2]; + std::ofstream output_file(fileName, std::ios::binary); + uint32_t n1 = 105, n2 = 18090; + const int F_ARRAY_LEN = 5; + float fs[] = {1.024f, 2.39021f, 4.532f, 7.980232f, 6.222f}; + const int I_ARRAY_LEN = 7; + uint8_t arr[] = {12, 13, 14, 33, 222, 183, 99}; + float f1 = 9283.1237f; + + output_file.write((const char*) &n1, sizeof(uint32_t)); + output_file.write((const char*) &n2, sizeof(uint32_t)); + output_file.write((const char*) &f1, sizeof(float)); + output_file.write((const char*) fs, F_ARRAY_LEN * sizeof(float)); + output_file.write((const char*) arr, I_ARRAY_LEN * sizeof(uint8_t)); + + output_file.close(); + + std::ifstream input_file(fileName, std::ios::binary | std::ios::ate); + size_t size = input_file.tellg(); + input_file.seekg(0); + + auto data = new char[size]; + input_file.read(data, size); + + diskann::cout << "Read " << size << " bytes from file" << std::endl; + + uint32_t int1, int2; + float f2; + float* fs1 = new float[F_ARRAY_LEN]; + uint8_t* arr1 = new uint8_t[I_ARRAY_LEN]; + memset(fs1, 0, sizeof(float) * F_ARRAY_LEN); + memset(arr1, 0, sizeof(uint8_t) * I_ARRAY_LEN); + + // std::ifstream reader(argv[1], std::ios::binary); + // reader.read((char *)&int1, sizeof(int)); + // reader.read((char*) &int2, sizeof(int)); + // reader.read((char*) &f2, sizeof(float)); + // reader.read((char*) fs1, sizeof(float) * F_ARRAY_LEN); + + ContentBuf cb(data, size); + std::basic_istream reader(&cb); + + reader.read((char*) &int1, sizeof(uint32_t)); + reader.read((char*) &int2, sizeof(uint32_t)); + reader.read((char*) &f2, sizeof(float)); + reader.read((char*) fs1, F_ARRAY_LEN * sizeof(float)); + reader.read((char*) arr1, I_ARRAY_LEN * sizeof(uint8_t)); + + diskann::cout << int1 << "," << int2 << "," << f2 << std::endl; + for (int i = 0; i < F_ARRAY_LEN; i++) { + diskann::cout << fs1[i] << ","; + } + diskann::cout << std::endl; + for (int i = 0; i < I_ARRAY_LEN; i++) { + diskann::cout << std::to_string(arr1[i]) << ","; + } + diskann::cout << std::endl; + + assert(int1 == n1); + assert(int2 == n2); + assert(abs(f1 - f2) < 0.0001); +} + +void testSubstringImpl(int argc, char** argv) { + assert(argc >= 4); + + std::string s1(argv[2]); + std::string s2(argv[3]); + assert(s1[0] == s2[0]); // at least they share one char in common! + + size_t index = -1; + size_t compareLen = s1.length() <= s2.length() ? s1.length() : s2.length(); + for (size_t i = 0; i < compareLen; i++) { + if (s1[i] != s2[i]) { + index = i - 1; + break; + } + } + diskann::cout << "Common substring at 0 is:" << s1.substr(0, index); +} + +// DISKANN_DLLIMPORT std::basic_ostream diskann::cout; +// DISKANN_DLLIMPORT std::basic_ostream diskann::cerr; + +int main(int argc, char** argv) { + if (argc < 2) { + diskann::cout + << std::string("Usage: ") << argv[0] << " [arguments]" + << std::endl + << "Modes: 1 for file concat. Args " + << std::endl + << " 2 for test unique_ptr assignment. No args." << std::endl + << " 3 for comparing distance computations (AVX and " + "normal). No args." + << std::endl + << " 4 for testing our streambuf() implemntation. No args" + << std::endl + << " 5 for testing membuf implementation. Args " + << std::endl + << " 6 for testing substring implementation. Args " + << std::endl + << " 7 for testing OpenMP thread pool implementation. No args." + << std::endl; + return -1; + } + + int mode = atoi(argv[1]); + + std::cout << "Testing" << std::endl; + + switch (mode) { + case 1: + concat_files(argc, argv); + return 0; + break; + case 2: + uniquePtrAssignment(); + break; + case 3: + compareDistanceComputationsInt(); + compareDistanceComputationsFloat(); + break; + case 4: + testStreamBufImpl(); + break; + case 5: + testMemBufImpl(argc, argv); + break; + case 6: + testSubstringImpl(argc, argv); + break; + case 7: + testOpenMPThreadPool(); + break; + default: + diskann::cout << "Don't know what to do with mode parameter: " << argv[1] + << std::endl; + } +} diff --git a/tests/utils/test_partitioning.cpp b/tests/utils/test_partitioning.cpp new file mode 100644 index 0000000000..3fff7af268 --- /dev/null +++ b/tests/utils/test_partitioning.cpp @@ -0,0 +1,35 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#include +#include +#include "partition_and_pq.h" + +// DEPRECATED: NEED TO REPROGRAM + +int main(int argc, char** argv) { + auto s = std::chrono::high_resolution_clock::now(); + + if (argc != 8) { + diskann::cout << argv[0] + << " format: data type base_set train_set " + "num_clusters " + "max_reps prefix_for_working_directory k_base " + << std::endl; + exit(-1); + } + size_t num_clusters = std::atoi(argv[4]); + size_t max_reps = std::atoi(argv[5]); + size_t k_base = std::atoi(argv[7]); + if (std::string(argv[1]) == std::string("float")) + partition(argv[2], argv[3], num_clusters, max_reps, argv[6], k_base); + else if (std::string(argv[1]) == std::string("int8")) + partition(argv[2], argv[3], num_clusters, max_reps, argv[6], + k_base); + else if (std::string(argv[1]) == std::string("uint8")) + partition(argv[2], argv[3], num_clusters, max_reps, argv[6], + k_base); + else + diskann::cout << "unsupported data format. use float/int8/uint8" + << std::endl; +} diff --git a/tests/utils/test_pq.cpp b/tests/utils/test_pq.cpp new file mode 100644 index 0000000000..2369688450 --- /dev/null +++ b/tests/utils/test_pq.cpp @@ -0,0 +1,27 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#include +#include +#include "partition_and_pq.h" +#include "util.h" + +// DEPRECATED: NEED TO REPROGRAM + +int main(int argc, char** argv) { + if (argc != 6) { + diskann::cout << argv[0] + << " format: train_set " + "num_clusters_per_chunk number_chunks " + "max_reps prefix_for_working_directory " + << std::endl; + exit(-1); + } + + size_t num_centers = (size_t) strtol(argv[2], NULL, 10); + size_t num_chunks = (size_t) strtol(argv[3], NULL, 10); + size_t max_reps = (size_t) strtol(argv[4], NULL, 10); + + generate_pq_pivots(argv[1], num_centers, num_chunks, max_reps, argv[5]); + return 0; +} diff --git a/tests/utils/test_pq_from_pivots.cpp b/tests/utils/test_pq_from_pivots.cpp new file mode 100644 index 0000000000..12f98e7bc3 --- /dev/null +++ b/tests/utils/test_pq_from_pivots.cpp @@ -0,0 +1,25 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#include +#include + +#include "partition_and_pq.h" +#include "util.h" + +// DEPRECATED: NEED TO REPROGRAM + +int main(int argc, char** argv) { + if (argc != 5) { + diskann::cout << argv[0] + << "format: base_set " + "num_clusters_per_chunk number_chunks " + "prefix_for_working_file " + << std::endl; + exit(-1); + } + + size_t num_centers = (size_t) strtol(argv[2], NULL, 10); + size_t num_chunks = (size_t) strtol(argv[3], NULL, 10); + generate_pq_data_from_pivots(argv[1], num_centers, num_chunks, argv[4]); +} diff --git a/tests/utils/tsv_to_bin.cpp b/tests/utils/tsv_to_bin.cpp index 111a6bb55d..1151c0a177 100644 --- a/tests/utils/tsv_to_bin.cpp +++ b/tests/utils/tsv_to_bin.cpp @@ -4,13 +4,10 @@ #include #include "utils.h" -template -void block_convert(std::ifstream& reader, std::ofstream& writer, _u64 npts, - _u64 ndims) { - auto read_buf = new T[4 * npts * (ndims + 1)]; - - auto cursor = read_buf; - T val; +void block_convert(std::ifstream& reader, std::ofstream& writer, + float* read_buf, _u64 npts, _u64 ndims) { + auto cursor = read_buf; + float val; for (_u64 i = 0; i < npts; i++) { for (_u64 d = 0; d < ndims; ++d) { @@ -19,54 +16,42 @@ void block_convert(std::ifstream& reader, std::ofstream& writer, _u64 npts, cursor++; } } - writer.write((char*) read_buf, npts * ndims * sizeof(T)); - delete[] read_buf; + writer.write((char*) read_buf, npts * ndims * sizeof(float)); } int main(int argc, char** argv) { - if (argc != 6) { - std::cout << argv[0] - << " input_filename.tsv output_filename.bin dim num_pts>" + if (argc != 5) { + diskann::cout << argv[0] + << " input_filename.tsv output_filename.bin dim num_pts>" << std::endl; exit(-1); } - if (std::string(argv[1]) != std::string("float") && - std::string(argv[1]) != std::string("int8") && - std::string(argv[1]) != std::string("uint8")) { - std::cout << "Unsupported type. float, int8 and uint8 types are supported." - << std::endl; - } - - _u64 ndims = atoi(argv[4]); - _u64 npts = atoi(argv[5]); + _u64 ndims = atoi(argv[3]); + _u64 npts = atoi(argv[4]); - std::ifstream reader(argv[2], std::ios::binary | std::ios::ate); + std::ifstream reader(argv[1], std::ios::binary | std::ios::ate); // _u64 fsize = reader.tellg(); reader.seekg(0, std::ios::beg); reader.seekg(0, std::ios::beg); _u64 blk_size = 131072; _u64 nblks = ROUND_UP(npts, blk_size) / blk_size; - std::cout << "# blks: " << nblks << std::endl; - std::ofstream writer(argv[3], std::ios::binary); - auto npts_s32 = (_u32) npts; - auto ndims_s32 = (_u32) ndims; - writer.write((char*) &npts_s32, sizeof(_u32)); - writer.write((char*) &ndims_s32, sizeof(_u32)); - + diskann::cout << "# blks: " << nblks << std::endl; + std::ofstream writer(argv[2], std::ios::binary); + int npts_s32 = (_s32) npts; + int ndims_s32 = (_s32) ndims; + writer.write((char*) &npts_s32, sizeof(_s32)); + writer.write((char*) &ndims_s32, sizeof(_s32)); + float* read_buf = new float[npts * (ndims + 1)]; for (_u64 i = 0; i < nblks; i++) { _u64 cblk_size = std::min(npts - i * blk_size, blk_size); - if (std::string(argv[1]) == std::string("float")) { - block_convert(reader, writer, cblk_size, ndims); - } else if (std::string(argv[1]) == std::string("int8")) { - block_convert(reader, writer, cblk_size, ndims); - } else if (std::string(argv[1]) == std::string("uint8")) { - block_convert(reader, writer, cblk_size, ndims); - } - std::cout << "Block #" << i << " written" << std::endl; + block_convert(reader, writer, read_buf, cblk_size, ndims); + diskann::cout << "Block #" << i << " written" << std::endl; } + delete[] read_buf; + reader.close(); writer.close(); } diff --git a/tests/utils/uint32_to_uint8.cpp b/tests/utils/uint32_to_uint8.cpp index 7c131180fd..21ce241e77 100644 --- a/tests/utils/uint32_to_uint8.cpp +++ b/tests/utils/uint32_to_uint8.cpp @@ -6,7 +6,8 @@ int main(int argc, char** argv) { if (argc != 3) { - std::cout << argv[0] << " input_uint32_bin output_int8_bin" << std::endl; + diskann::cout << argv[0] << " input_uint32_bin output_int8_bin" + << std::endl; exit(-1); } diff --git a/tests/utils/uint8_to_float.cpp b/tests/utils/uint8_to_float.cpp new file mode 100644 index 0000000000..04ab57864a --- /dev/null +++ b/tests/utils/uint8_to_float.cpp @@ -0,0 +1,22 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#include +#include "utils.h" + +int main(int argc, char** argv) { + if (argc != 3) { + diskann::cout << argv[0] << " input_uint8_bin output_float_bin" + << std::endl; + exit(-1); + } + + uint8_t* input; + size_t npts, nd; + diskann::load_bin(argv[1], input, npts, nd); + float* output = new float[npts * nd]; + diskann::convert_types(input, output, npts, nd); + diskann::save_bin(argv[2], output, npts, nd); + delete[] output; + delete[] input; +} diff --git a/tests/utils/update_metadata.cpp b/tests/utils/update_metadata.cpp new file mode 100644 index 0000000000..d57a4a14ba --- /dev/null +++ b/tests/utils/update_metadata.cpp @@ -0,0 +1,56 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#include +#include "utils.h" + +void block_convert(std::ifstream& reader, std::ofstream& writer, + float* read_buf, uint8_t* write_buf, _u64 npts, _u64 ndims) { + reader.read((char*) read_buf, + npts * (ndims * sizeof(float) + sizeof(unsigned))); + for (_u64 i = 0; i < npts; i++) { + memcpy(write_buf + i * ndims, (read_buf + i * (ndims + 1)) + 1, + ndims * sizeof(float)); + } + writer.write((char*) write_buf, npts * ndims * sizeof(float)); +} + +int main(int argc, char** argv) { + if (argc != 3) { + diskann::cout << argv[0] << " input_fvecs output_bin" << std::endl; + exit(-1); + } + std::ifstream reader(argv[1], std::ios::binary | std::ios::ate); + _u64 fsize = reader.tellg(); + reader.seekg(0, std::ios::beg); + + unsigned ndims_u32; + reader.read((char*) &ndims_u32, sizeof(unsigned)); + reader.seekg(0, std::ios::beg); + _u64 ndims = (_u64) ndims_u32; + _u64 npts = fsize / ((ndims + 1) * sizeof(float)); + diskann::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims + << std::endl; + + _u64 blk_size = 131072; + _u64 nblks = ROUND_UP(npts, blk_size) / blk_size; + diskann::cout << "# blks: " << nblks << std::endl; + std::ofstream writer(argv[2], std::ios::binary); + int npts_s32 = (_s32) npts; + int ndims_s32 = (_s32) ndims; + writer.write((char*) &npts_s32, sizeof(_s32)); + writer.write((char*) &ndims_s32, sizeof(_s32)); + float* read_buf = new float[npts * (ndims + 1)]; + uint8_t* write_buf = new uint8_t[npts * ndims]; + for (_u64 i = 0; i < nblks; i++) { + _u64 cblk_size = std::min(npts - i * blk_size, blk_size); + block_convert(reader, writer, read_buf, write_buf, cblk_size, ndims); + diskann::cout << "Block #" << i << " written" << std::endl; + } + + delete[] read_buf; + delete[] write_buf; + + reader.close(); + writer.close(); +}