diff --git a/CMakeLists.txt b/CMakeLists.txt index 3d3d2b860..a2fe163da 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -55,6 +55,8 @@ include_directories(${PROJECT_SOURCE_DIR}/include) if(NOT PYBIND) set(DISKANN_RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS ON) +elseif(MSVS) + set(DISKANN_RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS OFF) endif() # It's necessary to include tcmalloc headers only if calling into MallocExtension interface. # For using tcmalloc in DiskANN tools, it's enough to just link with tcmalloc. diff --git a/apps/build_memory_index.cpp b/apps/build_memory_index.cpp index 544e42dee..feb5dcecc 100644 --- a/apps/build_memory_index.cpp +++ b/apps/build_memory_index.cpp @@ -24,7 +24,7 @@ namespace po = boost::program_options; int main(int argc, char **argv) { - std::string data_type, dist_fn, data_path, index_path_prefix, label_file, universal_label, label_type; + std::string data_type, dist_fn, data_path, index_path_prefix, label_file, codebook_path, universal_label, label_type; uint32_t num_threads, R, L, Lf, build_PQ_bytes; float alpha; bool use_pq_build, use_opq; @@ -59,13 +59,14 @@ int main(int argc, char **argv) program_options_utils::GRAPH_BUILD_ALPHA); optional_configs.add_options()("build_PQ_bytes", po::value(&build_PQ_bytes)->default_value(0), program_options_utils::BUIlD_GRAPH_PQ_BYTES); + optional_configs.add_options()("codebook_path", po::value(&codebook_path)->default_value(""), + program_options_utils::CODEBOOK_PATH); optional_configs.add_options()("use_opq", po::bool_switch()->default_value(false), program_options_utils::USE_OPQ); optional_configs.add_options()("label_file", po::value(&label_file)->default_value(""), program_options_utils::LABEL_FILE); optional_configs.add_options()("universal_label", po::value(&universal_label)->default_value(""), program_options_utils::UNIVERSAL_LABEL); - optional_configs.add_options()("FilteredLbuild", po::value(&Lf)->default_value(0), program_options_utils::FILTERED_LBUILD); optional_configs.add_options()("label_type", po::value(&label_type)->default_value("uint"), @@ -146,6 +147,7 @@ int main(int argc, char **argv) .is_use_opq(use_opq) .is_pq_dist_build(use_pq_build) .with_num_pq_chunks(build_PQ_bytes) + .with_pq_codebook_path(codebook_path) .build(); auto index_factory = diskann::IndexFactory(config); diff --git a/apps/search_memory_index.cpp b/apps/search_memory_index.cpp index 1a9acc285..8f387d272 100644 --- a/apps/search_memory_index.cpp +++ b/apps/search_memory_index.cpp @@ -27,7 +27,9 @@ namespace po = boost::program_options; template int search_memory_index(diskann::Metric &metric, const std::string &index_path, const std::string &result_path_prefix, - const std::string &query_file, const std::string &truthset_file, const uint32_t num_threads, + const std::string &query_file, const std::string &truthset_file, + const std::string &codebook_file, const bool use_pq_build, const bool use_opq, + const uint32_t pq_num_chunks, const uint32_t num_threads, const uint32_t recall_at, const bool print_all_recalls, const std::vector &Lvec, const bool dynamic, const bool tags, const bool show_qps_per_thread, const std::vector &query_filters, const float fail_if_recall_below) @@ -82,12 +84,16 @@ int search_memory_index(diskann::Metric &metric, const std::string &index_path, .is_dynamic_index(dynamic) .is_enable_tags(tags) .is_concurrent_consolidate(false) - .is_pq_dist_build(false) - .is_use_opq(false) - .with_num_pq_chunks(0) + .is_pq_dist_build(use_pq_build) + .is_use_opq(use_opq) + .with_num_pq_chunks(pq_num_chunks) .with_num_frozen_pts(num_frozen_pts) + .with_pq_codebook_path(codebook_file) .build(); + std::cout << "******************** Attach Debugger ********************" << std::endl; + Sleep(60000); + auto index_factory = diskann::IndexFactory(config); auto index = index_factory.create_instance(); index->load(index_path.c_str(), num_threads, *(std::max_element(Lvec.begin(), Lvec.end()))); @@ -278,10 +284,10 @@ int search_memory_index(diskann::Metric &metric, const std::string &index_path, int main(int argc, char **argv) { std::string data_type, dist_fn, index_path_prefix, result_path, query_file, gt_file, filter_label, label_type, - query_filters_file; - uint32_t num_threads, K; + query_filters_file, codebook_path; + uint32_t num_threads, K, build_PQ_bytes; std::vector Lvec; - bool print_all_recalls, dynamic, tags, show_qps_per_thread; + bool print_all_recalls, dynamic, tags, show_qps_per_thread, use_pq_build, use_opq; float fail_if_recall_below = 0.0f; po::options_description desc{ @@ -331,6 +337,12 @@ int main(int argc, char **argv) optional_configs.add_options()("fail_if_recall_below", po::value(&fail_if_recall_below)->default_value(0.0f), program_options_utils::FAIL_IF_RECALL_BELOW); + optional_configs.add_options()("build_PQ_bytes", po::value(&build_PQ_bytes)->default_value(0), + program_options_utils::BUIlD_GRAPH_PQ_BYTES); + optional_configs.add_options()("codebook_path", po::value(&codebook_path)->default_value(""), + program_options_utils::CODEBOOK_PATH); + optional_configs.add_options()("use_opq", po::bool_switch()->default_value(false), + program_options_utils::USE_OPQ); // Output controls po::options_description output_controls("Output controls"); @@ -352,6 +364,8 @@ int main(int argc, char **argv) return 0; } po::notify(vm); + use_pq_build = (build_PQ_bytes > 0); + use_opq = vm["use_opq"].as(); } catch (const std::exception &ex) { @@ -420,18 +434,21 @@ int main(int argc, char **argv) if (data_type == std::string("int8")) { return search_memory_index( - metric, index_path_prefix, result_path, query_file, gt_file, num_threads, K, print_all_recalls, + metric, index_path_prefix, result_path, query_file, gt_file, codebook_path, use_pq_build, use_opq, + build_PQ_bytes, num_threads, K, print_all_recalls, Lvec, dynamic, tags, show_qps_per_thread, query_filters, fail_if_recall_below); } else if (data_type == std::string("uint8")) { return search_memory_index( - metric, index_path_prefix, result_path, query_file, gt_file, num_threads, K, print_all_recalls, + metric, index_path_prefix, result_path, query_file, gt_file, codebook_path, use_pq_build, use_opq, + build_PQ_bytes, num_threads, K, print_all_recalls, Lvec, dynamic, tags, show_qps_per_thread, query_filters, fail_if_recall_below); } else if (data_type == std::string("float")) { return search_memory_index(metric, index_path_prefix, result_path, query_file, gt_file, + codebook_path, use_pq_build, use_opq, build_PQ_bytes, num_threads, K, print_all_recalls, Lvec, dynamic, tags, show_qps_per_thread, query_filters, fail_if_recall_below); } @@ -446,18 +463,21 @@ int main(int argc, char **argv) if (data_type == std::string("int8")) { return search_memory_index(metric, index_path_prefix, result_path, query_file, gt_file, + codebook_path, use_pq_build, use_opq, build_PQ_bytes, num_threads, K, print_all_recalls, Lvec, dynamic, tags, show_qps_per_thread, query_filters, fail_if_recall_below); } else if (data_type == std::string("uint8")) { return search_memory_index(metric, index_path_prefix, result_path, query_file, gt_file, + codebook_path, use_pq_build, use_opq, build_PQ_bytes, num_threads, K, print_all_recalls, Lvec, dynamic, tags, show_qps_per_thread, query_filters, fail_if_recall_below); } else if (data_type == std::string("float")) { return search_memory_index(metric, index_path_prefix, result_path, query_file, gt_file, + codebook_path, use_pq_build, use_opq, build_PQ_bytes, num_threads, K, print_all_recalls, Lvec, dynamic, tags, show_qps_per_thread, query_filters, fail_if_recall_below); } diff --git a/apps/test_insert_deletes_consolidate.cpp b/apps/test_insert_deletes_consolidate.cpp index 97aed1864..505858e53 100644 --- a/apps/test_insert_deletes_consolidate.cpp +++ b/apps/test_insert_deletes_consolidate.cpp @@ -150,7 +150,7 @@ void build_incremental_index(const std::string &data_path, diskann::IndexWritePa uint32_t num_start_pts, size_t points_per_checkpoint, size_t checkpoints_per_snapshot, const std::string &save_path, size_t points_to_delete_from_beginning, size_t start_deletes_after, bool concurrent, const std::string &label_file, - const std::string &universal_label) + const std::string &universal_label, size_t num_pq_chunks, const std::string& pq_pivot_file) { size_t dim, aligned_dim; size_t num_points; @@ -161,7 +161,7 @@ void build_incremental_index(const std::string &data_path, diskann::IndexWritePa using LabelT = uint32_t; size_t current_point_offset = points_to_skip; - const size_t last_point_threshold = points_to_skip + max_points_to_insert; + size_t last_point_threshold = points_to_skip + max_points_to_insert; bool enable_tags = true; using TagT = uint32_t; @@ -182,6 +182,9 @@ void build_incremental_index(const std::string &data_path, diskann::IndexWritePa .is_filtered(has_labels) .with_num_frozen_pts(num_start_pts) .is_concurrent_consolidate(concurrent) + .with_pq_codebook_path(pq_pivot_file) + .is_pq_dist_build(!pq_pivot_file.empty()) + .with_num_pq_chunks(num_pq_chunks) .build(); diskann::IndexFactory index_factory = diskann::IndexFactory(index_config); @@ -206,6 +209,7 @@ void build_incremental_index(const std::string &data_path, diskann::IndexWritePa if (points_to_skip + max_points_to_insert > num_points) { max_points_to_insert = num_points - points_to_skip; + last_point_threshold = num_points; std::cerr << "WARNING: Reducing max_points_to_insert to " << max_points_to_insert << " points since the data file has only that many" << std::endl; } @@ -377,11 +381,11 @@ int main(int argc, char **argv) uint32_t num_threads, R, L, num_start_pts; float alpha, start_point_norm; size_t points_to_skip, max_points_to_insert, beginning_index_size, points_per_checkpoint, checkpoints_per_snapshot, - points_to_delete_from_beginning, start_deletes_after; + points_to_delete_from_beginning, start_deletes_after, num_pq_chunks; bool concurrent; // label options - std::string label_file, label_type, universal_label; + std::string label_file, label_type, universal_label, pq_pivot_file; std::uint32_t Lf, unique_labels_supported; po::options_description desc{program_options_utils::make_program_description("test_insert_deletes_consolidate", @@ -449,6 +453,11 @@ int main(int argc, char **argv) optional_configs.add_options()("unique_labels_supported", po::value(&unique_labels_supported)->default_value(0), "Number of unique labels supported by the dynamic index."); + optional_configs.add_options()("pq_pivot_file", po::value(&pq_pivot_file)->default_value(""), + "The file stored pq pivot info."); + optional_configs.add_options()("num_pq_chunks", po::value(&num_pq_chunks)->default_value(0), + "Number of PQ chunks to use."); + optional_configs.add_options()( "num_start_points", @@ -503,21 +512,27 @@ int main(int argc, char **argv) .with_filter_list_size(Lf) .build(); + + std::cout << "********** Attach Debugger Test insert delete consolidate **********" << std::endl; + if (data_type == std::string("int8")) build_incremental_index( data_path, params, points_to_skip, max_points_to_insert, beginning_index_size, start_point_norm, num_start_pts, points_per_checkpoint, checkpoints_per_snapshot, index_path_prefix, - points_to_delete_from_beginning, start_deletes_after, concurrent, label_file, universal_label); + points_to_delete_from_beginning, start_deletes_after, concurrent, + label_file, universal_label, num_pq_chunks, pq_pivot_file); else if (data_type == std::string("uint8")) build_incremental_index( data_path, params, points_to_skip, max_points_to_insert, beginning_index_size, start_point_norm, num_start_pts, points_per_checkpoint, checkpoints_per_snapshot, index_path_prefix, - points_to_delete_from_beginning, start_deletes_after, concurrent, label_file, universal_label); + points_to_delete_from_beginning, start_deletes_after, concurrent, + label_file, universal_label, num_pq_chunks, pq_pivot_file); else if (data_type == std::string("float")) build_incremental_index(data_path, params, points_to_skip, max_points_to_insert, beginning_index_size, start_point_norm, num_start_pts, points_per_checkpoint, checkpoints_per_snapshot, index_path_prefix, points_to_delete_from_beginning, - start_deletes_after, concurrent, label_file, universal_label); + start_deletes_after, concurrent, label_file, universal_label, + num_pq_chunks, pq_pivot_file); else std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; } diff --git a/include/fixed_chunk_pq_table.h b/include/fixed_chunk_pq_table.h new file mode 100644 index 000000000..431dc2577 --- /dev/null +++ b/include/fixed_chunk_pq_table.h @@ -0,0 +1,47 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#pragma once + +#include "common_includes.h" + +#ifdef EXEC_ENV_OLS +#include "content_buf.h" +#include "memory_mapped_files.h" +#endif + +namespace diskann +{ + class FixedChunkPQTable + { + public: + FixedChunkPQTable(); + virtual ~FixedChunkPQTable(); + + #ifdef EXEC_ENV_OLS + void load_pq_centroid_bin(MemoryMappedFiles &files, const char *pq_table_file, size_t num_chunks); + #else + void load_pq_centroid_bin(const char *pq_table_file, size_t num_chunks); + #endif + + void preprocess_query(float *query_vec); + // assumes pre-processed query + void populate_chunk_distances(const float *query_vec, float *dist_vec); + float l2_distance(const float *query_vec, uint8_t *base_vec); + float inner_product(const float *query_vec, uint8_t *base_vec); + // assumes no rotation is involved + template + void inflate_vector(InputType *base_vec, OutputType *out_vec) const; + + void populate_chunk_inner_products(const float *query_vec, float *dist_vec); + + float *tables = nullptr; // pq_tables = float array of size [256 * ndims] + uint64_t ndims = 0; // ndims = true dimension of vectors + uint64_t n_chunks = 0; + bool use_rotation = false; + uint32_t *chunk_offsets = nullptr; + float *centroid = nullptr; + float *tables_tr = nullptr; // same as pq_tables, but col-major + float *rotmat_tr = nullptr; + }; +} // namespace diskann diff --git a/include/index.h b/include/index.h index b9bf4f384..661354ad5 100644 --- a/include/index.h +++ b/include/index.h @@ -55,18 +55,37 @@ template clas public: // Constructor for Bulk operations and for creating the index object solely // for loading a prexisting index. - DISKANN_DLLEXPORT Index(const IndexConfig &index_config, std::shared_ptr> data_store, +#ifdef EXEC_ENV_OLS + DISKANN_DLLEXPORT Index(const IndexConfig &index_config, MemoryMappedFiles &files, + std::shared_ptr> data_store, std::unique_ptr graph_store, std::shared_ptr> pq_data_store = nullptr); +#else + DISKANN_DLLEXPORT Index(const IndexConfig &index_config, std::shared_ptr> data_store, + std::unique_ptr graph_store, + std::shared_ptr> pq_data_store = nullptr); +#endif // Constructor for incremental index +#ifdef EXEC_ENV_OLS + DISKANN_DLLEXPORT Index(Metric m, const size_t dim, const size_t max_points, + const std::shared_ptr index_parameters, + const std::shared_ptr index_search_params, MemoryMappedFiles &files, + const std::string &codebook_path = "", const size_t num_frozen_pts = 0, + const bool dynamic_index = false, const bool enable_tags = false, + const bool concurrent_consolidate = false, const bool pq_dist_build = false, + const size_t num_pq_chunks = 0, const bool use_opq = false, + const bool filtered_index = false); +#else DISKANN_DLLEXPORT Index(Metric m, const size_t dim, const size_t max_points, const std::shared_ptr index_parameters, const std::shared_ptr index_search_params, - const size_t num_frozen_pts = 0, const bool dynamic_index = false, - const bool enable_tags = false, const bool concurrent_consolidate = false, - const bool pq_dist_build = false, const size_t num_pq_chunks = 0, - const bool use_opq = false, const bool filtered_index = false); + const std::string &codebook_path = "", const size_t num_frozen_pts = 0, + const bool dynamic_index = false, const bool enable_tags = false, + const bool concurrent_consolidate = false, const bool pq_dist_build = false, + const size_t num_pq_chunks = 0, const bool use_opq = false, + const bool filtered_index = false); +#endif DISKANN_DLLEXPORT ~Index(); @@ -186,6 +205,9 @@ template clas // memory should be allocated for vec before calling this function DISKANN_DLLEXPORT int get_vector_by_tag(TagT &tag, T *vec); + // memory should be allocated for vec before calling this function + DISKANN_DLLEXPORT int get_pq_vector_by_tag(TagT &tag, T *vec); + DISKANN_DLLEXPORT void print_status(); DISKANN_DLLEXPORT void count_nodes_at_bfs_levels(); @@ -403,12 +425,7 @@ template clas bool _pq_dist = false; bool _use_opq = false; size_t _num_pq_chunks = 0; - // REFACTOR - // uint8_t *_pq_data = nullptr; - std::shared_ptr> _pq_distance_fn = nullptr; std::shared_ptr> _pq_data_store = nullptr; - bool _pq_generated = false; - FixedChunkPQTable _pq_table; // // Data structures, locks and flags for dynamic indexing and tags diff --git a/include/index_config.h b/include/index_config.h index 452498b01..4078463cd 100644 --- a/include/index_config.h +++ b/include/index_config.h @@ -35,6 +35,7 @@ struct IndexConfig std::string label_type; std::string tag_type; std::string data_type; + std::string pq_codebook_path; // Params for building index std::shared_ptr index_write_params; @@ -46,13 +47,15 @@ struct IndexConfig size_t max_points, size_t num_pq_chunks, size_t num_frozen_points, bool dynamic_index, bool enable_tags, bool pq_dist_build, bool concurrent_consolidate, bool use_opq, bool filtered_index, std::string &data_type, const std::string &tag_type, const std::string &label_type, + const std::string &pq_codebook_path, std::shared_ptr index_write_params, std::shared_ptr index_search_params) : data_strategy(data_strategy), graph_strategy(graph_strategy), metric(metric), dimension(dimension), max_points(max_points), dynamic_index(dynamic_index), enable_tags(enable_tags), pq_dist_build(pq_dist_build), concurrent_consolidate(concurrent_consolidate), use_opq(use_opq), filtered_index(filtered_index), num_pq_chunks(num_pq_chunks), num_frozen_pts(num_frozen_points), label_type(label_type), tag_type(tag_type), - data_type(data_type), index_write_params(index_write_params), index_search_params(index_search_params) + data_type(data_type), pq_codebook_path(pq_codebook_path), index_write_params(index_write_params), + index_search_params(index_search_params) { } @@ -160,6 +163,11 @@ class IndexConfigBuilder return *this; } + IndexConfigBuilder &with_pq_codebook_path(const std::string &pq_codebook_path) + { + this->_pq_codebook_path = pq_codebook_path; + return *this; + } IndexConfigBuilder &with_index_write_params(IndexWriteParameters &index_write_params) { this->_index_write_params = std::make_shared(index_write_params); @@ -219,8 +227,8 @@ class IndexConfigBuilder return IndexConfig(_data_strategy, _graph_strategy, _metric, _dimension, _max_points, _num_pq_chunks, _num_frozen_pts, _dynamic_index, _enable_tags, _pq_dist_build, _concurrent_consolidate, - _use_opq, _filtered_index, _data_type, _tag_type, _label_type, _index_write_params, - _index_search_params); + _use_opq, _filtered_index, _data_type, _tag_type, _label_type, _pq_codebook_path, + _index_write_params, _index_search_params); } IndexConfigBuilder(const IndexConfigBuilder &) = delete; @@ -247,6 +255,8 @@ class IndexConfigBuilder std::string _label_type{"uint32"}; std::string _tag_type{"uint32"}; std::string _data_type; + std::string _pq_codebook_path; + std::shared_ptr _index_write_params; std::shared_ptr _index_search_params; diff --git a/include/index_factory.h b/include/index_factory.h index 80bc40dba..ba2597e9d 100644 --- a/include/index_factory.h +++ b/include/index_factory.h @@ -23,6 +23,7 @@ class IndexFactory // flavours. template DISKANN_DLLEXPORT static std::shared_ptr> construct_pq_datastore(DataStoreStrategy strategy, + const std::string &codebook_path, size_t num_points, size_t dimension, Metric m, size_t num_pq_chunks, bool use_opq); diff --git a/include/pq.h b/include/pq.h index 3e6119f22..8b82836dc 100644 --- a/include/pq.h +++ b/include/pq.h @@ -8,44 +8,6 @@ namespace diskann { -class FixedChunkPQTable -{ - float *tables = nullptr; // pq_tables = float array of size [256 * ndims] - uint64_t ndims = 0; // ndims = true dimension of vectors - uint64_t n_chunks = 0; - bool use_rotation = false; - uint32_t *chunk_offsets = nullptr; - float *centroid = nullptr; - float *tables_tr = nullptr; // same as pq_tables, but col-major - float *rotmat_tr = nullptr; - - public: - FixedChunkPQTable(); - - virtual ~FixedChunkPQTable(); - -#ifdef EXEC_ENV_OLS - void load_pq_centroid_bin(MemoryMappedFiles &files, const char *pq_table_file, size_t num_chunks); -#else - void load_pq_centroid_bin(const char *pq_table_file, size_t num_chunks); -#endif - - uint32_t get_num_chunks(); - - void preprocess_query(float *query_vec); - - // assumes pre-processed query - void populate_chunk_distances(const float *query_vec, float *dist_vec); - - float l2_distance(const float *query_vec, uint8_t *base_vec); - - float inner_product(const float *query_vec, uint8_t *base_vec); - - // assumes no rotation is involved - void inflate_vector(uint8_t *base_vec, float *out_vec); - - void populate_chunk_inner_products(const float *query_vec, float *dist_vec); -}; void aggregate_coords(const std::vector &ids, const uint8_t *all_coords, const uint64_t ndims, uint8_t *out); diff --git a/include/pq_common.h b/include/pq_common.h index c6a3a5739..af395438a 100644 --- a/include/pq_common.h +++ b/include/pq_common.h @@ -3,6 +3,8 @@ #include #include +#include "ann_exception.h" + #define NUM_PQ_BITS 8 #define NUM_PQ_CENTROIDS (1 << NUM_PQ_BITS) #define MAX_OPQ_ITERS 20 @@ -14,11 +16,23 @@ namespace diskann { inline std::string get_quantized_vectors_filename(const std::string &prefix, bool use_opq, uint32_t num_chunks) { + if (num_chunks == 0) + { + throw ANNException("Must set num_chunks before calling get_quantized_vectors_filename", -1, + __FUNCSIG__, __FILE__, __LINE__); + } + return prefix + (use_opq ? "_opq" : "pq") + std::to_string(num_chunks) + "_compressed.bin"; } inline std::string get_pivot_data_filename(const std::string &prefix, bool use_opq, uint32_t num_chunks) { + if (num_chunks == 0) + { + throw ANNException("Must set num_chunks before calling get_pivot_data_filename", -1, __FUNCSIG__, __FILE__, + __LINE__); + } + return prefix + (use_opq ? "_opq" : "pq") + std::to_string(num_chunks) + "_pivots.bin"; } diff --git a/include/pq_data_store.h b/include/pq_data_store.h index 7c0cb5fe0..e708cc6e7 100644 --- a/include/pq_data_store.h +++ b/include/pq_data_store.h @@ -2,7 +2,7 @@ #include #include "distance.h" #include "quantized_distance.h" -#include "pq.h" +#include "fixed_chunk_pq_table.h" #include "abstract_data_store.h" namespace diskann @@ -13,8 +13,14 @@ template class PQDataStore : public AbstractDataStore { public: +#ifdef EXEC_ENV_OLS + PQDataStore(size_t dim, location_t num_points, size_t num_pq_chunks, std::unique_ptr> distance_fn, + std::unique_ptr> pq_distance_fn, MemoryMappedFiles &files, + const std::string &codebook_path); +#else PQDataStore(size_t dim, location_t num_points, size_t num_pq_chunks, std::unique_ptr> distance_fn, - std::unique_ptr> pq_distance_fn); + std::unique_ptr> pq_distance_fn, const std::string &codebook_path); +#endif PQDataStore(const PQDataStore &) = delete; PQDataStore &operator=(const PQDataStore &) = delete; ~PQDataStore(); @@ -85,6 +91,7 @@ template class PQDataStore : public AbstractDataStore private: uint8_t *_quantized_data = nullptr; size_t _num_chunks = 0; + size_t _aligned_dim; // REFACTOR TODO: Doing this temporarily before refactoring OPQ into // its own class. Remove later. @@ -92,6 +99,7 @@ template class PQDataStore : public AbstractDataStore Metric _distance_metric; std::unique_ptr> _distance_fn = nullptr; + std::string _pq_pivot_file_path; std::unique_ptr> _pq_distance_fn = nullptr; }; } // namespace diskann diff --git a/include/pq_flash_index.h b/include/pq_flash_index.h index ba5258e18..7b42670b0 100644 --- a/include/pq_flash_index.h +++ b/include/pq_flash_index.h @@ -10,6 +10,7 @@ #include "parameters.h" #include "percentile_stats.h" #include "pq.h" +#include "fixed_chunk_pq_table.h" #include "utils.h" #include "windows_customizations.h" #include "scratch.h" diff --git a/include/pq_l2_distance.h b/include/pq_l2_distance.h index e6fc6e41b..4208838d0 100644 --- a/include/pq_l2_distance.h +++ b/include/pq_l2_distance.h @@ -1,4 +1,5 @@ #pragma once + #include "quantized_distance.h" namespace diskann @@ -23,21 +24,18 @@ template class PQL2Distance : public QuantizedDistance virtual bool is_opq() const override; - virtual std::string get_quantized_vectors_filename(const std::string &prefix) const override; - virtual std::string get_pivot_data_filename(const std::string &prefix) const override; - virtual std::string get_rotation_matrix_suffix(const std::string &pq_pivots_filename) const override; - #ifdef EXEC_ENV_OLS - virtual void load_pivot_data(MemoryMappedFiles &files, const std::string &pq_table_file, - size_t num_chunks) override; + virtual void load_pivot_data(MemoryMappedFiles &files, const std::string &pq_table_file) override; #else - virtual void load_pivot_data(const std::string &pq_table_file, size_t num_chunks) override; + virtual void load_pivot_data(const std::string &pq_table_file) override; #endif // Number of chunks in the PQ table. Depends on the compression level used. // Has to be < ndim virtual uint32_t get_num_chunks() const override; + virtual const FixedChunkPQTable &get_pq_table() const override; + // Preprocess the query by computing chunk distances from the query vector to // various centroids. Since we don't want this class to do scratch management, // we will take a PQScratch object which can come either from Index class or @@ -71,17 +69,8 @@ template class PQL2Distance : public QuantizedDistance protected: // assumes pre-processed query virtual void prepopulate_chunkwise_distances(const float *query_vec, float *dist_vec); - - // assumes no rotation is involved - // virtual void inflate_vector(uint8_t *base_vec, float *out_vec); - - float *_tables = nullptr; // pq_tables = float array of size [256 * ndims] - uint64_t _ndims = 0; // ndims = true dimension of vectors + FixedChunkPQTable _pq_table; uint64_t _num_chunks = 0; bool _is_opq = false; - uint32_t *_chunk_offsets = nullptr; - float *_centroid = nullptr; - float *_tables_tr = nullptr; // same as pq_tables, but col-major - float *_rotmat_tr = nullptr; }; } // namespace diskann diff --git a/include/program_options_utils.hpp b/include/program_options_utils.hpp index 2be60595b..b15ad916a 100644 --- a/include/program_options_utils.hpp +++ b/include/program_options_utils.hpp @@ -70,6 +70,7 @@ const char *BUIlD_GRAPH_PQ_BYTES = "Number of PQ bytes to build the index; 0 for const char *USE_OPQ = "Use Optimized Product Quantization (OPQ)."; const char *LABEL_FILE = "Input label file in txt format for Filtered Index build. The file should contain comma " "separated filters for each node with each line corresponding to a graph node"; +const char *CODEBOOK_PATH = "Path for Codebook/piviot file to use when building PQ"; const char *UNIVERSAL_LABEL = "Universal label, Use only in conjunction with label file for filtered index build. If a " "graph node has all the labels against it, we can assign a special universal filter to the " diff --git a/include/quantized_distance.h b/include/quantized_distance.h index cc4aea929..d2306c263 100644 --- a/include/quantized_distance.h +++ b/include/quantized_distance.h @@ -1,8 +1,14 @@ #pragma once -#include #include #include + #include "abstract_scratch.h" +#include "fixed_chunk_pq_table.h" + +#ifdef EXEC_ENV_OLS +#include "content_buf.h" +#include "memory_mapped_files.h" +#endif namespace diskann { @@ -17,23 +23,23 @@ template class QuantizedDistance virtual ~QuantizedDistance() = default; virtual bool is_opq() const = 0; - virtual std::string get_quantized_vectors_filename(const std::string &prefix) const = 0; - virtual std::string get_pivot_data_filename(const std::string &prefix) const = 0; - virtual std::string get_rotation_matrix_suffix(const std::string &pq_pivots_filename) const = 0; // Loading the PQ centroid table need not be part of the abstract class. // However, we want to indicate that this function will change once we have a // file reader hierarchy, so leave it here as-is. #ifdef EXEC_ENV_OLS - virtual void load_pivot_data(MemoryMappedFiles &files, const std::String &pq_table_file, size_t num_chunks) = 0; + virtual void load_pivot_data(MemoryMappedFiles &files, const std::String &pq_table_file) = 0; #else - virtual void load_pivot_data(const std::string &pq_table_file, size_t num_chunks) = 0; + virtual void load_pivot_data(const std::string &pq_table_file) = 0; #endif // Number of chunks in the PQ table. Depends on the compression level used. // Has to be < ndim virtual uint32_t get_num_chunks() const = 0; + // Return the pq_table used for quantized distance calculation. + virtual const FixedChunkPQTable &get_pq_table() const = 0; + // Preprocess the query by computing chunk distances from the query vector to // various centroids. Since we don't want this class to do scratch management, // we will take a PQScratch object which can come either from Index class or diff --git a/include/utils.h b/include/utils.h index d3af5c3a9..45747d975 100644 --- a/include/utils.h +++ b/include/utils.h @@ -834,7 +834,7 @@ void convert_types(const InType *srcmat, OutType *destmat, size_t npts, size_t d { for (uint64_t j = 0; j < dim; j++) { - destmat[i * dim + j] = (OutType)srcmat[i * dim + j]; + destmat[i * dim + j] = static_cast(srcmat[i * dim + j]); } } } diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index cbca26440..e3bc82b74 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -13,7 +13,7 @@ else() linux_aligned_file_reader.cpp math_utils.cpp natural_number_map.cpp in_mem_data_store.cpp in_mem_graph_store.cpp natural_number_set.cpp memory_mapper.cpp partition.cpp pq.cpp - pq_flash_index.cpp scratch.cpp logger.cpp utils.cpp filter_utils.cpp index_factory.cpp abstract_index.cpp pq_l2_distance.cpp pq_data_store.cpp) + pq_flash_index.cpp scratch.cpp logger.cpp utils.cpp filter_utils.cpp index_factory.cpp abstract_index.cpp pq_l2_distance.cpp pq_data_store.cpp fixed_chunk_pq_table.cpp) if (RESTAPI) list(APPEND CPP_SOURCES restapi/search_wrapper.cpp restapi/server.cpp) endif() diff --git a/src/disk_utils.cpp b/src/disk_utils.cpp index 016560217..b35a53d42 100644 --- a/src/disk_utils.cpp +++ b/src/disk_utils.cpp @@ -650,7 +650,7 @@ int build_merged_vamana_index(std::string base_file, diskann::Metric compareMetr .build(); using TagT = uint32_t; diskann::Index _index(compareMetric, base_dim, base_num, - std::make_shared(paras), nullptr, + std::make_shared(paras), nullptr, "", defaults::NUM_FROZEN_POINTS_STATIC, false, false, false, build_pq_bytes > 0, build_pq_bytes, use_opq, use_filters); if (!use_filters) @@ -721,7 +721,7 @@ int build_merged_vamana_index(std::string base_file, diskann::Metric compareMetr get_bin_metadata(shard_base_file, shard_base_pts, shard_base_dim); diskann::Index _index(compareMetric, shard_base_dim, shard_base_pts, - std::make_shared(low_degree_params), nullptr, + std::make_shared(low_degree_params), nullptr, "", defaults::NUM_FROZEN_POINTS_STATIC, false, false, false, build_pq_bytes > 0, build_pq_bytes, use_opq); if (!use_filters) diff --git a/src/dll/CMakeLists.txt b/src/dll/CMakeLists.txt index 096d1b76e..e2890e696 100644 --- a/src/dll/CMakeLists.txt +++ b/src/dll/CMakeLists.txt @@ -1,7 +1,7 @@ #Copyright(c) Microsoft Corporation.All rights reserved. #Licensed under the MIT license. -add_library(${PROJECT_NAME} SHARED dllmain.cpp ../abstract_data_store.cpp ../partition.cpp ../pq.cpp ../pq_flash_index.cpp ../logger.cpp ../utils.cpp +add_library(${PROJECT_NAME} SHARED dllmain.cpp ../abstract_data_store.cpp ../partition.cpp ../fixed_chunk_pq_table.cpp ../pq.cpp ../pq_flash_index.cpp ../logger.cpp ../utils.cpp ../windows_aligned_file_reader.cpp ../distance.cpp ../pq_l2_distance.cpp ../memory_mapper.cpp ../index.cpp ../in_mem_data_store.cpp ../pq_data_store.cpp ../in_mem_graph_store.cpp ../math_utils.cpp ../disk_utils.cpp ../filter_utils.cpp ../ann_exception.cpp ../natural_number_set.cpp ../natural_number_map.cpp ../scratch.cpp ../index_factory.cpp ../abstract_index.cpp) diff --git a/src/fixed_chunk_pq_table.cpp b/src/fixed_chunk_pq_table.cpp new file mode 100644 index 000000000..67ecabbc7 --- /dev/null +++ b/src/fixed_chunk_pq_table.cpp @@ -0,0 +1,272 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#include "fixed_chunk_pq_table.h" +#include "pq_common.h" +#include "utils.h" + +namespace diskann +{ + FixedChunkPQTable::FixedChunkPQTable() + { + } + + FixedChunkPQTable::~FixedChunkPQTable() + { + #ifndef EXEC_ENV_OLS + if (tables != nullptr) + delete[] tables; + if (tables_tr != nullptr) + delete[] tables_tr; + if (chunk_offsets != nullptr) + delete[] chunk_offsets; + if (centroid != nullptr) + delete[] centroid; + if (rotmat_tr != nullptr) + delete[] rotmat_tr; + #endif + } + + #ifdef EXEC_ENV_OLS + void FixedChunkPQTable::load_pq_centroid_bin(MemoryMappedFiles &files, const char *pq_table_file, size_t num_chunks) + { + #else + void FixedChunkPQTable::load_pq_centroid_bin(const char *pq_table_file, size_t num_chunks) + { + #endif + + uint64_t nr, nc; + std::string rotmat_file = std::string(pq_table_file) + "_rotation_matrix.bin"; + + #ifdef EXEC_ENV_OLS + size_t *file_offset_data; // since load_bin only sets the pointer, no need + // to delete. + diskann::load_bin(files, pq_table_file, file_offset_data, nr, nc); + #else + std::unique_ptr file_offset_data; + diskann::load_bin(pq_table_file, file_offset_data, nr, nc); + #endif + + bool use_old_filetype = false; + + if (nr != 4 && nr != 5) + { + diskann::cout << "Error reading pq_pivots file " << pq_table_file + << ". Offsets dont contain correct metadata, # offsets = " << nr << ", but expecting " << 4 + << " or " << 5; + throw diskann::ANNException("Error reading pq_pivots file at offsets data.", -1, __FUNCSIG__, __FILE__, + __LINE__); + } + + if (nr == 4) + { + diskann::cout << "Offsets: " << file_offset_data[0] << " " << file_offset_data[1] << " " << file_offset_data[2] + << " " << file_offset_data[3] << std::endl; + } + else if (nr == 5) + { + use_old_filetype = true; + 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; + } + else + { + throw diskann::ANNException("Wrong number of offsets in pq_pivots", -1, __FUNCSIG__, __FILE__, __LINE__); + } + + #ifdef EXEC_ENV_OLS + diskann::load_bin(files, pq_table_file, tables, nr, nc, file_offset_data[0]); + #else + diskann::load_bin(pq_table_file, tables, nr, nc, file_offset_data[0]); + #endif + + if ((nr != NUM_PQ_CENTROIDS)) + { + diskann::cout << "Error reading pq_pivots file " << pq_table_file << ". 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; + + #ifdef EXEC_ENV_OLS + diskann::load_bin(files, pq_table_file, centroid, nr, nc, file_offset_data[1]); + #else + diskann::load_bin(pq_table_file, centroid, nr, nc, file_offset_data[1]); + #endif + + if ((nr != this->ndims) || (nc != 1)) + { + diskann::cerr << "Error reading centroids from pq_pivots file " << pq_table_file << ". 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__); + } + + int chunk_offsets_index = 2; + if (use_old_filetype) + { + chunk_offsets_index = 3; + } + #ifdef EXEC_ENV_OLS + diskann::load_bin(files, pq_table_file, chunk_offsets, nr, nc, file_offset_data[chunk_offsets_index]); + #else + diskann::load_bin(pq_table_file, chunk_offsets, nr, nc, file_offset_data[chunk_offsets_index]); + #endif + + if (nc != 1 || (nr != num_chunks + 1 && num_chunks != 0)) + { + diskann::cerr << "Error loading chunk offsets file. numc: " << nc << " (should be 1). numr: " << nr + << " (should be " << num_chunks + 1 << " or 0 if we need to infer)" << std::endl; + throw diskann::ANNException("Error loading chunk offsets file", -1, __FUNCSIG__, __FILE__, __LINE__); + } + + this->n_chunks = nr - 1; + diskann::cout << "Loaded PQ Pivots: #ctrs: " << NUM_PQ_CENTROIDS << ", #dims: " << this->ndims + << ", #chunks: " << this->n_chunks << std::endl; + + #ifdef EXEC_ENV_OLS + if (files.fileExists(rotmat_file)) + { + diskann::load_bin(files, rotmat_file, (float *&)rotmat_tr, nr, nc); + #else + if (file_exists(rotmat_file)) + { + diskann::load_bin(rotmat_file, rotmat_tr, nr, nc); + #endif + if (nr != this->ndims || nc != this->ndims) + { + diskann::cerr << "Error loading rotation matrix file" << std::endl; + throw diskann::ANNException("Error loading rotation matrix file", -1, __FUNCSIG__, __FILE__, __LINE__); + } + use_rotation = true; + } + + // alloc and compute transpose + tables_tr = new float[256 * this->ndims]; + for (size_t i = 0; i < 256; i++) + { + for (size_t j = 0; j < this->ndims; j++) + { + tables_tr[j * 256 + i] = tables[i * this->ndims + j]; + } + } + } + + void FixedChunkPQTable::preprocess_query(float *query_vec) + { + for (uint32_t d = 0; d < ndims; d++) + { + query_vec[d] -= centroid[d]; + } + std::vector tmp(ndims, 0); + if (use_rotation) + { + for (uint32_t d = 0; d < ndims; d++) + { + for (uint32_t d1 = 0; d1 < ndims; d1++) + { + tmp[d] += query_vec[d1] * rotmat_tr[d1 * ndims + d]; + } + } + std::memcpy(query_vec, tmp.data(), ndims * sizeof(float)); + } + } + + // assumes pre-processed query + void FixedChunkPQTable::populate_chunk_distances(const float *query_vec, float *dist_vec) + { + memset(dist_vec, 0, 256 * n_chunks * sizeof(float)); + // chunk wise distance computation + for (size_t chunk = 0; chunk < n_chunks; chunk++) + { + // sum (q-c)^2 for the dimensions associated with this chunk + float *chunk_dists = dist_vec + (256 * chunk); + for (size_t j = chunk_offsets[chunk]; j < chunk_offsets[chunk + 1]; j++) + { + const float *centers_dim_vec = tables_tr + (256 * j); + for (size_t idx = 0; idx < 256; idx++) + { + double diff = centers_dim_vec[idx] - (query_vec[j]); + chunk_dists[idx] += (float)(diff * diff); + } + } + } + } + + float FixedChunkPQTable::l2_distance(const float *query_vec, uint8_t *base_vec) + { + float res = 0; + for (size_t chunk = 0; chunk < n_chunks; chunk++) + { + for (size_t j = chunk_offsets[chunk]; j < chunk_offsets[chunk + 1]; j++) + { + const float *centers_dim_vec = tables_tr + (256 * j); + float diff = centers_dim_vec[base_vec[chunk]] - (query_vec[j]); + res += diff * diff; + } + } + return res; + } + + float FixedChunkPQTable::inner_product(const float *query_vec, uint8_t *base_vec) + { + float res = 0; + for (size_t chunk = 0; chunk < n_chunks; chunk++) + { + for (size_t j = chunk_offsets[chunk]; j < chunk_offsets[chunk + 1]; j++) + { + const float *centers_dim_vec = tables_tr + (256 * j); + float diff = centers_dim_vec[base_vec[chunk]] * query_vec[j]; // assumes centroid is 0 to + // prevent translation errors + res += diff; + } + } + return -res; // returns negative value to simulate distances (max -> min + // conversion) + } + + // assumes no rotation is involved + template + void FixedChunkPQTable::inflate_vector(InputType *base_vec, OutputType *out_vec) const + { + for (size_t chunk = 0; chunk < n_chunks; chunk++) + { + for (size_t j = chunk_offsets[chunk]; j < chunk_offsets[chunk + 1]; j++) + { + const float *centers_dim_vec = tables_tr + (256 * j); + out_vec[j] = static_cast (centers_dim_vec[static_cast(base_vec[chunk])] + centroid[j]); + } + } + } + + void FixedChunkPQTable::populate_chunk_inner_products(const float *query_vec, float *dist_vec) + { + memset(dist_vec, 0, 256 * n_chunks * sizeof(float)); + // chunk wise distance computation + for (size_t chunk = 0; chunk < n_chunks; chunk++) + { + // sum (q-c)^2 for the dimensions associated with this chunk + float *chunk_dists = dist_vec + (256 * chunk); + for (size_t j = chunk_offsets[chunk]; j < chunk_offsets[chunk + 1]; j++) + { + const float *centers_dim_vec = tables_tr + (256 * j); + for (size_t idx = 0; idx < 256; idx++) + { + double prod = centers_dim_vec[idx] * query_vec[j]; // assumes that we are not + // shifting the vectors to + // mean zero, i.e., centroid + // array should be all zeros + chunk_dists[idx] -= (float)prod; // returning negative to keep the search code + // clean (max inner product vs min distance) + } + } + } + } + + template void FixedChunkPQTable::inflate_vector(uint8_t *base_vec, float *out_vec) const; + template void FixedChunkPQTable::inflate_vector(uint8_t *base_vec, uint8_t *out_vec) const; + template void FixedChunkPQTable::inflate_vector(int8_t *base_vec, int8_t *out_vec) const; + template void FixedChunkPQTable::inflate_vector(float *base_vec, float *out_vec) const; + } // namespace diskann diff --git a/src/index.cpp b/src/index.cpp index bf93344fa..1c8383b8d 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -29,10 +29,21 @@ namespace diskann { // Initialize an index with metric m, load the data of type T with filename // (bin), and initialize max_points + +#ifdef EXEC_ENV_OLS +template +Index::Index(const IndexConfig &index_config, MemoryMappedFiles &files, + std::shared_ptr> data_store, + std::unique_ptr graph_store, + std::shared_ptr> pq_data_store) +#else template -Index::Index(const IndexConfig &index_config, std::shared_ptr> data_store, +Index::Index(const IndexConfig &index_config, + std::shared_ptr> data_store, std::unique_ptr graph_store, std::shared_ptr> pq_data_store) +#endif + : _dist_metric(index_config.metric), _dim(index_config.dimension), _max_points(index_config.max_points), _num_frozen_pts(index_config.num_frozen_pts), _dynamic_index(index_config.dynamic_index), _enable_tags(index_config.enable_tags), _indexingMaxC(DEFAULT_MAXC), _query_scratch(nullptr), @@ -45,17 +56,21 @@ Index::Index(const IndexConfig &index_config, std::shared_ptr::Index(const IndexConfig &index_config, std::shared_ptr(total_internal_points); if (_enable_tags) { @@ -111,44 +121,58 @@ Index::Index(const IndexConfig &index_config, std::shared_ptr index_parameters, + const std::shared_ptr index_search_params, MemoryMappedFiles &files, + const std::string &codebook_path = "", const size_t num_frozen_pts = 0, + const bool dynamic_index = false, const bool enable_tags = false, + const bool concurrent_consolidate = false, const bool pq_dist_build = false, + const size_t num_pq_chunks = 0, const bool use_opq = false, const bool filtered_index = false) +#else template Index::Index(Metric m, const size_t dim, const size_t max_points, const std::shared_ptr index_parameters, - const std::shared_ptr index_search_params, const size_t num_frozen_pts, + const std::shared_ptr index_search_params, + const std::string &codebook_path, const size_t num_frozen_pts, const bool dynamic_index, const bool enable_tags, const bool concurrent_consolidate, const bool pq_dist_build, const size_t num_pq_chunks, const bool use_opq, const bool filtered_index) - : Index( - IndexConfigBuilder() - .with_metric(m) - .with_dimension(dim) - .with_max_points(max_points) - .with_index_write_params(index_parameters) - .with_index_search_params(index_search_params) - .with_num_frozen_pts(num_frozen_pts) - .is_dynamic_index(dynamic_index) - .is_enable_tags(enable_tags) - .is_concurrent_consolidate(concurrent_consolidate) - .is_pq_dist_build(pq_dist_build) - .with_num_pq_chunks(num_pq_chunks) - .is_use_opq(use_opq) - .is_filtered(filtered_index) - .with_data_type(diskann_type_to_name()) - .build(), - IndexFactory::construct_datastore(DataStoreStrategy::MEMORY, - (max_points == 0 ? (size_t)1 : max_points) + - (dynamic_index && num_frozen_pts == 0 ? (size_t)1 : num_frozen_pts), - dim, m), - IndexFactory::construct_graphstore(GraphStoreStrategy::MEMORY, - (max_points == 0 ? (size_t)1 : max_points) + - (dynamic_index && num_frozen_pts == 0 ? (size_t)1 : num_frozen_pts), - (size_t)((index_parameters == nullptr ? 0 : index_parameters->max_degree) * - defaults::GRAPH_SLACK_FACTOR * 1.05))) +#endif + : Index(IndexConfigBuilder() + .with_metric(m) + .with_dimension(dim) + .with_max_points(max_points) + .with_index_write_params(index_parameters) + .with_index_search_params(index_search_params) + .with_num_frozen_pts(num_frozen_pts) + .is_dynamic_index(dynamic_index) + .is_enable_tags(enable_tags) + .is_concurrent_consolidate(concurrent_consolidate) + .is_pq_dist_build(pq_dist_build) + .with_num_pq_chunks(num_pq_chunks) + .is_use_opq(use_opq) + .is_filtered(filtered_index) + .with_data_type(diskann_type_to_name()) + .with_pq_codebook_path(codebook_path) + .build(), +#ifdef EXEC_ENV_OLS + files, +#endif + IndexFactory::construct_datastore( + DataStoreStrategy::MEMORY, + max_points + (dynamic_index && num_frozen_pts == 0 ? (size_t)1 : num_frozen_pts), dim, m), + IndexFactory::construct_graphstore( + GraphStoreStrategy::MEMORY, + max_points + (dynamic_index && num_frozen_pts == 0 ? (size_t)1 : num_frozen_pts), + (size_t)((index_parameters == nullptr ? 0 : index_parameters->max_degree) * + defaults::GRAPH_SLACK_FACTOR * 1.05))) { if (_pq_dist) { - _pq_data_store = IndexFactory::construct_pq_datastore(DataStoreStrategy::MEMORY, max_points + num_frozen_pts, - dim, m, num_pq_chunks, use_opq); + _pq_data_store = IndexFactory::construct_pq_datastore( + DataStoreStrategy::MEMORY, codebook_path, max_points + num_frozen_pts, dim, m, num_pq_chunks, use_opq); } else { @@ -237,6 +261,8 @@ template size_t Indexsave(data_file, (location_t)(_nd + _num_frozen_pts)); } @@ -730,6 +756,22 @@ template int Index return 0; } +template int Index::get_pq_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 " << get_tag_string(tag) << " does not exist" << std::endl; + return -1; + } + + location_t location = _tag_to_location[tag]; + _pq_data_store->get_vector(location, vec); + + return 0; +} + + template uint32_t Index::calculate_entry_point() { // REFACTOR TODO: This function does not support multi-threaded calculation of medoid. @@ -803,9 +845,7 @@ std::pair Index::iterate_to_fixed_point( assert(id_scratch.size() == 0); T *aligned_query = scratch->aligned_query(); - float *pq_dists = nullptr; - _pq_data_store->preprocess_query(aligned_query, scratch); if (expanded_nodes.size() > 0 || id_scratch.size() > 0) @@ -986,7 +1026,7 @@ void Index::search_for_point_and_prune(int location, uint32_t L if (!use_filter) { - _data_store->get_vector(location, scratch->aligned_query()); + _pq_data_store->get_vector(location, scratch->aligned_query()); iterate_to_fixed_point(scratch, Lindex, init_ids, false, unused_filter_label, false); } else @@ -1001,7 +1041,7 @@ void Index::search_for_point_and_prune(int location, uint32_t L if (_dynamic_index) tl.unlock(); - _data_store->get_vector(location, scratch->aligned_query()); + _pq_data_store->get_vector(location, scratch->aligned_query()); iterate_to_fixed_point(scratch, filteredLindex, filter_specific_start_nodes, true, _location_to_labels[location], false); @@ -1015,7 +1055,7 @@ void Index::search_for_point_and_prune(int location, uint32_t L // clear scratch for finding unfiltered candidates scratch->clear(); - _data_store->get_vector(location, scratch->aligned_query()); + _pq_data_store->get_vector(location, scratch->aligned_query()); iterate_to_fixed_point(scratch, Lindex, init_ids, false, unused_filter_label, false); for (auto unfiltered_neighbour : scratch->pool()) @@ -1582,11 +1622,6 @@ void Index::build(const T *data, const size_t num_points_to_loa { throw ANNException("Do not call build with 0 points", -1, __FUNCSIG__, __FILE__, __LINE__); } - if (_pq_dist) - { - throw ANNException("ERROR: DO not use this build interface with PQ distance", -1, __FUNCSIG__, __FILE__, - __LINE__); - } std::unique_lock ul(_update_lock); @@ -1595,6 +1630,7 @@ void Index::build(const T *data, const size_t num_points_to_loa _nd = num_points_to_load; _data_store->populate_data(data, (location_t)num_points_to_load); + _pq_data_store->populate_data(data, (location_t)num_points_to_load); } build_with_data_populated(tags); @@ -2164,10 +2200,8 @@ size_t Index::search_with_tags(const T *query, const uint64_t K std::shared_lock ul(_update_lock); const std::vector init_ids = get_init_ids(); + _pq_data_store->preprocess_query(query, scratch); - //_distance->preprocess_query(query, _data_store->get_dims(), - // scratch->aligned_query()); - _data_store->preprocess_query(query, scratch); if (!use_filters) { const std::vector unused_filter_label; @@ -2198,7 +2232,7 @@ size_t Index::search_with_tags(const T *query, const uint64_t K if (res_vectors.size() > 0) { - _data_store->get_vector(node.id, res_vectors[pos]); + _pq_data_store->get_vector(node.id, res_vectors[pos]); } if (distances != nullptr) @@ -2792,6 +2826,12 @@ template void Indexresize((location_t)new_internal_points); + + if (_pq_dist) + { + _pq_data_store->resize((location_t)new_internal_points); + } + _graph_store->resize_graph(new_internal_points); _locks = std::vector(new_internal_points); @@ -2902,6 +2942,12 @@ int Index::insert_point(const T *point, const TagT tag, const s _label_to_start_id[label] = (uint32_t)fz_location; _location_to_labels[fz_location] = {label}; _data_store->set_vector((location_t)fz_location, point); + + if (_pq_dist) + { + _pq_data_store->set_vector((location_t)fz_location, point); + } + _frozen_pts_used++; } } @@ -2964,6 +3010,11 @@ int Index::insert_point(const T *point, const TagT tag, const s _data_store->set_vector(location, point); // update datastore + if (_pq_dist) + { + _pq_data_store->set_vector(location, point); // Update PQDataStore + } + // Find and add appropriate graph edges ScratchStoreManager> manager(_query_scratch); auto scratch = manager.scratch_space(); diff --git a/src/index_factory.cpp b/src/index_factory.cpp index 35790f8d6..92614608c 100644 --- a/src/index_factory.cpp +++ b/src/index_factory.cpp @@ -23,10 +23,6 @@ void IndexFactory::check_config() if (_config->pq_dist_build) { - if (_config->dynamic_index) - throw ANNException("ERROR: Dynamic Indexing not supported with PQ distance based " - "index construction", - -1, __FUNCSIG__, __FILE__, __LINE__); if (_config->metric == diskann::Metric::INNER_PRODUCT) throw ANNException("ERROR: Inner product metrics not yet supported " "with PQ distance " @@ -94,7 +90,9 @@ std::unique_ptr IndexFactory::construct_graphstore(const Gra } template -std::shared_ptr> IndexFactory::construct_pq_datastore(DataStoreStrategy strategy, size_t num_points, +std::shared_ptr> IndexFactory::construct_pq_datastore(DataStoreStrategy strategy, + const std::string &codebook_path, + size_t num_points, size_t dimension, Metric m, size_t num_pq_chunks, bool use_opq) { @@ -107,7 +105,8 @@ std::shared_ptr> IndexFactory::construct_pq_datastore(DataStoreSt case DataStoreStrategy::MEMORY: distance_fn.reset(construct_inmem_distance_fn(m)); return std::make_shared>(dimension, (location_t)(num_points), num_pq_chunks, - std::move(distance_fn), std::move(quantized_distance_fn)); + std::move(distance_fn), std::move(quantized_distance_fn), + codebook_path); default: // REFACTOR TODO: We do support diskPQ - so we may need to add a new class for SSDPQDataStore! break; @@ -127,7 +126,7 @@ std::unique_ptr IndexFactory::create_instance() if (_config->data_strategy == DataStoreStrategy::MEMORY && _config->pq_dist_build) { pq_data_store = - construct_pq_datastore(_config->data_strategy, num_points + _config->num_frozen_pts, dim, + construct_pq_datastore(_config->data_strategy, _config->pq_codebook_path, num_points + _config->num_frozen_pts, dim, _config->metric, _config->num_pq_chunks, _config->use_opq); } else diff --git a/src/pq.cpp b/src/pq.cpp index d2b545c79..8e6518bc8 100644 --- a/src/pq.cpp +++ b/src/pq.cpp @@ -15,269 +15,6 @@ namespace diskann { -FixedChunkPQTable::FixedChunkPQTable() -{ -} - -FixedChunkPQTable::~FixedChunkPQTable() -{ -#ifndef EXEC_ENV_OLS - if (tables != nullptr) - delete[] tables; - if (tables_tr != nullptr) - delete[] tables_tr; - if (chunk_offsets != nullptr) - delete[] chunk_offsets; - if (centroid != nullptr) - delete[] centroid; - if (rotmat_tr != nullptr) - delete[] rotmat_tr; -#endif -} - -#ifdef EXEC_ENV_OLS -void FixedChunkPQTable::load_pq_centroid_bin(MemoryMappedFiles &files, const char *pq_table_file, size_t num_chunks) -{ -#else -void FixedChunkPQTable::load_pq_centroid_bin(const char *pq_table_file, size_t num_chunks) -{ -#endif - - uint64_t nr, nc; - std::string rotmat_file = std::string(pq_table_file) + "_rotation_matrix.bin"; - -#ifdef EXEC_ENV_OLS - size_t *file_offset_data; // since load_bin only sets the pointer, no need - // to delete. - diskann::load_bin(files, pq_table_file, file_offset_data, nr, nc); -#else - std::unique_ptr file_offset_data; - diskann::load_bin(pq_table_file, file_offset_data, nr, nc); -#endif - - bool use_old_filetype = false; - - if (nr != 4 && nr != 5) - { - diskann::cout << "Error reading pq_pivots file " << pq_table_file - << ". Offsets dont contain correct metadata, # offsets = " << nr << ", but expecting " << 4 - << " or " << 5; - throw diskann::ANNException("Error reading pq_pivots file at offsets data.", -1, __FUNCSIG__, __FILE__, - __LINE__); - } - - if (nr == 4) - { - diskann::cout << "Offsets: " << file_offset_data[0] << " " << file_offset_data[1] << " " << file_offset_data[2] - << " " << file_offset_data[3] << std::endl; - } - else if (nr == 5) - { - use_old_filetype = true; - 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; - } - else - { - throw diskann::ANNException("Wrong number of offsets in pq_pivots", -1, __FUNCSIG__, __FILE__, __LINE__); - } - -#ifdef EXEC_ENV_OLS - - diskann::load_bin(files, pq_table_file, tables, nr, nc, file_offset_data[0]); -#else - diskann::load_bin(pq_table_file, tables, nr, nc, file_offset_data[0]); -#endif - - if ((nr != NUM_PQ_CENTROIDS)) - { - diskann::cout << "Error reading pq_pivots file " << pq_table_file << ". 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; - -#ifdef EXEC_ENV_OLS - diskann::load_bin(files, pq_table_file, centroid, nr, nc, file_offset_data[1]); -#else - diskann::load_bin(pq_table_file, centroid, nr, nc, file_offset_data[1]); -#endif - - if ((nr != this->ndims) || (nc != 1)) - { - diskann::cerr << "Error reading centroids from pq_pivots file " << pq_table_file << ". 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__); - } - - int chunk_offsets_index = 2; - if (use_old_filetype) - { - chunk_offsets_index = 3; - } -#ifdef EXEC_ENV_OLS - diskann::load_bin(files, pq_table_file, chunk_offsets, nr, nc, file_offset_data[chunk_offsets_index]); -#else - diskann::load_bin(pq_table_file, chunk_offsets, nr, nc, file_offset_data[chunk_offsets_index]); -#endif - - if (nc != 1 || (nr != num_chunks + 1 && num_chunks != 0)) - { - diskann::cerr << "Error loading chunk offsets file. numc: " << nc << " (should be 1). numr: " << nr - << " (should be " << num_chunks + 1 << " or 0 if we need to infer)" << std::endl; - throw diskann::ANNException("Error loading chunk offsets file", -1, __FUNCSIG__, __FILE__, __LINE__); - } - - this->n_chunks = nr - 1; - diskann::cout << "Loaded PQ Pivots: #ctrs: " << NUM_PQ_CENTROIDS << ", #dims: " << this->ndims - << ", #chunks: " << this->n_chunks << std::endl; - -#ifdef EXEC_ENV_OLS - if (files.fileExists(rotmat_file)) - { - diskann::load_bin(files, rotmat_file, (float *&)rotmat_tr, nr, nc); -#else - if (file_exists(rotmat_file)) - { - diskann::load_bin(rotmat_file, rotmat_tr, nr, nc); -#endif - if (nr != this->ndims || nc != this->ndims) - { - diskann::cerr << "Error loading rotation matrix file" << std::endl; - throw diskann::ANNException("Error loading rotation matrix file", -1, __FUNCSIG__, __FILE__, __LINE__); - } - use_rotation = true; - } - - // alloc and compute transpose - tables_tr = new float[256 * this->ndims]; - for (size_t i = 0; i < 256; i++) - { - for (size_t j = 0; j < this->ndims; j++) - { - tables_tr[j * 256 + i] = tables[i * this->ndims + j]; - } - } -} - -uint32_t FixedChunkPQTable::get_num_chunks() -{ - return static_cast(n_chunks); -} - -void FixedChunkPQTable::preprocess_query(float *query_vec) -{ - for (uint32_t d = 0; d < ndims; d++) - { - query_vec[d] -= centroid[d]; - } - std::vector tmp(ndims, 0); - if (use_rotation) - { - for (uint32_t d = 0; d < ndims; d++) - { - for (uint32_t d1 = 0; d1 < ndims; d1++) - { - tmp[d] += query_vec[d1] * rotmat_tr[d1 * ndims + d]; - } - } - std::memcpy(query_vec, tmp.data(), ndims * sizeof(float)); - } -} - -// assumes pre-processed query -void FixedChunkPQTable::populate_chunk_distances(const float *query_vec, float *dist_vec) -{ - memset(dist_vec, 0, 256 * n_chunks * sizeof(float)); - // chunk wise distance computation - for (size_t chunk = 0; chunk < n_chunks; chunk++) - { - // sum (q-c)^2 for the dimensions associated with this chunk - float *chunk_dists = dist_vec + (256 * chunk); - for (size_t j = chunk_offsets[chunk]; j < chunk_offsets[chunk + 1]; j++) - { - const float *centers_dim_vec = tables_tr + (256 * j); - for (size_t idx = 0; idx < 256; idx++) - { - double diff = centers_dim_vec[idx] - (query_vec[j]); - chunk_dists[idx] += (float)(diff * diff); - } - } - } -} - -float FixedChunkPQTable::l2_distance(const float *query_vec, uint8_t *base_vec) -{ - float res = 0; - for (size_t chunk = 0; chunk < n_chunks; chunk++) - { - for (size_t j = chunk_offsets[chunk]; j < chunk_offsets[chunk + 1]; j++) - { - const float *centers_dim_vec = tables_tr + (256 * j); - float diff = centers_dim_vec[base_vec[chunk]] - (query_vec[j]); - res += diff * diff; - } - } - return res; -} - -float FixedChunkPQTable::inner_product(const float *query_vec, uint8_t *base_vec) -{ - float res = 0; - for (size_t chunk = 0; chunk < n_chunks; chunk++) - { - for (size_t j = chunk_offsets[chunk]; j < chunk_offsets[chunk + 1]; j++) - { - const float *centers_dim_vec = tables_tr + (256 * j); - float diff = centers_dim_vec[base_vec[chunk]] * query_vec[j]; // assumes centroid is 0 to - // prevent translation errors - res += diff; - } - } - return -res; // returns negative value to simulate distances (max -> min - // conversion) -} - -// assumes no rotation is involved -void FixedChunkPQTable::inflate_vector(uint8_t *base_vec, float *out_vec) -{ - for (size_t chunk = 0; chunk < n_chunks; chunk++) - { - for (size_t j = chunk_offsets[chunk]; j < chunk_offsets[chunk + 1]; j++) - { - const float *centers_dim_vec = tables_tr + (256 * j); - out_vec[j] = centers_dim_vec[base_vec[chunk]] + centroid[j]; - } - } -} - -void FixedChunkPQTable::populate_chunk_inner_products(const float *query_vec, float *dist_vec) -{ - memset(dist_vec, 0, 256 * n_chunks * sizeof(float)); - // chunk wise distance computation - for (size_t chunk = 0; chunk < n_chunks; chunk++) - { - // sum (q-c)^2 for the dimensions associated with this chunk - float *chunk_dists = dist_vec + (256 * chunk); - for (size_t j = chunk_offsets[chunk]; j < chunk_offsets[chunk + 1]; j++) - { - const float *centers_dim_vec = tables_tr + (256 * j); - for (size_t idx = 0; idx < 256; idx++) - { - double prod = centers_dim_vec[idx] * query_vec[j]; // assumes that we are not - // shifting the vectors to - // mean zero, i.e., centroid - // array should be all zeros - chunk_dists[idx] -= (float)prod; // returning negative to keep the search code - // clean (max inner product vs min distance) - } - } - } -} - void aggregate_coords(const std::vector &ids, const uint8_t *all_coords, const size_t ndims, uint8_t *out) { for (size_t i = 0; i < ids.size(); i++) diff --git a/src/pq_data_store.cpp b/src/pq_data_store.cpp index c47c16705..a5ff51b8f 100644 --- a/src/pq_data_store.cpp +++ b/src/pq_data_store.cpp @@ -1,9 +1,9 @@ #include +#include "pq_common.h" #include "pq_data_store.h" #include "pq.h" #include "pq_scratch.h" -#include "utils.h" #include "distance.h" namespace diskann @@ -12,18 +12,45 @@ namespace diskann // REFACTOR TODO: Assuming that num_pq_chunks is known already. Must verify if // this is true. template + +#ifdef EXEC_ENV_OLS +PQDataStore::PQDataStore(size_t dim, location_t num_points, size_t num_pq_chunks, + std::unique_ptr> distance_fn, + std::unique_ptr> pq_distance_fn, + MemoryMappedFiles &files, + const std::string &codebook_path) +#else PQDataStore::PQDataStore(size_t dim, location_t num_points, size_t num_pq_chunks, std::unique_ptr> distance_fn, - std::unique_ptr> pq_distance_fn) - : AbstractDataStore(num_points, dim), _quantized_data(nullptr), _num_chunks(num_pq_chunks), - _distance_metric(distance_fn->get_metric()) + std::unique_ptr> pq_distance_fn, + const std::string &codebook_path) +#endif + : AbstractDataStore(num_points, num_pq_chunks), _num_chunks(num_pq_chunks), + _pq_pivot_file_path(codebook_path) { if (num_pq_chunks > dim) { throw diskann::ANNException("ERROR: num_pq_chunks > dim", -1, __FUNCSIG__, __FILE__, __LINE__); } + _distance_fn = std::move(distance_fn); _pq_distance_fn = std::move(pq_distance_fn); + + _aligned_dim = ROUND_UP(num_pq_chunks, _distance_fn->get_required_alignment()); + alloc_aligned(((void **)&_quantized_data), this->_capacity * _aligned_dim * sizeof(data_t), 8 * sizeof(data_t)); + std::memset(_quantized_data, 0, this->_capacity * _aligned_dim * sizeof(data_t)); + +#ifdef EXEC_ENV_OLS + if (!codebook_path.empty()) + { + _pq_distance_fn->load_pivot_data(files, codebook_path); + } +#else + if (!codebook_path.empty()) + { + _pq_distance_fn->load_pivot_data(codebook_path.c_str()); + } +#endif } template PQDataStore::~PQDataStore() @@ -46,13 +73,17 @@ template size_t PQDataStore::save(const std::string &f template size_t PQDataStore::get_aligned_dim() const { - return this->get_dims(); + return _aligned_dim; } // Populate quantized data from regular data. template void PQDataStore::populate_data(const data_t *vectors, const location_t num_pts) { - throw std::logic_error("Not implemented yet"); + memset(_quantized_data, 0, _aligned_dim * sizeof(data_t) * num_pts); + for (location_t i = 0; i < num_pts; i++) + { + std::memmove(_quantized_data + i * _aligned_dim, vectors + i * this->_dim, this->_dim * sizeof(data_t)); + } } template void PQDataStore::populate_data(const std::string &filename, const size_t offset) @@ -69,10 +100,14 @@ template void PQDataStore::populate_data(const std::st double p_val = std::min(1.0, ((double)MAX_PQ_TRAINING_SET_SIZE / (double)file_num_points)); - auto pivots_file = _pq_distance_fn->get_pivot_data_filename(filename); - auto compressed_file = _pq_distance_fn->get_quantized_vectors_filename(filename); + auto pivots_file = _pq_pivot_file_path.empty() + ? get_pivot_data_filename(filename, _use_opq, static_cast(_num_chunks)) + : _pq_pivot_file_path; + + auto compressed_file = get_quantized_vectors_filename(filename, _use_opq, static_cast(_num_chunks)); - generate_quantized_data(filename, pivots_file, compressed_file, _distance_metric, p_val, _num_chunks, + generate_quantized_data(filename, pivots_file, compressed_file, _distance_fn->get_metric(), p_val, + _num_chunks, _pq_distance_fn->is_opq()); // REFACTOR TODO: Not sure of the alignment. Just copying from index.cpp @@ -84,22 +119,23 @@ template void PQDataStore::populate_data(const std::st "EXEC_ENV_OLS is defined.", -1, __FUNCSIG__, __FILE__, __LINE__); #else - _pq_distance_fn->load_pivot_data(pivots_file.c_str(), _num_chunks); + _pq_distance_fn->load_pivot_data(pivots_file); #endif } template void PQDataStore::extract_data_to_bin(const std::string &filename, const location_t num_pts) { - throw std::logic_error("Not implemented yet"); + diskann::save_bin(filename, _quantized_data, this->capacity(), _num_chunks, 0); } -template void PQDataStore::get_vector(const location_t i, data_t *target) const +template void PQDataStore::get_vector(const location_t i, data_t *dest) const { // REFACTOR TODO: Should we inflate the compressed vector here? if (i < this->capacity()) { - throw std::logic_error("Not implemented yet."); + const FixedChunkPQTable &pq_table = _pq_distance_fn->get_pq_table(); + pq_table.inflate_vector((data_t *)(_quantized_data + i * _aligned_dim), dest); } else { @@ -108,11 +144,37 @@ template void PQDataStore::get_vector(const location_t throw diskann::ANNException(ss.str(), -1); } } -template void PQDataStore::set_vector(const location_t i, const data_t *const vector) +template void PQDataStore::set_vector(const location_t loc, const data_t *const vector) { - // REFACTOR TODO: Should we accept a normal vector and compress here? - // memcpy (_data + i * _num_chunks, vector, _num_chunks * sizeof(data_t)); - throw std::logic_error("Not implemented yet"); + if (_pq_distance_fn == nullptr) + { + throw diskann::ANNException("PQ distance is not loaded, cannot set vector for PQDataStore.", -1); + } + + const FixedChunkPQTable &pq_table = _pq_distance_fn->get_pq_table(); + + if (pq_table.tables == nullptr) + { + throw diskann::ANNException("PQ table is not loaded for PQ distance, cannot set vector for PQDataStore.", -1); + } + + uint64_t full_dimension = pq_table.ndims; + uint64_t num_chunks = _num_chunks; + + std::vector vector_float(full_dimension); + diskann::convert_types(vector, vector_float.data(), 1, full_dimension); + std::vector compressed_vector(num_chunks * sizeof(data_t)); + std::vector compressed_vector_T(num_chunks); + + generate_pq_data_from_pivots_simplified(vector_float.data(), 1, pq_table.tables, 256 * full_dimension, + full_dimension, + num_chunks, compressed_vector); + + diskann::convert_types(compressed_vector.data(), compressed_vector_T.data(), 1, num_chunks); + + size_t offset_in_data = loc * _aligned_dim; + memset(_quantized_data + offset_in_data, 0, _aligned_dim * sizeof(data_t)); + memcpy(_quantized_data + offset_in_data, compressed_vector_T.data(), this->_dim * sizeof(data_t)); } template void PQDataStore::prefetch_vector(const location_t loc) @@ -123,10 +185,43 @@ template void PQDataStore::prefetch_vector(const locat template void PQDataStore::move_vectors(const location_t old_location_start, const location_t new_location_start, - const location_t num_points) + const location_t num_locations) { - // REFACTOR TODO: Moving vectors is only for in-mem fresh. - throw std::logic_error("Not implemented yet"); + if (num_locations == 0 || old_location_start == new_location_start) + { + return; + } + + // The [start, end) interval which will contain obsolete points to be + // cleared. + uint32_t mem_clear_loc_start = old_location_start; + uint32_t mem_clear_loc_end_limit = old_location_start + num_locations; + + if (new_location_start < old_location_start) + { + // If ranges are overlapping, make sure not to clear the newly copied + // data. + if (mem_clear_loc_start < new_location_start + num_locations) + { + // Clear only after the end of the new range. + mem_clear_loc_start = new_location_start + num_locations; + } + } + else + { + // If ranges are overlapping, make sure not to clear the newly copied + // data. + if (mem_clear_loc_end_limit > new_location_start) + { + // Clear only up to the beginning of the new range. + mem_clear_loc_end_limit = new_location_start; + } + } + + // Use memmove to handle overlapping ranges. + copy_vectors(old_location_start, new_location_start, num_locations); + memset(_quantized_data + _aligned_dim * mem_clear_loc_start, 0, + sizeof(data_t) * _aligned_dim * (mem_clear_loc_end_limit - mem_clear_loc_start)); } template @@ -158,12 +253,12 @@ void PQDataStore::preprocess_query(const data_t *aligned_query, Abstract template float PQDataStore::get_distance(const data_t *query, const location_t loc) const { - throw std::logic_error("Not implemented yet"); + throw diskann::ANNException("get_distance(const data_t *query, const location_t loc) hasn't been implemented for PQDataStore", -1); } template float PQDataStore::get_distance(const location_t loc1, const location_t loc2) const { - throw std::logic_error("Not implemented yet"); + throw diskann::ANNException("get_distance(const location_t loc1, const location_t loc2) hasn't been implemented for PQDataStore", -1); } template @@ -211,7 +306,7 @@ template location_t PQDataStore::calculate_medoid() co template size_t PQDataStore::get_alignment_factor() const { - return 1; + return _distance_fn->get_required_alignment(); } template Distance *PQDataStore::get_dist_fn() const @@ -225,26 +320,69 @@ template location_t PQDataStore::load_impl(const std:: { aligned_free(_quantized_data); } - auto quantized_vectors_file = _pq_distance_fn->get_quantized_vectors_filename(file_prefix); + auto quantized_vectors_file = + get_quantized_vectors_filename(file_prefix, _use_opq, static_cast(_num_chunks)); size_t num_points; load_aligned_bin(quantized_vectors_file, _quantized_data, num_points, _num_chunks, _num_chunks); this->_capacity = (location_t)num_points; - auto pivots_file = _pq_distance_fn->get_pivot_data_filename(file_prefix); - _pq_distance_fn->load_pivot_data(pivots_file, _num_chunks); + auto pivots_file = get_pivot_data_filename(file_prefix, _use_opq, static_cast(_num_chunks)); + _pq_distance_fn->load_pivot_data(pivots_file); return this->_capacity; } template location_t PQDataStore::expand(const location_t new_size) { - throw std::logic_error("Not implemented yet"); + if (new_size == this->capacity()) + { + return this->capacity(); + } + else if (new_size < this->capacity()) + { + std::stringstream ss; + ss << "Cannot 'expand' datastore when new capacity (" << new_size << ") < existing capacity(" + << this->capacity() << ")" << std::endl; + throw diskann::ANNException(ss.str(), -1); + } +#ifndef _WINDOWS + data_t *new_data; + alloc_aligned((void **)&new_data, new_size * _aligned_dim * sizeof(data_t), 8 * sizeof(data_t)); + memcpy(new_data, _quantized_data, this->capacity() * _aligned_dim * sizeof(data_t)); + aligned_free(_quantized_data); + _quantized_data = new_data; +#else + realloc_aligned((void **)&_quantized_data, new_size * _aligned_dim * sizeof(data_t), 8 * sizeof(data_t)); +#endif + this->_capacity = new_size; + return this->_capacity; } template location_t PQDataStore::shrink(const location_t new_size) { - throw std::logic_error("Not implemented yet"); + if (new_size == this->capacity()) + { + return this->capacity(); + } + else if (new_size > this->capacity()) + { + std::stringstream ss; + ss << "Cannot 'shrink' datastore when new capacity (" << new_size << ") > existing capacity(" + << this->capacity() << ")" << std::endl; + throw diskann::ANNException(ss.str(), -1); + } +#ifndef _WINDOWS + data_t *new_data; + alloc_aligned((void **)&new_data, new_size * _aligned_dim * sizeof(data_t), 8 * sizeof(data_t)); + memcpy(new_data, _quantized_data, new_size * _aligned_dim * sizeof(data_t)); + aligned_free(_quantized_data); + _quantized_data = new_data; +#else + realloc_aligned((void **)&_quantized_data, new_size * _aligned_dim * sizeof(data_t), 8 * sizeof(data_t)); +#endif + this->_capacity = new_size; + return this->_capacity; } #ifdef EXEC_ENV_OLS diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index bba101bb3..a82d9a10e 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -1015,7 +1015,7 @@ int PQFlashIndex::load_from_separate_paths(uint32_t num_threads, cons // chunk_offsets file the correct value _disk_pq_table.load_pq_centroid_bin(disk_pq_pivots_path.c_str(), 0); #endif - _disk_pq_n_chunks = _disk_pq_table.get_num_chunks(); + _disk_pq_n_chunks = _disk_pq_table.n_chunks; _disk_bytes_per_point = _disk_pq_n_chunks * sizeof(uint8_t); // revising disk_bytes_per_point since DISK PQ is used. diskann::cout << "Disk index uses PQ data compressed down to " << _disk_pq_n_chunks << " bytes per point." diff --git a/src/pq_l2_distance.cpp b/src/pq_l2_distance.cpp index c08744c35..2c7e9ffd8 100644 --- a/src/pq_l2_distance.cpp +++ b/src/pq_l2_distance.cpp @@ -16,18 +16,6 @@ PQL2Distance::PQL2Distance(uint32_t num_chunks, bool use_opq) : _num_chu template PQL2Distance::~PQL2Distance() { -#ifndef EXEC_ENV_OLS - if (_tables != nullptr) - delete[] _tables; - if (_chunk_offsets != nullptr) - delete[] _chunk_offsets; - if (_centroid != nullptr) - delete[] _centroid; - if (_rotmat_tr != nullptr) - delete[] _rotmat_tr; -#endif - if (_tables_tr != nullptr) - delete[] _tables_tr; } template bool PQL2Distance::is_opq() const @@ -35,165 +23,31 @@ template bool PQL2Distance::is_opq() const return this->_is_opq; } -template -std::string PQL2Distance::get_quantized_vectors_filename(const std::string &prefix) const -{ - if (_num_chunks == 0) - { - throw diskann::ANNException("Must set num_chunks before calling get_quantized_vectors_filename", -1, - __FUNCSIG__, __FILE__, __LINE__); - } - return diskann::get_quantized_vectors_filename(prefix, _is_opq, (uint32_t)_num_chunks); -} -template std::string PQL2Distance::get_pivot_data_filename(const std::string &prefix) const -{ - if (_num_chunks == 0) - { - throw diskann::ANNException("Must set num_chunks before calling get_pivot_data_filename", -1, __FUNCSIG__, - __FILE__, __LINE__); - } - return diskann::get_pivot_data_filename(prefix, _is_opq, (uint32_t)_num_chunks); -} -template -std::string PQL2Distance::get_rotation_matrix_suffix(const std::string &pq_pivots_filename) const -{ - return diskann::get_rotation_matrix_suffix(pq_pivots_filename); -} - #ifdef EXEC_ENV_OLS template -void PQL2Distance::load_pivot_data(MemoryMappedFiles &files, const std::string &pq_table_file, - size_t num_chunks) +void PQL2Distance::load_pivot_data(MemoryMappedFiles &files, const std::string &pq_table_file) { + _pq_table.load_pq_centroid_bin(files, pq_table_file.c_str(), _num_chunks); +} #else template -void PQL2Distance::load_pivot_data(const std::string &pq_table_file, size_t num_chunks) +void PQL2Distance::load_pivot_data(const std::string &pq_table_file) { -#endif - uint64_t nr, nc; - // std::string rotmat_file = get_opq_rot_matrix_filename(pq_table_file, - // false); - -#ifdef EXEC_ENV_OLS - size_t *file_offset_data; // since load_bin only sets the pointer, no need - // to delete. - diskann::load_bin(files, pq_table_file, file_offset_data, nr, nc); -#else - std::unique_ptr file_offset_data; - diskann::load_bin(pq_table_file, file_offset_data, nr, nc); -#endif - - bool use_old_filetype = false; - - if (nr != 4 && nr != 5) - { - diskann::cout << "Error reading pq_pivots file " << pq_table_file - << ". Offsets dont contain correct metadata, # offsets = " << nr << ", but expecting " << 4 - << " or " << 5; - throw diskann::ANNException("Error reading pq_pivots file at offsets data.", -1, __FUNCSIG__, __FILE__, - __LINE__); - } - - if (nr == 4) - { - diskann::cout << "Offsets: " << file_offset_data[0] << " " << file_offset_data[1] << " " << file_offset_data[2] - << " " << file_offset_data[3] << std::endl; - } - else if (nr == 5) - { - use_old_filetype = true; - 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; - } - else - { - throw diskann::ANNException("Wrong number of offsets in pq_pivots", -1, __FUNCSIG__, __FILE__, __LINE__); - } - -#ifdef EXEC_ENV_OLS - diskann::load_bin(files, pq_table_file, tables, nr, nc, file_offset_data[0]); -#else - diskann::load_bin(pq_table_file, _tables, nr, nc, file_offset_data[0]); -#endif - - if ((nr != NUM_PQ_CENTROIDS)) - { - diskann::cout << "Error reading pq_pivots file " << pq_table_file << ". 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; - -#ifdef EXEC_ENV_OLS - diskann::load_bin(files, pq_table_file, centroid, nr, nc, file_offset_data[1]); -#else - diskann::load_bin(pq_table_file, _centroid, nr, nc, file_offset_data[1]); -#endif - - if ((nr != this->_ndims) || (nc != 1)) - { - diskann::cerr << "Error reading centroids from pq_pivots file " << pq_table_file << ". 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__); - } - - int chunk_offsets_index = 2; - if (use_old_filetype) - { - chunk_offsets_index = 3; - } -#ifdef EXEC_ENV_OLS - diskann::load_bin(files, pq_table_file, chunk_offsets, nr, nc, file_offset_data[chunk_offsets_index]); -#else - diskann::load_bin(pq_table_file, _chunk_offsets, nr, nc, file_offset_data[chunk_offsets_index]); -#endif - - if (nc != 1 || (nr != num_chunks + 1 && num_chunks != 0)) - { - diskann::cerr << "Error loading chunk offsets file. numc: " << nc << " (should be 1). numr: " << nr - << " (should be " << num_chunks + 1 << " or 0 if we need to infer)" << std::endl; - throw diskann::ANNException("Error loading chunk offsets file", -1, __FUNCSIG__, __FILE__, __LINE__); - } - - this->_num_chunks = nr - 1; - diskann::cout << "Loaded PQ Pivots: #ctrs: " << NUM_PQ_CENTROIDS << ", #dims: " << this->_ndims - << ", #chunks: " << this->_num_chunks << std::endl; - - // For OPQ there will be a rotation matrix to load. - if (this->_is_opq) - { - std::string rotmat_file = get_rotation_matrix_suffix(pq_table_file); -#ifdef EXEC_ENV_OLS - diskann::load_bin(files, rotmat_file, (float *&)rotmat_tr, nr, nc); -#else - diskann::load_bin(rotmat_file, _rotmat_tr, nr, nc); -#endif - if (nr != this->_ndims || nc != this->_ndims) - { - diskann::cerr << "Error loading rotation matrix file" << std::endl; - throw diskann::ANNException("Error loading rotation matrix file", -1, __FUNCSIG__, __FILE__, __LINE__); - } - } - - // alloc and compute transpose - _tables_tr = new float[256 * this->_ndims]; - for (size_t i = 0; i < 256; i++) - { - for (size_t j = 0; j < this->_ndims; j++) - { - _tables_tr[j * 256 + i] = _tables[i * this->_ndims + j]; - } - } + _pq_table.load_pq_centroid_bin(pq_table_file.c_str(), _num_chunks); } +#endif template uint32_t PQL2Distance::get_num_chunks() const { return static_cast(_num_chunks); } +template +const FixedChunkPQTable & PQL2Distance::get_pq_table() const +{ + return _pq_table; +} + // REFACTOR: Instead of doing half the work in the caller and half in this // function, we let this function // do all of the work, making it easier for the caller. @@ -207,21 +61,21 @@ void PQL2Distance::preprocess_query(const data_t *aligned_query, uint32_ } scratch.initialize(dim, aligned_query); - for (uint32_t d = 0; d < _ndims; d++) + for (uint32_t d = 0; d < _pq_table.ndims; d++) { - scratch.rotated_query[d] -= _centroid[d]; + scratch.rotated_query[d] -= _pq_table.centroid[d]; } - std::vector tmp(_ndims, 0); + std::vector tmp(_pq_table.ndims, 0); if (_is_opq) { - for (uint32_t d = 0; d < _ndims; d++) + for (uint32_t d = 0; d < _pq_table.ndims; d++) { - for (uint32_t d1 = 0; d1 < _ndims; d1++) + for (uint32_t d1 = 0; d1 < _pq_table.ndims; d1++) { - tmp[d] += scratch.rotated_query[d1] * _rotmat_tr[d1 * _ndims + d]; + tmp[d] += scratch.rotated_query[d1] * _pq_table.rotmat_tr[d1 * _pq_table.ndims + d]; } } - std::memcpy(scratch.rotated_query, tmp.data(), _ndims * sizeof(float)); + std::memcpy(scratch.rotated_query, tmp.data(), _pq_table.ndims * sizeof(float)); } this->prepopulate_chunkwise_distances(scratch.rotated_query, scratch.aligned_pqtable_dist_scratch); } @@ -246,9 +100,9 @@ template float PQL2Distance::brute_force_distance(cons float res = 0; for (size_t chunk = 0; chunk < _num_chunks; chunk++) { - for (size_t j = _chunk_offsets[chunk]; j < _chunk_offsets[chunk + 1]; j++) + for (size_t j = _pq_table.chunk_offsets[chunk]; j < _pq_table.chunk_offsets[chunk + 1]; j++) { - const float *centers_dim_vec = _tables_tr + (256 * j); + const float *centers_dim_vec = _pq_table.tables + (256 * j); float diff = centers_dim_vec[base_vec[chunk]] - (query_vec[j]); res += diff * diff; } @@ -265,9 +119,9 @@ void PQL2Distance::prepopulate_chunkwise_distances(const float *query_ve { // sum (q-c)^2 for the dimensions associated with this chunk float *chunk_dists = dist_vec + (256 * chunk); - for (size_t j = _chunk_offsets[chunk]; j < _chunk_offsets[chunk + 1]; j++) + for (size_t j = _pq_table.chunk_offsets[chunk]; j < _pq_table.chunk_offsets[chunk + 1]; j++) { - const float *centers_dim_vec = _tables_tr + (256 * j); + const float *centers_dim_vec = _pq_table.tables_tr + (256 * j); for (size_t idx = 0; idx < 256; idx++) { double diff = centers_dim_vec[idx] - (query_vec[j]);