From 29bbef70c7214fa31fc8ea0ae6af08c4d24d69a8 Mon Sep 17 00:00:00 2001 From: harsha vardhan simhadri Date: Tue, 28 Mar 2023 01:19:47 -0700 Subject: [PATCH 01/69] gi# This is a combination of 2 commits. remove _u, _s typedefs --- include/cached_io.h | 4 +- include/common_includes.h | 1 + include/concurrent_queue.h | 4 +- include/disk_utils.h | 13 +- include/index.h | 31 +-- include/natural_number_map.h | 6 +- include/pq.h | 31 +-- include/pq_flash_index.h | 108 ++++----- include/scratch.h | 8 +- include/types.h | 12 + include/utils.h | 125 +++++----- src/CMakeLists.txt | 2 +- src/disk_utils.cpp | 218 +++++++++--------- src/distance.cpp | 3 +- src/dll/CMakeLists.txt | 3 +- src/index.cpp | 238 ++++++++++--------- src/math_utils.cpp | 18 +- src/partition.cpp | 70 +++--- src/pq.cpp | 142 ++++++------ src/pq_flash_index.cpp | 267 ++++++++++++---------- src/scratch.cpp | 12 +- src/utils.cpp | 37 +-- src/windows_aligned_file_reader.cpp | 7 +- tests/build_disk_index.cpp | 9 +- tests/build_memory_index.cpp | 16 +- tests/build_stitched_index.cpp | 183 +++++++-------- tests/range_search_disk_index.cpp | 28 +-- tests/search_disk_index.cpp | 22 +- tests/search_memory_index.cpp | 10 +- tests/test_insert_deletes_consolidate.cpp | 9 +- tests/test_streaming_scenario.cpp | 17 +- tests/utils/bin_to_tsv.cpp | 23 +- tests/utils/calculate_recall.cpp | 4 +- tests/utils/compute_groundtruth.cpp | 53 ++--- tests/utils/float_bin_to_int8.cpp | 28 +-- tests/utils/fvecs_to_bin.cpp | 48 ++-- tests/utils/fvecs_to_bvecs.cpp | 22 +- tests/utils/generate_synthetic_labels.cpp | 5 +- tests/utils/int8_to_float_scale.cpp | 28 +-- tests/utils/ivecs_to_bin.cpp | 37 +-- tests/utils/merge_shards.cpp | 10 +- tests/utils/rand_data_gen.cpp | 55 ++--- tests/utils/stats_label_data.cpp | 24 +- tests/utils/tsv_to_bin.cpp | 40 ++-- tests/utils/vector_analysis.cpp | 32 +-- 45 files changed, 1073 insertions(+), 990 deletions(-) create mode 100644 include/types.h diff --git a/include/cached_io.h b/include/cached_io.h index a41c03431..daef2f2f7 100644 --- a/include/cached_io.h +++ b/include/cached_io.h @@ -95,8 +95,8 @@ class cached_ifstream reader.read(cache_buf, cache_size); cur_off = 0; } - // note that if size_left < cache_size, then cur_off = cache_size, so - // subsequent reads will all be directly from file + // note that if size_left < cache_size, then cur_off = cache_size, + // so subsequent reads will all be directly from file } } diff --git a/include/common_includes.h b/include/common_includes.h index 96de5a46c..e1a51bdec 100644 --- a/include/common_includes.h +++ b/include/common_includes.h @@ -23,4 +23,5 @@ #include #include #include +#include #include diff --git a/include/concurrent_queue.h b/include/concurrent_queue.h index 405b90617..1e57bbf0f 100644 --- a/include/concurrent_queue.h +++ b/include/concurrent_queue.h @@ -88,8 +88,8 @@ template class ConcurrentQueue { T ret = this->q.front(); this->q.pop(); - // diskann::cout << "thread_id: " << std::this_thread::get_id() << ", - // ctx: " + // diskann::cout << "thread_id: " << std::this_thread::get_id() << + // ", ctx: " // << ret.ctx << "\n"; lk.unlock(); return ret; diff --git a/include/disk_utils.h b/include/disk_utils.h index ff0619e74..0497da71a 100644 --- a/include/disk_utils.h +++ b/include/disk_utils.h @@ -63,7 +63,7 @@ DISKANN_DLLEXPORT T *load_warmup(const std::string &cache_warmup_file, uint64_t DISKANN_DLLEXPORT int merge_shards(const std::string &vamana_prefix, const std::string &vamana_suffix, const std::string &idmaps_prefix, const std::string &idmaps_suffix, - const _u64 nshards, unsigned max_degree, const std::string &output_vamana, + const uint64_t nshards, unsigned max_degree, const std::string &output_vamana, const std::string &medoids_file, bool use_filters = false, const std::string &labels_to_medoids_file = std::string("")); @@ -82,20 +82,21 @@ DISKANN_DLLEXPORT int build_merged_vamana_index(std::string base_file, diskann:: bool use_filters = false, const std::string &label_file = std::string(""), const std::string &labels_to_medoids_file = std::string(""), - const std::string &universal_label = "", const _u32 Lf = 0); + const std::string &universal_label = "", const uint32_t Lf = 0); 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); + T *tuning_sample, uint64_t tuning_sample_num, + uint64_t tuning_sample_aligned_dim, uint32_t L, uint32_t nthreads, + uint32_t start_bw = 2); template DISKANN_DLLEXPORT int build_disk_index( const char *dataFilePath, const char *indexFilePath, const char *indexBuildParameters, diskann::Metric _compareMetric, bool use_opq = false, bool use_filters = false, const std::string &label_file = std::string(""), // default is empty string for no label_file - const std::string &universal_label = "", const _u32 filter_threshold = 0, - const _u32 Lf = 0); // default is empty string for no universal label + const std::string &universal_label = "", const uint32_t filter_threshold = 0, + const uint32_t Lf = 0); // default is empty string for no universal label template DISKANN_DLLEXPORT void create_disk_layout(const std::string base_file, const std::string mem_index_file, diff --git a/include/index.h b/include/index.h index 962f758e3..f266d9837 100644 --- a/include/index.h +++ b/include/index.h @@ -26,7 +26,7 @@ namespace diskann { -inline double estimate_ram_usage(_u64 size, _u32 dim, _u32 datasize, _u32 degree) +inline double estimate_ram_usage(size_t size, uint32_t dim, uint32_t datasize, uint32_t degree) { double size_of_data = ((double)size) * ROUND_UP(dim, 8) * datasize; double size_of_graph = ((double)size) * degree * sizeof(unsigned) * GRAPH_SLACK_FACTOR; @@ -226,7 +226,8 @@ template clas std::unordered_map load_label_map(const std::string &map_file); - // Returns the locations of start point and frozen points suitable for use with iterate_to_fixed_point. + // Returns the locations of start point and frozen points suitable for use + // with iterate_to_fixed_point. std::vector get_init_ids(); std::pair iterate_to_fixed_point(const T *node_coords, const unsigned Lindex, @@ -234,14 +235,15 @@ template clas InMemQueryScratch *scratch, bool use_filter, const std::vector &filters, bool search_invocation); - void search_for_point_and_prune(int location, _u32 Lindex, std::vector &pruned_list, - InMemQueryScratch *scratch, bool use_filter = false, _u32 filteredLindex = 0); + void search_for_point_and_prune(int location, uint32_t Lindex, std::vector &pruned_list, + InMemQueryScratch *scratch, bool use_filter = false, + uint32_t filteredLindex = 0); void prune_neighbors(const unsigned location, std::vector &pool, std::vector &pruned_list, InMemQueryScratch *scratch); - void prune_neighbors(const unsigned location, std::vector &pool, const _u32 range, - const _u32 max_candidate_size, const float alpha, std::vector &pruned_list, + void prune_neighbors(const unsigned location, std::vector &pool, const uint32_t range, + const uint32_t max_candidate_size, const float alpha, std::vector &pruned_list, InMemQueryScratch *scratch); // Prunes candidates in @pool to a shorter list @result @@ -251,7 +253,8 @@ template clas const tsl::robin_set *const delete_set_ptr = nullptr); // add reverse links from all the visited nodes to node n. - void inter_insert(unsigned n, std::vector &pruned_list, const _u32 range, InMemQueryScratch *scratch); + void inter_insert(unsigned n, std::vector &pruned_list, const uint32_t range, + InMemQueryScratch *scratch); void inter_insert(unsigned n, std::vector &pruned_list, InMemQueryScratch *scratch); @@ -288,10 +291,10 @@ template clas // Do not call without acquiring appropriate locks // call public member functions save and load to invoke these. - DISKANN_DLLEXPORT _u64 save_graph(std::string filename); - DISKANN_DLLEXPORT _u64 save_data(std::string filename); - DISKANN_DLLEXPORT _u64 save_tags(std::string filename); - DISKANN_DLLEXPORT _u64 save_delete_list(const std::string &filename); + DISKANN_DLLEXPORT size_t save_graph(std::string filename); + DISKANN_DLLEXPORT size_t save_data(std::string filename); + DISKANN_DLLEXPORT size_t save_tags(std::string filename); + DISKANN_DLLEXPORT size_t save_delete_list(const std::string &filename); #ifdef EXEC_ENV_OLS DISKANN_DLLEXPORT size_t load_graph(AlignedFileReader &reader, size_t expected_num_points); DISKANN_DLLEXPORT size_t load_data(AlignedFileReader &reader); @@ -350,8 +353,8 @@ template clas std::vector> _pts_to_labels; tsl::robin_set _labels; std::string _labels_file; - std::unordered_map _label_to_medoid_id; - std::unordered_map<_u32, _u32> _medoid_counts; + std::unordered_map _label_to_medoid_id; + std::unordered_map _medoid_counts; bool _use_universal_label = false; LabelT _universal_label = 0; uint32_t _filterIndexingQueueSize; @@ -370,7 +373,7 @@ template clas bool _pq_dist = false; bool _use_opq = false; size_t _num_pq_chunks = 0; - _u8 *_pq_data = nullptr; + char *_pq_data = nullptr; bool _pq_generated = false; FixedChunkPQTable _pq_table; diff --git a/include/natural_number_map.h b/include/natural_number_map.h index dd4b09e41..3bbdcb640 100644 --- a/include/natural_number_map.h +++ b/include/natural_number_map.h @@ -35,9 +35,9 @@ template class natural_number_map struct position { size_t _key; - // The number of keys that were enumerated when iterating through the map - // so far. Used to early-terminate enumeration when ithere - // are no more entries in the map. + // The number of keys that were enumerated when iterating through the + // map so far. Used to early-terminate enumeration when ithere are no + // more entries in the map. size_t _keys_already_enumerated; // Returns whether it's valid to access the element at this position in diff --git a/include/pq.h b/include/pq.h index 9f6e50ee4..ff36af7ac 100644 --- a/include/pq.h +++ b/include/pq.h @@ -17,10 +17,10 @@ namespace diskann class FixedChunkPQTable { float *tables = nullptr; // pq_tables = float array of size [256 * ndims] - _u64 ndims = 0; // ndims = true dimension of vectors - _u64 n_chunks = 0; + size_t ndims = 0; // ndims = true dimension of vectors + size_t n_chunks = 0; bool use_rotation = false; - _u32 *chunk_offsets = nullptr; + uint32_t *chunk_offsets = nullptr; float *centroid = nullptr; float *tables_tr = nullptr; // same as pq_tables, but col-major float *rotmat_tr = nullptr; @@ -36,19 +36,19 @@ class FixedChunkPQTable void load_pq_centroid_bin(const char *pq_table_file, size_t num_chunks); #endif - _u32 get_num_chunks(); + 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, _u8 *base_vec); + float l2_distance(const float *query_vec, char *base_vec); - float inner_product(const float *query_vec, _u8 *base_vec); + float inner_product(const float *query_vec, char *base_vec); // assumes no rotation is involved - void inflate_vector(_u8 *base_vec, float *out_vec); + void inflate_vector(char *base_vec, float *out_vec); void populate_chunk_inner_products(const float *query_vec, float *dist_vec); }; @@ -57,16 +57,17 @@ template struct PQScratch { float *aligned_pqtable_dist_scratch = nullptr; // MUST BE AT LEAST [256 * NCHUNKS] float *aligned_dist_scratch = nullptr; // MUST BE AT LEAST diskann MAX_DEGREE - _u8 *aligned_pq_coord_scratch = nullptr; // MUST BE AT LEAST [N_CHUNKS * MAX_DEGREE] + char *aligned_pq_coord_scratch = nullptr; // MUST BE AT LEAST [N_CHUNKS * MAX_DEGREE] float *rotated_query = nullptr; float *aligned_query_float = nullptr; PQScratch(size_t graph_degree, size_t aligned_dim) { diskann::alloc_aligned((void **)&aligned_pq_coord_scratch, - (_u64)graph_degree * (_u64)MAX_PQ_CHUNKS * sizeof(_u8), 256); - diskann::alloc_aligned((void **)&aligned_pqtable_dist_scratch, 256 * (_u64)MAX_PQ_CHUNKS * sizeof(float), 256); - diskann::alloc_aligned((void **)&aligned_dist_scratch, (_u64)graph_degree * sizeof(float), 256); + (size_t)graph_degree * (size_t)MAX_PQ_CHUNKS * sizeof(char), 256); + diskann::alloc_aligned((void **)&aligned_pqtable_dist_scratch, 256 * (size_t)MAX_PQ_CHUNKS * sizeof(float), + 256); + diskann::alloc_aligned((void **)&aligned_dist_scratch, (size_t)graph_degree * sizeof(float), 256); diskann::alloc_aligned((void **)&aligned_query_float, aligned_dim * sizeof(float), 8 * sizeof(float)); diskann::alloc_aligned((void **)&rotated_query, aligned_dim * sizeof(float), 8 * sizeof(float)); @@ -86,15 +87,15 @@ template struct PQScratch } }; -void aggregate_coords(const std::vector &ids, const _u8 *all_coords, const _u64 ndims, _u8 *out); +void aggregate_coords(const std::vector &ids, const char *all_coords, const size_t ndims, char *out); -void pq_dist_lookup(const _u8 *pq_ids, const _u64 n_pts, const _u64 pq_nchunks, const float *pq_dists, +void pq_dist_lookup(const char *pq_ids, const size_t n_pts, const size_t pq_nchunks, const float *pq_dists, std::vector &dists_out); // Need to replace calls to these with calls to vector& based functions above -void aggregate_coords(const unsigned *ids, const _u64 n_ids, const _u8 *all_coords, const _u64 ndims, _u8 *out); +void aggregate_coords(const unsigned *ids, const size_t n_ids, const char *all_coords, const size_t ndims, char *out); -void pq_dist_lookup(const _u8 *pq_ids, const _u64 n_pts, const _u64 pq_nchunks, const float *pq_dists, +void pq_dist_lookup(const char *pq_ids, const size_t n_pts, const size_t pq_nchunks, const float *pq_dists, float *dists_out); DISKANN_DLLEXPORT int generate_pq_pivots(const float *const train_data, size_t num_train, unsigned dim, diff --git a/include/pq_flash_index.h b/include/pq_flash_index.h index f166219ac..4f1815105 100644 --- a/include/pq_flash_index.h +++ b/include/pq_flash_index.h @@ -48,45 +48,47 @@ template class PQFlashIndex #ifdef EXEC_ENV_OLS DISKANN_DLLEXPORT void 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, + uint64_t l_search, uint64_t beamwidth, + uint64_t num_nodes_to_cache, uint32_t nthreads, std::vector &node_list); #else - DISKANN_DLLEXPORT void generate_cache_list_from_sample_queries(std::string sample_bin, _u64 l_search, - _u64 beamwidth, _u64 num_nodes_to_cache, + DISKANN_DLLEXPORT void generate_cache_list_from_sample_queries(std::string sample_bin, uint64_t l_search, + uint64_t beamwidth, uint64_t num_nodes_to_cache, uint32_t num_threads, std::vector &node_list); #endif - DISKANN_DLLEXPORT void cache_bfs_levels(_u64 num_nodes_to_cache, std::vector &node_list, + DISKANN_DLLEXPORT void cache_bfs_levels(uint64_t num_nodes_to_cache, std::vector &node_list, const bool shuffle = false); - 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, + DISKANN_DLLEXPORT void cached_beam_search(const T *query, const uint64_t k_search, const uint64_t l_search, + uint64_t *res_ids, float *res_dists, const uint64_t beam_width, const bool use_reorder_data = false, QueryStats *stats = nullptr); - 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, const bool use_filter, - const LabelT &filter_label, const bool use_reorder_data = false, - QueryStats *stats = nullptr); - - 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, const _u32 io_limit, + DISKANN_DLLEXPORT void cached_beam_search(const T *query, const uint64_t k_search, const uint64_t l_search, + uint64_t *res_ids, float *res_dists, const uint64_t beam_width, + const bool use_filter, const LabelT &filter_label, const bool use_reorder_data = false, QueryStats *stats = nullptr); - 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, const bool use_filter, - const LabelT &filter_label, const _u32 io_limit, - const bool use_reorder_data = false, QueryStats *stats = nullptr); + DISKANN_DLLEXPORT void cached_beam_search(const T *query, const uint64_t k_search, const uint64_t l_search, + uint64_t *res_ids, float *res_dists, const uint64_t beam_width, + const uint32_t io_limit, const bool use_reorder_data = false, + QueryStats *stats = nullptr); + + DISKANN_DLLEXPORT void cached_beam_search(const T *query, const uint64_t k_search, const uint64_t l_search, + uint64_t *res_ids, float *res_dists, const uint64_t beam_width, + const bool use_filter, const LabelT &filter_label, + const uint32_t io_limit, const bool use_reorder_data = false, + QueryStats *stats = nullptr); DISKANN_DLLEXPORT LabelT get_converted_label(const std::string &filter_label); - DISKANN_DLLEXPORT _u32 range_search(const T *query1, const double range, const _u64 min_l_search, - const _u64 max_l_search, std::vector<_u64> &indices, - std::vector &distances, const _u64 min_beam_width, - QueryStats *stats = nullptr); + DISKANN_DLLEXPORT uint32_t range_search(const T *query1, const double range, const uint64_t min_l_search, + const uint64_t max_l_search, std::vector &indices, + std::vector &distances, const uint64_t min_beam_width, + QueryStats *stats = nullptr); - DISKANN_DLLEXPORT _u64 get_data_dim(); + DISKANN_DLLEXPORT uint64_t get_data_dim(); std::shared_ptr &reader; @@ -94,15 +96,15 @@ template class PQFlashIndex protected: DISKANN_DLLEXPORT void use_medoids_data_as_centroids(); - DISKANN_DLLEXPORT void setup_thread_data(_u64 nthreads, _u64 visited_reserve = 4096); + DISKANN_DLLEXPORT void setup_thread_data(uint64_t nthreads, uint64_t visited_reserve = 4096); DISKANN_DLLEXPORT void set_universal_label(const LabelT &label); private: - DISKANN_DLLEXPORT inline bool point_has_label(_u32 point_id, _u32 label_id); + DISKANN_DLLEXPORT inline bool point_has_label(uint32_t point_id, uint32_t label_id); std::unordered_map load_label_map(const std::string &map_file); DISKANN_DLLEXPORT void parse_label_file(const std::string &map_file, size_t &num_pts_labels); - DISKANN_DLLEXPORT void get_label_file_metadata(std::string map_file, _u32 &num_pts, _u32 &num_total_labels); + DISKANN_DLLEXPORT void get_label_file_metadata(std::string map_file, uint32_t &num_pts, uint32_t &num_total_labels); DISKANN_DLLEXPORT inline int32_t get_filter_number(const LabelT &filter_label); // index info @@ -111,10 +113,10 @@ template class PQFlashIndex // nnbrs of node `i`: *(unsigned*) (buf) // nbrs of node `i`: ((unsigned*)buf) + 1 - _u64 max_node_len = 0, nnodes_per_sector = 0, max_degree = 0; + uint64_t max_node_len = 0, nnodes_per_sector = 0, max_degree = 0; // Data used for searching with re-order vectors - _u64 ndims_reorder_vecs = 0, reorder_data_start_sector = 0, nvecs_per_sector = 0; + uint64_t ndims_reorder_vecs = 0, reorder_data_start_sector = 0, nvecs_per_sector = 0; diskann::Metric metric = diskann::Metric::L2; @@ -123,25 +125,25 @@ template class PQFlashIndex float max_base_norm = 0.0f; // data info - _u64 num_points = 0; - _u64 num_frozen_points = 0; - _u64 frozen_location = 0; - _u64 data_dim = 0; - _u64 disk_data_dim = 0; // will be different from data_dim only if we use - // PQ for disk data (very large dimensionality) - _u64 aligned_dim = 0; - _u64 disk_bytes_per_point = 0; + uint64_t num_points = 0; + uint64_t num_frozen_points = 0; + uint64_t frozen_location = 0; + uint64_t data_dim = 0; + uint64_t disk_data_dim = 0; // will be different from data_dim only if we use + // PQ for disk data (very large dimensionality) + uint64_t aligned_dim = 0; + uint64_t disk_bytes_per_point = 0; std::string disk_index_file; - std::vector> node_visit_counter; + std::vector> node_visit_counter; // PQ data // n_chunks = # of chunks ndims is split into - // data: _u8 * n_chunks + // data: char * n_chunks // chunk_size = chunk size of each dimension chunk // pq_tables = float* [[2^8 * [chunk_size]] * n_chunks] - _u8 *data = nullptr; - _u64 n_chunks; + char *data = nullptr; + uint64_t n_chunks; FixedChunkPQTable pq_table; // distance comparator @@ -150,7 +152,7 @@ template class PQFlashIndex // for very large datasets: we use PQ even for the disk resident index bool use_disk_index_pq = false; - _u64 disk_pq_n_chunks = 0; + uint64_t disk_pq_n_chunks = 0; FixedChunkPQTable disk_pq_table; // medoid/start info @@ -167,32 +169,32 @@ template class PQFlashIndex // nhood_cache unsigned *nhood_cache_buf = nullptr; - tsl::robin_map<_u32, std::pair<_u32, _u32 *>> nhood_cache; + tsl::robin_map> nhood_cache; // coord_cache T *coord_cache_buf = nullptr; - tsl::robin_map<_u32, T *> coord_cache; + tsl::robin_map coord_cache; // thread-specific scratch ConcurrentQueue *> thread_data; - _u64 max_nthreads; + uint64_t max_nthreads; bool load_flag = false; bool count_visited_nodes = false; bool reorder_data_exists = false; - _u64 reoreder_data_offset = 0; + uint64_t reoreder_data_offset = 0; // filter support - _u32 *_pts_to_label_offsets = nullptr; - _u32 *_pts_to_labels = nullptr; + uint32_t *_pts_to_label_offsets = nullptr; + uint32_t *_pts_to_labels = nullptr; tsl::robin_set _labels; - std::unordered_map _filter_to_medoid_id; + std::unordered_map _filter_to_medoid_id; bool _use_universal_label; - _u32 _universal_filter_num; + uint32_t _universal_filter_num; std::vector _filter_list; - tsl::robin_set<_u32> _dummy_pts; - tsl::robin_set<_u32> _has_dummy_pts; - tsl::robin_map<_u32, _u32> _dummy_to_real_map; - tsl::robin_map<_u32, std::vector<_u32>> _real_to_dummy_map; + tsl::robin_set _dummy_pts; + tsl::robin_set _has_dummy_pts; + tsl::robin_map _dummy_to_real_map; + tsl::robin_map> _real_to_dummy_map; std::unordered_map _label_map; #ifdef EXEC_ENV_OLS diff --git a/include/scratch.h b/include/scratch.h index 92348ecc5..bd44fbb07 100644 --- a/include/scratch.h +++ b/include/scratch.h @@ -22,7 +22,7 @@ // SSD Index related limits #define MAX_GRAPH_DEGREE 512 #define MAX_N_CMPS 16384 -#define SECTOR_LEN (_u64)4096 +#define SECTOR_LEN (size_t)4096 #define MAX_N_SECTOR_READS 128 namespace diskann @@ -152,16 +152,16 @@ template class SSDQueryScratch { public: T *coord_scratch = nullptr; // MUST BE AT LEAST [MAX_N_CMPS * data_dim] - _u64 coord_idx = 0; // index of next [data_dim] scratch to use + size_t coord_idx = 0; // index of next [data_dim] scratch to use char *sector_scratch = nullptr; // MUST BE AT LEAST [MAX_N_SECTOR_READS * SECTOR_LEN] - _u64 sector_idx = 0; // index of next [SECTOR_LEN] scratch to use + size_t sector_idx = 0; // index of next [SECTOR_LEN] scratch to use T *aligned_query_T = nullptr; PQScratch *_pq_scratch; - tsl::robin_set<_u64> visited; + tsl::robin_set visited; NeighborPriorityQueue retset; std::vector full_retset; diff --git a/include/types.h b/include/types.h new file mode 100644 index 000000000..ea04cd34d --- /dev/null +++ b/include/types.h @@ -0,0 +1,12 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#pragma once + +#include +#include + +namespace diskann +{ +typedef uint32_t location_t; +} // namespace diskann \ No newline at end of file diff --git a/include/utils.h b/include/utils.h index b5d5874c8..728fe9a82 100644 --- a/include/utils.h +++ b/include/utils.h @@ -2,6 +2,7 @@ // Licensed under the MIT license. #pragma once + #include #include "common_includes.h" @@ -25,16 +26,13 @@ typedef int FileHandle; #include "ann_exception.h" #include "windows_customizations.h" #include "tsl/robin_set.h" +#include "types.h" #ifdef EXEC_ENV_OLS #include "content_buf.h" #include "memory_mapped_files.h" #endif -#include -#include -#include - // taken from // https://github.com/Microsoft/BLAS-on-flash/blob/master/include/utils.h // round up X to the nearest multiple of Y @@ -53,7 +51,7 @@ typedef int FileHandle; 4096 // all metadata of individual sub-component files is written in first // 4KB for unified files -#define BUFFER_SIZE_FOR_CACHED_IO (_u64)1024 * (_u64)1048576 +#define BUFFER_SIZE_FOR_CACHED_IO (uint64_t)1024 * (uint64_t)1048576 inline bool file_exists(const std::string &name, bool dirCheck = false) { @@ -63,8 +61,8 @@ inline bool file_exists(const std::string &name, bool dirCheck = false) val = stat(name.c_str(), &buffer); #else // It is the 21st century but Windows API still thinks in 32-bit terms. - // Turns out calling stat() on a file > 4GB results in errno = 132 (OVERFLOW). - // How silly is this!? So calling _stat64() + // Turns out calling stat() on a file > 4GB results in errno = 132 + // (OVERFLOW). How silly is this!? So calling _stat64() struct _stat64 buffer; val = _stat64(name.c_str(), &buffer); #endif @@ -92,15 +90,6 @@ inline bool file_exists(const std::string &name, bool dirCheck = false) } } -typedef uint64_t _u64; -typedef int64_t _s64; -typedef uint32_t _u32; -typedef int32_t _s32; -typedef uint16_t _u16; -typedef int16_t _s16; -typedef uint8_t _u8; -typedef int8_t _s8; - inline void open_file_to_write(std::ofstream &writer, const std::string &filename) { writer.exceptions(std::ofstream::failbit | std::ofstream::badbit); @@ -122,12 +111,12 @@ inline void open_file_to_write(std::ofstream &writer, const std::string &filenam } } -inline _u64 get_file_size(const std::string &fname) +inline uint64_t 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(); + uint64_t end_pos = reader.tellg(); reader.close(); return end_pos; } @@ -146,7 +135,8 @@ inline int delete_file(const std::string &fileName) if (rc != 0) { diskann::cerr << "Could not delete file: " << fileName - << " even though it exists. This might indicate a permissions issue. " + << " even though it exists. This might indicate a permissions " + "issue. " "If you see this message, please contact the diskann team." << std::endl; } @@ -161,7 +151,7 @@ inline int delete_file(const std::string &fileName) inline void convert_labels_string_to_int(const std::string &inFileName, const std::string &outFileName, const std::string &mapFileName, const std::string &unv_label) { - std::unordered_map string_int_map; + std::unordered_map string_int_map; std::ofstream label_writer(outFileName); std::ifstream label_reader(inFileName); if (unv_label != "") @@ -170,14 +160,14 @@ inline void convert_labels_string_to_int(const std::string &inFileName, const st while (std::getline(label_reader, line)) { std::istringstream new_iss(line); - std::vector<_u32> lbls; + std::vector lbls; while (getline(new_iss, token, ',')) { token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); if (string_int_map.find(token) == string_int_map.end()) { - _u32 nextId = (_u32)string_int_map.size() + 1; + uint32_t nextId = (uint32_t)string_int_map.size() + 1; string_int_map[token] = nextId; } lbls.push_back(string_int_map[token]); @@ -442,7 +432,7 @@ inline void wait_for_keystroke() 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; + uint64_t read_blk_size = 64 * 1024 * 1024; cached_ifstream reader(bin_file, read_blk_size); diskann::cout << "Reading truthset file " << bin_file.c_str() << " ..." << std::endl; size_t actual_file_size = reader.get_file_size(); @@ -490,9 +480,9 @@ inline void load_truthset(const std::string &bin_file, uint32_t *&ids, float *&d } inline void prune_truthset_for_range(const std::string &bin_file, float range, - std::vector> &groundtruth, size_t &npts) + std::vector> &groundtruth, size_t &npts) { - _u64 read_blk_size = 64 * 1024 * 1024; + uint64_t read_blk_size = 64 * 1024 * 1024; cached_ifstream reader(bin_file, read_blk_size); diskann::cout << "Reading truthset file " << bin_file.c_str() << "... " << std::endl; size_t actual_file_size = reader.get_file_size(); @@ -501,8 +491,8 @@ inline void prune_truthset_for_range(const std::string &bin_file, float range, reader.read((char *)&npts_i32, sizeof(int)); reader.read((char *)&dim_i32, sizeof(int)); npts = (unsigned)npts_i32; - _u64 dim = (unsigned)dim_i32; - _u32 *ids; + uint64_t dim = (unsigned)dim_i32; + uint32_t *ids; float *dists; diskann::cout << "Metadata: #pts = " << npts << ", #dims = " << dim << "... " << std::endl; @@ -536,10 +526,10 @@ inline void prune_truthset_for_range(const std::string &bin_file, float range, float min_dist = std::numeric_limits::max(); float max_dist = 0; groundtruth.resize(npts); - for (_u32 i = 0; i < npts; i++) + for (uint32_t i = 0; i < npts; i++) { groundtruth[i].clear(); - for (_u32 j = 0; j < dim; j++) + for (uint32_t j = 0; j < dim; j++) { if (dists[i * dim + j] <= range) { @@ -555,23 +545,24 @@ inline void prune_truthset_for_range(const std::string &bin_file, float range, delete[] dists; } -inline void load_range_truthset(const std::string &bin_file, std::vector> &groundtruth, _u64 >_num) +inline void load_range_truthset(const std::string &bin_file, std::vector> &groundtruth, + uint64_t >_num) { - _u64 read_blk_size = 64 * 1024 * 1024; + uint64_t read_blk_size = 64 * 1024 * 1024; cached_ifstream reader(bin_file, read_blk_size); diskann::cout << "Reading truthset file " << bin_file.c_str() << "... " << std::flush; size_t actual_file_size = reader.get_file_size(); - int npts_u32, total_u32; - reader.read((char *)&npts_u32, sizeof(int)); - reader.read((char *)&total_u32, sizeof(int)); + int nptsuint32_t, totaluint32_t; + reader.read((char *)&nptsuint32_t, sizeof(int)); + reader.read((char *)&totaluint32_t, sizeof(int)); - gt_num = (_u64)npts_u32; - _u64 total_res = (_u64)total_u32; + gt_num = (uint64_t)nptsuint32_t; + uint64_t total_res = (uint64_t)totaluint32_t; diskann::cout << "Metadata: #pts = " << gt_num << ", #total_results = " << total_res << "..." << std::endl; - size_t expected_file_size = 2 * sizeof(_u32) + gt_num * sizeof(_u32) + total_res * sizeof(_u32); + size_t expected_file_size = 2 * sizeof(uint32_t) + gt_num * sizeof(uint32_t) + total_res * sizeof(uint32_t); if (actual_file_size != expected_file_size) { @@ -583,26 +574,26 @@ inline void load_range_truthset(const std::string &bin_file, std::vector gt_count(gt_num); + std::vector gt_count(gt_num); - reader.read((char *)gt_count.data(), sizeof(_u32) * gt_num); + reader.read((char *)gt_count.data(), sizeof(uint32_t) * gt_num); - std::vector<_u32> gt_stats(gt_count); + std::vector gt_stats(gt_count); std::sort(gt_stats.begin(), gt_stats.end()); std::cout << "GT count percentiles:" << std::endl; - for (_u32 p = 0; p < 100; p += 5) + for (uint32_t p = 0; p < 100; p += 5) std::cout << "percentile " << p << ": " << gt_stats[static_cast(std::floor((p / 100.0) * gt_num))] << std::endl; std::cout << "percentile 100" << ": " << gt_stats[gt_num - 1] << std::endl; - for (_u32 i = 0; i < gt_num; i++) + for (uint32_t i = 0; i < gt_num; i++) { groundtruth[i].clear(); groundtruth[i].resize(gt_count[i]); if (gt_count[i] != 0) - reader.read((char *)groundtruth[i].data(), sizeof(_u32) * gt_count[i]); + reader.read((char *)groundtruth[i].data(), sizeof(uint32_t) * gt_count[i]); } } @@ -639,8 +630,8 @@ DISKANN_DLLEXPORT double calculate_recall(unsigned num_queries, unsigned *gold_s const tsl::robin_set &active_tags); DISKANN_DLLEXPORT double calculate_range_search_recall(unsigned num_queries, - std::vector> &groundtruth, - std::vector> &our_results); + std::vector> &groundtruth, + std::vector> &our_results); template inline void load_bin(const std::string &bin_file, std::unique_ptr &data, size_t &npts, size_t &dim, @@ -778,7 +769,7 @@ template void convert_types(const InType *srcmat, OutType *destmat, size_t npts, size_t dim) { #pragma omp parallel for schedule(static, 65536) - for (int64_t i = 0; i < (_s64)npts; i++) + for (int64_t i = 0; i < (int64_t)npts; i++) { for (uint64_t j = 0; j < dim; j++) { @@ -800,17 +791,17 @@ template float prepare_base_for_inner_products(const std::string in std::cout << "Pre-processing base file by adding extra coordinate" << std::endl; std::ifstream in_reader(in_file.c_str(), std::ios::binary); std::ofstream out_writer(out_file.c_str(), std::ios::binary); - _u64 npts, in_dims, out_dims; + uint64_t npts, in_dims, out_dims; float max_norm = 0; - _u32 npts32, dims32; + uint32_t npts32, dims32; in_reader.read((char *)&npts32, sizeof(uint32_t)); in_reader.read((char *)&dims32, sizeof(uint32_t)); npts = npts32; in_dims = dims32; out_dims = in_dims + 1; - _u32 outdims32 = (_u32)out_dims; + uint32_t outdims32 = (uint32_t)out_dims; out_writer.write((char *)&npts32, sizeof(uint32_t)); out_writer.write((char *)&outdims32, sizeof(uint32_t)); @@ -821,19 +812,19 @@ template float prepare_base_for_inner_products(const std::string in std::unique_ptr out_block_data = std::make_unique(block_size * out_dims); std::memset(out_block_data.get(), 0, sizeof(float) * block_size * out_dims); - _u64 num_blocks = DIV_ROUND_UP(npts, block_size); + uint64_t num_blocks = DIV_ROUND_UP(npts, block_size); std::vector norms(npts, 0); - for (_u64 b = 0; b < num_blocks; b++) + for (uint64_t b = 0; b < num_blocks; b++) { - _u64 start_id = b * block_size; - _u64 end_id = (b + 1) * block_size < npts ? (b + 1) * block_size : npts; - _u64 block_pts = end_id - start_id; + uint64_t start_id = b * block_size; + uint64_t end_id = (b + 1) * block_size < npts ? (b + 1) * block_size : npts; + uint64_t block_pts = end_id - start_id; in_reader.read((char *)in_block_data.get(), block_pts * in_dims * sizeof(T)); - for (_u64 p = 0; p < block_pts; p++) + for (uint64_t p = 0; p < block_pts; p++) { - for (_u64 j = 0; j < in_dims; j++) + for (uint64_t j = 0; j < in_dims; j++) { norms[start_id + p] += in_block_data[p * in_dims + j] * in_block_data[p * in_dims + j]; } @@ -843,16 +834,16 @@ template float prepare_base_for_inner_products(const std::string in max_norm = std::sqrt(max_norm); - in_reader.seekg(2 * sizeof(_u32), std::ios::beg); - for (_u64 b = 0; b < num_blocks; b++) + in_reader.seekg(2 * sizeof(uint32_t), std::ios::beg); + for (uint64_t b = 0; b < num_blocks; b++) { - _u64 start_id = b * block_size; - _u64 end_id = (b + 1) * block_size < npts ? (b + 1) * block_size : npts; - _u64 block_pts = end_id - start_id; + uint64_t start_id = b * block_size; + uint64_t end_id = (b + 1) * block_size < npts ? (b + 1) * block_size : npts; + uint64_t block_pts = end_id - start_id; in_reader.read((char *)in_block_data.get(), block_pts * in_dims * sizeof(T)); - for (_u64 p = 0; p < block_pts; p++) + for (uint64_t p = 0; p < block_pts; p++) { - for (_u64 j = 0; j < in_dims; j++) + for (uint64_t j = 0; j < in_dims; j++) { out_block_data[p * out_dims + j] = in_block_data[p * in_dims + j] / max_norm; } @@ -874,13 +865,13 @@ template void save_Tvecs(const char *filename, T *data, size_t npts // create cached ofstream with 64MB cache cached_ofstream writer(fname, 64 * 1048576); - unsigned dims_u32 = (unsigned)ndims; + unsigned dimsuint32_t = (unsigned)ndims; // start writing for (uint64_t i = 0; i < npts; i++) { // write dims in u32 - writer.write((char *)&dims_u32, sizeof(unsigned)); + writer.write((char *)&dimsuint32_t, sizeof(unsigned)); // get cur point in data T *cur_pt = data + i * ndims; @@ -894,7 +885,7 @@ inline uint64_t save_data_in_base_dimensions(const std::string &filename, T *dat 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); + uint64_t 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)); @@ -952,7 +943,7 @@ inline void prefetch_vector_l2(const char *vec, size_t vecsize) } // NOTE: Implementation in utils.cpp. -void block_convert(std::ofstream &writr, std::ifstream &readr, float *read_buf, _u64 npts, _u64 ndims); +void block_convert(std::ofstream &writr, std::ifstream &readr, float *read_buf, uint64_t npts, uint64_t ndims); DISKANN_DLLEXPORT void normalize_data_file(const std::string &inFileName, const std::string &outFileName); diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index d993ccc69..1cebbf6bd 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -7,7 +7,7 @@ if(MSVC) add_subdirectory(dll) else() #file(GLOB CPP_SOURCES *.cpp) - set(CPP_SOURCES ann_exception.cpp disk_utils.cpp distance.cpp index.cpp + set(CPP_SOURCES ann_exception.cpp disk_utils.cpp distance.cpp index.cpp linux_aligned_file_reader.cpp math_utils.cpp natural_number_map.cpp natural_number_set.cpp memory_mapper.cpp partition.cpp pq.cpp pq_flash_index.cpp scratch.cpp logger.cpp utils.cpp) diff --git a/src/disk_utils.cpp b/src/disk_utils.cpp index aaf3dd71d..fa15348f9 100644 --- a/src/disk_utils.cpp +++ b/src/disk_utils.cpp @@ -24,9 +24,9 @@ namespace diskann 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); + std::unique_ptr metadata; + uint64_t nr, nc; + diskann::load_bin(index_file, metadata, nr, nc); if (nc != 1) { std::stringstream stream; @@ -34,9 +34,9 @@ void add_new_file_to_single_index(std::string index_file, std::string new_file) throw diskann::ANNException(stream.str(), -1); } size_t index_ending_offset = metadata[nr - 1]; - _u64 read_blk_size = 64 * 1024 * 1024; + uint64_t read_blk_size = 64 * 1024 * 1024; cached_ofstream writer(index_file, read_blk_size); - _u64 check_file_size = get_file_size(index_file); + uint64_t check_file_size = get_file_size(index_file); if (check_file_size != index_ending_offset) { std::stringstream stream; @@ -57,7 +57,7 @@ void add_new_file_to_single_index(std::string index_file, std::string new_file) 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++) + for (uint64_t 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; @@ -68,12 +68,12 @@ void add_new_file_to_single_index(std::string index_file, std::string new_file) // writer.close(); delete[] dump; - std::vector<_u64> new_meta; - for (_u64 i = 0; i < nr; i++) + std::vector new_meta; + for (uint64_t 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); + diskann::save_bin(index_file, new_meta.data(), new_meta.size(), 1); } double get_memory_budget(double search_ram_budget) @@ -96,7 +96,7 @@ double get_memory_budget(const std::string &mem_budget_str) size_t calculate_num_pq_chunks(double final_index_ram_limit, size_t points_num, uint32_t dim, const std::vector ¶m_list) { - size_t num_pq_chunks = (size_t)(std::floor)(_u64(final_index_ram_limit / (double)points_num)); + size_t num_pq_chunks = (size_t)(std::floor)(uint64_t(final_index_ram_limit / (double)points_num)); diskann::cout << "Calculated num_pq_chunks :" << num_pq_chunks << std::endl; if (param_list.size() >= 6) { @@ -239,27 +239,27 @@ void read_idmap(const std::string &fname, std::vector &ivecs) } int merge_shards(const std::string &vamana_prefix, const std::string &vamana_suffix, const std::string &idmaps_prefix, - const std::string &idmaps_suffix, const _u64 nshards, unsigned max_degree, + const std::string &idmaps_suffix, const uint64_t nshards, unsigned max_degree, const std::string &output_vamana, const std::string &medoids_file, bool use_filters, const std::string &labels_to_medoids_file) { // Read ID maps std::vector vamana_names(nshards); std::vector> idmaps(nshards); - for (_u64 shard = 0; shard < nshards; shard++) + for (size_t shard = 0; shard < nshards; shard++) { vamana_names[shard] = vamana_prefix + std::to_string(shard) + vamana_suffix; read_idmap(idmaps_prefix + std::to_string(shard) + idmaps_suffix, idmaps[shard]); } // find max node id - _u64 nnodes = 0; - _u64 nelems = 0; + size_t nnodes = 0; + size_t nelems = 0; for (auto &idmap : idmaps) { for (auto &id : idmap) { - nnodes = std::max(nnodes, (_u64)id); + nnodes = std::max(nnodes, (size_t)id); } nelems += idmap.size(); } @@ -269,13 +269,13 @@ int merge_shards(const std::string &vamana_prefix, const std::string &vamana_suf // compute inverse map: node -> shards std::vector> node_shard; node_shard.reserve(nelems); - for (_u64 shard = 0; shard < nshards; shard++) + for (size_t shard = 0; shard < nshards; shard++) { diskann::cout << "Creating inverse map -- shard #" << shard << std::endl; - for (_u64 idx = 0; idx < idmaps[shard].size(); idx++) + for (size_t idx = 0; idx < idmaps[shard].size(); idx++) { - _u64 node_id = idmaps[shard][idx]; - node_shard.push_back(std::make_pair((_u32)node_id, (_u32)shard)); + size_t node_id = idmaps[shard][idx]; + node_shard.push_back(std::make_pair((uint32_t)node_id, (uint32_t)shard)); } } std::sort(node_shard.begin(), node_shard.end(), [](const auto &left, const auto &right) { @@ -287,9 +287,9 @@ int merge_shards(const std::string &vamana_prefix, const std::string &vamana_suf // combined file if (use_filters) { - std::unordered_map> global_label_to_medoids; + std::unordered_map> global_label_to_medoids; - for (_u64 i = 0; i < nshards; i++) + for (size_t i = 0; i < nshards; i++) { std::ifstream mapping_reader; std::string map_file = vamana_names[i] + "_labels_to_medoids.txt"; @@ -301,9 +301,9 @@ int merge_shards(const std::string &vamana_prefix, const std::string &vamana_suf while (std::getline(mapping_reader, line)) { std::istringstream iss(line); - _u32 cnt = 0; - _u32 medoid; - _u32 label; + uint32_t cnt = 0; + uint32_t medoid; + uint32_t label; while (std::getline(iss, token, ',')) { token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); @@ -329,7 +329,7 @@ int merge_shards(const std::string &vamana_prefix, const std::string &vamana_suf { mapping_writer << iter.first << ", "; auto &vec = iter.second; - for (_u32 idx = 0; idx < vec.size() - 1; idx++) + for (uint32_t idx = 0; idx < vec.size() - 1; idx++) { mapping_writer << vec[idx] << ", "; } @@ -340,7 +340,7 @@ int merge_shards(const std::string &vamana_prefix, const std::string &vamana_suf // create cached vamana readers std::vector vamana_readers(nshards); - for (_u64 i = 0; i < nshards; i++) + for (size_t i = 0; i < nshards; i++) { vamana_readers[i].open(vamana_names[i], BUFFER_SIZE_FOR_CACHED_IO); size_t expected_file_size; @@ -348,8 +348,8 @@ int merge_shards(const std::string &vamana_prefix, const std::string &vamana_suf } size_t vamana_metadata_size = - sizeof(_u64) + sizeof(_u32) + sizeof(_u32) + sizeof(_u64); // expected file size + max degree + medoid_id + - // frozen_point info + sizeof(uint64_t) + sizeof(uint32_t) + sizeof(uint32_t) + sizeof(uint64_t); // expected file size + max degree + + // medoid_id + frozen_point info // create cached vamana writers cached_ofstream merged_vamana_writer(output_vamana, BUFFER_SIZE_FOR_CACHED_IO); @@ -374,20 +374,20 @@ int merge_shards(const std::string &vamana_prefix, const std::string &vamana_suf merged_vamana_writer.write((char *)&output_width, sizeof(unsigned)); std::ofstream medoid_writer(medoids_file.c_str(), std::ios::binary); - _u32 nshards_u32 = (_u32)nshards; - _u32 one_val = 1; + uint32_t nshards_u32 = (uint32_t)nshards; + uint32_t one_val = 1; medoid_writer.write((char *)&nshards_u32, sizeof(uint32_t)); medoid_writer.write((char *)&one_val, sizeof(uint32_t)); - _u64 vamana_index_frozen = 0; // as of now the functionality to merge many overlapping vamana - // indices is supported only for bulk indices without frozen point. - // Hence the final index will also not have any frozen points. - for (_u64 shard = 0; shard < nshards; shard++) + uint64_t vamana_index_frozen = 0; // as of now the functionality to merge many overlapping vamana + // indices is supported only for bulk indices without frozen point. + // Hence the final index will also not have any frozen points. + for (uint64_t 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)); + vamana_readers[shard].read((char *)&vamana_index_frozen, sizeof(uint64_t)); assert(vamana_index_frozen == false); // rename medoid medoid = idmaps[shard][medoid]; @@ -397,7 +397,7 @@ int merge_shards(const std::string &vamana_prefix, const std::string &vamana_suf if (shard == (nshards - 1)) //--> uncomment if running hierarchical merged_vamana_writer.write((char *)&medoid, sizeof(unsigned)); } - merged_vamana_writer.write((char *)&merged_index_frozen, sizeof(_u64)); + merged_vamana_writer.write((char *)&merged_index_frozen, sizeof(uint64_t)); medoid_writer.close(); diskann::cout << "Starting merge" << std::endl; @@ -446,7 +446,7 @@ int merge_shards(const std::string &vamana_prefix, const std::string &vamana_suf if (shard_nnbrs > 0) vamana_readers[shard_id].read((char *)shard_nhood.data(), shard_nnbrs * sizeof(unsigned)); // rename nodes - for (_u64 j = 0; j < shard_nnbrs; j++) + for (uint64_t j = 0; j < shard_nnbrs; j++) { if (nhood_set[idmaps[shard_id][shard_nhood[j]]] == 0) { @@ -488,32 +488,32 @@ int merge_shards(const std::string &vamana_prefix, const std::string &vamana_suf the new nodes at the end. The dummy map contains the real graph id of the new nodes added to the graph */ template -void breakup_dense_points(const std::string data_file, const std::string labels_file, _u32 density, +void breakup_dense_points(const std::string data_file, const std::string labels_file, uint32_t density, const std::string out_data_file, const std::string out_labels_file, const std::string out_metadata_file) { std::string token, line; std::ifstream labels_stream(labels_file); T *data; - _u64 npts, ndims; + uint64_t npts, ndims; diskann::load_bin(data_file, data, npts, ndims); - std::unordered_map<_u32, _u32> dummy_pt_ids; - _u32 next_dummy_id = (_u32)npts; + std::unordered_map dummy_pt_ids; + uint32_t next_dummy_id = (uint32_t)npts; - _u32 point_cnt = 0; + uint32_t point_cnt = 0; std::vector> labels_per_point; labels_per_point.resize(npts); - _u32 dense_pts = 0; + uint32_t dense_pts = 0; if (labels_stream.is_open()) { while (getline(labels_stream, line)) { std::stringstream iss(line); - _u32 lbl_cnt = 0; - _u32 label_host = point_cnt; + uint32_t lbl_cnt = 0; + uint32_t label_host = point_cnt; while (getline(iss, token, ',')) { if (lbl_cnt == density) @@ -522,7 +522,7 @@ void breakup_dense_points(const std::string data_file, const std::string labels_ dense_pts++; label_host = next_dummy_id; labels_per_point.resize(next_dummy_id + 1); - dummy_pt_ids[next_dummy_id] = (_u32)point_cnt; + dummy_pt_ids[next_dummy_id] = (uint32_t)point_cnt; next_dummy_id++; lbl_cnt = 0; } @@ -543,9 +543,9 @@ void breakup_dense_points(const std::string data_file, const std::string labels_ diskann::cout << labels_per_point.size() << " is the new number of points" << std::endl; std::ofstream label_writer(out_labels_file); assert(label_writer.is_open()); - for (_u32 i = 0; i < labels_per_point.size(); i++) + for (uint32_t i = 0; i < labels_per_point.size(); i++) { - for (_u32 j = 0; j < (labels_per_point[i].size() - 1); j++) + for (uint32_t j = 0; j < (labels_per_point[i].size() - 1); j++) { label_writer << labels_per_point[i][j] << ","; } @@ -579,11 +579,11 @@ void extract_shard_labels(const std::string &in_label_file, const std::string &s // point in labels file diskann::cout << "Extracting labels for shard" << std::endl; - _u32 *ids = nullptr; - _u64 num_ids, tmp_dim; + uint32_t *ids = nullptr; + uint64_t num_ids, tmp_dim; diskann::load_bin(shard_ids_bin, ids, num_ids, tmp_dim); - _u32 counter = 0, shard_counter = 0; + uint32_t counter = 0, shard_counter = 0; std::string cur_line; std::ifstream label_reader(in_label_file); @@ -616,7 +616,7 @@ int build_merged_vamana_index(std::string base_file, diskann::Metric compareMetr std::string medoids_file, std::string centroids_file, size_t build_pq_bytes, bool use_opq, bool use_filters, const std::string &label_file, const std::string &labels_to_medoids_file, const std::string &universal_label, - const _u32 Lf) + const uint32_t Lf) { size_t base_num, base_dim; diskann::get_bin_metadata(base_file, base_num, base_dim); @@ -659,7 +659,8 @@ int build_merged_vamana_index(std::string base_file, diskann::Metric compareMetr if (use_filters) { - // need to copy the labels_to_medoids file to the specified input file + // need to copy the labels_to_medoids file to the specified input + // file std::remove(labels_to_medoids_file.c_str()); std::string mem_labels_to_medoid_file = mem_index_path + "_labels_to_medoids.txt"; copy_file(mem_labels_to_medoid_file, labels_to_medoids_file); @@ -707,7 +708,7 @@ int build_merged_vamana_index(std::string base_file, diskann::Metric compareMetr paras.Set("saturate_graph", 0); paras.Set("save_path", shard_index_file); - _u64 shard_base_dim, shard_base_pts; + uint64_t 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_dim, shard_base_pts, false, false, false, @@ -781,8 +782,8 @@ int build_merged_vamana_index(std::string base_file, diskann::Metric compareMetr // 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) + uint64_t tuning_sample_num, uint64_t tuning_sample_aligned_dim, uint32_t L, + uint32_t nthreads, uint32_t start_bw) { uint32_t cur_bw = start_bw; double max_qps = 0; @@ -797,7 +798,7 @@ uint32_t optimize_beamwidth(std::unique_ptr> &p 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++) + for (int64_t 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), @@ -838,8 +839,8 @@ void create_disk_layout(const std::string base_file, const std::string mem_index 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; + uint64_t read_blk_size = 64 * 1024 * 1024; + uint64_t write_blk_size = read_blk_size; cached_ifstream base_reader(base_file, read_blk_size); base_reader.read((char *)&npts, sizeof(uint32_t)); base_reader.read((char *)&ndims, sizeof(uint32_t)); @@ -865,8 +866,9 @@ void create_disk_layout(const std::string base_file, const std::string mem_index reorder_data_reader.read((char *)&npts_reorder_file, sizeof(unsigned)); reorder_data_reader.read((char *)&ndims_reorder_file, sizeof(unsigned)); if (npts_reorder_file != npts) - throw ANNException("Mismatch in num_points between reorder data file and base file", -1, __FUNCSIG__, - __FILE__, __LINE__); + throw ANNException("Mismatch in num_points between reorder " + "data file and base file", + -1, __FUNCSIG__, __FILE__, __LINE__); if (reorder_data_file_size != 8 + sizeof(float) * (size_t)npts_reorder_file * (size_t)ndims_reorder_file) throw ANNException("Discrepancy in reorder data file size ", -1, __FUNCSIG__, __FILE__, __LINE__); } @@ -896,18 +898,18 @@ void create_disk_layout(const std::string base_file, const std::string mem_index throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } - _u64 vamana_frozen_num = false, vamana_frozen_loc = 0; + uint64_t 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)); + vamana_reader.read((char *)&vamana_frozen_num, sizeof(uint64_t)); // compute - _u64 medoid, max_node_len, nnodes_per_sector; - npts_64 = (_u64)npts; - medoid = (_u64)medoid_u32; + uint64_t medoid, max_node_len, nnodes_per_sector; + npts_64 = (uint64_t)npts; + medoid = (uint64_t)medoid_u32; if (vamana_frozen_num == 1) vamana_frozen_loc = medoid; - max_node_len = (((_u64)width_u32 + 1) * sizeof(unsigned)) + (ndims_64 * sizeof(T)); + max_node_len = (((uint64_t)width_u32 + 1) * sizeof(unsigned)) + (ndims_64 * sizeof(T)); nnodes_per_sector = SECTOR_LEN / max_node_len; diskann::cout << "medoid: " << medoid << "B" << std::endl; @@ -921,18 +923,18 @@ void create_disk_layout(const std::string base_file, const std::string mem_index 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_reorder_sectors = 0; - _u64 n_data_nodes_per_sector = 0; + uint64_t n_sectors = ROUND_UP(npts_64, nnodes_per_sector) / nnodes_per_sector; + uint64_t n_reorder_sectors = 0; + uint64_t n_data_nodes_per_sector = 0; if (append_reorder_data) { n_data_nodes_per_sector = SECTOR_LEN / (ndims_reorder_file * sizeof(float)); n_reorder_sectors = ROUND_UP(npts_64, n_data_nodes_per_sector) / n_data_nodes_per_sector; } - _u64 disk_index_file_size = (n_sectors + n_reorder_sectors + 1) * SECTOR_LEN; + uint64_t disk_index_file_size = (n_sectors + n_reorder_sectors + 1) * SECTOR_LEN; - std::vector<_u64> output_file_meta; + std::vector output_file_meta; output_file_meta.push_back(npts_64); output_file_meta.push_back(ndims_64); output_file_meta.push_back(medoid); @@ -940,7 +942,7 @@ void create_disk_layout(const std::string base_file, const std::string mem_index 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((_u64)append_reorder_data); + output_file_meta.push_back((uint64_t)append_reorder_data); if (append_reorder_data) { output_file_meta.push_back(n_sectors + 1); @@ -953,15 +955,15 @@ void create_disk_layout(const std::string base_file, const std::string mem_index std::unique_ptr cur_node_coords = std::make_unique(ndims_64); diskann::cout << "# sectors: " << n_sectors << std::endl; - _u64 cur_node_id = 0; - for (_u64 sector = 0; sector < n_sectors; sector++) + uint64_t cur_node_id = 0; + for (uint64_t sector = 0; sector < n_sectors; sector++) { if (sector % 100000 == 0) { diskann::cout << "Sector #" << sector << "written" << std::endl; } memset(sector_buf.get(), 0, SECTOR_LEN); - for (_u64 sector_node_id = 0; sector_node_id < nnodes_per_sector && cur_node_id < npts_64; sector_node_id++) + for (uint64_t sector_node_id = 0; sector_node_id < nnodes_per_sector && cur_node_id < npts_64; sector_node_id++) { memset(node_buf.get(), 0, max_node_len); // read cur node's nnbrs @@ -979,7 +981,7 @@ void create_disk_layout(const std::string base_file, const std::string mem_index } // write coords of node first - // T *node_coords = data + ((_u64) ndims_64 * cur_node_id); + // T *node_coords = data + ((uint64_t) ndims_64 * cur_node_id); base_reader.read((char *)cur_node_coords.get(), sizeof(T) * ndims_64); memcpy(node_buf.get(), cur_node_coords.get(), ndims_64 * sizeof(T)); @@ -1007,7 +1009,7 @@ void create_disk_layout(const std::string base_file, const std::string mem_index auto vec_len = ndims_reorder_file * sizeof(float); std::unique_ptr vec_buf = std::make_unique(vec_len); - for (_u64 sector = 0; sector < n_reorder_sectors; sector++) + for (uint64_t sector = 0; sector < n_reorder_sectors; sector++) { if (sector % 100000 == 0) { @@ -1016,7 +1018,7 @@ void create_disk_layout(const std::string base_file, const std::string mem_index memset(sector_buf.get(), 0, SECTOR_LEN); - for (_u64 sector_node_id = 0; sector_node_id < n_data_nodes_per_sector && sector_node_id < npts_64; + for (uint64_t sector_node_id = 0; sector_node_id < n_data_nodes_per_sector && sector_node_id < npts_64; sector_node_id++) { memset(vec_buf.get(), 0, vec_len); @@ -1030,14 +1032,14 @@ void create_disk_layout(const std::string base_file, const std::string mem_index } } diskann_writer.close(); - diskann::save_bin<_u64>(output_file, output_file_meta.data(), output_file_meta.size(), 1, 0); + diskann::save_bin(output_file, output_file_meta.data(), output_file_meta.size(), 1, 0); diskann::cout << "Output disk index file written to " << output_file << std::endl; } template int build_disk_index(const char *dataFilePath, const char *indexFilePath, const char *indexBuildParameters, diskann::Metric compareMetric, bool use_opq, bool use_filters, const std::string &label_file, - const std::string &universal_label, const _u32 filter_threshold, const _u32 Lf) + const std::string &universal_label, const uint32_t filter_threshold, const uint32_t Lf) { std::stringstream parser; parser << std::string(indexBuildParameters); @@ -1165,7 +1167,7 @@ int build_disk_index(const char *dataFilePath, const char *indexFilePath, const std::cerr << "Not building index. Please provide more RAM budget" << std::endl; return -1; } - _u32 num_threads = (_u32)atoi(param_list[4].c_str()); + uint32_t num_threads = (uint32_t)atoi(param_list[4].c_str()); if (num_threads != 0) { @@ -1193,8 +1195,8 @@ int build_disk_index(const char *dataFilePath, const char *indexFilePath, const dummy_remap_file = index_prefix_path + "_dummy_remap.txt"; breakup_dense_points(data_file_to_use, labels_file_to_use, filter_threshold, augmented_data_file, augmented_labels_file, - dummy_remap_file); // RKNOTE: This has large memory footprint, need - // to make this streaming + dummy_remap_file); // RKNOTE: This has large memory footprint, + // need to make this streaming data_file_to_use = augmented_data_file; labels_file_to_use = augmented_labels_file; } @@ -1211,7 +1213,7 @@ int build_disk_index(const char *dataFilePath, const char *indexFilePath, const generate_disk_quantized_data(data_file_to_use, disk_pq_pivots_path, disk_pq_compressed_vectors_path, compareMetric, p_val, disk_pq_dims); } - size_t num_pq_chunks = (size_t)(std::floor)(_u64(final_index_ram_limit / points_num)); + size_t num_pq_chunks = (size_t)(std::floor)(uint64_t(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; @@ -1245,10 +1247,10 @@ int build_disk_index(const char *dataFilePath, const char *indexFilePath, const else { if (!reorder_data) - diskann::create_disk_layout<_u8>(disk_pq_compressed_vectors_path, mem_index_path, disk_index_path); + diskann::create_disk_layout(disk_pq_compressed_vectors_path, mem_index_path, disk_index_path); else - diskann::create_disk_layout<_u8>(disk_pq_compressed_vectors_path, mem_index_path, disk_index_path, - data_file_to_use.c_str()); + diskann::create_disk_layout(disk_pq_compressed_vectors_path, mem_index_path, disk_index_path, + data_file_to_use.c_str()); } diskann::cout << timer.elapsed_seconds_for_step("generating disk layout") << std::endl; @@ -1315,91 +1317,91 @@ template DISKANN_DLLEXPORT float *load_warmup(MemoryMappedFiles &files, c 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); + uint64_t tuning_sample_num, uint64_t 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); + uint64_t tuning_sample_num, uint64_t 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); + std::unique_ptr> &pFlashIndex, float *tuning_sample, + uint64_t tuning_sample_num, uint64_t 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, int8_t *tuning_sample, - _u64 tuning_sample_num, _u64 tuning_sample_aligned_dim, uint32_t L, uint32_t nthreads, uint32_t start_bw); + uint64_t tuning_sample_num, uint64_t 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); + uint64_t tuning_sample_num, uint64_t 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); + std::unique_ptr> &pFlashIndex, float *tuning_sample, + uint64_t tuning_sample_num, uint64_t tuning_sample_aligned_dim, uint32_t L, uint32_t nthreads, uint32_t start_bw); template DISKANN_DLLEXPORT int build_disk_index(const char *dataFilePath, const char *indexFilePath, const char *indexBuildParameters, diskann::Metric compareMetric, bool use_opq, bool use_filters, const std::string &label_file, const std::string &universal_label, - const _u32 filter_threshold, const _u32 Lf); + const uint32_t filter_threshold, const uint32_t Lf); template DISKANN_DLLEXPORT int build_disk_index(const char *dataFilePath, const char *indexFilePath, const char *indexBuildParameters, diskann::Metric compareMetric, bool use_opq, bool use_filters, const std::string &label_file, const std::string &universal_label, - const _u32 filter_threshold, const _u32 Lf); + const uint32_t filter_threshold, const uint32_t Lf); template DISKANN_DLLEXPORT int build_disk_index(const char *dataFilePath, const char *indexFilePath, const char *indexBuildParameters, diskann::Metric compareMetric, bool use_opq, bool use_filters, const std::string &label_file, const std::string &universal_label, - const _u32 filter_threshold, const _u32 Lf); + const uint32_t filter_threshold, const uint32_t Lf); // LabelT = uint16 template DISKANN_DLLEXPORT int build_disk_index(const char *dataFilePath, const char *indexFilePath, const char *indexBuildParameters, diskann::Metric compareMetric, bool use_opq, bool use_filters, const std::string &label_file, const std::string &universal_label, - const _u32 filter_threshold, const _u32 Lf); + const uint32_t filter_threshold, const uint32_t Lf); template DISKANN_DLLEXPORT int build_disk_index(const char *dataFilePath, const char *indexFilePath, const char *indexBuildParameters, diskann::Metric compareMetric, bool use_opq, bool use_filters, const std::string &label_file, const std::string &universal_label, - const _u32 filter_threshold, const _u32 Lf); + const uint32_t filter_threshold, const uint32_t Lf); template DISKANN_DLLEXPORT int build_disk_index(const char *dataFilePath, const char *indexFilePath, const char *indexBuildParameters, diskann::Metric compareMetric, bool use_opq, bool use_filters, const std::string &label_file, const std::string &universal_label, - const _u32 filter_threshold, const _u32 Lf); + const uint32_t filter_threshold, const uint32_t Lf); 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, size_t build_pq_bytes, bool use_opq, bool use_filters, const std::string &label_file, - const std::string &labels_to_medoids_file, const std::string &universal_label, const _u32 Lf); + const std::string &labels_to_medoids_file, const std::string &universal_label, const uint32_t Lf); 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, size_t build_pq_bytes, bool use_opq, bool use_filters, const std::string &label_file, - const std::string &labels_to_medoids_file, const std::string &universal_label, const _u32 Lf); + const std::string &labels_to_medoids_file, const std::string &universal_label, const uint32_t Lf); 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, size_t build_pq_bytes, bool use_opq, bool use_filters, const std::string &label_file, - const std::string &labels_to_medoids_file, const std::string &universal_label, const _u32 Lf); + const std::string &labels_to_medoids_file, const std::string &universal_label, const uint32_t Lf); // Label=16_t 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, size_t build_pq_bytes, bool use_opq, bool use_filters, const std::string &label_file, - const std::string &labels_to_medoids_file, const std::string &universal_label, const _u32 Lf); + const std::string &labels_to_medoids_file, const std::string &universal_label, const uint32_t Lf); 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, size_t build_pq_bytes, bool use_opq, bool use_filters, const std::string &label_file, - const std::string &labels_to_medoids_file, const std::string &universal_label, const _u32 Lf); + const std::string &labels_to_medoids_file, const std::string &universal_label, const uint32_t Lf); 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, size_t build_pq_bytes, bool use_opq, bool use_filters, const std::string &label_file, - const std::string &labels_to_medoids_file, const std::string &universal_label, const _u32 Lf); + const std::string &labels_to_medoids_file, const std::string &universal_label, const uint32_t Lf); }; // namespace diskann diff --git a/src/distance.cpp b/src/distance.cpp index 15d30c8cb..ecf6597ae 100644 --- a/src/distance.cpp +++ b/src/distance.cpp @@ -616,7 +616,8 @@ template <> diskann::Distance *get_distance_function(diskann::Metric m) if (m == diskann::Metric::L2) { #ifdef _WINDOWS - diskann::cout << "WARNING: AVX/AVX2 distance function not defined for Uint8. Using " + diskann::cout << "WARNING: AVX/AVX2 distance function not defined for Uint8. " + "Using " "slow version. " "Contact gopalsr@microsoft.com if you need AVX/AVX2 support." << std::endl; diff --git a/src/dll/CMakeLists.txt b/src/dll/CMakeLists.txt index f9abe1e15..ca62a8f0e 100644 --- a/src/dll/CMakeLists.txt +++ b/src/dll/CMakeLists.txt @@ -2,7 +2,8 @@ #Licensed under the MIT license. add_library(${PROJECT_NAME} SHARED dllmain.cpp ../partition.cpp ../pq.cpp ../pq_flash_index.cpp ../logger.cpp ../utils.cpp - ../windows_aligned_file_reader.cpp ../distance.cpp ../memory_mapper.cpp ../index.cpp ../math_utils.cpp ../disk_utils.cpp + ../windows_aligned_file_reader.cpp ../distance.cpp ../memory_mapper.cpp ../index.cpp + ../math_utils.cpp ../disk_utils.cpp ../ann_exception.cpp ../natural_number_set.cpp ../natural_number_map.cpp ../scratch.cpp) set(TARGET_DIR "$<$:${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}>$<$:${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}>") diff --git a/src/index.cpp b/src/index.cpp index 17a285992..41d5e7e95 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -69,7 +69,8 @@ Index::Index(Metric m, const size_t dim, const size_t max_point "index construction", -1, __FUNCSIG__, __FILE__, __LINE__); if (m == diskann::Metric::INNER_PRODUCT) - throw ANNException("ERROR: Inner product metrics not yet supported with PQ distance " + throw ANNException("ERROR: Inner product metrics not yet supported " + "with PQ distance " "base index", -1, __FUNCSIG__, __FILE__, __LINE__); } @@ -172,7 +173,7 @@ void Index::initialize_query_scratch(uint32_t num_threads, uint } } -template _u64 Index::save_tags(std::string tags_file) +template size_t Index::save_tags(std::string tags_file) { if (!_enable_tags) { @@ -181,7 +182,7 @@ template _u64 Index _u64 Index _u64 Index::save_data(std::string data_file) +template size_t Index::save_data(std::string data_file) { - // Note: at this point, either _nd == _max_points or any frozen points have been - // temporarily moved to _nd, so _nd + _num_frozen_points is the valid location limit. + // Note: at this point, either _nd == _max_points or any frozen points have + // been temporarily moved to _nd, so _nd + _num_frozen_points is the valid + // location limit. return save_data_in_base_dimensions(data_file, _data, _nd + _num_frozen_pts, _dim, _aligned_dim); } // 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) +template size_t Index::save_graph(std::string graph_file) { std::ofstream out; open_file_to_write(out, graph_file); - _u64 file_offset = 0; // we will use this if we want + size_t file_offset = 0; // we will use this if we want out.seekp(file_offset, out.beg); - _u64 index_size = 24; - _u32 max_degree = 0; + size_t index_size = 24; + uint32_t max_degree = 0; out.write((char *)&index_size, sizeof(uint64_t)); out.write((char *)&_max_observed_degree, sizeof(unsigned)); unsigned ep_u32 = _start; out.write((char *)&ep_u32, sizeof(unsigned)); - out.write((char *)&_num_frozen_pts, sizeof(_u64)); - // Note: at this point, either _nd == _max_points or any frozen points have been - // temporarily moved to _nd, so _nd + _num_frozen_points is the valid location limit. + out.write((char *)&_num_frozen_pts, sizeof(size_t)); + // Note: at this point, either _nd == _max_points or any frozen points have + // been temporarily moved to _nd, so _nd + _num_frozen_points is the valid + // location limit. 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)); + max_degree = _final_graph[i].size() > max_degree ? (uint32_t)_final_graph[i].size() : max_degree; + index_size += (size_t)(sizeof(unsigned) * (GK + 1)); } out.seekp(file_offset, out.beg); out.write((char *)&index_size, sizeof(uint64_t)); - out.write((char *)&max_degree, sizeof(_u32)); + out.write((char *)&max_degree, sizeof(uint32_t)); out.close(); return index_size; // number of bytes written } template -_u64 Index::save_delete_list(const std::string &filename) +size_t Index::save_delete_list(const std::string &filename) { if (_delete_set->size() == 0) { return 0; } - std::unique_ptr<_u32[]> delete_list = std::make_unique<_u32[]>(_delete_set->size()); - _u32 i = 0; + std::unique_ptr delete_list = std::make_unique(_delete_set->size()); + uint32_t i = 0; for (auto &del : *_delete_set) { delete_list[i++] = del; } - return save_bin<_u32>(filename, delete_list.get(), _delete_set->size(), 1); + return save_bin(filename, delete_list.get(), _delete_set->size(), 1); } template @@ -321,9 +324,9 @@ void Index::save(const char *filename, bool compact_before_save { std::ofstream label_writer(std::string(filename) + "_labels.txt"); assert(label_writer.is_open()); - for (_u32 i = 0; i < _pts_to_labels.size(); i++) + for (uint32_t i = 0; i < _pts_to_labels.size(); i++) { - for (_u32 j = 0; j < (_pts_to_labels[i].size() - 1); j++) + for (uint32_t j = 0; j < (_pts_to_labels[i].size() - 1); j++) { label_writer << _pts_to_labels[i][j] << ","; } @@ -360,7 +363,8 @@ void Index::save(const char *filename, bool compact_before_save << std::endl; } - // If frozen points were temporarily compacted to _nd, move back to _max_points. + // If frozen points were temporarily compacted to _nd, move back to + // _max_points. reposition_frozen_point_to_end(); diskann::cout << "Time taken for save: " << timer.elapsed() / 1000000.0 << "s." << std::endl; @@ -407,7 +411,7 @@ size_t Index::load_tags(const std::string tag_filename) const size_t num_data_points = file_num_points - _num_frozen_pts; _location_to_tag.reserve(num_data_points); _tag_to_location.reserve(num_data_points); - for (_u32 i = 0; i < (_u32)num_data_points; i++) + for (uint32_t i = 0; i < (uint32_t)num_data_points; i++) { TagT tag = *(tag_data + i); if (_delete_set->find(i) == _delete_set->end()) @@ -480,13 +484,13 @@ template size_t Index::load_delete_set(const std::string &filename) { #endif - std::unique_ptr<_u32[]> delete_list; - _u64 npts, ndim; + std::unique_ptr delete_list; + size_t npts, ndim; #ifdef EXEC_ENV_OLS - diskann::load_bin<_u32>(reader, delete_list, npts, ndim); + diskann::load_bin(reader, delete_list, npts, ndim); #else - diskann::load_bin<_u32>(filename, delete_list, npts, ndim); + diskann::load_bin(filename, delete_list, npts, ndim); #endif assert(ndim == 1); for (uint32_t i = 0; i < npts; i++) @@ -522,7 +526,8 @@ void Index::load(const char *filename, uint32_t num_threads, ui if (!_save_as_one_file) { - // For DLVS Store, we will not support saving the index in multiple files. + // For DLVS Store, we will not support saving the index in multiple + // files. #ifndef EXEC_ENV_OLS std::string data_file = std::string(filename) + ".data"; std::string tags_file = std::string(filename) + ".tags"; @@ -576,8 +581,8 @@ void Index::load(const char *filename, uint32_t num_threads, ui while (std::getline(medoid_stream, line)) { std::istringstream iss(line); - _u32 cnt = 0; - _u32 medoid = 0; + uint32_t cnt = 0; + uint32_t medoid = 0; LabelT label; while (std::getline(iss, token, ',')) { @@ -638,16 +643,16 @@ size_t Index::get_graph_num_frozen_points(const std::string &gr { size_t expected_file_size; unsigned max_observed_degree, start; - _u64 file_frozen_pts; + size_t file_frozen_pts; std::ifstream in; in.exceptions(std::ios::badbit | std::ios::failbit); in.open(graph_file, std::ios::binary); - in.read((char *)&expected_file_size, sizeof(_u64)); + in.read((char *)&expected_file_size, sizeof(size_t)); in.read((char *)&max_observed_degree, sizeof(unsigned)); in.read((char *)&start, sizeof(unsigned)); - in.read((char *)&file_frozen_pts, sizeof(_u64)); + in.read((char *)&file_frozen_pts, sizeof(size_t)); return file_frozen_pts; } @@ -664,29 +669,29 @@ size_t Index::load_graph(std::string filename, size_t expected_ { #endif size_t expected_file_size; - _u64 file_frozen_pts; + size_t file_frozen_pts; #ifdef EXEC_ENV_OLS - int header_size = 2 * sizeof(_u64) + 2 * sizeof(unsigned); + int header_size = 2 * sizeof(size_t) + 2 * sizeof(unsigned); std::unique_ptr header = std::make_unique(header_size); read_array(reader, header.get(), header_size); - expected_file_size = *((_u64 *)header.get()); - _max_observed_degree = *((_u32 *)(header.get() + sizeof(_u64))); - _start = *((_u32 *)(header.get() + sizeof(_u64) + sizeof(unsigned))); - file_frozen_pts = *((_u64 *)(header.get() + sizeof(_u64) + sizeof(unsigned) + sizeof(unsigned))); + expected_file_size = *((size_t *)header.get()); + _max_observed_degree = *((uint32_t *)(header.get() + sizeof(size_t))); + _start = *((uint32_t *)(header.get() + sizeof(size_t) + sizeof(unsigned))); + file_frozen_pts = *((size_t *)(header.get() + sizeof(size_t) + sizeof(unsigned) + sizeof(unsigned))); #else - _u64 file_offset = 0; // will need this for single file format support + size_t file_offset = 0; // will need this for single file format support std::ifstream in; in.exceptions(std::ios::badbit | std::ios::failbit); in.open(filename, std::ios::binary); in.seekg(file_offset, in.beg); - in.read((char *)&expected_file_size, sizeof(_u64)); + in.read((char *)&expected_file_size, sizeof(size_t)); in.read((char *)&_max_observed_degree, sizeof(unsigned)); in.read((char *)&_start, sizeof(unsigned)); - in.read((char *)&file_frozen_pts, sizeof(_u64)); - _u64 vamana_metadata_size = sizeof(_u64) + sizeof(_u32) + sizeof(_u32) + sizeof(_u64); + in.read((char *)&file_frozen_pts, sizeof(size_t)); + size_t vamana_metadata_size = sizeof(size_t) + sizeof(uint32_t) + sizeof(uint32_t) + sizeof(size_t); #endif diskann::cout << "From graph header, expected_file_size: " << expected_file_size @@ -732,18 +737,18 @@ size_t Index::load_graph(std::string filename, size_t expected_ _max_points = expected_max_points; } #ifdef EXEC_ENV_OLS - _u32 nodes_read = 0; - _u64 cc = 0; - _u64 graph_offset = header_size; + uint32_t nodes_read = 0; + size_t cc = 0; + size_t graph_offset = header_size; while (nodes_read < expected_num_points) { - _u32 k; + uint32_t k; read_value(reader, k, graph_offset); - graph_offset += sizeof(_u32); - std::vector<_u32> tmp(k); + graph_offset += sizeof(uint32_t); + std::vector tmp(k); tmp.reserve(k); read_array(reader, tmp.data(), k, graph_offset); - graph_offset += k * sizeof(_u32); + graph_offset += k * sizeof(uint32_t); cc += k; _final_graph[nodes_read].swap(tmp); nodes_read++; @@ -776,7 +781,7 @@ size_t Index::load_graph(std::string filename, size_t expected_ tmp.reserve(k); in.read((char *)tmp.data(), k * sizeof(unsigned)); _final_graph[nodes_read - 1].swap(tmp); - bytes_read += sizeof(uint32_t) * ((_u64)k + 1); + bytes_read += sizeof(uint32_t) * ((size_t)k + 1); if (nodes_read % 10000000 == 0) diskann::cout << "." << std::flush; if (k > _max_range_of_loaded_graph) @@ -829,7 +834,7 @@ template unsigned Index Index::iterate_to_fixed_point( float *query_float; float *query_rotated; float *pq_dists; - _u8 *pq_coord_scratch; + char *pq_coord_scratch; // Intialize PQ related scratch to use PQ based distances if (_pq_dist) { @@ -1119,9 +1124,10 @@ std::pair Index::iterate_to_fixed_point( } template -void Index::search_for_point_and_prune(int location, _u32 Lindex, std::vector &pruned_list, +void Index::search_for_point_and_prune(int location, uint32_t Lindex, + std::vector &pruned_list, InMemQueryScratch *scratch, bool use_filter, - _u32 filteredLindex) + uint32_t filteredLindex) { const std::vector init_ids = get_init_ids(); const std::vector unused_filter_label; @@ -1133,7 +1139,7 @@ void Index::search_for_point_and_prune(int location, _u32 Linde } else { - std::vector<_u32> filter_specific_start_nodes; + std::vector filter_specific_start_nodes; for (auto &x : _pts_to_labels[location]) filter_specific_start_nodes.emplace_back(_label_to_medoid_id[x]); iterate_to_fixed_point(_data + _aligned_dim * location, filteredLindex, filter_specific_start_nodes, scratch, @@ -1198,8 +1204,8 @@ void Index::occlude_list(const unsigned location, std::vector::max(); - // Add the entry to the result if its not been deleted, and doesn't add - // a self loop + // Add the entry to the result if its not been deleted, and doesn't + // add a self loop if (delete_set_ptr == nullptr || delete_set_ptr->find(iter->id) == delete_set_ptr->end()) { if (iter->id != location) @@ -1218,8 +1224,8 @@ void Index::occlude_list(const unsigned location, std::vectorid; - _u32 b = iter2->id; + uint32_t a = iter->id; + uint32_t b = iter2->id; for (auto &x : _pts_to_labels[b]) { if (std::find(_pts_to_labels[a].begin(), _pts_to_labels[a].end(), x) == _pts_to_labels[a].end()) @@ -1264,8 +1270,8 @@ void Index::prune_neighbors(const unsigned location, std::vecto } template -void Index::prune_neighbors(const unsigned location, std::vector &pool, const _u32 range, - const _u32 max_candidate_size, const float alpha, +void Index::prune_neighbors(const unsigned location, std::vector &pool, const uint32_t range, + const uint32_t max_candidate_size, const float alpha, std::vector &pruned_list, InMemQueryScratch *scratch) { if (pool.size() == 0) @@ -1306,7 +1312,7 @@ void Index::prune_neighbors(const unsigned location, std::vecto } template -void Index::inter_insert(unsigned n, std::vector &pruned_list, const _u32 range, +void Index::inter_insert(unsigned n, std::vector &pruned_list, const uint32_t range, InMemQueryScratch *scratch) { const auto &src_pool = pruned_list; @@ -1325,7 +1331,7 @@ void Index::inter_insert(unsigned n, std::vector &pru auto &des_pool = _final_graph[des]; if (std::find(des_pool.begin(), des_pool.end(), n) == des_pool.end()) { - if (des_pool.size() < (_u64)(GRAPH_SLACK_FACTOR * range)) + if (des_pool.size() < (size_t)(GRAPH_SLACK_FACTOR * range)) { des_pool.emplace_back(n); prune_needed = false; @@ -1423,7 +1429,7 @@ template void Index void Index void Index _indexingRange) @@ -1506,7 +1512,7 @@ void Index::prune_all_nbrs(const Parameters ¶meters) diskann::Timer timer; #pragma omp parallel for - for (_s64 node = 0; node < (_s64)(_max_points + _num_frozen_pts); node++) + for (int64_t node = 0; node < (int64_t)(_max_points + _num_frozen_pts); node++) { if ((size_t)node < _nd || (size_t)node >= _max_points) { @@ -1530,7 +1536,7 @@ void Index::prune_all_nbrs(const Parameters ¶meters) } } - prune_neighbors((_u32)node, dummy_pool, range, maxc, alpha, new_out_neighbors, scratch); + prune_neighbors((uint32_t)node, dummy_pool, range, maxc, alpha, new_out_neighbors, scratch); _final_graph[node].clear(); for (auto id : new_out_neighbors) _final_graph[node].emplace_back(id); @@ -1771,11 +1777,12 @@ void Index::build(const char *filename, const size_t num_points generate_quantized_data(std::string(filename), pq_pivots_file, pq_compressed_file, _dist_metric, p_val, _num_pq_chunks, _use_opq); - copy_aligned_data_from_file<_u8>(pq_compressed_file.c_str(), _pq_data, file_num_points, _num_pq_chunks, - _num_pq_chunks); + copy_aligned_data_from_file(pq_compressed_file.c_str(), _pq_data, file_num_points, _num_pq_chunks, + _num_pq_chunks); #ifdef EXEC_ENV_OLS - throw ANNException("load_pq_centroid_bin should not be called when EXEC_ENV_OLS is defined.", -1, __FUNCSIG__, - __FILE__, __LINE__); + throw ANNException("load_pq_centroid_bin should not be called when " + "EXEC_ENV_OLS is defined.", + -1, __FUNCSIG__, __FILE__, __LINE__); #else _pq_table.load_pq_centroid_bin(pq_pivots_file.c_str(), _num_pq_chunks); #endif @@ -1947,7 +1954,7 @@ void Index::build_filtered_index(const char *filename, const st num_points_labels); // determines medoid for each label and // identifies the points to label mapping - std::unordered_map> label_to_points; + std::unordered_map> label_to_points; for (int lbl = 0; lbl < _labels.size(); lbl++) { @@ -1955,8 +1962,8 @@ void Index::build_filtered_index(const char *filename, const st std::advance(itr, lbl); auto &x = *itr; - std::vector<_u32> labeled_points; - for (_u32 point_id = 0; point_id < num_points_to_load; point_id++) + std::vector labeled_points; + for (uint32_t point_id = 0; point_id < num_points_to_load; point_id++) { bool pt_has_lbl = std::find(_pts_to_labels[point_id].begin(), _pts_to_labels[point_id].end(), x) != _pts_to_labels[point_id].end(); @@ -1973,17 +1980,17 @@ void Index::build_filtered_index(const char *filename, const st label_to_points[x] = labeled_points; } - _u32 num_cands = 25; + uint32_t num_cands = 25; for (auto itr = _labels.begin(); itr != _labels.end(); itr++) { - _u32 best_medoid_count = std::numeric_limits<_u32>::max(); + uint32_t best_medoid_count = std::numeric_limits::max(); auto &curr_label = *itr; - _u32 best_medoid; + uint32_t best_medoid; auto labeled_points = label_to_points[curr_label]; - for (_u32 cnd = 0; cnd < num_cands; cnd++) + for (uint32_t cnd = 0; cnd < num_cands; cnd++) { - _u32 cur_cnd = labeled_points[rand() % labeled_points.size()]; - _u32 cur_cnt = std::numeric_limits<_u32>::max(); + uint32_t cur_cnd = labeled_points[rand() % labeled_points.size()]; + uint32_t cur_cnt = std::numeric_limits::max(); if (_medoid_counts.find(cur_cnd) == _medoid_counts.end()) { _medoid_counts[cur_cnd] = 0; @@ -2037,7 +2044,7 @@ std::pair Index::search(const T *query, con NeighborPriorityQueue &best_L_nodes = scratch->best_l_nodes(); size_t pos = 0; - for (int i = 0; i < best_L_nodes.size(); ++i) + for (size_t i = 0; i < best_L_nodes.size(); ++i) { if (best_L_nodes[i].id < _max_points) { @@ -2114,7 +2121,7 @@ std::pair Index::search_with_filters(const auto best_L_nodes = scratch->best_l_nodes(); size_t pos = 0; - for (int i = 0; i < best_L_nodes.size(); ++i) + for (size_t i = 0; i < best_L_nodes.size(); ++i) { if (best_L_nodes[i].id < _max_points) { @@ -2408,9 +2415,9 @@ consolidation_report Index::consolidate_deletes(const Parameter unsigned num_calls_to_process_delete = 0; diskann::Timer timer; #pragma omp parallel for num_threads(num_threads) schedule(dynamic, 8192) reduction(+ : num_calls_to_process_delete) - for (_s64 loc = 0; loc < (_s64)_max_points; loc++) + for (int64_t loc = 0; loc < (int64_t)_max_points; loc++) { - if (old_delete_set->find((_u32)loc) == old_delete_set->end() && !_empty_slots.is_in_set((_u32)loc)) + if (old_delete_set->find((uint32_t)loc) == old_delete_set->end() && !_empty_slots.is_in_set((uint32_t)loc)) { ScratchStoreManager> manager(_query_scratch); auto scratch = manager.scratch_space(); @@ -2418,7 +2425,7 @@ consolidation_report Index::consolidate_deletes(const Parameter num_calls_to_process_delete += 1; } } - for (_s64 loc = _max_points; loc < (_s64)(_max_points + _num_frozen_pts); loc++) + for (int64_t loc = _max_points; loc < (int64_t)(_max_points + _num_frozen_pts); loc++) { ScratchStoreManager> manager(_query_scratch); auto scratch = manager.scratch_space(); @@ -2452,7 +2459,7 @@ template void Index 0) { reposition_points(_max_points, _nd, _num_frozen_pts); - _start = (_u32)_nd; + _start = (uint32_t)_nd; } } @@ -2478,11 +2485,11 @@ template void Index new_location = std::vector(_max_points + _num_frozen_pts, (_u32)UINT32_MAX); + std::vector new_location = std::vector(_max_points + _num_frozen_pts, UINT32_MAX); - _u32 new_counter = 0; - std::set<_u32> empty_locations; - for (_u32 old_location = 0; old_location < _max_points; old_location++) + uint32_t new_counter = 0; + std::set empty_locations; + for (uint32_t old_location = 0; old_location < _max_points; old_location++) { if (_location_to_tag.contains(old_location)) { @@ -2494,7 +2501,7 @@ template void Index void Index::release_locations(const tsl::robin_set for (auto location : locations) { if (_empty_slots.is_in_set(location)) - throw ANNException("Trying to release location, but location already in empty slots", -1, __FUNCSIG__, - __FILE__, __LINE__); + throw ANNException("Trying to release location, but location " + "already in empty slots", + -1, __FUNCSIG__, __FILE__, __LINE__); _empty_slots.insert(location); _nd--; @@ -2642,9 +2650,9 @@ void Index::reposition_points(unsigned old_location_start, unsi return; } - // Update pointers to the moved nodes. Note: the computation is correct even when - // new_location_start < old_location_start given the C++ unsigned integer arithmetic - // rules. + // Update pointers to the moved nodes. Note: the computation is correct even + // when new_location_start < old_location_start given the C++ unsigned + // integer arithmetic rules. const unsigned location_delta = new_location_start - old_location_start; for (unsigned i = 0; i < _max_points + _num_frozen_pts; i++) @@ -2652,11 +2660,13 @@ void Index::reposition_points(unsigned old_location_start, unsi if (loc >= old_location_start && loc < old_location_start + num_locations) loc += location_delta; - // The [start, end) interval which will contain obsolete points to be cleared. + // The [start, end) interval which will contain obsolete points to be + // cleared. unsigned mem_clear_loc_start = old_location_start; unsigned mem_clear_loc_end_limit = old_location_start + num_locations; - // Move the adjacency lists. Make sure that overlapping ranges are handled correctly. + // Move the adjacency lists. Make sure that overlapping ranges are handled + // correctly. if (new_location_start < old_location_start) { // New location before the old location: copy the entries in order @@ -2667,7 +2677,8 @@ void Index::reposition_points(unsigned old_location_start, unsi _final_graph[new_location_start + loc_offset].swap(_final_graph[old_location_start + loc_offset]); } - // If ranges are overlapping, make sure not to clear the newly copied data. + // 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. @@ -2684,7 +2695,8 @@ void Index::reposition_points(unsigned old_location_start, unsi _final_graph[new_location_start + loc_offset - 1u].swap(_final_graph[old_location_start + loc_offset - 1u]); } - // If ranges are overlapping, make sure not to clear the newly copied data. + // 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. @@ -2710,8 +2722,8 @@ template void Index void Index::resize(size_t new_max_points) @@ -2734,8 +2746,8 @@ template void Index::insert_point(const T *point, const TagT tag) assert(_has_built); if (tag == static_cast(0)) { - throw diskann::ANNException("Do not insert point with tag 0. That is reserved for points hidden " + throw diskann::ANNException("Do not insert point with tag 0. That is " + "reserved for points hidden " "from the user.", -1, __FUNCSIG__, __FILE__, __LINE__); } @@ -2796,8 +2809,9 @@ int Index::insert_point(const T *point, const TagT tag) location = reserve_location(); if (location == -1) { - throw diskann::ANNException("Cannot reserve location even after expanding graph. Terminating.", -1, - __FUNCSIG__, __FILE__, __LINE__); + throw diskann::ANNException("Cannot reserve location even after " + "expanding graph. Terminating.", + -1, __FUNCSIG__, __FILE__, __LINE__); } #else return -1; @@ -2848,7 +2862,7 @@ int Index::insert_point(const T *point, const TagT tag) LockGuard guard(_locks[location]); _final_graph[location].clear(); - _final_graph[location].reserve((_u64)(_indexingRange * GRAPH_SLACK_FACTOR * 1.05)); + _final_graph[location].reserve((size_t)(_indexingRange * GRAPH_SLACK_FACTOR * 1.05)); for (auto link : pruned_list) { diff --git a/src/math_utils.cpp b/src/math_utils.cpp index 151e6a97d..abc9a240d 100644 --- a/src/math_utils.cpp +++ b/src/math_utils.cpp @@ -27,7 +27,7 @@ float calc_distance(float *vec_1, float *vec_2, size_t dim) void compute_vecs_l2sq(float *vecs_l2sq, float *data, const size_t num_points, const size_t dim) { #pragma omp parallel for schedule(static, 8192) - for (int64_t n_iter = 0; n_iter < (_s64)num_points; n_iter++) + for (int64_t n_iter = 0; n_iter < (uint64_t)num_points; n_iter++) { vecs_l2sq[n_iter] = cblas_snrm2((MKL_INT)dim, (data + (n_iter * dim)), 1); vecs_l2sq[n_iter] *= vecs_l2sq[n_iter]; @@ -96,7 +96,7 @@ void compute_closest_centers_in_block(const float *const data, const size_t num_ if (k == 1) { #pragma omp parallel for schedule(static, 8192) - for (int64_t i = 0; i < (_s64)num_points; i++) + for (int64_t i = 0; i < (uint64_t)num_points; i++) { float min = std::numeric_limits::max(); float *current = dist_matrix + (i * num_centers); @@ -113,7 +113,7 @@ void compute_closest_centers_in_block(const float *const data, const size_t num_ else { #pragma omp parallel for schedule(static, 8192) - for (int64_t i = 0; i < (_s64)num_points; i++) + for (int64_t i = 0; i < (uint64_t)num_points; i++) { std::priority_queue top_k_queue; float *current = dist_matrix + (i * num_centers); @@ -182,7 +182,7 @@ void compute_closest_centers(float *data, size_t num_points, size_t dim, float * #pragma omp parallel for schedule(static, 1) for (int64_t j = cur_blk * PAR_BLOCK_SIZE; - j < std::min((_s64)num_points, (_s64)((cur_blk + 1) * PAR_BLOCK_SIZE)); j++) + j < std::min((uint64_t)num_points, (uint64_t)((cur_blk + 1) * PAR_BLOCK_SIZE)); j++) { for (size_t l = 0; l < k; l++) { @@ -212,7 +212,7 @@ void process_residuals(float *data_load, size_t num_points, size_t dim, float *c 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 (int64_t n_iter = 0; n_iter < (uint64_t)num_points; n_iter++) { for (size_t d_iter = 0; d_iter < dim; d_iter++) { @@ -259,7 +259,7 @@ float lloyds_iter(float *data, size_t num_points, size_t dim, float *centers, si memset(centers, 0, sizeof(float) * (size_t)num_centers * (size_t)dim); #pragma omp parallel for schedule(static, 1) - for (int64_t c = 0; c < (_s64)num_centers; ++c) + for (int64_t c = 0; c < (uint64_t)num_centers; ++c) { float *center = centers + (size_t)c * (size_t)dim; double *cluster_sum = new double[dim]; @@ -290,7 +290,7 @@ float lloyds_iter(float *data, size_t num_points, size_t dim, float *centers, si std::vector residuals(nchunks * BUF_PAD, 0.0); #pragma omp parallel for schedule(static, 32) - for (int64_t chunk = 0; chunk < (_s64)nchunks; ++chunk) + for (int64_t chunk = 0; chunk < (uint64_t)nchunks; ++chunk) for (size_t d = chunk * CHUNK_SIZE; d < num_points && d < (chunk + 1) * CHUNK_SIZE; ++d) residuals[chunk * BUF_PAD] += math_utils::calc_distance(data + (d * dim), centers + (size_t)closest_center[d] * (size_t)dim, dim); @@ -405,7 +405,7 @@ void kmeanspp_selecting_pivots(float *data, size_t num_points, size_t dim, float float *dist = new float[num_points]; #pragma omp parallel for schedule(static, 8192) - for (int64_t i = 0; i < (_s64)num_points; i++) + for (int64_t i = 0; i < (uint64_t)num_points; i++) { dist[i] = math_utils::calc_distance(data + i * dim, data + init_id * dim, dim); } @@ -446,7 +446,7 @@ void kmeanspp_selecting_pivots(float *data, size_t num_points, size_t dim, float std::memcpy(pivot_data + num_picked * dim, data + tmp_pivot * dim, dim * sizeof(float)); #pragma omp parallel for schedule(static, 8192) - for (int64_t i = 0; i < (_s64)num_points; i++) + for (int64_t i = 0; i < (uint64_t)num_points; i++) { dist[i] = (std::min)(dist[i], math_utils::calc_distance(data + i * dim, data + tmp_pivot * dim, dim)); } diff --git a/src/partition.cpp b/src/partition.cpp index ece453e87..7b88badf0 100644 --- a/src/partition.cpp +++ b/src/partition.cpp @@ -33,7 +33,7 @@ template void gen_random_slice(const std::string base_file, const std::string output_prefix, double sampling_rate) { - _u64 read_blk_size = 64 * 1024 * 1024; + uint64_t read_blk_size = 64 * 1024 * 1024; cached_ifstream base_reader(base_file.c_str(), read_blk_size); std::ofstream sample_writer(std::string(output_prefix + "_data.bin").c_str(), std::ios::binary); std::ofstream sample_id_writer(std::string(output_prefix + "_ids.bin").c_str(), std::ios::binary); @@ -68,7 +68,7 @@ void gen_random_slice(const std::string base_file, const std::string output_pref if (sample < sampling_rate) { sample_writer.write((char *)cur_row.get(), sizeof(T) * nd); - uint32_t cur_i_u32 = (_u32)i; + uint32_t cur_i_u32 = (uint32_t)i; sample_id_writer.write((char *)&cur_i_u32, sizeof(uint32_t)); num_sampled_pts_u32++; } @@ -100,7 +100,7 @@ void gen_random_slice(const std::string data_file, double p_val, float *&sampled std::vector> sampled_vectors; // amount to read in one shot - _u64 read_blk_size = 64 * 1024 * 1024; + uint64_t read_blk_size = 64 * 1024 * 1024; // create cached reader + writer cached_ifstream base_reader(data_file.c_str(), read_blk_size); @@ -193,7 +193,7 @@ int estimate_cluster_sizes(float *test_data_float, size_t num_test, float *pivot } size_t block_size = num_test <= BLOCK_SIZE ? num_test : BLOCK_SIZE; - _u32 *block_closest_centers = new _u32[block_size * k_base]; + uint32_t *block_closest_centers = new uint32_t[block_size * k_base]; float *block_data_float; size_t num_blocks = DIV_ROUND_UP(num_test, block_size); @@ -222,7 +222,7 @@ int estimate_cluster_sizes(float *test_data_float, size_t num_test, float *pivot diskann::cout << "Estimated cluster sizes: "; for (size_t i = 0; i < num_centers; i++) { - _u32 cur_shard_count = (_u32)shard_counts[i]; + uint32_t cur_shard_count = (uint32_t)shard_counts[i]; cluster_sizes.push_back((size_t)cur_shard_count); diskann::cout << cur_shard_count << " "; } @@ -236,12 +236,12 @@ template int shard_data_into_clusters(const std::string data_file, float *pivots, const size_t num_centers, const size_t dim, const size_t k_base, std::string prefix_path) { - _u64 read_blk_size = 64 * 1024 * 1024; - // _u64 write_blk_size = 64 * 1024 * 1024; + uint64_t read_blk_size = 64 * 1024 * 1024; + // uint64_t write_blk_size = 64 * 1024 * 1024; // create cached reader + writer cached_ifstream base_reader(data_file, read_blk_size); - _u32 npts32; - _u32 basedim32; + uint32_t npts32; + uint32_t basedim32; base_reader.read((char *)&npts32, sizeof(uint32_t)); base_reader.read((char *)&basedim32, sizeof(uint32_t)); size_t num_points = npts32; @@ -254,8 +254,8 @@ int shard_data_into_clusters(const std::string data_file, float *pivots, const s std::unique_ptr shard_counts = std::make_unique(num_centers); std::vector shard_data_writer(num_centers); std::vector shard_idmap_writer(num_centers); - _u32 dummy_size = 0; - _u32 const_one = 1; + uint32_t dummy_size = 0; + uint32_t const_one = 1; for (size_t i = 0; i < num_centers; i++) { @@ -271,7 +271,7 @@ int shard_data_into_clusters(const std::string data_file, float *pivots, const s } 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); + std::unique_ptr block_closest_centers = std::make_unique(block_size * k_base); std::unique_ptr block_data_T = std::make_unique(block_size * dim); std::unique_ptr block_data_float = std::make_unique(block_size * dim); @@ -306,7 +306,7 @@ int shard_data_into_clusters(const std::string data_file, float *pivots, const s diskann::cout << "Actual shard sizes: " << std::flush; for (size_t i = 0; i < num_centers; i++) { - _u32 cur_shard_count = (_u32)shard_counts[i]; + uint32_t cur_shard_count = (uint32_t)shard_counts[i]; total_count += cur_shard_count; diskann::cout << cur_shard_count << " "; shard_data_writer[i].seekp(0); @@ -328,12 +328,12 @@ template int shard_data_into_clusters_only_ids(const std::string data_file, float *pivots, const size_t num_centers, const size_t dim, const size_t k_base, std::string prefix_path) { - _u64 read_blk_size = 64 * 1024 * 1024; - // _u64 write_blk_size = 64 * 1024 * 1024; + uint64_t read_blk_size = 64 * 1024 * 1024; + // uint64_t write_blk_size = 64 * 1024 * 1024; // create cached reader + writer cached_ifstream base_reader(data_file, read_blk_size); - _u32 npts32; - _u32 basedim32; + uint32_t npts32; + uint32_t basedim32; base_reader.read((char *)&npts32, sizeof(uint32_t)); base_reader.read((char *)&basedim32, sizeof(uint32_t)); size_t num_points = npts32; @@ -346,8 +346,8 @@ int shard_data_into_clusters_only_ids(const std::string data_file, float *pivots std::unique_ptr shard_counts = std::make_unique(num_centers); std::vector shard_idmap_writer(num_centers); - _u32 dummy_size = 0; - _u32 const_one = 1; + uint32_t dummy_size = 0; + uint32_t const_one = 1; for (size_t i = 0; i < num_centers; i++) { @@ -359,7 +359,7 @@ int shard_data_into_clusters_only_ids(const std::string data_file, float *pivots } 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); + std::unique_ptr block_closest_centers = std::make_unique(block_size * k_base); std::unique_ptr block_data_T = std::make_unique(block_size * dim); std::unique_ptr block_data_float = std::make_unique(block_size * dim); @@ -393,7 +393,7 @@ int shard_data_into_clusters_only_ids(const std::string data_file, float *pivots diskann::cout << "Actual shard sizes: " << std::flush; for (size_t i = 0; i < num_centers; i++) { - _u32 cur_shard_count = (_u32)shard_counts[i]; + uint32_t cur_shard_count = (uint32_t)shard_counts[i]; total_count += cur_shard_count; diskann::cout << cur_shard_count << " "; shard_idmap_writer[i].seekp(0); @@ -409,29 +409,29 @@ int shard_data_into_clusters_only_ids(const std::string data_file, float *pivots template int retrieve_shard_data_from_ids(const std::string data_file, std::string idmap_filename, std::string data_filename) { - _u64 read_blk_size = 64 * 1024 * 1024; - // _u64 write_blk_size = 64 * 1024 * 1024; + uint64_t read_blk_size = 64 * 1024 * 1024; + // uint64_t write_blk_size = 64 * 1024 * 1024; // create cached reader + writer cached_ifstream base_reader(data_file, read_blk_size); - _u32 npts32; - _u32 basedim32; + uint32_t npts32; + uint32_t basedim32; base_reader.read((char *)&npts32, sizeof(uint32_t)); base_reader.read((char *)&basedim32, sizeof(uint32_t)); size_t num_points = npts32; size_t dim = basedim32; - _u32 dummy_size = 0; + uint32_t dummy_size = 0; std::ofstream shard_data_writer(data_filename.c_str(), std::ios::binary); shard_data_writer.write((char *)&dummy_size, sizeof(uint32_t)); shard_data_writer.write((char *)&basedim32, sizeof(uint32_t)); - _u32 *shard_ids; - _u64 shard_size, tmp; - diskann::load_bin<_u32>(idmap_filename, shard_ids, shard_size, tmp); + uint32_t *shard_ids; + uint64_t shard_size, tmp; + diskann::load_bin(idmap_filename, shard_ids, shard_size, tmp); - _u32 cur_pos = 0; - _u32 num_written = 0; + uint32_t cur_pos = 0; + uint32_t num_written = 0; std::cout << "Shard has " << shard_size << " points" << std::endl; size_t block_size = num_points <= BLOCK_SIZE ? num_points : BLOCK_SIZE; @@ -495,7 +495,8 @@ int partition(const std::string data_file, const float sampling_rate, size_t num // kmeans_partitioning on training data - // cur_file = cur_file + "_kmeans_partitioning-" + std::to_string(num_parts); + // cur_file = cur_file + "_kmeans_partitioning-" + + // std::to_string(num_parts); output_file = cur_file + "_centroids.bin"; pivot_data = new float[num_parts * train_dim]; @@ -544,7 +545,8 @@ int partition_with_ram_budget(const std::string data_file, const double sampling // kmeans_partitioning on training data - // cur_file = cur_file + "_kmeans_partitioning-" + std::to_string(num_parts); + // cur_file = cur_file + "_kmeans_partitioning-" + + // std::to_string(num_parts); output_file = cur_file + "_centroids.bin"; while (!fit_in_ram) @@ -572,7 +574,7 @@ int partition_with_ram_budget(const std::string data_file, const double sampling { // to account for the fact that p is the size of the shard over the // testing sample. - p = (_u64)(p / sampling_rate); + p = (uint64_t)(p / sampling_rate); double cur_shard_ram_estimate = diskann::estimate_ram_usage(p, train_dim, sizeof(T), graph_degree); if (cur_shard_ram_estimate > max_ram_usage) diff --git a/src/pq.cpp b/src/pq.cpp index 2eb0a3cb0..20263f882 100644 --- a/src/pq.cpp +++ b/src/pq.cpp @@ -41,16 +41,16 @@ void FixedChunkPQTable::load_pq_centroid_bin(const char *pq_table_file, size_t n { #endif - _u64 nr, nc; + size_t nr, nc; std::string rotmat_file = std::string(pq_table_file) + "_rotation_matrix.bin"; #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_table_file, file_offset_data, nr, nc); + 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<_u64[]> file_offset_data; - diskann::load_bin<_u64>(pq_table_file, file_offset_data, nr, nc); + std::unique_ptr file_offset_data; + diskann::load_bin(pq_table_file, file_offset_data, nr, nc); #endif bool use_old_filetype = false; @@ -150,32 +150,32 @@ void FixedChunkPQTable::load_pq_centroid_bin(const char *pq_table_file, size_t n // alloc and compute transpose tables_tr = new float[256 * this->ndims]; - for (_u64 i = 0; i < 256; i++) + for (size_t i = 0; i < 256; i++) { - for (_u64 j = 0; j < this->ndims; j++) + for (size_t j = 0; j < this->ndims; j++) { tables_tr[j * 256 + i] = tables[i * this->ndims + j]; } } } -_u32 FixedChunkPQTable::get_num_chunks() +uint32_t FixedChunkPQTable::get_num_chunks() { - return static_cast<_u32>(n_chunks); + return static_cast(n_chunks); } void FixedChunkPQTable::preprocess_query(float *query_vec) { - for (_u32 d = 0; d < ndims; d++) + for (uint32_t d = 0; d < ndims; d++) { query_vec[d] -= centroid[d]; } std::vector tmp(ndims, 0); if (use_rotation) { - for (_u32 d = 0; d < ndims; d++) + for (uint32_t d = 0; d < ndims; d++) { - for (_u32 d1 = 0; d1 < ndims; d1++) + for (uint32_t d1 = 0; d1 < ndims; d1++) { tmp[d] += query_vec[d1] * rotmat_tr[d1 * ndims + d]; } @@ -189,14 +189,14 @@ void FixedChunkPQTable::populate_chunk_distances(const float *query_vec, float * { memset(dist_vec, 0, 256 * n_chunks * sizeof(float)); // chunk wise distance computation - for (_u64 chunk = 0; chunk < n_chunks; chunk++) + 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 (_u64 j = chunk_offsets[chunk]; j < chunk_offsets[chunk + 1]; j++) + for (size_t j = chunk_offsets[chunk]; j < chunk_offsets[chunk + 1]; j++) { const float *centers_dim_vec = tables_tr + (256 * j); - for (_u64 idx = 0; idx < 256; idx++) + for (size_t idx = 0; idx < 256; idx++) { double diff = centers_dim_vec[idx] - (query_vec[j]); chunk_dists[idx] += (float)(diff * diff); @@ -205,12 +205,12 @@ void FixedChunkPQTable::populate_chunk_distances(const float *query_vec, float * } } -float FixedChunkPQTable::l2_distance(const float *query_vec, _u8 *base_vec) +float FixedChunkPQTable::l2_distance(const float *query_vec, char *base_vec) { float res = 0; - for (_u64 chunk = 0; chunk < n_chunks; chunk++) + for (size_t chunk = 0; chunk < n_chunks; chunk++) { - for (_u64 j = chunk_offsets[chunk]; j < chunk_offsets[chunk + 1]; j++) + 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]); @@ -220,12 +220,12 @@ float FixedChunkPQTable::l2_distance(const float *query_vec, _u8 *base_vec) return res; } -float FixedChunkPQTable::inner_product(const float *query_vec, _u8 *base_vec) +float FixedChunkPQTable::inner_product(const float *query_vec, char *base_vec) { float res = 0; - for (_u64 chunk = 0; chunk < n_chunks; chunk++) + for (size_t chunk = 0; chunk < n_chunks; chunk++) { - for (_u64 j = chunk_offsets[chunk]; j < chunk_offsets[chunk + 1]; j++) + 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 @@ -238,11 +238,11 @@ float FixedChunkPQTable::inner_product(const float *query_vec, _u8 *base_vec) } // assumes no rotation is involved -void FixedChunkPQTable::inflate_vector(_u8 *base_vec, float *out_vec) +void FixedChunkPQTable::inflate_vector(char *base_vec, float *out_vec) { - for (_u64 chunk = 0; chunk < n_chunks; chunk++) + for (size_t chunk = 0; chunk < n_chunks; chunk++) { - for (_u64 j = chunk_offsets[chunk]; j < chunk_offsets[chunk + 1]; j++) + 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]; @@ -254,35 +254,35 @@ void FixedChunkPQTable::populate_chunk_inner_products(const float *query_vec, fl { memset(dist_vec, 0, 256 * n_chunks * sizeof(float)); // chunk wise distance computation - for (_u64 chunk = 0; chunk < n_chunks; chunk++) + 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 (_u64 j = chunk_offsets[chunk]; j < chunk_offsets[chunk + 1]; j++) + for (size_t j = chunk_offsets[chunk]; j < chunk_offsets[chunk + 1]; j++) { const float *centers_dim_vec = tables_tr + (256 * j); - for (_u64 idx = 0; idx < 256; idx++) + 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) + 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 _u8 *all_coords, const _u64 ndims, _u8 *out) +void aggregate_coords(const std::vector &ids, const char *all_coords, const size_t ndims, char *out) { - for (_u64 i = 0; i < ids.size(); i++) + for (size_t i = 0; i < ids.size(); i++) { - memcpy(out + i * ndims, all_coords + ids[i] * ndims, ndims * sizeof(_u8)); + memcpy(out + i * ndims, all_coords + ids[i] * ndims, ndims * sizeof(char)); } } -void pq_dist_lookup(const _u8 *pq_ids, const _u64 n_pts, const _u64 pq_nchunks, const float *pq_dists, +void pq_dist_lookup(const char *pq_ids, const size_t n_pts, const size_t pq_nchunks, const float *pq_dists, std::vector &dists_out) { //_mm_prefetch((char*) dists_out, _MM_HINT_T0); @@ -291,47 +291,49 @@ void pq_dist_lookup(const _u8 *pq_ids, const _u64 n_pts, const _u64 pq_nchunks, _mm_prefetch((char *)(pq_ids + 128), _MM_HINT_T0); dists_out.clear(); dists_out.resize(n_pts, 0); - for (_u64 chunk = 0; chunk < pq_nchunks; chunk++) + for (size_t chunk = 0; chunk < pq_nchunks; chunk++) { const float *chunk_dists = pq_dists + 256 * chunk; if (chunk < pq_nchunks - 1) { _mm_prefetch((char *)(chunk_dists + 256), _MM_HINT_T0); } - for (_u64 idx = 0; idx < n_pts; idx++) + for (size_t idx = 0; idx < n_pts; idx++) { - _u8 pq_centerid = pq_ids[pq_nchunks * idx + chunk]; + char pq_centerid = pq_ids[pq_nchunks * idx + chunk]; dists_out[idx] += chunk_dists[pq_centerid]; } } } -// Need to replace calls to these functions with calls to vector& based functions above -void aggregate_coords(const unsigned *ids, const _u64 n_ids, const _u8 *all_coords, const _u64 ndims, _u8 *out) +// Need to replace calls to these functions with calls to vector& based +// functions above +void aggregate_coords(const unsigned *ids, const size_t n_ids, const char *all_coords, const size_t ndims, char *out) { - for (_u64 i = 0; i < n_ids; i++) + for (size_t i = 0; i < n_ids; i++) { - memcpy(out + i * ndims, all_coords + ids[i] * ndims, ndims * sizeof(_u8)); + memcpy(out + i * ndims, all_coords + ids[i] * ndims, ndims * sizeof(char)); } } -void pq_dist_lookup(const _u8 *pq_ids, const _u64 n_pts, const _u64 pq_nchunks, const float *pq_dists, float *dists_out) +void pq_dist_lookup(const char *pq_ids, const size_t n_pts, const size_t pq_nchunks, const float *pq_dists, + float *dists_out) { _mm_prefetch((char *)dists_out, _MM_HINT_T0); _mm_prefetch((char *)pq_ids, _MM_HINT_T0); _mm_prefetch((char *)(pq_ids + 64), _MM_HINT_T0); _mm_prefetch((char *)(pq_ids + 128), _MM_HINT_T0); memset(dists_out, 0, n_pts * sizeof(float)); - for (_u64 chunk = 0; chunk < pq_nchunks; chunk++) + for (size_t chunk = 0; chunk < pq_nchunks; chunk++) { const float *chunk_dists = pq_dists + 256 * chunk; if (chunk < pq_nchunks - 1) { _mm_prefetch((char *)(chunk_dists + 256), _MM_HINT_T0); } - for (_u64 idx = 0; idx < n_pts; idx++) + for (size_t idx = 0; idx < n_pts; idx++) { - _u8 pq_centerid = pq_ids[pq_nchunks * idx + chunk]; + char pq_centerid = pq_ids[pq_nchunks * idx + chunk]; dists_out[idx] += chunk_dists[pq_centerid]; } } @@ -460,7 +462,7 @@ int generate_pq_pivots(const float *const passed_train_data, size_t num_train, u << chunk_offsets[i + 1] << ")" << std::endl; #pragma omp parallel for schedule(static, 65536) - for (int64_t j = 0; j < (_s64)num_train; j++) + for (int64_t j = 0; j < (int64_t)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)); @@ -486,7 +488,7 @@ int generate_pq_pivots(const float *const passed_train_data, size_t num_train, u 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(), chunk_offsets.data(), chunk_offsets.size(), 1, cumul_bytes[2]); - diskann::save_bin<_u64>(pq_pivots_path.c_str(), cumul_bytes.data(), cumul_bytes.size(), 1, 0); + diskann::save_bin(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; @@ -599,10 +601,10 @@ int generate_opq_pivots(const float *passed_train_data, size_t num_train, unsign rotmat_tr.reset(new float[dim * dim]); std::memset(rotmat_tr.get(), 0, dim * dim * sizeof(float)); - for (_u32 d1 = 0; d1 < dim; d1++) + for (uint32_t d1 = 0; d1 < dim; d1++) *(rotmat_tr.get() + d1 * dim + d1) = 1; - for (_u32 rnd = 0; rnd < MAX_OPQ_ITERS; rnd++) + for (uint32_t rnd = 0; rnd < MAX_OPQ_ITERS; rnd++) { // rotate the training data using the current rotation matrix cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, (MKL_INT)num_train, (MKL_INT)dim, (MKL_INT)dim, 1.0f, @@ -624,7 +626,7 @@ int generate_opq_pivots(const float *passed_train_data, size_t num_train, unsign << chunk_offsets[i + 1] << ")" << std::endl; #pragma omp parallel for schedule(static, 65536) - for (int64_t j = 0; j < (_s64)num_train; j++) + for (int64_t j = 0; j < (int64_t)num_train; j++) { std::memcpy(cur_data.get() + j * cur_chunk_size, rotated_train_data.get() + j * dim + chunk_offsets[i], cur_chunk_size * sizeof(float)); @@ -644,7 +646,7 @@ int generate_opq_pivots(const float *passed_train_data, size_t num_train, unsign } } - _u32 num_lloyds_iters = 8; + uint32_t num_lloyds_iters = 8; kmeans::run_lloyds(cur_data.get(), num_train, cur_chunk_size, cur_pivot_data.get(), num_centers, num_lloyds_iters, NULL, closest_center.get()); @@ -654,10 +656,10 @@ int generate_opq_pivots(const float *passed_train_data, size_t num_train, unsign cur_pivot_data.get() + j * cur_chunk_size, cur_chunk_size * sizeof(float)); } - for (_u64 j = 0; j < num_train; j++) + for (size_t j = 0; j < num_train; j++) { std::memcpy(rotated_and_quantized_train_data.get() + j * dim + chunk_offsets[i], - cur_pivot_data.get() + (_u64)closest_center[j] * cur_chunk_size, + cur_pivot_data.get() + (size_t)closest_center[j] * cur_chunk_size, cur_chunk_size * sizeof(float)); } } @@ -670,7 +672,7 @@ int generate_opq_pivots(const float *passed_train_data, size_t num_train, unsign // compute the SVD of the correlation matrix to help determine the new // rotation matrix - _u32 errcode = + uint32_t errcode = LAPACKE_sgesdd(LAPACK_ROW_MAJOR, 'A', (MKL_INT)dim, (MKL_INT)dim, correlation_matrix.get(), (MKL_INT)dim, singular_values.get(), Umat.get(), (MKL_INT)dim, Vmat_T.get(), (MKL_INT)dim); @@ -694,7 +696,7 @@ int generate_opq_pivots(const float *passed_train_data, size_t num_train, unsign diskann::save_bin(opq_pivots_path.c_str(), centroid.get(), (size_t)dim, 1, cumul_bytes[1]); cumul_bytes[3] = cumul_bytes[2] + diskann::save_bin(opq_pivots_path.c_str(), chunk_offsets.data(), chunk_offsets.size(), 1, cumul_bytes[2]); - diskann::save_bin<_u64>(opq_pivots_path.c_str(), cumul_bytes.data(), cumul_bytes.size(), 1, 0); + diskann::save_bin(opq_pivots_path.c_str(), cumul_bytes.data(), cumul_bytes.size(), 1, 0); diskann::cout << "Saved opq pivot data to " << opq_pivots_path << " of size " << cumul_bytes[cumul_bytes.size() - 1] << "B." << std::endl; @@ -714,10 +716,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, bool use_opq) { - _u64 read_blk_size = 64 * 1024 * 1024; + size_t read_blk_size = 64 * 1024 * 1024; cached_ifstream base_reader(data_file, read_blk_size); - _u32 npts32; - _u32 basedim32; + uint32_t npts32; + uint32_t basedim32; base_reader.read((char *)&npts32, sizeof(uint32_t)); base_reader.read((char *)&basedim32, sizeof(uint32_t)); size_t num_points = npts32; @@ -737,10 +739,10 @@ int generate_pq_data_from_pivots(const std::string data_file, unsigned num_cente } else { - _u64 nr, nc; - std::unique_ptr<_u64[]> file_offset_data; + size_t nr, nc; + std::unique_ptr file_offset_data; - diskann::load_bin<_u64>(pq_pivots_path.c_str(), file_offset_data, nr, nc, 0); + diskann::load_bin(pq_pivots_path.c_str(), file_offset_data, nr, nc, 0); if (nr != 4) { @@ -796,7 +798,7 @@ int generate_pq_data_from_pivots(const std::string data_file, unsigned num_cente } std::ofstream compressed_file_writer(pq_compressed_vectors_path, std::ios::binary); - _u32 num_pq_chunks_u32 = num_pq_chunks; + uint32_t 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)); @@ -812,8 +814,9 @@ int generate_pq_data_from_pivots(const std::string data_file, unsigned num_cente std::memset(block_inflated_base.get(), 0, block_size * dim * sizeof(float)); #endif - std::unique_ptr<_u32[]> block_compressed_base = std::make_unique<_u32[]>(block_size * (_u64)num_pq_chunks); - std::memset(block_compressed_base.get(), 0, block_size * (_u64)num_pq_chunks * sizeof(uint32_t)); + std::unique_ptr block_compressed_base = + std::make_unique(block_size * (size_t)num_pq_chunks); + std::memset(block_compressed_base.get(), 0, block_size * (size_t)num_pq_chunks * sizeof(uint32_t)); std::unique_ptr block_data_T = std::make_unique(block_size * dim); std::unique_ptr block_data_float = std::make_unique(block_size * dim); @@ -850,7 +853,8 @@ int generate_pq_data_from_pivots(const std::string data_file, unsigned num_cente if (use_opq) { - // rotate the current block with the trained rotation matrix before PQ + // rotate the current block with the trained rotation matrix before + // PQ cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, (MKL_INT)cur_blk_size, (MKL_INT)dim, (MKL_INT)dim, 1.0f, block_data_float.get(), (MKL_INT)dim, rotmat_tr.get(), (MKL_INT)dim, 0.0f, block_data_tmp.get(), (MKL_INT)dim); @@ -868,14 +872,14 @@ int generate_pq_data_from_pivots(const std::string data_file, unsigned num_cente std::unique_ptr closest_center = std::make_unique(cur_blk_size); #pragma omp parallel for schedule(static, 8192) - for (int64_t j = 0; j < (_s64)cur_blk_size; j++) + for (int64_t j = 0; j < (int64_t)cur_blk_size; j++) { for (uint64_t k = 0; k < cur_chunk_size; k++) cur_data[j * cur_chunk_size + k] = block_data_float[j * dim + chunk_offsets[i] + k]; } #pragma omp parallel for schedule(static, 1) - for (int64_t j = 0; j < (_s64)num_centers; j++) + for (int64_t j = 0; j < (int64_t)num_centers; j++) { std::memcpy(cur_pivot_data.get() + j * cur_chunk_size, full_pivot_data.get() + j * dim + chunk_offsets[i], cur_chunk_size * sizeof(float)); @@ -885,7 +889,7 @@ int generate_pq_data_from_pivots(const std::string data_file, unsigned num_cente num_centers, 1, closest_center.get()); #pragma omp parallel for schedule(static, 8192) - for (int64_t j = 0; j < (_s64)cur_blk_size; j++) + for (int64_t j = 0; j < (int64_t)cur_blk_size; j++) { block_compressed_base[j * num_pq_chunks + i] = closest_center[j]; #ifdef SAVE_INFLATED_PQ @@ -978,7 +982,7 @@ void generate_quantized_data(const std::string data_file_to_use, const std::stri } else { - generate_opq_pivots(train_data, train_size, (_u32)train_dim, NUM_PQ_CENTROIDS, (_u32)num_pq_chunks, + generate_opq_pivots(train_data, train_size, (uint32_t)train_dim, NUM_PQ_CENTROIDS, (uint32_t)num_pq_chunks, pq_pivots_path, make_zero_mean); } generate_pq_data_from_pivots(data_file_to_use.c_str(), NUM_PQ_CENTROIDS, (uint32_t)num_pq_chunks, pq_pivots_path, diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index ad7809ccb..6101f65c8 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -13,27 +13,28 @@ #include "linux_aligned_file_reader.h" #endif -#define READ_U64(stream, val) stream.read((char *)&val, sizeof(_u64)) -#define READ_U32(stream, val) stream.read((char *)&val, sizeof(_u32)) +#define READ_U64(stream, val) stream.read((char *)&val, sizeof(uint64_t)) +#define READ_U32(stream, val) stream.read((char *)&val, sizeof(uint32_t)) #define READ_UNSIGNED(stream, val) stream.read((char *)&val, sizeof(unsigned)) // sector # on disk where node_id is present with in the graph part -#define NODE_SECTOR_NO(node_id) (((_u64)(node_id)) / nnodes_per_sector + 1) +#define NODE_SECTOR_NO(node_id) (((uint64_t)(node_id)) / nnodes_per_sector + 1) // obtains region of sector containing node -#define OFFSET_TO_NODE(sector_buf, node_id) ((char *)sector_buf + (((_u64)node_id) % nnodes_per_sector) * max_node_len) +#define OFFSET_TO_NODE(sector_buf, node_id) \ + ((char *)sector_buf + (((uint64_t)node_id) % nnodes_per_sector) * max_node_len) -// returns region of `node_buf` containing [NNBRS][NBR_ID(_u32)] +// returns region of `node_buf` containing [NNBRS][NBR_ID(uint32_t)] #define OFFSET_TO_NODE_NHOOD(node_buf) (unsigned *)((char *)node_buf + disk_bytes_per_point) // returns region of `node_buf` containing [COORD(T)] #define OFFSET_TO_NODE_COORDS(node_buf) (T *)(node_buf) // sector # beyond the end of graph where data for id is present for reordering -#define VECTOR_SECTOR_NO(id) (((_u64)(id)) / nvecs_per_sector + reorder_data_start_sector) +#define VECTOR_SECTOR_NO(id) (((uint64_t)(id)) / nvecs_per_sector + reorder_data_start_sector) // sector # beyond the end of graph where data for id is present for reordering -#define VECTOR_SECTOR_OFFSET(id) ((((_u64)(id)) % nvecs_per_sector) * data_dim * sizeof(float)) +#define VECTOR_SECTOR_OFFSET(id) ((((uint64_t)(id)) % nvecs_per_sector) * data_dim * sizeof(float)) namespace diskann { @@ -101,12 +102,12 @@ template PQFlashIndex::~PQFlashIndex() } template -void PQFlashIndex::setup_thread_data(_u64 nthreads, _u64 visited_reserve) +void PQFlashIndex::setup_thread_data(uint64_t nthreads, uint64_t visited_reserve) { diskann::cout << "Setting up thread-specific contexts for nthreads: " << nthreads << std::endl; // omp parallel for to generate unique thread IDs #pragma omp parallel for num_threads((int)nthreads) - for (_s64 thread = 0; thread < (_s64)nthreads; thread++) + for (int64_t thread = 0; thread < (int64_t)nthreads; thread++) { #pragma omp critical { @@ -122,7 +123,7 @@ void PQFlashIndex::setup_thread_data(_u64 nthreads, _u64 visited_rese 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(); + uint64_t num_cached_nodes = node_list.size(); // borrow thread data ScratchStoreManager> manager(this->thread_data); @@ -132,20 +133,20 @@ template void PQFlashIndex::load_cache_ nhood_cache_buf = new unsigned[num_cached_nodes * (max_degree + 1)]; memset(nhood_cache_buf, 0, num_cached_nodes * (max_degree + 1)); - _u64 coord_cache_buf_len = num_cached_nodes * aligned_dim; + uint64_t coord_cache_buf_len = num_cached_nodes * aligned_dim; diskann::alloc_aligned((void **)&coord_cache_buf, coord_cache_buf_len * sizeof(T), 8 * sizeof(T)); memset(coord_cache_buf, 0, coord_cache_buf_len * sizeof(T)); size_t BLOCK_SIZE = 8; size_t num_blocks = DIV_ROUND_UP(num_cached_nodes, BLOCK_SIZE); - for (_u64 block = 0; block < num_blocks; block++) + for (uint64_t block = 0; block < num_blocks; block++) { - _u64 start_idx = block * BLOCK_SIZE; - _u64 end_idx = (std::min)(num_cached_nodes, (block + 1) * BLOCK_SIZE); + uint64_t start_idx = block * BLOCK_SIZE; + uint64_t end_idx = (std::min)(num_cached_nodes, (block + 1) * BLOCK_SIZE); std::vector read_reqs; - std::vector> nhoods; - for (_u64 node_idx = start_idx; node_idx < end_idx; node_idx++) + std::vector> nhoods; + for (uint64_t node_idx = start_idx; node_idx < end_idx; node_idx++) { AlignedRead read; char *buf = nullptr; @@ -159,8 +160,8 @@ template void PQFlashIndex::load_cache_ reader->read(read_reqs, ctx); - _u64 node_idx = start_idx; - for (_u32 i = 0; i < read_reqs.size(); i++) + uint64_t node_idx = start_idx; + for (uint32_t i = 0; i < read_reqs.size(); i++) { #if defined(_WINDOWS) && defined(USE_BING_INFRA) // this block is to handle failed reads in // production settings @@ -181,7 +182,7 @@ template void PQFlashIndex::load_cache_ auto nnbrs = *node_nhood; unsigned *nbrs = node_nhood + 1; - std::pair<_u32, unsigned *> cnhood; + std::pair cnhood; cnhood.first = nnbrs; cnhood.second = nhood_cache_buf + node_idx * (max_degree + 1); memcpy(cnhood.second, nbrs, nnbrs * sizeof(unsigned)); @@ -196,14 +197,14 @@ template void PQFlashIndex::load_cache_ #ifdef EXEC_ENV_OLS 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, + uint64_t l_search, uint64_t beamwidth, + uint64_t num_nodes_to_cache, uint32_t nthreads, std::vector &node_list) { #else template -void PQFlashIndex::generate_cache_list_from_sample_queries(std::string sample_bin, _u64 l_search, - _u64 beamwidth, _u64 num_nodes_to_cache, +void PQFlashIndex::generate_cache_list_from_sample_queries(std::string sample_bin, uint64_t l_search, + uint64_t beamwidth, uint64_t num_nodes_to_cache, uint32_t nthreads, std::vector &node_list) { @@ -211,13 +212,13 @@ void PQFlashIndex::generate_cache_list_from_sample_queries(std::strin this->count_visited_nodes = true; this->node_visit_counter.clear(); this->node_visit_counter.resize(this->num_points); - for (_u32 i = 0; i < node_visit_counter.size(); i++) + for (uint32_t i = 0; i < node_visit_counter.size(); i++) { this->node_visit_counter[i].first = i; this->node_visit_counter[i].second = 0; } - _u64 sample_num, sample_dim, sample_aligned_dim; + uint64_t sample_num, sample_dim, sample_aligned_dim; T *samples; #ifdef EXEC_ENV_OLS @@ -241,18 +242,20 @@ void PQFlashIndex::generate_cache_list_from_sample_queries(std::strin 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++) + for (int64_t 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_dists.data() + (i * 1), beamwidth); } std::sort(this->node_visit_counter.begin(), node_visit_counter.end(), - [](std::pair<_u32, _u32> &left, std::pair<_u32, _u32> &right) { return left.second > right.second; }); + [](std::pair &left, std::pair &right) { + return left.second > right.second; + }); node_list.clear(); node_list.shrink_to_fit(); node_list.reserve(num_nodes_to_cache); - for (_u64 i = 0; i < num_nodes_to_cache; i++) + for (uint64_t i = 0; i < num_nodes_to_cache; i++) { node_list.push_back(this->node_visit_counter[i].first); } @@ -262,7 +265,7 @@ void PQFlashIndex::generate_cache_list_from_sample_queries(std::strin } template -void PQFlashIndex::cache_bfs_levels(_u64 num_nodes_to_cache, std::vector &node_list, +void PQFlashIndex::cache_bfs_levels(uint64_t num_nodes_to_cache, std::vector &node_list, const bool shuffle) { std::random_device rng; @@ -271,7 +274,7 @@ void PQFlashIndex::cache_bfs_levels(_u64 num_nodes_to_cache, std::vec tsl::robin_set node_set; // Do not cache more than 10% of the nodes in the index - _u64 tenp_nodes = (_u64)(std::round(this->num_points * 0.1)); + uint64_t tenp_nodes = (uint64_t)(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 @@ -289,12 +292,12 @@ void PQFlashIndex::cache_bfs_levels(_u64 num_nodes_to_cache, std::vec cur_level = std::make_unique>(); prev_level = std::make_unique>(); - for (_u64 miter = 0; miter < num_medoids; miter++) + for (uint64_t miter = 0; miter < num_medoids; miter++) { cur_level->insert(medoids[miter]); } - _u64 lvl = 1; + uint64_t lvl = 1; uint64_t prev_node_set_size = 0; while ((node_set.size() + cur_level->size() < num_nodes_to_cache) && cur_level->size() != 0) { @@ -331,7 +334,7 @@ void PQFlashIndex::cache_bfs_levels(_u64 num_nodes_to_cache, std::vec 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> nhoods; + std::vector> nhoods; for (size_t cur_pt = start; cur_pt < end; cur_pt++) { char *buf = nullptr; @@ -348,7 +351,7 @@ void PQFlashIndex::cache_bfs_levels(_u64 num_nodes_to_cache, std::vec reader->read(read_reqs, ctx); // process each nhood buf - for (_u32 i = 0; i < read_reqs.size(); i++) + for (uint32_t i = 0; i < read_reqs.size(); i++) { #if defined(_WINDOWS) && defined(USE_BING_INFRA) // this block is to handle read failures in // production settings @@ -362,10 +365,10 @@ void PQFlashIndex::cache_bfs_levels(_u64 num_nodes_to_cache, std::vec // insert node coord into coord_cache char *node_buf = OFFSET_TO_NODE(nhood.second, nhood.first); unsigned *node_nhood = OFFSET_TO_NODE_NHOOD(node_buf); - _u64 nnbrs = (_u64)*node_nhood; + uint64_t nnbrs = (uint64_t)*node_nhood; unsigned *nbrs = node_nhood + 1; // explore next level - for (_u64 j = 0; j < nnbrs && !finish_flag; j++) + for (uint64_t j = 0; j < nnbrs && !finish_flag; j++) { if (node_set.find(nbrs[j]) == node_set.end()) { @@ -440,7 +443,7 @@ template void PQFlashIndex::use_medoids } else { - disk_pq_table.inflate_vector((_u8 *)medoid_coords, (centroid_data + cur_m * aligned_dim)); + disk_pq_table.inflate_vector((char *)medoid_coords, (centroid_data + cur_m * aligned_dim)); } aligned_free(medoid_buf); @@ -452,7 +455,7 @@ template inline int32_t PQFlashIndex::get_filter_number(const LabelT &filter_label) { int idx = -1; - for (_u32 i = 0; i < _filter_list.size(); i++) + for (uint32_t i = 0; i < _filter_list.size(); i++) { if (_filter_list[i] == filter_label) { @@ -498,7 +501,8 @@ LabelT PQFlashIndex::get_converted_label(const std::string &filter_la } template -void PQFlashIndex::get_label_file_metadata(std::string map_file, _u32 &num_pts, _u32 &num_total_labels) +void PQFlashIndex::get_label_file_metadata(std::string map_file, uint32_t &num_pts, + uint32_t &num_total_labels) { std::ifstream infile(map_file); std::string line, token; @@ -523,12 +527,12 @@ void PQFlashIndex::get_label_file_metadata(std::string map_file, _u32 } template -inline bool PQFlashIndex::point_has_label(_u32 point_id, _u32 label_id) +inline bool PQFlashIndex::point_has_label(uint32_t point_id, uint32_t label_id) { - _u32 start_vec = _pts_to_label_offsets[point_id]; - _u32 num_lbls = _pts_to_labels[start_vec]; + uint32_t start_vec = _pts_to_label_offsets[point_id]; + uint32_t num_lbls = _pts_to_labels[start_vec]; bool ret_val = false; - for (_u32 i = 0; i < num_lbls; i++) + for (uint32_t i = 0; i < num_lbls; i++) { if (_pts_to_labels[start_vec + 1 + i] == label_id) { @@ -549,23 +553,23 @@ void PQFlashIndex::parse_label_file(const std::string &label_file, si } std::string line, token; - _u32 line_cnt = 0; + uint32_t line_cnt = 0; - _u32 num_pts_in_label_file; - _u32 num_total_labels; + uint32_t num_pts_in_label_file; + uint32_t num_total_labels; get_label_file_metadata(label_file, num_pts_in_label_file, num_total_labels); - _pts_to_label_offsets = new _u32[num_pts_in_label_file]; - _pts_to_labels = new _u32[num_pts_in_label_file + num_total_labels]; - _u32 counter = 0; + _pts_to_label_offsets = new uint32_t[num_pts_in_label_file]; + _pts_to_labels = new uint32_t[num_pts_in_label_file + num_total_labels]; + uint32_t counter = 0; while (std::getline(infile, line)) { std::istringstream iss(line); - std::vector<_u32> lbls(0); + std::vector lbls(0); _pts_to_label_offsets[line_cnt] = counter; - _u32 &num_lbls_in_cur_pt = _pts_to_labels[counter]; + uint32_t &num_lbls_in_cur_pt = _pts_to_labels[counter]; num_lbls_in_cur_pt = 0; counter++; getline(iss, token, '\t'); @@ -612,7 +616,7 @@ template void PQFlashIndex::set_univers else { _use_universal_label = true; - _universal_filter_num = (_u32)temp_filter_num; + _universal_filter_num = (uint32_t)temp_filter_num; } } @@ -685,9 +689,9 @@ int PQFlashIndex::load_from_separate_paths(uint32_t num_threads, cons size_t npts_u64, nchunks_u64; #ifdef EXEC_ENV_OLS - diskann::load_bin<_u8>(files, pq_compressed_vectors, this->data, npts_u64, nchunks_u64); + diskann::load_bin(files, pq_compressed_vectors, this->data, npts_u64, nchunks_u64); #else - diskann::load_bin<_u8>(pq_compressed_vectors, this->data, npts_u64, nchunks_u64); + diskann::load_bin(pq_compressed_vectors, this->data, npts_u64, nchunks_u64); #endif this->num_points = npts_u64; @@ -709,15 +713,15 @@ int PQFlashIndex::load_from_separate_paths(uint32_t num_threads, cons while (std::getline(medoid_stream, line)) { std::istringstream iss(line); - _u32 cnt = 0; - _u32 medoid; + uint32_t cnt = 0; + uint32_t medoid; LabelT label; while (std::getline(iss, token, ',')) { if (cnt == 0) label = std::stoul(token); else - medoid = (_u32)stoul(token); + medoid = (uint32_t)stoul(token); cnt++; } _filter_to_medoid_id[label] = medoid; @@ -748,15 +752,15 @@ int PQFlashIndex::load_from_separate_paths(uint32_t num_threads, cons while (std::getline(dummy_map_stream, line)) { std::istringstream iss(line); - _u32 cnt = 0; - _u32 dummy_id; - _u32 real_id; + uint32_t cnt = 0; + uint32_t dummy_id; + uint32_t real_id; while (std::getline(iss, token, ',')) { if (cnt == 0) - dummy_id = (_u32)stoul(token); + dummy_id = (uint32_t)stoul(token); else - real_id = (_u32)stoul(token); + real_id = (uint32_t)stoul(token); cnt++; } _dummy_pts.insert(dummy_id); @@ -764,7 +768,7 @@ int PQFlashIndex::load_from_separate_paths(uint32_t num_threads, cons _dummy_to_real_map[dummy_id] = real_id; if (_real_to_dummy_map.find(real_id) == _real_to_dummy_map.end()) - _real_to_dummy_map[real_id] = std::vector<_u32>(); + _real_to_dummy_map[real_id] = std::vector(); _real_to_dummy_map[real_id].emplace_back(dummy_id); } @@ -805,7 +809,8 @@ int PQFlashIndex::load_from_separate_paths(uint32_t num_threads, cons 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_bytes_per_point = disk_pq_n_chunks * sizeof(_u8); // revising disk_bytes_per_point since DISK PQ is used. + 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." << std::endl; } @@ -828,13 +833,13 @@ int PQFlashIndex::load_from_separate_paths(uint32_t num_threads, cons std::ifstream index_metadata(disk_index_file, std::ios::binary); #endif - _u32 nr, nc; // metadata itself is stored as bin format (nr is number of - // metadata, nc should be 1) + uint32_t nr, nc; // metadata itself is stored as bin format (nr is number of + // metadata, nc should be 1) READ_U32(index_metadata, nr); READ_U32(index_metadata, nc); - _u64 disk_nnodes; - _u64 disk_ndims; // can be disk PQ dim if disk_PQ is set to true + uint64_t disk_nnodes; + uint64_t disk_ndims; // can be disk PQ dim if disk_PQ is set to true READ_U64(index_metadata, disk_nnodes); READ_U64(index_metadata, disk_ndims); @@ -863,7 +868,7 @@ int PQFlashIndex::load_from_separate_paths(uint32_t num_threads, cons // setting up concept of frozen points in disk index for streaming-DiskANN READ_U64(index_metadata, this->num_frozen_points); - _u64 file_frozen_id; + uint64_t file_frozen_id; READ_U64(index_metadata, file_frozen_id); if (this->num_frozen_points == 1) this->frozen_location = file_frozen_id; @@ -878,8 +883,9 @@ int PQFlashIndex::load_from_separate_paths(uint32_t num_threads, cons { if (this->use_disk_index_pq == false) { - throw ANNException("Reordering is designed for used with disk PQ compression option", -1, __FUNCSIG__, - __FILE__, __LINE__); + throw ANNException("Reordering is designed for used with disk PQ " + "compression option", + -1, __FUNCSIG__, __FILE__, __LINE__); } READ_U64(index_metadata, this->reorder_data_start_sector); READ_U64(index_metadata, this->ndims_reorder_vecs); @@ -950,7 +956,8 @@ int PQFlashIndex::load_from_separate_paths(uint32_t num_threads, cons if (aligned_tmp_dim != aligned_dim || num_centroids != num_medoids) { std::stringstream stream; - stream << "Error loading centroids data file. Expected bin format of " + 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."; @@ -963,7 +970,7 @@ int PQFlashIndex::load_from_separate_paths(uint32_t num_threads, cons { num_medoids = 1; medoids = new uint32_t[1]; - medoids[0] = (_u32)(medoid_id_on_file); + medoids[0] = (uint32_t)(medoid_id_on_file); use_medoids_data_as_centroids(); } @@ -971,7 +978,7 @@ int PQFlashIndex::load_from_separate_paths(uint32_t num_threads, cons if (file_exists(norm_file) && metric == diskann::Metric::INNER_PRODUCT) { - _u64 dumr, dumc; + uint64_t dumr, dumc; float *norm_val; diskann::load_bin(norm_file, norm_val, dumr, dumc); this->max_base_norm = norm_val[0]; @@ -1018,39 +1025,41 @@ bool getNextCompletedRequest(const IOContext &ctx, size_t size, int &completedIn #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, +void PQFlashIndex::cached_beam_search(const T *query1, const uint64_t k_search, const uint64_t l_search, + uint64_t *indices, float *distances, const uint64_t beam_width, const bool use_reorder_data, QueryStats *stats) { - cached_beam_search(query1, k_search, l_search, indices, distances, beam_width, std::numeric_limits<_u32>::max(), + cached_beam_search(query1, k_search, l_search, indices, distances, beam_width, std::numeric_limits::max(), use_reorder_data, stats); } 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, +void PQFlashIndex::cached_beam_search(const T *query1, const uint64_t k_search, const uint64_t l_search, + uint64_t *indices, float *distances, const uint64_t beam_width, const bool use_filter, const LabelT &filter_label, const bool use_reorder_data, QueryStats *stats) { cached_beam_search(query1, k_search, l_search, indices, distances, beam_width, use_filter, filter_label, - std::numeric_limits<_u32>::max(), use_reorder_data, stats); + std::numeric_limits::max(), use_reorder_data, stats); } 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, - const _u32 io_limit, const bool use_reorder_data, QueryStats *stats) +void PQFlashIndex::cached_beam_search(const T *query1, const uint64_t k_search, const uint64_t l_search, + uint64_t *indices, float *distances, const uint64_t beam_width, + const uint32_t io_limit, const bool use_reorder_data, + QueryStats *stats) { LabelT dummy_filter = 0; cached_beam_search(query1, k_search, l_search, indices, distances, beam_width, false, dummy_filter, - std::numeric_limits<_u32>::max(), use_reorder_data, stats); + std::numeric_limits::max(), use_reorder_data, stats); } 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, - const bool use_filter, const LabelT &filter_label, const _u32 io_limit, - const bool use_reorder_data, QueryStats *stats) +void PQFlashIndex::cached_beam_search(const T *query1, const uint64_t k_search, const uint64_t l_search, + uint64_t *indices, float *distances, const uint64_t beam_width, + const bool use_filter, const LabelT &filter_label, + const uint32_t io_limit, const bool use_reorder_data, + QueryStats *stats) { int32_t filter_num = 0; if (use_filter) @@ -1118,12 +1127,12 @@ void PQFlashIndex::cached_beam_search(const T *query1, const _u64 k_s // pointers to buffers for data T *data_buf = query_scratch->coord_scratch; - _u64 &data_buf_idx = query_scratch->coord_idx; + uint64_t &data_buf_idx = query_scratch->coord_idx; _mm_prefetch((char *)data_buf, _MM_HINT_T1); // sector scratch char *sector_scratch = query_scratch->sector_scratch; - _u64 §or_scratch_idx = query_scratch->sector_idx; + uint64_t §or_scratch_idx = query_scratch->sector_idx; // query <-> PQ chunk centers distances pq_table.preprocess_query(query_rotated); // center the query and rotate if @@ -1133,25 +1142,26 @@ void PQFlashIndex::cached_beam_search(const T *query1, const _u64 k_s // query <-> neighbor list float *dist_scratch = pq_query_scratch->aligned_dist_scratch; - _u8 *pq_coord_scratch = pq_query_scratch->aligned_pq_coord_scratch; + char *pq_coord_scratch = pq_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 uint64_t n_ids, + float *dists_out) { diskann::aggregate_coords(ids, n_ids, this->data, this->n_chunks, pq_coord_scratch); diskann::pq_dist_lookup(pq_coord_scratch, n_ids, this->n_chunks, pq_dists, dists_out); }; Timer query_timer, io_timer, cpu_timer; - tsl::robin_set<_u64> &visited = query_scratch->visited; + tsl::robin_set &visited = query_scratch->visited; NeighborPriorityQueue &retset = query_scratch->retset; retset.reserve(l_search); std::vector &full_retset = query_scratch->full_retset; - _u32 best_medoid = 0; + uint32_t best_medoid = 0; float best_dist = (std::numeric_limits::max)(); if (!use_filter) { - for (_u64 cur_m = 0; cur_m < num_medoids; cur_m++) + for (uint64_t cur_m = 0; cur_m < num_medoids; cur_m++) { float cur_expanded_dist = dist_cmp_float->compare(query_float, centroid_data + aligned_dim * cur_m, (unsigned)aligned_dim); @@ -1198,7 +1208,7 @@ void PQFlashIndex::cached_beam_search(const T *query1, const _u64 k_s cached_nhoods.clear(); sector_scratch_idx = 0; // find new beam - _u32 num_seen = 0; + uint32_t num_seen = 0; while (retset.has_unexpanded_node() && frontier.size() < beam_width && num_seen < beam_width) { auto nbr = retset.closest_unexpanded(); @@ -1218,7 +1228,7 @@ void PQFlashIndex::cached_beam_search(const T *query1, const _u64 k_s } if (this->count_visited_nodes) { - reinterpret_cast &>(this->node_visit_counter[nbr.id].second).fetch_add(1); + reinterpret_cast &>(this->node_visit_counter[nbr.id].second).fetch_add(1); } } @@ -1227,10 +1237,10 @@ void PQFlashIndex::cached_beam_search(const T *query1, const _u64 k_s { if (stats != nullptr) stats->n_hops++; - for (_u64 i = 0; i < frontier.size(); i++) + for (uint64_t i = 0; i < frontier.size(); i++) { auto id = frontier[i]; - std::pair<_u32, char *> fnhood; + std::pair fnhood; fnhood.first = id; fnhood.second = sector_scratch + sector_scratch_idx * SECTOR_LEN; sector_scratch_idx++; @@ -1245,7 +1255,8 @@ void PQFlashIndex::cached_beam_search(const T *query1, const _u64 k_s } io_timer.reset(); #ifdef USE_BING_INFRA - reader->read(frontier_read_reqs, ctx, true); // async reader windows. + reader->read(frontier_read_reqs, ctx, + true); // async reader windows. #else reader->read(frontier_read_reqs, ctx); // synchronous IO linux #endif @@ -1268,14 +1279,14 @@ void PQFlashIndex::cached_beam_search(const T *query1, const _u64 k_s else { if (metric == diskann::Metric::INNER_PRODUCT) - cur_expanded_dist = disk_pq_table.inner_product(query_float, (_u8 *)node_fp_coords_copy); + cur_expanded_dist = disk_pq_table.inner_product(query_float, (char *)node_fp_coords_copy); else cur_expanded_dist = disk_pq_table.l2_distance( // disk_pq does not support OPQ yet - query_float, (_u8 *)node_fp_coords_copy); + query_float, (char *)node_fp_coords_copy); } full_retset.push_back(Neighbor((unsigned)cached_nhood.first, cur_expanded_dist)); - _u64 nnbrs = cached_nhood.second.first; + uint64_t nnbrs = cached_nhood.second.first; unsigned *node_nbrs = cached_nhood.second.second; // compute node_nbrs <-> query dists in PQ space @@ -1288,7 +1299,7 @@ void PQFlashIndex::cached_beam_search(const T *query1, const _u64 k_s } // process prefetched nhood - for (_u64 m = 0; m < nnbrs; ++m) + for (uint64_t m = 0; m < nnbrs; ++m) { unsigned id = node_nbrs[m]; if (visited.insert(id).second) @@ -1309,8 +1320,8 @@ void PQFlashIndex::cached_beam_search(const T *query1, const _u64 k_s // process each frontier nhood - compute distances to unvisited nodes int completedIndex = -1; long requestCount = static_cast(frontier_read_reqs.size()); - // If we issued read requests and if a read is complete or there are reads - // in wait state, then enter the while loop. + // If we issued read requests and if a read is complete or there are + // reads in wait state, then enter the while loop. while (requestCount > 0 && getNextCompletedRequest(ctx, requestCount, completedIndex)) { assert(completedIndex >= 0); @@ -1322,7 +1333,7 @@ void PQFlashIndex::cached_beam_search(const T *query1, const _u64 k_s #endif char *node_disk_buf = OFFSET_TO_NODE(frontier_nhood.second, frontier_nhood.first); unsigned *node_buf = OFFSET_TO_NODE_NHOOD(node_disk_buf); - _u64 nnbrs = (_u64)(*node_buf); + uint64_t nnbrs = (uint64_t)(*node_buf); T *node_fp_coords = OFFSET_TO_NODE_COORDS(node_disk_buf); // assert(data_buf_idx < MAX_N_CMPS); if (data_buf_idx == MAX_N_CMPS) @@ -1339,9 +1350,9 @@ void PQFlashIndex::cached_beam_search(const T *query1, const _u64 k_s else { if (metric == diskann::Metric::INNER_PRODUCT) - cur_expanded_dist = disk_pq_table.inner_product(query_float, (_u8 *)node_fp_coords_copy); + cur_expanded_dist = disk_pq_table.inner_product(query_float, (char *)node_fp_coords_copy); else - cur_expanded_dist = disk_pq_table.l2_distance(query_float, (_u8 *)node_fp_coords_copy); + cur_expanded_dist = disk_pq_table.l2_distance(query_float, (char *)node_fp_coords_copy); } full_retset.push_back(Neighbor(frontier_nhood.first, cur_expanded_dist)); unsigned *node_nbrs = (node_buf + 1); @@ -1356,7 +1367,7 @@ void PQFlashIndex::cached_beam_search(const T *query1, const _u64 k_s cpu_timer.reset(); // process prefetch-ed nhood - for (_u64 m = 0; m < nnbrs; ++m) + for (uint64_t m = 0; m < nnbrs; ++m) { unsigned id = node_nbrs[m]; if (visited.insert(id).second) @@ -1394,7 +1405,8 @@ void PQFlashIndex::cached_beam_search(const T *query1, const _u64 k_s { if (!(this->reorder_data_exists)) { - throw ANNException("Requested use of reordering data which does not exist in index " + throw ANNException("Requested use of reordering data which does " + "not exist in index " "file", -1, __FUNCSIG__, __FILE__, __LINE__); } @@ -1438,7 +1450,7 @@ void PQFlashIndex::cached_beam_search(const T *query1, const _u64 k_s } // copy k_search values - for (_u64 i = 0; i < k_search; i++) + for (uint64_t i = 0; i < k_search; i++) { indices[i] = full_retset[i].id; @@ -1454,8 +1466,8 @@ void PQFlashIndex::cached_beam_search(const T *query1, const _u64 k_s { // flip the sign to convert min to max distances[i] = (-distances[i]); - // rescale to revert back to original norms (cancelling the effect of - // base and query pre-processing) + // rescale to revert back to original norms (cancelling the + // effect of base and query pre-processing) if (max_base_norm != 0) distances[i] *= (max_base_norm * query_norm); } @@ -1476,25 +1488,26 @@ void PQFlashIndex::cached_beam_search(const T *query1, const _u64 k_s // indices and distances need to be pre-allocated of size l_search and the // return value is the number of matching hits. template -_u32 PQFlashIndex::range_search(const T *query1, const double range, const _u64 min_l_search, - const _u64 max_l_search, std::vector<_u64> &indices, - std::vector &distances, const _u64 min_beam_width, QueryStats *stats) +uint32_t PQFlashIndex::range_search(const T *query1, const double range, const uint64_t min_l_search, + const uint64_t max_l_search, std::vector &indices, + std::vector &distances, const uint64_t min_beam_width, + QueryStats *stats) { - _u32 res_count = 0; + uint32_t res_count = 0; bool stop_flag = false; - _u32 l_search = min_l_search; // starting size of the candidate list + uint32_t l_search = min_l_search; // starting size of the candidate list while (!stop_flag) { indices.resize(l_search); distances.resize(l_search); - _u64 cur_bw = min_beam_width > (l_search / 5) ? min_beam_width : l_search / 5; + uint64_t cur_bw = min_beam_width > (l_search / 5) ? min_beam_width : l_search / 5; cur_bw = (cur_bw > 100) ? 100 : cur_bw; for (auto &x : distances) x = std::numeric_limits::max(); this->cached_beam_search(query1, l_search, l_search, indices.data(), distances.data(), cur_bw, false, stats); - for (_u32 i = 0; i < l_search; i++) + for (uint32_t i = 0; i < l_search; i++) { if (distances[i] > (float)range) { @@ -1504,7 +1517,7 @@ _u32 PQFlashIndex::range_search(const T *query1, const double range, else if (i == l_search - 1) res_count = l_search; } - if (res_count < (_u32)(l_search / 2.0)) + if (res_count < (uint32_t)(l_search / 2.0)) stop_flag = true; l_search = l_search * 2; if (l_search > max_l_search) @@ -1515,7 +1528,7 @@ _u32 PQFlashIndex::range_search(const T *query1, const double range, return res_count; } -template _u64 PQFlashIndex::get_data_dim() +template uint64_t PQFlashIndex::get_data_dim() { return data_dim; } @@ -1544,11 +1557,11 @@ template char *PQFlashIndex::getHeaderB #endif // instantiations -template class PQFlashIndex<_u8>; -template class PQFlashIndex<_s8>; +template class PQFlashIndex; +template class PQFlashIndex; template class PQFlashIndex; -template class PQFlashIndex<_u8, uint16_t>; -template class PQFlashIndex<_s8, uint16_t>; +template class PQFlashIndex; +template class PQFlashIndex; template class PQFlashIndex; } // namespace diskann diff --git a/src/scratch.cpp b/src/scratch.cpp index c83ecac9f..7b14b14fa 100644 --- a/src/scratch.cpp +++ b/src/scratch.cpp @@ -95,10 +95,10 @@ template void SSDQueryScratch::reset() template SSDQueryScratch::SSDQueryScratch(size_t aligned_dim, size_t visited_reserve) { - _u64 coord_alloc_size = ROUND_UP(MAX_N_CMPS * aligned_dim, 256); + size_t coord_alloc_size = ROUND_UP(MAX_N_CMPS * aligned_dim, 256); diskann::alloc_aligned((void **)&coord_scratch, coord_alloc_size, 256); - diskann::alloc_aligned((void **)§or_scratch, (_u64)MAX_N_SECTOR_READS * (_u64)SECTOR_LEN, SECTOR_LEN); + diskann::alloc_aligned((void **)§or_scratch, (size_t)MAX_N_SECTOR_READS * (size_t)SECTOR_LEN, SECTOR_LEN); diskann::alloc_aligned((void **)&aligned_query_T, aligned_dim * sizeof(T), 8 * sizeof(T)); _pq_scratch = new PQScratch(MAX_GRAPH_DEGREE, aligned_dim); @@ -132,11 +132,11 @@ template DISKANN_DLLEXPORT class InMemQueryScratch; template DISKANN_DLLEXPORT class InMemQueryScratch; template DISKANN_DLLEXPORT class InMemQueryScratch; -template DISKANN_DLLEXPORT class SSDQueryScratch<_u8>; -template DISKANN_DLLEXPORT class SSDQueryScratch<_s8>; +template DISKANN_DLLEXPORT class SSDQueryScratch; +template DISKANN_DLLEXPORT class SSDQueryScratch; template DISKANN_DLLEXPORT class SSDQueryScratch; -template DISKANN_DLLEXPORT class SSDThreadData<_u8>; -template DISKANN_DLLEXPORT class SSDThreadData<_s8>; +template DISKANN_DLLEXPORT class SSDThreadData; +template DISKANN_DLLEXPORT class SSDThreadData; template DISKANN_DLLEXPORT class SSDThreadData; } // namespace diskann \ No newline at end of file diff --git a/src/utils.cpp b/src/utils.cpp index ec25dead3..c70f884ef 100644 --- a/src/utils.cpp +++ b/src/utils.cpp @@ -73,20 +73,20 @@ bool AvxSupportedCPU = false; namespace diskann { -void block_convert(std::ofstream &writr, std::ifstream &readr, float *read_buf, _u64 npts, _u64 ndims) +void block_convert(std::ofstream &writr, std::ifstream &readr, float *read_buf, size_t npts, size_t ndims) { readr.read((char *)read_buf, npts * ndims * sizeof(float)); - _u32 ndims_u32 = (_u32)ndims; + uint32_t ndims_u32 = (uint32_t)ndims; #pragma omp parallel for - for (_s64 i = 0; i < (_s64)npts; i++) + for (int64_t i = 0; i < (int64_t)npts; i++) { float norm_pt = std::numeric_limits::epsilon(); - for (_u32 dim = 0; dim < ndims_u32; dim++) + for (uint32_t 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++) + for (uint32_t dim = 0; dim < ndims_u32; dim++) { *(read_buf + i * ndims + dim) = *(read_buf + i * ndims + dim) / norm_pt; } @@ -100,24 +100,25 @@ void normalize_data_file(const std::string &inFileName, const std::string &outFi 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)); + readr.read((char *)&npts_s32, sizeof(int32_t)); + readr.read((char *)&ndims_s32, sizeof(int32_t)); - writr.write((char *)&npts_s32, sizeof(_s32)); - writr.write((char *)&ndims_s32, sizeof(_s32)); + writr.write((char *)&npts_s32, sizeof(int32_t)); + writr.write((char *)&ndims_s32, sizeof(int32_t)); - _u64 npts = (_u64)npts_s32, ndims = (_u64)ndims_s32; + size_t npts = (size_t)npts_s32; + size_t ndims = (size_t)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; + size_t blk_size = 131072; + size_t 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++) + for (size_t i = 0; i < nblks; i++) { - _u64 cblk_size = std::min(npts - i * blk_size, blk_size); + size_t cblk_size = std::min(npts - i * blk_size, blk_size); block_convert(writr, readr, read_buf, cblk_size, ndims); } delete[] read_buf; @@ -148,8 +149,8 @@ double calculate_recall(unsigned num_queries, unsigned *gold_std, float *gs_dist gt.insert(gt_vec, gt_vec + tie_breaker); res.insert(res_vec, - res_vec + recall_at); // change to recall_at for recall k@k or - // dim_or for k@dim_or + res_vec + recall_at); // change to recall_at for recall k@k + // or dim_or for k@dim_or unsigned cur_recall = 0; for (auto &v : gt) { @@ -222,8 +223,8 @@ double calculate_recall(unsigned num_queries, unsigned *gold_std, float *gs_dist return ((double)(total_recall / (num_queries))) * ((double)(100.0 / recall_at)); } -double calculate_range_search_recall(unsigned num_queries, std::vector> &groundtruth, - std::vector> &our_results) +double calculate_range_search_recall(unsigned num_queries, std::vector> &groundtruth, + std::vector> &our_results) { double total_recall = 0; std::set gt, res; diff --git a/src/windows_aligned_file_reader.cpp b/src/windows_aligned_file_reader.cpp index 953b01641..a02929eee 100644 --- a/src/windows_aligned_file_reader.cpp +++ b/src/windows_aligned_file_reader.cpp @@ -153,11 +153,14 @@ void WindowsAlignedFileReader::read(std::vector &read_reqs, IOConte DWORD error = GetLastError(); if (error != WAIT_TIMEOUT) { - diskann::cerr << "GetQueuedCompletionStatus() failed with error = " << error << std::endl; + diskann::cerr << "GetQueuedCompletionStatus() failed " + "with error = " + << error << std::endl; throw diskann::ANNException("GetQueuedCompletionStatus failed with error: ", error, __FUNCSIG__, __FILE__, __LINE__); } - // no completion packet dequeued ==> sleep for 5us and try again + // no completion packet dequeued ==> sleep for 5us and try + // again std::this_thread::sleep_for(5us); } else diff --git a/tests/build_disk_index.cpp b/tests/build_disk_index.cpp index 025f74f4f..fb90502db 100644 --- a/tests/build_disk_index.cpp +++ b/tests/build_disk_index.cpp @@ -48,7 +48,8 @@ int main(int argc, char **argv) "Include full precision data in the index. Use only in " "conjuction with compressed data on SSD."); desc.add_options()("build_PQ_bytes", po::value(&build_PQ)->default_value(0), - "Number of PQ bytes to build the index; 0 for full precision build"); + "Number of PQ bytes to build the index; 0 for full " + "precision build"); desc.add_options()("use_opq", po::bool_switch()->default_value(false), "Use Optimized Product Quantization (OPQ)."); desc.add_options()("label_file", po::value(&label_file)->default_value(""), @@ -56,8 +57,10 @@ int main(int argc, char **argv) "The file should contain comma separated filters for each node " "with each line corresponding to a graph node"); desc.add_options()("universal_label", po::value(&universal_label)->default_value(""), - "Universal label, Use only in conjuction with label file for filtered " - "index build. If a graph node has all the labels against it, we can " + "Universal label, Use only in conjuction 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 point instead of comma " "separated filters for that point"); desc.add_options()("filtered_Lbuild,Lf", po::value(&Lf)->default_value(0), diff --git a/tests/build_memory_index.cpp b/tests/build_memory_index.cpp index c8066b3b7..df719d7e0 100644 --- a/tests/build_memory_index.cpp +++ b/tests/build_memory_index.cpp @@ -24,7 +24,7 @@ template int build_in_memory_index(const diskann::Metric &metric, const std::string &data_path, const unsigned R, const unsigned L, const float alpha, const std::string &save_path, const unsigned num_threads, const bool use_pq_build, const size_t num_pq_bytes, const bool use_opq, - const std::string &label_file, const std::string &universal_label, const _u32 Lf) + const std::string &label_file, const std::string &universal_label, const uint32_t Lf) { diskann::Parameters paras; paras.Set("R", R); @@ -37,7 +37,7 @@ int build_in_memory_index(const diskann::Metric &metric, const std::string &data std::string labels_file_to_use = save_path + "_label_formatted.txt"; std::string mem_labels_int_map_file = save_path + "_labels_map.txt"; - _u64 data_num, data_dim; + size_t data_num, data_dim; diskann::get_bin_metadata(data_path, data_num, data_dim); diskann::Index index(metric, data_dim, data_num, false, false, false, use_pq_build, num_pq_bytes, @@ -87,22 +87,26 @@ int main(int argc, char **argv) desc.add_options()("Lbuild,L", po::value(&L)->default_value(100), "Build complexity, higher value results in better graphs"); desc.add_options()("alpha", po::value(&alpha)->default_value(1.2f), - "alpha controls density and diameter of graph, set 1 for sparse graph, " + "alpha controls density and diameter of graph, set " + "1 for sparse graph, " "1.2 or 1.4 for denser graphs with lower diameter"); desc.add_options()("num_threads,T", po::value(&num_threads)->default_value(omp_get_num_procs()), "Number of threads used for building index (defaults to " "omp_get_num_procs())"); desc.add_options()("build_PQ_bytes", po::value(&build_PQ_bytes)->default_value(0), - "Number of PQ bytes to build the index; 0 for full precision build"); + "Number of PQ bytes to build the index; 0 for full precision " + "build"); desc.add_options()("use_opq", po::bool_switch()->default_value(false), - "Set true for OPQ compression while using PQ distance comparisons for " + "Set true for OPQ compression while using PQ " + "distance comparisons for " "building the index, and false for PQ compression"); desc.add_options()("label_file", po::value(&label_file)->default_value(""), "Input label file in txt format for Filtered Index search. " "The file should contain comma separated filters for each node " "with each line corresponding to a graph node"); desc.add_options()("universal_label", po::value(&universal_label)->default_value(""), - "Universal label, if using it, only in conjunction with labels_file"); + "Universal label, if using it, only in conjunction with " + "labels_file"); desc.add_options()("FilteredLbuild,Lf", po::value(&Lf)->default_value(0), "Build complexity for filtered points, higher value " "results in better graphs"); diff --git a/tests/build_stitched_index.cpp b/tests/build_stitched_index.cpp index 22c93d846..2a93a2dc9 100644 --- a/tests/build_stitched_index.cpp +++ b/tests/build_stitched_index.cpp @@ -30,9 +30,10 @@ typedef tsl::robin_set label_set; typedef std::string path; // structs for returning multiple items from a function -typedef std::tuple, tsl::robin_map, label_set> parse_label_file_return_values; -typedef std::tuple>, _u64> load_label_index_return_values; -typedef std::tuple>, _u64> stitch_indices_return_values; +typedef std::tuple, tsl::robin_map, label_set> + parse_label_file_return_values; +typedef std::tuple>, uint64_t> load_label_index_return_values; +typedef std::tuple>, uint64_t> stitch_indices_return_values; /* * Inline function to display progress bar. @@ -80,7 +81,8 @@ void handle_args(int argc, char **argv, std::string &data_type, path &input_data desc.add_options()("stitched_R", po::value(&stitched_R)->default_value(100), "Degree to prune final graph down to"); desc.add_options()("alpha", po::value(&alpha)->default_value(1.2f), - "alpha controls density and diameter of graph, set 1 for sparse graph, " + "alpha controls density and diameter of graph, set " + "1 for sparse graph, " "1.2 or 1.4 for denser graphs with lower diameter"); desc.add_options()("num_threads,T", po::value(&num_threads)->default_value(omp_get_num_procs()), "Number of threads used for building index (defaults to " @@ -90,7 +92,8 @@ void handle_args(int argc, char **argv, std::string &data_type, path &input_data desc.add_options()("universal_label", po::value(&universal_label)->default_value(""), "If a point comes with the specified universal label (and only the " "univ. " - "label), then the point is considered to have every possible label"); + "label), then the point is considered to have every possible " + "label"); po::variables_map vm; po::store(po::parse_command_line(argc, argv, desc), vm); @@ -131,10 +134,10 @@ parse_label_file_return_values parse_label_file(path label_data_path, std::strin // values to return std::vector point_ids_to_labels(line_cnt); - tsl::robin_map labels_to_number_of_points; + tsl::robin_map labels_to_number_of_points; label_set all_labels; - std::vector<_u32> points_with_universal_label; + std::vector points_with_universal_label; line_cnt = 0; while (std::getline(label_data_stream, line)) { @@ -142,7 +145,7 @@ parse_label_file_return_values parse_label_file(path label_data_path, std::strin label_set current_labels; // get point id - _u32 point_id = line_cnt; + uint32_t point_id = line_cnt; // parse comma separated labels bool current_universal_label_check = false; @@ -201,19 +204,19 @@ parse_label_file_return_values parse_label_file(path label_data_path, std::strin * input_data_path + "_" + label */ template -tsl::robin_map> generate_label_specific_vector_files( - path input_data_path, tsl::robin_map labels_to_number_of_points, +tsl::robin_map> generate_label_specific_vector_files( + path input_data_path, tsl::robin_map labels_to_number_of_points, std::vector point_ids_to_labels, label_set all_labels) { auto file_writing_timer = std::chrono::high_resolution_clock::now(); diskann::MemoryMapper input_data(input_data_path); char *input_start = input_data.getBuf(); - _u32 number_of_points, dimension; - std::memcpy(&number_of_points, input_start, sizeof(_u32)); - std::memcpy(&dimension, input_start + sizeof(_u32), sizeof(_u32)); - const _u32 VECTOR_SIZE = dimension * sizeof(T); - const size_t METADATA = 2 * sizeof(_u32); + uint32_t number_of_points, dimension; + std::memcpy(&number_of_points, input_start, sizeof(uint32_t)); + std::memcpy(&dimension, input_start + sizeof(uint32_t), sizeof(uint32_t)); + const uint32_t VECTOR_SIZE = dimension * sizeof(T); + const size_t METADATA = 2 * sizeof(uint32_t); if (number_of_points != point_ids_to_labels.size()) { std::cerr << "Error: number of points in labels file and data file differ." << std::endl; @@ -221,8 +224,8 @@ tsl::robin_map> generate_label_specific_vector_fi } tsl::robin_map label_to_iovec_map; - tsl::robin_map label_to_curr_iovec; - tsl::robin_map> label_id_to_orig_id; + tsl::robin_map label_to_curr_iovec; + tsl::robin_map> label_id_to_orig_id; // setup iovec list for each label for (const auto &lbl : all_labels) @@ -238,7 +241,7 @@ tsl::robin_map> generate_label_specific_vector_fi } // each point added to corresponding per-label iovec list - for (_u32 point_id = 0; point_id < number_of_points; point_id++) + for (uint32_t point_id = 0; point_id < number_of_points; point_id++) { char *curr_point = input_start + METADATA + (VECTOR_SIZE * point_id); iovec curr_iovec; @@ -258,7 +261,7 @@ tsl::robin_map> generate_label_specific_vector_fi { int label_input_data_fd; path curr_label_input_data_path(input_data_path + "_" + lbl); - _u32 curr_num_pts = labels_to_number_of_points[lbl]; + uint32_t curr_num_pts = labels_to_number_of_points[lbl]; label_input_data_fd = open(curr_label_input_data_path.c_str(), O_CREAT | O_WRONLY | O_TRUNC | O_APPEND, (mode_t)0644); @@ -266,8 +269,8 @@ tsl::robin_map> generate_label_specific_vector_fi throw; // write metadata - _u32 metadata[2] = {curr_num_pts, dimension}; - int return_value = write(label_input_data_fd, metadata, sizeof(_u32) * 2); + uint32_t metadata[2] = {curr_num_pts, dimension}; + int return_value = write(label_input_data_fd, metadata, sizeof(uint32_t) * 2); if (return_value == -1) { throw; @@ -308,17 +311,17 @@ tsl::robin_map> generate_label_specific_vector_fi // for use on systems without writev (i.e. Windows) template -tsl::robin_map> generate_label_specific_vector_files_compat( - path input_data_path, tsl::robin_map labels_to_number_of_points, +tsl::robin_map> generate_label_specific_vector_files_compat( + path input_data_path, tsl::robin_map labels_to_number_of_points, std::vector point_ids_to_labels, label_set all_labels) { auto file_writing_timer = std::chrono::high_resolution_clock::now(); std::ifstream input_data_stream(input_data_path); - _u32 number_of_points, dimension; - input_data_stream.read((char *)&number_of_points, sizeof(_u32)); - input_data_stream.read((char *)&dimension, sizeof(_u32)); - const _u32 VECTOR_SIZE = dimension * sizeof(T); + uint32_t number_of_points, dimension; + input_data_stream.read((char *)&number_of_points, sizeof(uint32_t)); + input_data_stream.read((char *)&dimension, sizeof(uint32_t)); + const uint32_t VECTOR_SIZE = dimension * sizeof(T); if (number_of_points != point_ids_to_labels.size()) { std::cerr << "Error: number of points in labels file and data file differ." << std::endl; @@ -326,12 +329,12 @@ tsl::robin_map> generate_label_specific_vector_fi } tsl::robin_map labels_to_vectors; - tsl::robin_map labels_to_curr_vector; - tsl::robin_map> label_id_to_orig_id; + tsl::robin_map labels_to_curr_vector; + tsl::robin_map> label_id_to_orig_id; for (const auto &lbl : all_labels) { - _u32 number_of_label_pts = labels_to_number_of_points[lbl]; + uint32_t number_of_label_pts = labels_to_number_of_points[lbl]; char *vectors = (char *)malloc(number_of_label_pts * VECTOR_SIZE); if (vectors == nullptr) { @@ -342,7 +345,7 @@ tsl::robin_map> generate_label_specific_vector_fi label_id_to_orig_id[lbl].reserve(number_of_label_pts); } - for (_u32 point_id = 0; point_id < number_of_points; point_id++) + for (uint32_t point_id = 0; point_id < number_of_points; point_id++) { char *curr_vector = (char *)malloc(VECTOR_SIZE); input_data_stream.read(curr_vector, VECTOR_SIZE); @@ -359,13 +362,13 @@ tsl::robin_map> generate_label_specific_vector_fi for (const auto &lbl : all_labels) { path curr_label_input_data_path(input_data_path + "_" + lbl); - _u32 number_of_label_pts = labels_to_number_of_points[lbl]; + uint32_t number_of_label_pts = labels_to_number_of_points[lbl]; std::ofstream label_file_stream; label_file_stream.exceptions(std::ios::badbit | std::ios::failbit); label_file_stream.open(curr_label_input_data_path, std::ios_base::binary); - label_file_stream.write((char *)&number_of_label_pts, sizeof(_u32)); - label_file_stream.write((char *)&dimension, sizeof(_u32)); + label_file_stream.write((char *)&number_of_label_pts, sizeof(uint32_t)); + label_file_stream.write((char *)&dimension, sizeof(uint32_t)); label_file_stream.write((char *)labels_to_vectors[lbl], number_of_label_pts * VECTOR_SIZE); label_file_stream.close(); @@ -437,33 +440,33 @@ void generate_label_indices(path input_data_path, path final_index_path_prefix, * * Returns both the graph index and the size of the file in bytes. */ -load_label_index_return_values load_label_index(path label_index_path, _u32 label_number_of_points) +load_label_index_return_values load_label_index(path label_index_path, uint32_t label_number_of_points) { std::ifstream label_index_stream; label_index_stream.exceptions(std::ios::badbit | std::ios::failbit); label_index_stream.open(label_index_path, std::ios::binary); - _u64 index_file_size, index_num_frozen_points; - _u32 index_max_observed_degree, index_entry_point; - const size_t INDEX_METADATA = 2 * sizeof(_u64) + 2 * sizeof(_u32); - label_index_stream.read((char *)&index_file_size, sizeof(_u64)); - label_index_stream.read((char *)&index_max_observed_degree, sizeof(_u32)); - label_index_stream.read((char *)&index_entry_point, sizeof(_u32)); - label_index_stream.read((char *)&index_num_frozen_points, sizeof(_u64)); + uint64_t index_file_size, index_num_frozen_points; + uint32_t index_max_observed_degree, index_entry_point; + const size_t INDEX_METADATA = 2 * sizeof(uint64_t) + 2 * sizeof(uint32_t); + label_index_stream.read((char *)&index_file_size, sizeof(uint64_t)); + label_index_stream.read((char *)&index_max_observed_degree, sizeof(uint32_t)); + label_index_stream.read((char *)&index_entry_point, sizeof(uint32_t)); + label_index_stream.read((char *)&index_num_frozen_points, sizeof(uint64_t)); size_t bytes_read = INDEX_METADATA; - std::vector> label_index(label_number_of_points); - _u32 nodes_read = 0; + std::vector> label_index(label_number_of_points); + uint32_t nodes_read = 0; while (bytes_read != index_file_size) { - _u32 current_node_num_neighbors; - label_index_stream.read((char *)¤t_node_num_neighbors, sizeof(_u32)); + uint32_t current_node_num_neighbors; + label_index_stream.read((char *)¤t_node_num_neighbors, sizeof(uint32_t)); nodes_read++; - std::vector<_u32> current_node_neighbors(current_node_num_neighbors); - label_index_stream.read((char *)current_node_neighbors.data(), current_node_num_neighbors * sizeof(_u32)); + std::vector current_node_neighbors(current_node_num_neighbors); + label_index_stream.read((char *)current_node_neighbors.data(), current_node_num_neighbors * sizeof(uint32_t)); label_index[nodes_read - 1].swap(current_node_neighbors); - bytes_read += sizeof(_u32) * (current_node_num_neighbors + 1); + bytes_read += sizeof(uint32_t) * (current_node_num_neighbors + 1); } return std::make_tuple(label_index, index_file_size); @@ -477,9 +480,10 @@ load_label_index_return_values load_label_index(path label_index_path, _u32 labe * 3. data (redundant for static indices) * 4. labels (redundant for static indices) */ -void save_full_index(path final_index_path_prefix, path input_data_path, _u64 final_index_size, - std::vector> stitched_graph, tsl::robin_map entry_points, - std::string universal_label, path label_data_path) +void save_full_index(path final_index_path_prefix, path input_data_path, uint64_t final_index_size, + std::vector> stitched_graph, + tsl::robin_map entry_points, std::string universal_label, + path label_data_path) { // aux. file 1 auto saving_index_timer = std::chrono::high_resolution_clock::now(); @@ -523,34 +527,34 @@ void save_full_index(path final_index_path_prefix, path input_data_path, _u64 fi } // main index - _u64 index_num_frozen_points = 0, index_num_edges = 0; - _u32 index_max_observed_degree = 0, index_entry_point = 0; - const size_t METADATA = 2 * sizeof(_u64) + 2 * sizeof(_u32); + uint64_t index_num_frozen_points = 0, index_num_edges = 0; + uint32_t index_max_observed_degree = 0, index_entry_point = 0; + const size_t METADATA = 2 * sizeof(uint64_t) + 2 * sizeof(uint32_t); for (auto &point_neighbors : stitched_graph) { - index_max_observed_degree = std::max(index_max_observed_degree, (_u32)point_neighbors.size()); + index_max_observed_degree = std::max(index_max_observed_degree, (uint32_t)point_neighbors.size()); } std::ofstream stitched_graph_writer; stitched_graph_writer.exceptions(std::ios::badbit | std::ios::failbit); stitched_graph_writer.open(final_index_path_prefix, std::ios_base::binary); - stitched_graph_writer.write((char *)&final_index_size, sizeof(_u64)); - stitched_graph_writer.write((char *)&index_max_observed_degree, sizeof(_u32)); - stitched_graph_writer.write((char *)&index_entry_point, sizeof(_u32)); - stitched_graph_writer.write((char *)&index_num_frozen_points, sizeof(_u64)); + stitched_graph_writer.write((char *)&final_index_size, sizeof(uint64_t)); + stitched_graph_writer.write((char *)&index_max_observed_degree, sizeof(uint32_t)); + stitched_graph_writer.write((char *)&index_entry_point, sizeof(uint32_t)); + stitched_graph_writer.write((char *)&index_num_frozen_points, sizeof(uint64_t)); size_t bytes_written = METADATA; - for (_u32 node_point = 0; node_point < stitched_graph.size(); node_point++) + for (uint32_t node_point = 0; node_point < stitched_graph.size(); node_point++) { - _u32 current_node_num_neighbors = stitched_graph[node_point].size(); - std::vector<_u32> current_node_neighbors = stitched_graph[node_point]; - stitched_graph_writer.write((char *)¤t_node_num_neighbors, sizeof(_u32)); - bytes_written += sizeof(_u32); + uint32_t current_node_num_neighbors = stitched_graph[node_point].size(); + std::vector current_node_neighbors = stitched_graph[node_point]; + stitched_graph_writer.write((char *)¤t_node_num_neighbors, sizeof(uint32_t)); + bytes_written += sizeof(uint32_t); for (const auto ¤t_node_neighbor : current_node_neighbors) { - stitched_graph_writer.write((char *)¤t_node_neighbor, sizeof(_u32)); - bytes_written += sizeof(_u32); + stitched_graph_writer.write((char *)¤t_node_neighbor, sizeof(uint32_t)); + bytes_written += sizeof(uint32_t); } index_num_edges += current_node_num_neighbors; } @@ -578,45 +582,46 @@ void save_full_index(path final_index_path_prefix, path input_data_path, _u64 fi */ template stitch_indices_return_values stitch_label_indices( - path final_index_path_prefix, _u32 total_number_of_points, label_set all_labels, - tsl::robin_map labels_to_number_of_points, tsl::robin_map &label_entry_points, - tsl::robin_map> label_id_to_orig_id_map) + path final_index_path_prefix, uint32_t total_number_of_points, label_set all_labels, + tsl::robin_map labels_to_number_of_points, + tsl::robin_map &label_entry_points, + tsl::robin_map> label_id_to_orig_id_map) { size_t final_index_size = 0; - std::vector> stitched_graph(total_number_of_points); + std::vector> stitched_graph(total_number_of_points); auto stitching_index_timer = std::chrono::high_resolution_clock::now(); for (const auto &lbl : all_labels) { path curr_label_index_path(final_index_path_prefix + "_" + lbl); - std::vector> curr_label_index; - _u64 curr_label_index_size; - _u32 curr_label_entry_point; + std::vector> curr_label_index; + uint64_t curr_label_index_size; + uint32_t curr_label_entry_point; std::tie(curr_label_index, curr_label_index_size) = load_label_index(curr_label_index_path, labels_to_number_of_points[lbl]); curr_label_entry_point = random(0, curr_label_index.size()); label_entry_points[lbl] = label_id_to_orig_id_map[lbl][curr_label_entry_point]; - for (_u32 node_point = 0; node_point < curr_label_index.size(); node_point++) + for (uint32_t node_point = 0; node_point < curr_label_index.size(); node_point++) { - _u32 original_point_id = label_id_to_orig_id_map[lbl][node_point]; + uint32_t original_point_id = label_id_to_orig_id_map[lbl][node_point]; for (auto &node_neighbor : curr_label_index[node_point]) { - _u32 original_neighbor_id = label_id_to_orig_id_map[lbl][node_neighbor]; - std::vector<_u32> curr_point_neighbors = stitched_graph[original_point_id]; + uint32_t original_neighbor_id = label_id_to_orig_id_map[lbl][node_neighbor]; + std::vector curr_point_neighbors = stitched_graph[original_point_id]; if (std::find(curr_point_neighbors.begin(), curr_point_neighbors.end(), original_neighbor_id) == curr_point_neighbors.end()) { stitched_graph[original_point_id].push_back(original_neighbor_id); - final_index_size += sizeof(_u32); + final_index_size += sizeof(uint32_t); } } } } - const size_t METADATA = 2 * sizeof(_u64) + 2 * sizeof(_u32); - final_index_size += (total_number_of_points * sizeof(_u32) + METADATA); + const size_t METADATA = 2 * sizeof(uint64_t) + 2 * sizeof(uint32_t); + final_index_size += (total_number_of_points * sizeof(uint32_t) + METADATA); std::chrono::duration stitching_index_time = std::chrono::high_resolution_clock::now() - stitching_index_timer; @@ -634,8 +639,8 @@ stitch_indices_return_values stitch_label_indices( */ template void prune_and_save(path final_index_path_prefix, path full_index_path_prefix, path input_data_path, - std::vector> stitched_graph, unsigned stitched_R, - tsl::robin_map label_entry_points, std::string universal_label, + std::vector> stitched_graph, unsigned stitched_R, + tsl::robin_map label_entry_points, std::string universal_label, path label_data_path, unsigned num_threads) { size_t dimension, number_of_label_points; @@ -710,15 +715,15 @@ int main(int argc, char **argv) // 2. parse label file and create necessary data structures std::vector point_ids_to_labels; - tsl::robin_map labels_to_number_of_points; + tsl::robin_map labels_to_number_of_points; label_set all_labels; std::tie(point_ids_to_labels, labels_to_number_of_points, all_labels) = parse_label_file(labels_file_to_use, universal_label); // 3. for each label, make a separate data file - tsl::robin_map> label_id_to_orig_id_map; - _u32 total_number_of_points = point_ids_to_labels.size(); + tsl::robin_map> label_id_to_orig_id_map; + uint32_t total_number_of_points = point_ids_to_labels.size(); #ifndef _WINDOWS if (data_type == "uint8") @@ -757,9 +762,9 @@ int main(int argc, char **argv) throw; // 5. "stitch" the indices together - std::vector> stitched_graph; - tsl::robin_map label_entry_points; - _u64 stitched_graph_size; + std::vector> stitched_graph; + tsl::robin_map label_entry_points; + uint64_t stitched_graph_size; if (data_type == "uint8") std::tie(stitched_graph, stitched_graph_size) = diff --git a/tests/range_search_disk_index.cpp b/tests/range_search_disk_index.cpp index 2120c2f6c..8a5a23404 100644 --- a/tests/range_search_disk_index.cpp +++ b/tests/range_search_disk_index.cpp @@ -66,7 +66,7 @@ int search_disk_index(diskann::Metric &metric, const std::string &index_path_pre // load query bin T *query = nullptr; - std::vector> groundtruth_ids; + std::vector> groundtruth_ids; size_t query_num, query_dim, query_aligned_dim, gt_num; diskann::load_aligned_bin(query_file, query, query_num, query_dim, query_aligned_dim); @@ -110,7 +110,8 @@ int search_disk_index(diskann::Metric &metric, const std::string &index_path_pre 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); + // 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(); @@ -129,7 +130,7 @@ int search_disk_index(diskann::Metric &metric, const std::string &index_path_pre } else { - warmup_num = (std::min)((_u32)150000, (_u32)15000 * num_threads); + warmup_num = (std::min)((uint32_t)150000, (uint32_t)15000 * num_threads); warmup_dim = query_dim; warmup_aligned_dim = query_aligned_dim; diskann::alloc_aligned(((void **)&warmup), warmup_num * warmup_aligned_dim * sizeof(T), 8 * sizeof(T)); @@ -150,7 +151,7 @@ int search_disk_index(diskann::Metric &metric, const std::string &index_path_pre 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++) + for (int64_t 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), @@ -183,7 +184,7 @@ int search_disk_index(diskann::Metric &metric, const std::string &index_path_pre for (uint32_t test_id = 0; test_id < Lvec.size(); test_id++) { - _u64 L = Lvec[test_id]; + uint64_t L = Lvec[test_id]; if (beamwidth <= 0) { @@ -200,15 +201,16 @@ int search_disk_index(diskann::Metric &metric, const std::string &index_path_pre auto s = std::chrono::high_resolution_clock::now(); #pragma omp parallel for schedule(dynamic, 1) - for (_s64 i = 0; i < (int64_t)query_num; i++) + for (int64_t i = 0; i < (int64_t)query_num; i++) { - std::vector<_u64> indices; + std::vector indices; std::vector distances; - _u32 res_count = _pFlashIndex->range_search(query + (i * query_aligned_dim), search_range, L, max_list_size, - indices, distances, optimized_beamwidth, stats + i); + uint32_t res_count = + _pFlashIndex->range_search(query + (i * query_aligned_dim), search_range, L, max_list_size, indices, + distances, optimized_beamwidth, stats + i); query_result_ids[test_id][i].reserve(res_count); query_result_ids[test_id][i].resize(res_count); - for (_u32 idx = 0; idx < res_count; idx++) + for (uint32_t idx = 0; idx < res_count; idx++) query_result_ids[test_id][i][idx] = indices[idx]; } auto e = std::chrono::high_resolution_clock::now(); @@ -233,9 +235,9 @@ int search_disk_index(diskann::Metric &metric, const std::string &index_path_pre { recall = diskann::calculate_range_search_recall(query_num, groundtruth_ids, query_result_ids[test_id]); - _u32 total_true_positive = 0; - _u32 total_positive = 0; - for (_u32 i = 0; i < query_num; i++) + uint32_t total_true_positive = 0; + uint32_t total_positive = 0; + for (uint32_t i = 0; i < query_num; i++) { total_true_positive += query_result_ids[test_id][i].size(); total_positive += groundtruth_ids[i].size(); diff --git a/tests/search_disk_index.cpp b/tests/search_disk_index.cpp index b1d04d7c6..641ccf893 100644 --- a/tests/search_disk_index.cpp +++ b/tests/search_disk_index.cpp @@ -50,16 +50,16 @@ template int search_disk_index(diskann::Metric &metric, const std::string &index_path_prefix, const std::string &result_output_prefix, const std::string &query_file, std::string >_file, const unsigned num_threads, const unsigned recall_at, const unsigned beamwidth, - const unsigned num_nodes_to_cache, const _u32 search_io_limit, const std::vector &Lvec, - const float fail_if_recall_below, const bool use_reorder_data = false, - const std::string &filter_label = "") + const unsigned num_nodes_to_cache, const uint32_t search_io_limit, + const std::vector &Lvec, const float fail_if_recall_below, + const bool use_reorder_data = false, const std::string &filter_label = "") { diskann::cout << "Search parameters: #threads: " << num_threads << ", "; if (beamwidth <= 0) diskann::cout << "beamwidth to be optimized for each L value" << std::flush; else diskann::cout << " beamwidth: " << beamwidth << std::flush; - if (search_io_limit == std::numeric_limits<_u32>::max()) + if (search_io_limit == std::numeric_limits::max()) diskann::cout << "." << std::endl; else diskann::cout << ", io_limit: " << search_io_limit << "." << std::endl; @@ -133,7 +133,7 @@ int search_disk_index(diskann::Metric &metric, const std::string &index_path_pre } else { - warmup_num = (std::min)((_u32)150000, (_u32)15000 * num_threads); + warmup_num = (std::min)((uint32_t)150000, (uint32_t)15000 * num_threads); warmup_dim = query_dim; warmup_aligned_dim = query_aligned_dim; diskann::alloc_aligned(((void **)&warmup), warmup_num * warmup_aligned_dim * sizeof(T), 8 * sizeof(T)); @@ -154,7 +154,7 @@ int search_disk_index(diskann::Metric &metric, const std::string &index_path_pre 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++) + for (int64_t 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), @@ -189,7 +189,7 @@ int search_disk_index(diskann::Metric &metric, const std::string &index_path_pre for (uint32_t test_id = 0; test_id < Lvec.size(); test_id++) { - _u64 L = Lvec[test_id]; + uint64_t L = Lvec[test_id]; if (L < recall_at) { @@ -215,7 +215,7 @@ int search_disk_index(diskann::Metric &metric, const std::string &index_path_pre auto s = std::chrono::high_resolution_clock::now(); #pragma omp parallel for schedule(dynamic, 1) - for (_s64 i = 0; i < (int64_t)query_num; i++) + for (int64_t i = 0; i < (int64_t)query_num; i++) { if (!filtered_search) { @@ -273,14 +273,14 @@ int search_disk_index(diskann::Metric &metric, const std::string &index_path_pre } diskann::cout << "Done searching. Now saving results " << std::endl; - _u64 test_id = 0; + uint64_t test_id = 0; for (auto L : Lvec) { if (L < recall_at) continue; std::string cur_result_path = result_output_prefix + "_" + std::to_string(L) + "_idx_uint32.bin"; - diskann::save_bin<_u32>(cur_result_path, query_result_ids[test_id].data(), query_num, recall_at); + diskann::save_bin(cur_result_path, query_result_ids[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); @@ -324,7 +324,7 @@ int main(int argc, char **argv) desc.add_options()("num_nodes_to_cache", po::value(&num_nodes_to_cache)->default_value(0), "Beamwidth for search"); desc.add_options()("search_io_limit", - po::value(&search_io_limit)->default_value(std::numeric_limits<_u32>::max()), + po::value(&search_io_limit)->default_value(std::numeric_limits::max()), "Max #IOs for search"); desc.add_options()("num_threads,T", po::value(&num_threads)->default_value(omp_get_num_procs()), "Number of threads used for building index (defaults to " diff --git a/tests/search_memory_index.cpp b/tests/search_memory_index.cpp index 4e093998d..ececf5f57 100644 --- a/tests/search_memory_index.cpp +++ b/tests/search_memory_index.cpp @@ -123,7 +123,7 @@ int search_memory_index(diskann::Metric &metric, const std::string &index_path, for (uint32_t test_id = 0; test_id < Lvec.size(); test_id++) { - _u64 L = Lvec[test_id]; + uint64_t L = Lvec[test_id]; if (L < recall_at) { diskann::cout << "Ignoring search with L:" << L << " since it's smaller than K:" << recall_at << std::endl; @@ -200,13 +200,13 @@ int search_memory_index(diskann::Metric &metric, const std::string &index_path, if (tags) { std::cout << std::setw(4) << L << std::setw(12) << displayed_qps << std::setw(20) << (float)mean_latency - << std::setw(15) << (float)latency_stats[(_u64)(0.999 * query_num)]; + << std::setw(15) << (float)latency_stats[(uint64_t)(0.999 * query_num)]; } else { std::cout << std::setw(4) << L << std::setw(12) << displayed_qps << std::setw(18) << avg_cmps << std::setw(20) << (float)mean_latency << std::setw(15) - << (float)latency_stats[(_u64)(0.999 * query_num)]; + << (float)latency_stats[(uint64_t)(0.999 * query_num)]; } for (float recall : recalls) { @@ -217,7 +217,7 @@ int search_memory_index(diskann::Metric &metric, const std::string &index_path, } std::cout << "Done searching. Now saving results " << std::endl; - _u64 test_id = 0; + uint64_t test_id = 0; for (auto L : Lvec) { if (L < recall_at) @@ -226,7 +226,7 @@ int search_memory_index(diskann::Metric &metric, const std::string &index_path, continue; } std::string cur_result_path = result_path_prefix + "_" + std::to_string(L) + "_idx_uint32.bin"; - diskann::save_bin<_u32>(cur_result_path, query_result_ids[test_id].data(), query_num, recall_at); + diskann::save_bin(cur_result_path, query_result_ids[test_id].data(), query_num, recall_at); test_id++; } diff --git a/tests/test_insert_deletes_consolidate.cpp b/tests/test_insert_deletes_consolidate.cpp index ef598c659..886265448 100644 --- a/tests/test_insert_deletes_consolidate.cpp +++ b/tests/test_insert_deletes_consolidate.cpp @@ -350,7 +350,8 @@ int main(int argc, char **argv) desc.add_options()("Lbuild,L", po::value(&L)->default_value(100), "Build complexity, higher value results in better graphs"); desc.add_options()("alpha", po::value(&alpha)->default_value(1.2f), - "alpha controls density and diameter of graph, set 1 for sparse graph, " + "alpha controls density and diameter of graph, set " + "1 for sparse graph, " "1.2 or 1.4 for denser graphs with lower diameter"); desc.add_options()("num_threads,T", po::value(&num_threads)->default_value(omp_get_num_procs()), "Number of threads used for building index (defaults to " @@ -373,7 +374,8 @@ int main(int argc, char **argv) desc.add_options()("start_point_norm", po::value(&start_point_norm)->default_value(0), "Set the start point to a random point on a sphere of this radius"); desc.add_options()("num_start_points", po::value(&num_start_pts)->default_value(0), - "Set the number of random start (frozen) points to use when inserting and searching"); + "Set the number of random start (frozen) points to use when " + "inserting and searching"); po::variables_map vm; po::store(po::parse_command_line(argc, argv, desc), vm); @@ -386,7 +388,8 @@ int main(int argc, char **argv) if (beginning_index_size == 0) if (start_point_norm == 0) { - std::cout << "When beginning_index_size is 0, use a start point with " + std::cout << "When beginning_index_size is 0, use a start " + "point with " "appropriate norm" << std::endl; return -1; diff --git a/tests/test_streaming_scenario.cpp b/tests/test_streaming_scenario.cpp index b1f655162..18adef67f 100644 --- a/tests/test_streaming_scenario.cpp +++ b/tests/test_streaming_scenario.cpp @@ -208,8 +208,9 @@ void build_incremental_index(const std::string &data_path, const unsigned L, con throw diskann::ANNException("num_points < max_points_to_insert", -1, __FUNCSIG__, __FILE__, __LINE__); if (max_points_to_insert < active_window + consolidate_interval) - throw diskann::ANNException("ERROR: max_points_to_insert < active_window + consolidate_interval", -1, - __FUNCSIG__, __FILE__, __LINE__); + throw diskann::ANNException("ERROR: max_points_to_insert < " + "active_window + consolidate_interval", + -1, __FUNCSIG__, __FILE__, __LINE__); if (consolidate_interval < max_points_to_insert / 1000) throw diskann::ANNException("ERROR: consolidate_interval is too small", -1, __FUNCSIG__, __FILE__, __LINE__); @@ -296,7 +297,8 @@ int main(int argc, char **argv) desc.add_options()("Lbuild,L", po::value(&L)->default_value(100), "Build complexity, higher value results in better graphs"); desc.add_options()("alpha", po::value(&alpha)->default_value(1.2f), - "alpha controls density and diameter of graph, set 1 for sparse graph, " + "alpha controls density and diameter of graph, set " + "1 for sparse graph, " "1.2 or 1.4 for denser graphs with lower diameter"); desc.add_options()("insert_threads", po::value(&insert_threads)->default_value(omp_get_num_procs() / 2), @@ -308,17 +310,20 @@ int main(int argc, char **argv) "the index (defaults to omp_get_num_procs()/2)"); desc.add_options()("max_points_to_insert", po::value(&max_points_to_insert)->default_value(0), - "The number of points from the file that the program streams over "); + "The number of points from the file that the program streams " + "over "); desc.add_options()("active_window", po::value(&active_window)->required(), "Program maintains an index over an active window of " "this size that slides through the data"); desc.add_options()("consolidate_interval", po::value(&consolidate_interval)->required(), - "The program simultaneously adds this number of points to the right of " + "The program simultaneously adds this number of points to the " + "right of " "the window while deleting the same number from the left"); desc.add_options()("start_point_norm", po::value(&start_point_norm)->required(), "Set the start point to a random point on a sphere of this radius"); desc.add_options()("num_start_points", po::value(&num_start_pts)->default_value(0), - "Set the number of random start (frozen) points to use when inserting and searching"); + "Set the number of random start (frozen) points to use when " + "inserting and searching"); po::variables_map vm; po::store(po::parse_command_line(argc, argv, desc), vm); diff --git a/tests/utils/bin_to_tsv.cpp b/tests/utils/bin_to_tsv.cpp index 813fa6e9f..7851bef6d 100644 --- a/tests/utils/bin_to_tsv.cpp +++ b/tests/utils/bin_to_tsv.cpp @@ -4,13 +4,14 @@ #include #include "utils.h" -template void block_convert(std::ofstream &writer, std::ifstream &reader, T *read_buf, _u64 npts, _u64 ndims) +template +void block_convert(std::ofstream &writer, std::ifstream &reader, T *read_buf, size_t npts, size_t ndims) { reader.read((char *)read_buf, npts * ndims * sizeof(float)); - for (_u64 i = 0; i < npts; i++) + for (size_t i = 0; i < npts; i++) { - for (_u64 d = 0; d < ndims; d++) + for (size_t d = 0; d < ndims; d++) { writer << read_buf[d + i * ndims]; if (d < ndims - 1) @@ -36,22 +37,22 @@ int main(int argc, char **argv) } std::ifstream reader(argv[2], std::ios::binary); - _u32 npts_u32; - _u32 ndims_u32; - reader.read((char *)&npts_u32, sizeof(_s32)); - reader.read((char *)&ndims_u32, sizeof(_s32)); + uint32_t npts_u32; + uint32_t ndims_u32; + reader.read((char *)&npts_u32, sizeof(uint32_t)); + reader.read((char *)&ndims_u32, sizeof(uint32_t)); size_t npts = npts_u32; size_t ndims = ndims_u32; std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims << std::endl; - _u64 blk_size = 131072; - _u64 nblks = ROUND_UP(npts, blk_size) / blk_size; + size_t blk_size = 131072; + size_t nblks = ROUND_UP(npts, blk_size) / blk_size; std::ofstream writer(argv[3]); char *read_buf = new char[blk_size * ndims * 4]; - for (_u64 i = 0; i < nblks; i++) + for (size_t i = 0; i < nblks; i++) { - _u64 cblk_size = std::min(npts - i * blk_size, blk_size); + size_t 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); else if (type_string == std::string("int8")) diff --git a/tests/utils/calculate_recall.cpp b/tests/utils/calculate_recall.cpp index 3307104ce..a67c7bfdc 100644 --- a/tests/utils/calculate_recall.cpp +++ b/tests/utils/calculate_recall.cpp @@ -31,7 +31,9 @@ int main(int argc, char **argv) if (points_num_gs != points_num_or) { - std::cout << "Error. Number of queries mismatch in ground truth and our results" << std::endl; + std::cout << "Error. Number of queries mismatch in ground truth and " + "our results" + << std::endl; return -1; } points_num = points_num_gs; diff --git a/tests/utils/compute_groundtruth.cpp b/tests/utils/compute_groundtruth.cpp index 231579d77..3c81c8bd6 100644 --- a/tests/utils/compute_groundtruth.cpp +++ b/tests/utils/compute_groundtruth.cpp @@ -142,28 +142,28 @@ void exact_knn(const size_t dim, const size_t k, points = new float[npoints * dim]; queries = new float[nqueries * dim]; #pragma omp parallel for schedule(static, 4096) - for (_s64 i = 0; i < (_s64)npoints; i++) + for (int64_t i = 0; i < (int64_t)npoints; i++) { float norm = std::sqrt(points_l2sq[i]); if (norm == 0) { norm = std::numeric_limits::epsilon(); } - for (_u32 j = 0; j < dim; j++) + for (uint32_t j = 0; j < dim; j++) { points[i * dim + j] = points_in[i * dim + j] / norm; } } #pragma omp parallel for schedule(static, 4096) - for (_s64 i = 0; i < (_s64)nqueries; i++) + for (int64_t i = 0; i < (int64_t)nqueries; i++) { float norm = std::sqrt(queries_l2sq[i]); if (norm == 0) { norm = std::numeric_limits::epsilon(); } - for (_u32 j = 0; j < dim; j++) + for (uint32_t j = 0; j < dim; j++) { queries[i * dim + j] = queries_in[i * dim + j] / norm; } @@ -186,7 +186,7 @@ void exact_knn(const size_t dim, const size_t k, size_t q_batch_size = (1 << 9); float *dist_matrix = new float[(size_t)q_batch_size * (size_t)npoints]; - for (_u64 b = 0; b < div_round_up(nqueries, q_batch_size); ++b) + for (size_t b = 0; b < div_round_up(nqueries, q_batch_size); ++b) { int64_t q_b = b * q_batch_size; int64_t q_e = ((b + 1) * q_batch_size > nqueries) ? nqueries : (b + 1) * q_batch_size; @@ -207,9 +207,9 @@ void exact_knn(const size_t dim, const size_t k, for (long long q = q_b; q < q_e; q++) { maxPQIFCS point_dist; - for (_u64 p = 0; p < k; p++) + for (size_t p = 0; p < k; p++) point_dist.emplace(p, dist_matrix[(ptrdiff_t)p + (ptrdiff_t)(q - q_b) * (ptrdiff_t)npoints]); - for (_u64 p = k; p < npoints; p++) + for (size_t p = k; p < npoints; p++) { if (point_dist.top().second > dist_matrix[(ptrdiff_t)p + (ptrdiff_t)(q - q_b) * (ptrdiff_t)npoints]) point_dist.emplace(p, dist_matrix[(ptrdiff_t)p + (ptrdiff_t)(q - q_b) * (ptrdiff_t)npoints]); @@ -257,7 +257,7 @@ template inline int get_num_parts(const char *filename) } template -inline void load_bin_as_float(const char *filename, float *&data, size_t &npts_u64, size_t &ndims_u64, int part_num) +inline void load_bin_as_float(const char *filename, float *&data, size_t &npts, size_t &ndims, int part_num) { std::ifstream reader; reader.exceptions(std::ios::failbit | std::ios::badbit); @@ -268,24 +268,24 @@ inline void load_bin_as_float(const char *filename, float *&data, size_t &npts_u reader.read((char *)&ndims_i32, sizeof(int)); uint64_t start_id = part_num * PARTSIZE; uint64_t end_id = (std::min)(start_id + PARTSIZE, (uint64_t)npts_i32); - npts_u64 = end_id - start_id; - ndims_u64 = (uint64_t)ndims_i32; - std::cout << "#pts in part = " << npts_u64 << ", #dims = " << ndims_u64 - << ", size = " << npts_u64 * ndims_u64 * sizeof(T) << "B" << std::endl; - - reader.seekg(start_id * ndims_u64 * sizeof(T) + 2 * sizeof(uint32_t), std::ios::beg); - T *data_T = new T[npts_u64 * ndims_u64]; - reader.read((char *)data_T, sizeof(T) * npts_u64 * ndims_u64); + npts = end_id - start_id; + ndims = (uint64_t)ndims_i32; + std::cout << "#pts in part = " << npts << ", #dims = " << ndims << ", size = " << npts * ndims * sizeof(T) << "B" + << std::endl; + + reader.seekg(start_id * ndims * sizeof(T) + 2 * sizeof(uint32_t), std::ios::beg); + T *data_T = new T[npts * ndims]; + reader.read((char *)data_T, sizeof(T) * npts * ndims); std::cout << "Finished reading part of the bin file." << std::endl; reader.close(); - data = aligned_malloc(npts_u64 * ndims_u64, ALIGNMENT); + data = aligned_malloc(npts * ndims, ALIGNMENT); #pragma omp parallel for schedule(dynamic, 32768) - for (int64_t i = 0; i < (int64_t)npts_u64; i++) + for (int64_t i = 0; i < (int64_t)npts; i++) { - for (int64_t j = 0; j < (int64_t)ndims_u64; j++) + for (int64_t j = 0; j < (int64_t)ndims; j++) { - float cur_val_float = (float)data_T[i * ndims_u64 + j]; - std::memcpy((char *)(data + i * ndims_u64 + j), (char *)&cur_val_float, sizeof(float)); + float cur_val_float = (float)data_T[i * ndims + j]; + std::memcpy((char *)(data + i * ndims + j), (char *)&cur_val_float, sizeof(float)); } } delete[] data_T; @@ -489,7 +489,7 @@ int aux_main(const std::string &base_file, const std::string &label_file, const int *closest_points_part = new int[nqueries * k]; float *dist_closest_points_part = new float[nqueries * k]; - _u32 part_k; + uint32_t part_k; if (filter_label == "") { part_k = k < npoints ? k : npoints; @@ -506,9 +506,9 @@ int aux_main(const std::string &base_file, const std::string &label_file, const } } - for (_u64 i = 0; i < nqueries; i++) + for (size_t i = 0; i < nqueries; i++) { - for (_u64 j = 0; j < part_k; j++) + for (size_t j = 0; j < part_k; j++) { if (tags_enabled) if (location_to_tag[closest_points_part[i * k + j] + start_id] == 0) @@ -532,7 +532,7 @@ int aux_main(const std::string &base_file, const std::string &label_file, const diskann::aligned_free(base_data); } - for (_u64 i = 0; i < nqueries; i++) + for (size_t i = 0; i < nqueries; i++) { std::vector> &cur_res = results[i]; std::sort(cur_res.begin(), cur_res.end(), custom_dist); @@ -592,7 +592,8 @@ int main(int argc, char **argv) desc.add_options()("filter_label", po::value(&filter_label)->default_value(""), "Input filter label if doing filtered groundtruth"); desc.add_options()("universal_label", po::value(&universal_label)->default_value(""), - "Universal label, if using it, only in conjunction with label_file"); + "Universal label, if using it, only in conjunction with " + "label_file"); desc.add_options()("gt_file", po::value(>_file)->required(), "File name for the writing ground truth in binary format"); diff --git a/tests/utils/float_bin_to_int8.cpp b/tests/utils/float_bin_to_int8.cpp index fd63cb353..a7632a6cf 100644 --- a/tests/utils/float_bin_to_int8.cpp +++ b/tests/utils/float_bin_to_int8.cpp @@ -4,14 +4,14 @@ #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) +void block_convert(std::ofstream &writer, int8_t *write_buf, std::ifstream &reader, float *read_buf, size_t npts, + size_t ndims, float bias, float scale) { reader.read((char *)read_buf, npts * ndims * sizeof(float)); - for (_u64 i = 0; i < npts; i++) + for (size_t i = 0; i < npts; i++) { - for (_u64 d = 0; d < ndims; d++) + for (size_t d = 0; d < ndims; d++) { write_buf[d + i * ndims] = (int8_t)((read_buf[d + i * ndims] - bias) * (254.0 / scale)); } @@ -28,16 +28,16 @@ int main(int argc, char **argv) } std::ifstream reader(argv[1], std::ios::binary); - _u32 npts_u32; - _u32 ndims_u32; - reader.read((char *)&npts_u32, sizeof(_s32)); - reader.read((char *)&ndims_u32, sizeof(_s32)); + uint32_t npts_u32; + uint32_t ndims_u32; + reader.read((char *)&npts_u32, sizeof(uint32_t)); + reader.read((char *)&ndims_u32, sizeof(uint32_t)); size_t npts = npts_u32; size_t ndims = ndims_u32; std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims << std::endl; - _u64 blk_size = 131072; - _u64 nblks = ROUND_UP(npts, blk_size) / blk_size; + size_t blk_size = 131072; + size_t nblks = ROUND_UP(npts, blk_size) / blk_size; std::ofstream writer(argv[2], std::ios::binary); auto read_buf = new float[blk_size * ndims]; @@ -45,12 +45,12 @@ int main(int argc, char **argv) float bias = atof(argv[3]); float scale = atof(argv[4]); - writer.write((char *)(&npts_u32), sizeof(_u32)); - writer.write((char *)(&ndims_u32), sizeof(_u32)); + writer.write((char *)(&npts_u32), sizeof(uint32_t)); + writer.write((char *)(&ndims_u32), sizeof(uint32_t)); - for (_u64 i = 0; i < nblks; i++) + for (size_t i = 0; i < nblks; i++) { - _u64 cblk_size = std::min(npts - i * blk_size, blk_size); + size_t cblk_size = std::min(npts - i * blk_size, blk_size); block_convert(writer, write_buf, reader, read_buf, cblk_size, ndims, bias, scale); std::cout << "Block #" << i << " written" << std::endl; } diff --git a/tests/utils/fvecs_to_bin.cpp b/tests/utils/fvecs_to_bin.cpp index 10e28076a..c8725898b 100644 --- a/tests/utils/fvecs_to_bin.cpp +++ b/tests/utils/fvecs_to_bin.cpp @@ -5,11 +5,11 @@ #include "utils.h" // Convert float types -void block_convert_float(std::ifstream &reader, std::ofstream &writer, float *read_buf, float *write_buf, _u64 npts, - _u64 ndims) +void block_convert_float(std::ifstream &reader, std::ofstream &writer, float *read_buf, float *write_buf, size_t npts, + size_t ndims) { reader.read((char *)read_buf, npts * (ndims * sizeof(float) + sizeof(unsigned))); - for (_u64 i = 0; i < npts; i++) + for (size_t i = 0; i < npts; i++) { memcpy(write_buf + i * ndims, (read_buf + i * (ndims + 1)) + 1, ndims * sizeof(float)); } @@ -17,16 +17,16 @@ void block_convert_float(std::ifstream &reader, std::ofstream &writer, float *re } // Convert byte types -void block_convert_byte(std::ifstream &reader, std::ofstream &writer, _u8 *read_buf, _u8 *write_buf, _u64 npts, - _u64 ndims) +void block_convert_byte(std::ifstream &reader, std::ofstream &writer, char *read_buf, char *write_buf, size_t npts, + size_t ndims) { - reader.read((char *)read_buf, npts * (ndims * sizeof(_u8) + sizeof(unsigned))); - for (_u64 i = 0; i < npts; i++) + reader.read((char *)read_buf, npts * (ndims * sizeof(char) + sizeof(unsigned))); + for (size_t i = 0; i < npts; i++) { memcpy(write_buf + i * ndims, (read_buf + i * (ndims + sizeof(unsigned))) + sizeof(unsigned), - ndims * sizeof(_u8)); + ndims * sizeof(char)); } - writer.write((char *)write_buf, npts * ndims * sizeof(_u8)); + writer.write((char *)write_buf, npts * ndims * sizeof(char)); } int main(int argc, char **argv) @@ -41,7 +41,7 @@ int main(int argc, char **argv) if (strcmp(argv[1], "uint8") == 0 || strcmp(argv[1], "int8") == 0) { - datasize = sizeof(_u8); + datasize = sizeof(char); } else if (strcmp(argv[1], "float") != 0) { @@ -50,32 +50,32 @@ int main(int argc, char **argv) } std::ifstream reader(argv[2], std::ios::binary | std::ios::ate); - _u64 fsize = reader.tellg(); + size_t 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 * datasize) + sizeof(unsigned)); + size_t ndims = (size_t)ndims_u32; + size_t npts = fsize / ((ndims * datasize) + sizeof(unsigned)); std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims << std::endl; - _u64 blk_size = 131072; - _u64 nblks = ROUND_UP(npts, blk_size) / blk_size; + size_t blk_size = 131072; + size_t nblks = ROUND_UP(npts, blk_size) / blk_size; std::cout << "# blks: " << nblks << std::endl; std::ofstream writer(argv[3], std::ios::binary); - _s32 npts_s32 = (_s32)npts; - _s32 ndims_s32 = (_s32)ndims; - writer.write((char *)&npts_s32, sizeof(_s32)); - writer.write((char *)&ndims_s32, sizeof(_s32)); + int npts_s32 = (int)npts; + int ndims_s32 = (int)ndims; + writer.write((char *)&npts_s32, sizeof(int)); + writer.write((char *)&ndims_s32, sizeof(int)); - _u64 chunknpts = std::min(npts, blk_size); - _u8 *read_buf = new _u8[chunknpts * ((ndims * datasize) + sizeof(unsigned))]; - _u8 *write_buf = new _u8[chunknpts * ndims * datasize]; + size_t chunknpts = std::min(npts, blk_size); + char *read_buf = new char[chunknpts * ((ndims * datasize) + sizeof(unsigned))]; + char *write_buf = new char[chunknpts * ndims * datasize]; - for (_u64 i = 0; i < nblks; i++) + for (size_t i = 0; i < nblks; i++) { - _u64 cblk_size = std::min(npts - i * blk_size, blk_size); + size_t cblk_size = std::min(npts - i * blk_size, blk_size); if (datasize == sizeof(float)) { block_convert_float(reader, writer, (float *)read_buf, (float *)write_buf, cblk_size, ndims); diff --git a/tests/utils/fvecs_to_bvecs.cpp b/tests/utils/fvecs_to_bvecs.cpp index 8324e9df7..c43bd1682 100644 --- a/tests/utils/fvecs_to_bvecs.cpp +++ b/tests/utils/fvecs_to_bvecs.cpp @@ -4,14 +4,14 @@ #include #include "utils.h" -void block_convert(std::ifstream &reader, std::ofstream &writer, float *read_buf, uint8_t *write_buf, _u64 npts, - _u64 ndims) +void block_convert(std::ifstream &reader, std::ofstream &writer, float *read_buf, uint8_t *write_buf, size_t npts, + size_t ndims) { reader.read((char *)read_buf, npts * (ndims * sizeof(float) + sizeof(unsigned))); - for (_u64 i = 0; i < npts; i++) + for (size_t i = 0; i < npts; i++) { memcpy(write_buf + i * (ndims + 4), read_buf + i * (ndims + 1), sizeof(unsigned)); - for (_u64 d = 0; d < ndims; d++) + for (size_t d = 0; d < ndims; d++) write_buf[i * (ndims + 4) + 4 + d] = (uint8_t)read_buf[i * (ndims + 1) + 1 + d]; } writer.write((char *)write_buf, npts * (ndims * 1 + 4)); @@ -25,25 +25,25 @@ int main(int argc, char **argv) exit(-1); } std::ifstream reader(argv[1], std::ios::binary | std::ios::ate); - _u64 fsize = reader.tellg(); + size_t 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)); + size_t ndims = (size_t)ndims_u32; + size_t npts = fsize / ((ndims + 1) * sizeof(float)); std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims << std::endl; - _u64 blk_size = 131072; - _u64 nblks = ROUND_UP(npts, blk_size) / blk_size; + size_t blk_size = 131072; + size_t nblks = ROUND_UP(npts, blk_size) / blk_size; std::cout << "# blks: " << nblks << std::endl; std::ofstream writer(argv[2], std::ios::binary); auto read_buf = new float[npts * (ndims + 1)]; auto write_buf = new uint8_t[npts * (ndims + 4)]; - for (_u64 i = 0; i < nblks; i++) + for (size_t i = 0; i < nblks; i++) { - _u64 cblk_size = std::min(npts - i * blk_size, blk_size); + size_t 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; } diff --git a/tests/utils/generate_synthetic_labels.cpp b/tests/utils/generate_synthetic_labels.cpp index c96d3361d..09b475260 100644 --- a/tests/utils/generate_synthetic_labels.cpp +++ b/tests/utils/generate_synthetic_labels.cpp @@ -91,7 +91,7 @@ class ZipfDistribution int main(int argc, char **argv) { std::string output_file, distribution_type; - _u64 num_labels, num_points; + size_t num_labels, num_points; try { @@ -104,7 +104,8 @@ int main(int argc, char **argv) desc.add_options()("num_labels,L", po::value(&num_labels)->required(), "Number of unique labels, up to 5000"); desc.add_options()("distribution_type,DT", po::value(&distribution_type)->default_value("random"), - "Distribution function for labels defaults to random"); + "Distribution function for labels defaults to " + "random"); po::variables_map vm; po::store(po::parse_command_line(argc, argv, desc), vm); diff --git a/tests/utils/int8_to_float_scale.cpp b/tests/utils/int8_to_float_scale.cpp index 42aa06d0e..2de1a3a56 100644 --- a/tests/utils/int8_to_float_scale.cpp +++ b/tests/utils/int8_to_float_scale.cpp @@ -4,14 +4,14 @@ #include #include "utils.h" -void block_convert(std::ofstream &writer, float *write_buf, std::ifstream &reader, int8_t *read_buf, _u64 npts, - _u64 ndims, float bias, float scale) +void block_convert(std::ofstream &writer, float *write_buf, std::ifstream &reader, int8_t *read_buf, size_t npts, + size_t ndims, float bias, float scale) { reader.read((char *)read_buf, npts * ndims * sizeof(int8_t)); - for (_u64 i = 0; i < npts; i++) + for (size_t i = 0; i < npts; i++) { - for (_u64 d = 0; d < ndims; d++) + for (size_t d = 0; d < ndims; d++) { write_buf[d + i * ndims] = (((float)read_buf[d + i * ndims] - bias) * scale); } @@ -28,16 +28,16 @@ int main(int argc, char **argv) } std::ifstream reader(argv[1], std::ios::binary); - _u32 npts_u32; - _u32 ndims_u32; - reader.read((char *)&npts_u32, sizeof(_s32)); - reader.read((char *)&ndims_u32, sizeof(_s32)); + uint32_t npts_u32; + uint32_t ndims_u32; + reader.read((char *)&npts_u32, sizeof(uint32_t)); + reader.read((char *)&ndims_u32, sizeof(uint32_t)); size_t npts = npts_u32; size_t ndims = ndims_u32; std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims << std::endl; - _u64 blk_size = 131072; - _u64 nblks = ROUND_UP(npts, blk_size) / blk_size; + size_t blk_size = 131072; + size_t nblks = ROUND_UP(npts, blk_size) / blk_size; std::ofstream writer(argv[2], std::ios::binary); auto read_buf = new int8_t[blk_size * ndims]; @@ -45,12 +45,12 @@ int main(int argc, char **argv) float bias = atof(argv[3]); float scale = atof(argv[4]); - writer.write((char *)(&npts_u32), sizeof(_u32)); - writer.write((char *)(&ndims_u32), sizeof(_u32)); + writer.write((char *)(&npts_u32), sizeof(uint32_t)); + writer.write((char *)(&ndims_u32), sizeof(uint32_t)); - for (_u64 i = 0; i < nblks; i++) + for (size_t i = 0; i < nblks; i++) { - _u64 cblk_size = std::min(npts - i * blk_size, blk_size); + size_t cblk_size = std::min(npts - i * blk_size, blk_size); block_convert(writer, write_buf, reader, read_buf, cblk_size, ndims, bias, scale); std::cout << "Block #" << i << " written" << std::endl; } diff --git a/tests/utils/ivecs_to_bin.cpp b/tests/utils/ivecs_to_bin.cpp index b7ee2304a..22fbe7737 100644 --- a/tests/utils/ivecs_to_bin.cpp +++ b/tests/utils/ivecs_to_bin.cpp @@ -4,14 +4,15 @@ #include #include "utils.h" -void block_convert(std::ifstream &reader, std::ofstream &writer, _u32 *read_buf, _u32 *write_buf, _u64 npts, _u64 ndims) +void block_convert(std::ifstream &reader, std::ofstream &writer, uint32_t *read_buf, uint32_t *write_buf, size_t npts, + size_t ndims) { - reader.read((char *)read_buf, npts * (ndims * sizeof(_u32) + sizeof(unsigned))); - for (_u64 i = 0; i < npts; i++) + reader.read((char *)read_buf, npts * (ndims * sizeof(uint32_t) + sizeof(unsigned))); + for (size_t i = 0; i < npts; i++) { - memcpy(write_buf + i * ndims, (read_buf + i * (ndims + 1)) + 1, ndims * sizeof(_u32)); + memcpy(write_buf + i * ndims, (read_buf + i * (ndims + 1)) + 1, ndims * sizeof(uint32_t)); } - writer.write((char *)write_buf, npts * ndims * sizeof(_u32)); + writer.write((char *)write_buf, npts * ndims * sizeof(uint32_t)); } int main(int argc, char **argv) @@ -22,29 +23,29 @@ int main(int argc, char **argv) exit(-1); } std::ifstream reader(argv[1], std::ios::binary | std::ios::ate); - _u64 fsize = reader.tellg(); + size_t 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(_u32)); + size_t ndims = (size_t)ndims_u32; + size_t npts = fsize / ((ndims + 1) * sizeof(uint32_t)); std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims << std::endl; - _u64 blk_size = 131072; - _u64 nblks = ROUND_UP(npts, blk_size) / blk_size; + size_t blk_size = 131072; + size_t nblks = ROUND_UP(npts, blk_size) / blk_size; std::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)); - _u32 *read_buf = new _u32[npts * (ndims + 1)]; - _u32 *write_buf = new _u32[npts * ndims]; - for (_u64 i = 0; i < nblks; i++) + int npts_s32 = (int)npts; + int ndims_s32 = (int)ndims; + writer.write((char *)&npts_s32, sizeof(int)); + writer.write((char *)&ndims_s32, sizeof(int)); + uint32_t *read_buf = new uint32_t[npts * (ndims + 1)]; + uint32_t *write_buf = new uint32_t[npts * ndims]; + for (size_t i = 0; i < nblks; i++) { - _u64 cblk_size = std::min(npts - i * blk_size, blk_size); + size_t 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; } diff --git a/tests/utils/merge_shards.cpp b/tests/utils/merge_shards.cpp index 5e14fb14d..106c15eef 100644 --- a/tests/utils/merge_shards.cpp +++ b/tests/utils/merge_shards.cpp @@ -19,8 +19,10 @@ int main(int argc, char **argv) if (argc != 9) { std::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] " + << " vamana_index_prefix[1] vamana_index_suffix[2] " + "idmaps_prefix[3] " + "idmaps_suffix[4] n_shards[5] max_degree[6] " + "output_vamana_path[7] " "output_medoids_path[8]" << std::endl; exit(-1); @@ -30,8 +32,8 @@ int main(int argc, char **argv) std::string vamana_suffix(argv[2]); std::string idmaps_prefix(argv[3]); std::string idmaps_suffix(argv[4]); - _u64 nshards = (_u64)std::atoi(argv[5]); - _u32 max_degree = (_u64)std::atoi(argv[6]); + uint64_t nshards = (uint64_t)std::atoi(argv[5]); + uint32_t max_degree = (uint64_t)std::atoi(argv[6]); std::string output_index(argv[7]); std::string output_medoids(argv[8]); diff --git a/tests/utils/rand_data_gen.cpp b/tests/utils/rand_data_gen.cpp index c4461137d..ea2e67478 100644 --- a/tests/utils/rand_data_gen.cpp +++ b/tests/utils/rand_data_gen.cpp @@ -11,7 +11,7 @@ namespace po = boost::program_options; -int block_write_float(std::ofstream &writer, _u64 ndims, _u64 npts, float norm) +int block_write_float(std::ofstream &writer, size_t ndims, size_t npts, float norm) { auto vec = new float[ndims]; @@ -19,14 +19,14 @@ int block_write_float(std::ofstream &writer, _u64 ndims, _u64 npts, float norm) std::mt19937 gen{rd()}; std::normal_distribution<> normal_rand{0, 1}; - for (_u64 i = 0; i < npts; i++) + for (size_t i = 0; i < npts; i++) { float sum = 0; - for (_u64 d = 0; d < ndims; ++d) + for (size_t d = 0; d < ndims; ++d) vec[d] = normal_rand(gen); - for (_u64 d = 0; d < ndims; ++d) + for (size_t d = 0; d < ndims; ++d) sum += vec[d] * vec[d]; - for (_u64 d = 0; d < ndims; ++d) + for (size_t d = 0; d < ndims; ++d) vec[d] = vec[d] * norm / std::sqrt(sum); writer.write((char *)vec, ndims * sizeof(float)); @@ -36,7 +36,7 @@ int block_write_float(std::ofstream &writer, _u64 ndims, _u64 npts, float norm) return 0; } -int block_write_int8(std::ofstream &writer, _u64 ndims, _u64 npts, float norm) +int block_write_int8(std::ofstream &writer, size_t ndims, size_t npts, float norm) { auto vec = new float[ndims]; auto vec_T = new int8_t[ndims]; @@ -45,17 +45,17 @@ int block_write_int8(std::ofstream &writer, _u64 ndims, _u64 npts, float norm) std::mt19937 gen{rd()}; std::normal_distribution<> normal_rand{0, 1}; - for (_u64 i = 0; i < npts; i++) + for (size_t i = 0; i < npts; i++) { float sum = 0; - for (_u64 d = 0; d < ndims; ++d) + for (size_t d = 0; d < ndims; ++d) vec[d] = normal_rand(gen); - for (_u64 d = 0; d < ndims; ++d) + for (size_t d = 0; d < ndims; ++d) sum += vec[d] * vec[d]; - for (_u64 d = 0; d < ndims; ++d) + for (size_t d = 0; d < ndims; ++d) vec[d] = vec[d] * norm / std::sqrt(sum); - for (_u64 d = 0; d < ndims; ++d) + for (size_t d = 0; d < ndims; ++d) { vec_T[d] = std::round(vec[d]); } @@ -68,7 +68,7 @@ int block_write_int8(std::ofstream &writer, _u64 ndims, _u64 npts, float norm) return 0; } -int block_write_uint8(std::ofstream &writer, _u64 ndims, _u64 npts, float norm) +int block_write_uint8(std::ofstream &writer, size_t ndims, size_t npts, float norm) { auto vec = new float[ndims]; auto vec_T = new int8_t[ndims]; @@ -77,17 +77,17 @@ int block_write_uint8(std::ofstream &writer, _u64 ndims, _u64 npts, float norm) std::mt19937 gen{rd()}; std::normal_distribution<> normal_rand{0, 1}; - for (_u64 i = 0; i < npts; i++) + for (size_t i = 0; i < npts; i++) { float sum = 0; - for (_u64 d = 0; d < ndims; ++d) + for (size_t d = 0; d < ndims; ++d) vec[d] = normal_rand(gen); - for (_u64 d = 0; d < ndims; ++d) + for (size_t d = 0; d < ndims; ++d) sum += vec[d] * vec[d]; - for (_u64 d = 0; d < ndims; ++d) + for (size_t d = 0; d < ndims; ++d) vec[d] = vec[d] * norm / std::sqrt(sum); - for (_u64 d = 0; d < ndims; ++d) + for (size_t d = 0; d < ndims; ++d) { vec_T[d] = 128 + std::round(vec[d]); } @@ -103,7 +103,7 @@ int block_write_uint8(std::ofstream &writer, _u64 ndims, _u64 npts, float norm) int main(int argc, char **argv) { std::string data_type, output_file; - _u64 ndims, npts; + size_t ndims, npts; float norm; try @@ -149,7 +149,8 @@ int main(int argc, char **argv) { if (norm > 127) { - std::cerr << "Error: for int8/uint8 datatypes, L2 norm can not be greater " + std::cerr << "Error: for int8/uint8 datatypes, L2 norm can not be " + "greater " "than 127" << std::endl; return -1; @@ -161,19 +162,19 @@ int main(int argc, char **argv) std::ofstream writer; writer.exceptions(std::ofstream::failbit | std::ofstream::badbit); writer.open(output_file, 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)); + auto npts_u32 = (uint32_t)npts; + auto ndims_u32 = (uint32_t)ndims; + writer.write((char *)&npts_u32, sizeof(uint32_t)); + writer.write((char *)&ndims_u32, sizeof(uint32_t)); - _u64 blk_size = 131072; - _u64 nblks = ROUND_UP(npts, blk_size) / blk_size; + size_t blk_size = 131072; + size_t nblks = ROUND_UP(npts, blk_size) / blk_size; std::cout << "# blks: " << nblks << std::endl; int ret = 0; - for (_u64 i = 0; i < nblks; i++) + for (size_t i = 0; i < nblks; i++) { - _u64 cblk_size = std::min(npts - i * blk_size, blk_size); + size_t cblk_size = std::min(npts - i * blk_size, blk_size); if (data_type == std::string("float")) { ret = block_write_float(writer, ndims, cblk_size, norm); diff --git a/tests/utils/stats_label_data.cpp b/tests/utils/stats_label_data.cpp index c5aabd5ff..549ae3e6a 100644 --- a/tests/utils/stats_label_data.cpp +++ b/tests/utils/stats_label_data.cpp @@ -28,26 +28,26 @@ #endif namespace po = boost::program_options; -void stats_analysis(const std::string labels_file, std::string univeral_label, _u32 density = 10) +void stats_analysis(const std::string labels_file, std::string univeral_label, uint32_t density = 10) { std::string token, line; std::ifstream labels_stream(labels_file); - std::unordered_map label_counts; + std::unordered_map label_counts; std::string label_with_max_points; - _u32 max_points = 0; + uint32_t max_points = 0; long long sum = 0; long long point_cnt = 0; float avg_labels_per_pt, avg_labels_per_pt_incl_0, mean_label_size, mean_label_size_incl_0; - std::vector<_u32> labels_per_point; - _u32 dense_pts = 0; + std::vector labels_per_point; + uint32_t dense_pts = 0; if (labels_stream.is_open()) { while (getline(labels_stream, line)) { point_cnt++; std::stringstream iss(line); - _u32 lbl_cnt = 0; + uint32_t lbl_cnt = 0; while (getline(iss, token, ',')) { lbl_cnt++; @@ -69,7 +69,7 @@ void stats_analysis(const std::string labels_file, std::string univeral_label, _ << " labels = " << (float)dense_pts / (float)labels_per_point.size() << std::endl; std::sort(labels_per_point.begin(), labels_per_point.end()); - std::vector> label_count_vec; + std::vector> label_count_vec; for (auto it = label_counts.begin(); it != label_counts.end(); it++) { @@ -84,14 +84,14 @@ void stats_analysis(const std::string labels_file, std::string univeral_label, _ } sort(label_count_vec.begin(), label_count_vec.end(), - [](const std::pair &lhs, const std::pair &rhs) { + [](const std::pair &lhs, const std::pair &rhs) { return lhs.second < rhs.second; }); for (float p = 0; p < 1; p += 0.05) { - std::cout << "Percentile " << (100 * p) << "\t" << label_count_vec[(_u32)(p * label_count_vec.size())].first - << " with count=" << label_count_vec[(_u32)(p * label_count_vec.size())].second << std::endl; + std::cout << "Percentile " << (100 * p) << "\t" << label_count_vec[(size_t)(p * label_count_vec.size())].first + << " with count=" << label_count_vec[(size_t)(p * label_count_vec.size())].second << std::endl; } std::cout << "Most common label " @@ -117,7 +117,7 @@ void stats_analysis(const std::string labels_file, std::string univeral_label, _ int main(int argc, char **argv) { std::string labels_file, universal_label; - _u32 density; + uint32_t density; po::options_description desc{"Arguments"}; try @@ -127,7 +127,7 @@ int main(int argc, char **argv) "path to labels data file."); desc.add_options()("universal_label", po::value(&universal_label)->required(), "Universal label used in labels file."); - desc.add_options()("density", po::value<_u32>(&density)->default_value(1), + desc.add_options()("density", po::value(&density)->default_value(1), "Number of labels each point in labels file, defaults to 1"); po::variables_map vm; po::store(po::parse_command_line(argc, argv, desc), vm); diff --git a/tests/utils/tsv_to_bin.cpp b/tests/utils/tsv_to_bin.cpp index 5aaa9a03c..c590a8f73 100644 --- a/tests/utils/tsv_to_bin.cpp +++ b/tests/utils/tsv_to_bin.cpp @@ -4,16 +4,16 @@ #include #include "utils.h" -void block_convert_float(std::ifstream &reader, std::ofstream &writer, _u64 npts, _u64 ndims) +void block_convert_float(std::ifstream &reader, std::ofstream &writer, size_t npts, size_t ndims) { auto read_buf = new float[npts * (ndims + 1)]; auto cursor = read_buf; float val; - for (_u64 i = 0; i < npts; i++) + for (size_t i = 0; i < npts; i++) { - for (_u64 d = 0; d < ndims; ++d) + for (size_t d = 0; d < ndims; ++d) { reader >> val; *cursor = val; @@ -24,16 +24,16 @@ void block_convert_float(std::ifstream &reader, std::ofstream &writer, _u64 npts delete[] read_buf; } -void block_convert_int8(std::ifstream &reader, std::ofstream &writer, _u64 npts, _u64 ndims) +void block_convert_int8(std::ifstream &reader, std::ofstream &writer, size_t npts, size_t ndims) { auto read_buf = new int8_t[npts * (ndims + 1)]; auto cursor = read_buf; int val; - for (_u64 i = 0; i < npts; i++) + for (size_t i = 0; i < npts; i++) { - for (_u64 d = 0; d < ndims; ++d) + for (size_t d = 0; d < ndims; ++d) { reader >> val; *cursor = (int8_t)val; @@ -44,16 +44,16 @@ void block_convert_int8(std::ifstream &reader, std::ofstream &writer, _u64 npts, delete[] read_buf; } -void block_convert_uint8(std::ifstream &reader, std::ofstream &writer, _u64 npts, _u64 ndims) +void block_convert_uint8(std::ifstream &reader, std::ofstream &writer, size_t npts, size_t ndims) { auto read_buf = new uint8_t[npts * (ndims + 1)]; auto cursor = read_buf; int val; - for (_u64 i = 0; i < npts; i++) + for (size_t i = 0; i < npts; i++) { - for (_u64 d = 0; d < ndims; ++d) + for (size_t d = 0; d < ndims; ++d) { reader >> val; *cursor = (uint8_t)val; @@ -81,26 +81,26 @@ int main(int argc, char **argv) std::cout << "Unsupported type. float, int8 and uint8 types are supported." << std::endl; } - _u64 ndims = atoi(argv[4]); - _u64 npts = atoi(argv[5]); + size_t ndims = atoi(argv[4]); + size_t npts = atoi(argv[5]); std::ifstream reader(argv[2], std::ios::binary | std::ios::ate); - // _u64 fsize = reader.tellg(); + // size_t 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; + size_t blk_size = 131072; + size_t 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)); + auto npts_u32 = (uint32_t)npts; + auto ndims_u32 = (uint32_t)ndims; + writer.write((char *)&npts_u32, sizeof(uint32_t)); + writer.write((char *)&ndims_u32, sizeof(uint32_t)); - for (_u64 i = 0; i < nblks; i++) + for (size_t i = 0; i < nblks; i++) { - _u64 cblk_size = std::min(npts - i * blk_size, blk_size); + size_t cblk_size = std::min(npts - i * blk_size, blk_size); if (std::string(argv[1]) == std::string("float")) { block_convert_float(reader, writer, cblk_size, ndims); diff --git a/tests/utils/vector_analysis.cpp b/tests/utils/vector_analysis.cpp index 99e5627b1..5e4cb9bf4 100644 --- a/tests/utils/vector_analysis.cpp +++ b/tests/utils/vector_analysis.cpp @@ -24,18 +24,18 @@ template int analyze_norm(std::string base_file) { std::cout << "Analyzing data norms" << std::endl; T *data; - _u64 npts, ndims; + size_t npts, ndims; diskann::load_bin(base_file, data, npts, ndims); std::vector norms(npts, 0); #pragma omp parallel for schedule(dynamic) - for (_s64 i = 0; i < (_s64)npts; i++) + for (int64_t i = 0; i < (int64_t)npts; i++) { - for (_u32 d = 0; d < ndims; d++) + for (size_t d = 0; d < ndims; d++) norms[i] += data[i * ndims + d] * data[i * ndims + d]; norms[i] = std::sqrt(norms[i]); } std::sort(norms.begin(), norms.end()); - for (_u32 p = 0; p < 100; p += 5) + for (int p = 0; p < 100; p += 5) std::cout << "percentile " << p << ": " << norms[std::floor((p / 100.0) * npts)] << std::endl; std::cout << "percentile 100" << ": " << norms[npts - 1] << std::endl; @@ -47,17 +47,17 @@ template int normalize_base(std::string base_file, std::string out_ { std::cout << "Normalizing base" << std::endl; T *data; - _u64 npts, ndims; + size_t npts, ndims; diskann::load_bin(base_file, data, npts, ndims); // std::vector norms(npts, 0); #pragma omp parallel for schedule(dynamic) - for (_s64 i = 0; i < (_s64)npts; i++) + for (int64_t i = 0; i < (int64_t)npts; i++) { float pt_norm = 0; - for (_u32 d = 0; d < ndims; d++) + for (size_t d = 0; d < ndims; d++) pt_norm += data[i * ndims + d] * data[i * ndims + d]; pt_norm = std::sqrt(pt_norm); - for (_u32 d = 0; d < ndims; d++) + for (size_t d = 0; d < ndims; d++) data[i * ndims + d] = data[i * ndims + d] / pt_norm; } diskann::save_bin(out_file, data, npts, ndims); @@ -69,14 +69,14 @@ template int augment_base(std::string base_file, std::string out_fi { std::cout << "Analyzing data norms" << std::endl; T *data; - _u64 npts, ndims; + size_t npts, ndims; diskann::load_bin(base_file, data, npts, ndims); std::vector norms(npts, 0); float max_norm = 0; #pragma omp parallel for schedule(dynamic) - for (_s64 i = 0; i < (_s64)npts; i++) + for (int64_t i = 0; i < (int64_t)npts; i++) { - for (_u32 d = 0; d < ndims; d++) + for (size_t d = 0; d < ndims; d++) norms[i] += data[i * ndims + d] * data[i * ndims + d]; max_norm = norms[i] > max_norm ? norms[i] : max_norm; } @@ -84,13 +84,13 @@ template int augment_base(std::string base_file, std::string out_fi max_norm = std::sqrt(max_norm); std::cout << "Max norm: " << max_norm << std::endl; T *new_data; - _u64 newdims = ndims + 1; + size_t newdims = ndims + 1; new_data = new T[npts * newdims]; - for (_u64 i = 0; i < npts; i++) + for (size_t i = 0; i < npts; i++) { if (prep_base) { - for (_u64 j = 0; j < ndims; j++) + for (size_t j = 0; j < ndims; j++) { new_data[i * newdims + j] = data[i * ndims + j] / max_norm; } @@ -104,7 +104,7 @@ template int augment_base(std::string base_file, std::string out_fi } else { - for (_u64 j = 0; j < ndims; j++) + for (size_t j = 0; j < ndims; j++) { new_data[i * newdims + j] = data[i * ndims + j] / std::sqrt(norms[i]); } @@ -120,7 +120,7 @@ template int augment_base(std::string base_file, std::string out_fi template int aux_main(char **argv) { std::string base_file(argv[2]); - _u32 option = atoi(argv[3]); + uint32_t option = atoi(argv[3]); if (option == 1) analyze_norm(base_file); else if (option == 2) From 0267012c504fca6af324f61622cd7a566658fe52 Mon Sep 17 00:00:00 2001 From: harsha vardhan simhadri Date: Tue, 28 Mar 2023 01:45:55 -0700 Subject: [PATCH 02/69] added some seed files --- include/abstract_data_store.h | 58 ++++++++++++++++ include/abstract_graph_store.h | 26 +++++++ include/in_mem_data_store.h | 58 ++++++++++++++++ include/in_mem_graph_store.h | 23 +++++++ src/CMakeLists.txt | 1 + src/dll/CMakeLists.txt | 2 +- src/in_mem_data_store.cpp | 119 +++++++++++++++++++++++++++++++++ src/in_mem_graph_store.cpp | 28 ++++++++ 8 files changed, 314 insertions(+), 1 deletion(-) create mode 100644 include/abstract_data_store.h create mode 100644 include/abstract_graph_store.h create mode 100644 include/in_mem_data_store.h create mode 100644 include/in_mem_graph_store.h create mode 100644 src/in_mem_data_store.cpp create mode 100644 src/in_mem_graph_store.cpp diff --git a/include/abstract_data_store.h b/include/abstract_data_store.h new file mode 100644 index 000000000..82d7d646a --- /dev/null +++ b/include/abstract_data_store.h @@ -0,0 +1,58 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#pragma once + +#include + +#include "types.h" + +namespace diskann +{ + +template class AbstractDataStore +{ + public: + AbstractDataStore(const location_t max_pts, const location_t num_frozen_pts, const size_t dim) + : _max_pts(max_pts), _num_pts(0), _num_frozen_pts(num_frozen_pts), _dim(dim) + { + } + + // Return number of points returned + virtual size_t load(const std::string &filename) = 0; + virtual void store(const std::string &filename) = 0; + + virtual data_t *get_vector(location_t i) = 0; + virtual data_t *get_vector_by_UID(id_t uid) = 0; + + virtual void set_vector(const location_t i, const data_t *const vector) = 0; + + location_t get_max_pts() + { + return _max_pts; + } + + location_t get_num_pts() + { + return _num_pts; + } + + location_t get_num_frozen_pts() + { + return _num_frozen_pts; + } + + size_t get_dims() + { + return _dim; + } + + protected: + location_t _max_pts; + location_t _num_pts; + location_t _num_frozen_pts; + + const size_t _dim; +}; + +} // namespace diskann \ No newline at end of file diff --git a/include/abstract_graph_store.h b/include/abstract_graph_store.h new file mode 100644 index 000000000..fe088d51b --- /dev/null +++ b/include/abstract_graph_store.h @@ -0,0 +1,26 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#pragma once + +#include +#include + +#include "types.h" + +namespace diskann +{ + +class AbstractGraphStore +{ + public: + AbstractGraphStore(const size_t max_pts); + + virtual int load(const std::string &index_path_prefix) = 0; + virtual int store(const std::string &index_path_prefix) = 0; + + virtual void get_adj_list(const location_t i, std::vector &neighbors) = 0; + virtual void set_adj_list(const location_t i, std::vector &neighbors) = 0; +}; + +} // namespace diskann \ No newline at end of file diff --git a/include/in_mem_data_store.h b/include/in_mem_data_store.h new file mode 100644 index 000000000..d948f74ee --- /dev/null +++ b/include/in_mem_data_store.h @@ -0,0 +1,58 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#include + +#include "tsl/robin_map.h" +#include "tsl/robin_set.h" +#include "tsl/sparse_map.h" +#include "boost/dynamic_bitset.hpp" + +#include "abstract_data_store.h" + +#include "natural_number_map.h" +#include "natural_number_set.h" + +namespace diskann +{ +template class InMemDataStore : public AbstractDataStore +{ + public: + InMemDataStore(const location_t max_pts, const location_t num_frozen_pts, const size_t dim); + ~InMemDataStore(); + + void load(const std::string &filename); + void store(const std::string &filename); + + data_t *get_vector(location_t i); + data_t *get_vector_by_UID(id_t uid); + + void set_vector(const location_t i, const data_t *const vector); + + protected: + location_t load_data(const std::string &filename); +#ifdef EXEC_ENV_OLS + location_t load_data(AlignedFileReader &reader); +#endif + + private: + data_t *_data = nullptr; + + const size_t _aligned_dim; + + // lazy_delete removes entry from _location_to_tag and _tag_to_location. If + // _location_to_tag does not resolve a location, infer that it was deleted. + tsl::sparse_map _tag_to_location; + natural_number_map _location_to_tag; + + // _empty_slots has unallocated slots and those freed by consolidate_delete. + // _delete_set has locations marked deleted by lazy_delete. Will not be + // immediately available for insert. consolidate_delete will release these + // slots to _empty_slots. + natural_number_set _empty_slots; + std::unique_ptr> _delete_set; + + std::shared_timed_mutex _lock; // Takes please of Index::_tag_lock +}; + +} // namespace diskann \ No newline at end of file diff --git a/include/in_mem_graph_store.h b/include/in_mem_graph_store.h new file mode 100644 index 000000000..31b502fd7 --- /dev/null +++ b/include/in_mem_graph_store.h @@ -0,0 +1,23 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#pragma once + +#include "abstract_graph_store.h" + +namespace diskann +{ + +class InMemGraphStore : public AbstractGraphStore +{ + public: + InMemGraphStore(const size_t max_pts); + + int load(const std::string &index_path_prefix); + int store(const std::string &index_path_prefix); + + void get_adj_list(const location_t i, std::vector &neighbors); + void set_adj_list(const location_t i, std::vector &neighbors); +}; + +} // namespace diskann \ No newline at end of file diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 1cebbf6bd..1556d2dcc 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -8,6 +8,7 @@ if(MSVC) else() #file(GLOB CPP_SOURCES *.cpp) set(CPP_SOURCES ann_exception.cpp disk_utils.cpp distance.cpp index.cpp + in_mem_graph_store.cpp in_mem_data_store.cpp linux_aligned_file_reader.cpp math_utils.cpp natural_number_map.cpp natural_number_set.cpp memory_mapper.cpp partition.cpp pq.cpp pq_flash_index.cpp scratch.cpp logger.cpp utils.cpp) diff --git a/src/dll/CMakeLists.txt b/src/dll/CMakeLists.txt index ca62a8f0e..fc611d435 100644 --- a/src/dll/CMakeLists.txt +++ b/src/dll/CMakeLists.txt @@ -3,7 +3,7 @@ add_library(${PROJECT_NAME} SHARED dllmain.cpp ../partition.cpp ../pq.cpp ../pq_flash_index.cpp ../logger.cpp ../utils.cpp ../windows_aligned_file_reader.cpp ../distance.cpp ../memory_mapper.cpp ../index.cpp - ../math_utils.cpp ../disk_utils.cpp + ../in_mem_graph_store.cpp ../in_mem_data_store.cpp ../math_utils.cpp ../disk_utils.cpp ../ann_exception.cpp ../natural_number_set.cpp ../natural_number_map.cpp ../scratch.cpp) set(TARGET_DIR "$<$:${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}>$<$:${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}>") diff --git a/src/in_mem_data_store.cpp b/src/in_mem_data_store.cpp new file mode 100644 index 000000000..4e39297a9 --- /dev/null +++ b/src/in_mem_data_store.cpp @@ -0,0 +1,119 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#include "in_mem_data_store.h" + +#include "utils.h" + +namespace diskann +{ + +template +InMemDataStore::InMemDataStore(const location_t max_pts, const location_t num_frozen_pts, + const size_t dim) + : AbstractDataStore(max_pts, num_frozen_pts, dim), _aligned_dim(ROUND_UP(dim, 8)) +{ + location_t total_internal_points = max_pts + num_frozen_pts; + alloc_aligned(((void **)&_data), total_internal_points * _aligned_dim * sizeof(data_t), 8 * sizeof(data_t)); + std::memset(_data, 0, total_internal_points * _aligned_dim * sizeof(data_t)); +} + +template InMemDataStore::~InMemDataStore() +{ + if (_data != nullptr) + { + aligned_free(this->_data); + } +} + +template void InMemDataStore::load(const std::string &filename) +{ +} + +template void InMemDataStore::store(const std::string &filename) +{ +} + +#ifdef EXEC_ENV_OLS +template location_t Index::load_data(AlignedFileReader &reader) +{ + size_t file_dim, file_num_points; + + diskann::get_bin_metadata(reader, file_num_points, file_dim); + + // since we are loading a new dataset, _empty_slots must be cleared + _empty_slots.clear(); + + if (file_dim != _dim) + { + std::stringstream stream; + stream << "ERROR: Driver requests loading " << _dim << " dimension," + << "but file has " << file_dim << " dimension." << std::endl; + diskann::cerr << stream.str() << std::endl; + aligned_free(_data); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } + + if (file_num_points > _max_points + _num_frozen_pts) + { + resize(file_num_points - _num_frozen_pts); + } + + return file_num_points; +} +#endif + +template +location_t InMemDataStore::load_data(const std::string &filename) +{ + size_t file_dim, file_num_points; + if (!file_exists(filename)) + { + std::stringstream stream; + stream << "ERROR: data file " << filename << " does not exist." << std::endl; + diskann::cerr << stream.str() << std::endl; + aligned_free(_data); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } + diskann::get_bin_metadata(filename, file_num_points, file_dim); + + // since we are loading a new dataset, _empty_slots must be cleared + _empty_slots.clear(); + + if (file_dim != this->_dim) + { + std::stringstream stream; + stream << "ERROR: Driver requests loading " << this->_dim << " dimension," + << "but file has " << file_dim << " dimension." << std::endl; + diskann::cerr << stream.str() << std::endl; + aligned_free(_data); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } + + if (file_num_points > this->get_max_points() + this->get_num_frozen_pts()) + { + resize(file_num_points - this->get_num_frozen_pts()); + } + + copy_aligned_data_from_file(filename.c_str(), _data, file_num_points, file_dim, _aligned_dim); + + return file_num_points; +} + +template data_t *InMemDataStore::get_vector(location_t i) +{ + return _data + i * _aligned_dim; +} + +template data_t *InMemDataStore::get_vector_by_UID(id_t uid) +{ + return get_vector(_tag_to_location(uid)); +} + +template +void InMemDataStore::set_vector(const location_t loc, const data_t *const vector) +{ + memcpy(_data + loc * _aligned_dim, vector, this->_dim * sizeof(data_t)); +} + +} // namespace diskann \ No newline at end of file diff --git a/src/in_mem_graph_store.cpp b/src/in_mem_graph_store.cpp new file mode 100644 index 000000000..df9933080 --- /dev/null +++ b/src/in_mem_graph_store.cpp @@ -0,0 +1,28 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#include "in_mem_graph_store.h" + +namespace diskann +{ + +InMemGraphStore::InMemGraphStore(const size_t max_pts) : AbstractGraphStore(max_pts) +{ +} + +int InMemGraphStore::load(const std::string &index_path_prefix) +{ +} +int InMemGraphStore::store(const std::string &index_path_prefix) +{ +} + +void InMemGraphStore::get_adj_list(const location_t i, std::vector &neighbors) +{ +} + +void InMemGraphStore::set_adj_list(const location_t i, std::vector &neighbors) +{ +} + +} // namespace diskann \ No newline at end of file From 7131d70edf3902eaf70b00e4ba7461d48b7742cf Mon Sep 17 00:00:00 2001 From: harsha vardhan simhadri Date: Fri, 31 Mar 2023 16:00:54 -0700 Subject: [PATCH 03/69] add seed files --- include/abstract_data_store.h | 58 ++++++++++++++++ include/abstract_graph_store.h | 26 +++++++ include/in_mem_data_store.h | 58 ++++++++++++++++ include/in_mem_graph_store.h | 23 +++++++ src/CMakeLists.txt | 1 + src/dll/CMakeLists.txt | 3 +- src/in_mem_data_store.cpp | 119 +++++++++++++++++++++++++++++++++ src/in_mem_graph_store.cpp | 28 ++++++++ 8 files changed, 315 insertions(+), 1 deletion(-) create mode 100644 include/abstract_data_store.h create mode 100644 include/abstract_graph_store.h create mode 100644 include/in_mem_data_store.h create mode 100644 include/in_mem_graph_store.h create mode 100644 src/in_mem_data_store.cpp create mode 100644 src/in_mem_graph_store.cpp diff --git a/include/abstract_data_store.h b/include/abstract_data_store.h new file mode 100644 index 000000000..82d7d646a --- /dev/null +++ b/include/abstract_data_store.h @@ -0,0 +1,58 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#pragma once + +#include + +#include "types.h" + +namespace diskann +{ + +template class AbstractDataStore +{ + public: + AbstractDataStore(const location_t max_pts, const location_t num_frozen_pts, const size_t dim) + : _max_pts(max_pts), _num_pts(0), _num_frozen_pts(num_frozen_pts), _dim(dim) + { + } + + // Return number of points returned + virtual size_t load(const std::string &filename) = 0; + virtual void store(const std::string &filename) = 0; + + virtual data_t *get_vector(location_t i) = 0; + virtual data_t *get_vector_by_UID(id_t uid) = 0; + + virtual void set_vector(const location_t i, const data_t *const vector) = 0; + + location_t get_max_pts() + { + return _max_pts; + } + + location_t get_num_pts() + { + return _num_pts; + } + + location_t get_num_frozen_pts() + { + return _num_frozen_pts; + } + + size_t get_dims() + { + return _dim; + } + + protected: + location_t _max_pts; + location_t _num_pts; + location_t _num_frozen_pts; + + const size_t _dim; +}; + +} // namespace diskann \ No newline at end of file diff --git a/include/abstract_graph_store.h b/include/abstract_graph_store.h new file mode 100644 index 000000000..fe088d51b --- /dev/null +++ b/include/abstract_graph_store.h @@ -0,0 +1,26 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#pragma once + +#include +#include + +#include "types.h" + +namespace diskann +{ + +class AbstractGraphStore +{ + public: + AbstractGraphStore(const size_t max_pts); + + virtual int load(const std::string &index_path_prefix) = 0; + virtual int store(const std::string &index_path_prefix) = 0; + + virtual void get_adj_list(const location_t i, std::vector &neighbors) = 0; + virtual void set_adj_list(const location_t i, std::vector &neighbors) = 0; +}; + +} // namespace diskann \ No newline at end of file diff --git a/include/in_mem_data_store.h b/include/in_mem_data_store.h new file mode 100644 index 000000000..d948f74ee --- /dev/null +++ b/include/in_mem_data_store.h @@ -0,0 +1,58 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#include + +#include "tsl/robin_map.h" +#include "tsl/robin_set.h" +#include "tsl/sparse_map.h" +#include "boost/dynamic_bitset.hpp" + +#include "abstract_data_store.h" + +#include "natural_number_map.h" +#include "natural_number_set.h" + +namespace diskann +{ +template class InMemDataStore : public AbstractDataStore +{ + public: + InMemDataStore(const location_t max_pts, const location_t num_frozen_pts, const size_t dim); + ~InMemDataStore(); + + void load(const std::string &filename); + void store(const std::string &filename); + + data_t *get_vector(location_t i); + data_t *get_vector_by_UID(id_t uid); + + void set_vector(const location_t i, const data_t *const vector); + + protected: + location_t load_data(const std::string &filename); +#ifdef EXEC_ENV_OLS + location_t load_data(AlignedFileReader &reader); +#endif + + private: + data_t *_data = nullptr; + + const size_t _aligned_dim; + + // lazy_delete removes entry from _location_to_tag and _tag_to_location. If + // _location_to_tag does not resolve a location, infer that it was deleted. + tsl::sparse_map _tag_to_location; + natural_number_map _location_to_tag; + + // _empty_slots has unallocated slots and those freed by consolidate_delete. + // _delete_set has locations marked deleted by lazy_delete. Will not be + // immediately available for insert. consolidate_delete will release these + // slots to _empty_slots. + natural_number_set _empty_slots; + std::unique_ptr> _delete_set; + + std::shared_timed_mutex _lock; // Takes please of Index::_tag_lock +}; + +} // namespace diskann \ No newline at end of file diff --git a/include/in_mem_graph_store.h b/include/in_mem_graph_store.h new file mode 100644 index 000000000..31b502fd7 --- /dev/null +++ b/include/in_mem_graph_store.h @@ -0,0 +1,23 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#pragma once + +#include "abstract_graph_store.h" + +namespace diskann +{ + +class InMemGraphStore : public AbstractGraphStore +{ + public: + InMemGraphStore(const size_t max_pts); + + int load(const std::string &index_path_prefix); + int store(const std::string &index_path_prefix); + + void get_adj_list(const location_t i, std::vector &neighbors); + void set_adj_list(const location_t i, std::vector &neighbors); +}; + +} // namespace diskann \ No newline at end of file diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index f8dda23c7..b4a6a4e89 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -9,6 +9,7 @@ else() #file(GLOB CPP_SOURCES *.cpp) set(CPP_SOURCES ann_exception.cpp disk_utils.cpp distance.cpp index.cpp 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) if (RESTAPI) diff --git a/src/dll/CMakeLists.txt b/src/dll/CMakeLists.txt index 1423d6c2c..0f511a8ae 100644 --- a/src/dll/CMakeLists.txt +++ b/src/dll/CMakeLists.txt @@ -2,7 +2,8 @@ #Licensed under the MIT license. add_library(${PROJECT_NAME} SHARED dllmain.cpp ../partition.cpp ../pq.cpp ../pq_flash_index.cpp ../logger.cpp ../utils.cpp - ../windows_aligned_file_reader.cpp ../distance.cpp ../memory_mapper.cpp ../index.cpp ../math_utils.cpp ../disk_utils.cpp ../filter_utils.cpp + ../windows_aligned_file_reader.cpp ../distance.cpp ../memory_mapper.cpp ../index.cpp + ../in_mem_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) set(TARGET_DIR "$<$:${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}>$<$:${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}>") diff --git a/src/in_mem_data_store.cpp b/src/in_mem_data_store.cpp new file mode 100644 index 000000000..4e39297a9 --- /dev/null +++ b/src/in_mem_data_store.cpp @@ -0,0 +1,119 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#include "in_mem_data_store.h" + +#include "utils.h" + +namespace diskann +{ + +template +InMemDataStore::InMemDataStore(const location_t max_pts, const location_t num_frozen_pts, + const size_t dim) + : AbstractDataStore(max_pts, num_frozen_pts, dim), _aligned_dim(ROUND_UP(dim, 8)) +{ + location_t total_internal_points = max_pts + num_frozen_pts; + alloc_aligned(((void **)&_data), total_internal_points * _aligned_dim * sizeof(data_t), 8 * sizeof(data_t)); + std::memset(_data, 0, total_internal_points * _aligned_dim * sizeof(data_t)); +} + +template InMemDataStore::~InMemDataStore() +{ + if (_data != nullptr) + { + aligned_free(this->_data); + } +} + +template void InMemDataStore::load(const std::string &filename) +{ +} + +template void InMemDataStore::store(const std::string &filename) +{ +} + +#ifdef EXEC_ENV_OLS +template location_t Index::load_data(AlignedFileReader &reader) +{ + size_t file_dim, file_num_points; + + diskann::get_bin_metadata(reader, file_num_points, file_dim); + + // since we are loading a new dataset, _empty_slots must be cleared + _empty_slots.clear(); + + if (file_dim != _dim) + { + std::stringstream stream; + stream << "ERROR: Driver requests loading " << _dim << " dimension," + << "but file has " << file_dim << " dimension." << std::endl; + diskann::cerr << stream.str() << std::endl; + aligned_free(_data); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } + + if (file_num_points > _max_points + _num_frozen_pts) + { + resize(file_num_points - _num_frozen_pts); + } + + return file_num_points; +} +#endif + +template +location_t InMemDataStore::load_data(const std::string &filename) +{ + size_t file_dim, file_num_points; + if (!file_exists(filename)) + { + std::stringstream stream; + stream << "ERROR: data file " << filename << " does not exist." << std::endl; + diskann::cerr << stream.str() << std::endl; + aligned_free(_data); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } + diskann::get_bin_metadata(filename, file_num_points, file_dim); + + // since we are loading a new dataset, _empty_slots must be cleared + _empty_slots.clear(); + + if (file_dim != this->_dim) + { + std::stringstream stream; + stream << "ERROR: Driver requests loading " << this->_dim << " dimension," + << "but file has " << file_dim << " dimension." << std::endl; + diskann::cerr << stream.str() << std::endl; + aligned_free(_data); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } + + if (file_num_points > this->get_max_points() + this->get_num_frozen_pts()) + { + resize(file_num_points - this->get_num_frozen_pts()); + } + + copy_aligned_data_from_file(filename.c_str(), _data, file_num_points, file_dim, _aligned_dim); + + return file_num_points; +} + +template data_t *InMemDataStore::get_vector(location_t i) +{ + return _data + i * _aligned_dim; +} + +template data_t *InMemDataStore::get_vector_by_UID(id_t uid) +{ + return get_vector(_tag_to_location(uid)); +} + +template +void InMemDataStore::set_vector(const location_t loc, const data_t *const vector) +{ + memcpy(_data + loc * _aligned_dim, vector, this->_dim * sizeof(data_t)); +} + +} // namespace diskann \ No newline at end of file diff --git a/src/in_mem_graph_store.cpp b/src/in_mem_graph_store.cpp new file mode 100644 index 000000000..df9933080 --- /dev/null +++ b/src/in_mem_graph_store.cpp @@ -0,0 +1,28 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#include "in_mem_graph_store.h" + +namespace diskann +{ + +InMemGraphStore::InMemGraphStore(const size_t max_pts) : AbstractGraphStore(max_pts) +{ +} + +int InMemGraphStore::load(const std::string &index_path_prefix) +{ +} +int InMemGraphStore::store(const std::string &index_path_prefix) +{ +} + +void InMemGraphStore::get_adj_list(const location_t i, std::vector &neighbors) +{ +} + +void InMemGraphStore::set_adj_list(const location_t i, std::vector &neighbors) +{ +} + +} // namespace diskann \ No newline at end of file From 73e5287c6166957584c1a8f302f4f3df2f0f9193 Mon Sep 17 00:00:00 2001 From: Gopal Srinivasa Date: Sun, 2 Apr 2023 14:08:48 +0530 Subject: [PATCH 04/69] New distance metric hierarchy --- include/distance.h | 63 +++++++++++++++++++++++++++++++++++++++++++++- src/distance.cpp | 28 +++++++++++++++++++++ 2 files changed, 90 insertions(+), 1 deletion(-) diff --git a/include/distance.h b/include/distance.h index e04be7ee2..39db1a0a7 100644 --- a/include/distance.h +++ b/include/distance.h @@ -1,6 +1,7 @@ #pragma once #include "windows_customizations.h" + namespace diskann { enum Metric @@ -14,8 +15,58 @@ enum Metric template class Distance { public: + // distance comparison function virtual float compare(const T *a, const T *b, uint32_t length) const = 0; - virtual ~Distance() + + // For MIPS, normalization => a new dimension gets added to the vectors. + // This function lets callers know if the normalization process + // changes the dimension. + virtual uint32_t new_dimension(uint32_t orig_dimension) const + { + return orig_dimension; + } + + // This is for efficiency. If no normalization is required, the callers + //can simply ignore the normalize_data_for_build() function. + virtual bool normalization_required() const + { + return false; + } + + //Check the normalization_required() function before calling this. + //Clients can call the function like this: + // + // if (metric->normalization_required()){ + // T* normalized_data_batch; + // if ( metric->new_dimension() == orig_dim ) { + // normalized_data_batch = nullptr; + // modify_data = true; + // } else { + // normalized_data_batch = new T[batch_size * metric->new_dimension()]; + // modify_data = false; + // } + // Split data into batches of batch_size and for each, call: + // metric->normalize_data_for_build(data_batch, batch_size, orig_dim + // normalized_data_batch, modify_data); + //The default implementation is correct but very inefficient! + virtual void normalize_data_for_build(const T *original_data, + const uint32_t num_points, const uint32_t orig_dim, + T *normalized_data, bool modify_data = true) + { + } + + //Invokes normalization for a single vector during search. The scratch space + //has to be created by the caller keeping track of the fact that normalization + //might change the dimension of the query vector. + virtual void normalize_vector_for_search(const T *query_vec, const uint32_t query_dim, + T* scratch_query) + { + memcpy(scratch_query, query_vec, query_dim * sizeof(T)); + } + + //Providing a default implementation for the virtual destructor because we don't + //expect most metric implementations to need it. + virtual ~Distance() { } }; @@ -137,6 +188,16 @@ class AVXNormalizedCosineDistanceFloat : public Distance // This will ensure that cosine is between -1 and 1. return 1.0f + _innerProduct.compare(a, b, length); } + DISKANN_DLLEXPORT virtual uint32_t new_dimension(uint32_t orig_dimension); + + DISKANN_DLLEXPORT virtual bool normalization_required() const; + + DISKANN_DLLEXPORT virtual void normalize_data_for_build(const float *original_data, const uint32_t num_points, + const uint32_t orig_dim, float *normalized_data, + bool modify_orig); + + DISKANN_DLLEXPORT virtual void normalize_vector_for_search(const float *query_vec, const uint32_t query_dim, + float* scratch_query_vector); }; template Distance *get_distance_function(Metric m); diff --git a/src/distance.cpp b/src/distance.cpp index ecf6597ae..aa206ef7d 100644 --- a/src/distance.cpp +++ b/src/distance.cpp @@ -522,6 +522,34 @@ float AVXDistanceInnerProductFloat::compare(const float *a, const float *b, uint return -result; } +uint32_t AVXNormalizedCosineDistanceFloat::new_dimension(uint32_t orig_dimension) +{ + return orig_dimension; +} +bool AVXNormalizedCosineDistanceFloat::normalization_required() const +{ + return true; +} +void AVXNormalizedCosineDistanceFloat::normalize_data_for_build(const float *original_data, const uint32_t num_points, + const uint32_t orig_dim, float *normalized_data, + bool modify_orig) +{ + if (!modify_orig) + { + throw diskann::ANNException("For cosine distance, modify_orig should be set to true for efficiency.", -1); + } + for (auto i = 0; i < num_points; i++) + { + normalize(original_data + i * orig_dim, orig_dim); + } +} +void AVXNormalizedCosineDistanceFloat::normalize_vector_for_search(const float *query_vec, const uint32_t query_dim, + float *scratch_space) +{ + normalize(query_vec, query_dim); +} + + // Get the right distance function for the given metric. template <> diskann::Distance *get_distance_function(diskann::Metric m) { From 8adc959810d1aa8411d5d6d0067fc2fd4cd0dfd7 Mon Sep 17 00:00:00 2001 From: Gopal Srinivasa Date: Mon, 3 Apr 2023 14:51:34 +0530 Subject: [PATCH 05/69] Refactoring changes --- include/abstract_data_store.h | 16 ++++--------- include/distance.h | 40 +++++++++++++++++++++------------ include/in_mem_data_store.h | 19 +++++++++++----- src/in_mem_data_store.cpp | 42 ++++++++++++++++------------------- 4 files changed, 62 insertions(+), 55 deletions(-) diff --git a/include/abstract_data_store.h b/include/abstract_data_store.h index 9307ac625..9ee4e8108 100644 --- a/include/abstract_data_store.h +++ b/include/abstract_data_store.h @@ -10,11 +10,11 @@ namespace diskann { -template class AbstractDataStore +template class AbstractDataStore { public: - AbstractDataStore(const location_t max_pts, const location_t num_frozen_pts, const size_t dim) - : _max_pts(max_pts), _num_pts(0), _num_frozen_pts(num_frozen_pts), _dim(dim) + AbstractDataStore(const location_t max_pts, const size_t dim) + : _max_pts(max_pts), _num_pts(0), _dim(dim) { } @@ -23,8 +23,6 @@ template class AbstractDataStore virtual void store(const std::string &filename) = 0; virtual data_t *get_vector(location_t i) = 0; - virtual data_t *get_vector_by_UID(id_t uid) = 0; - virtual void set_vector(const location_t i, const data_t *const vector) = 0; location_t get_max_pts() @@ -37,12 +35,7 @@ template class AbstractDataStore return _num_pts; } - //location_t get_num_frozen_pts() - //{ - // return _num_frozen_pts; - //} - - virtual void get_distance(const T *query, const location_t *locations, const uint32_t location_count, const shared_ptr> &metric, + virtual void get_distance(const T *query, const location_t *locations, const uint32_t location_count, float *distances) = 0; @@ -54,7 +47,6 @@ template class AbstractDataStore protected: location_t _max_pts; location_t _num_pts; - location_t _num_frozen_pts; const size_t _dim; }; diff --git a/include/distance.h b/include/distance.h index 39db1a0a7..ddf49e42b 100644 --- a/include/distance.h +++ b/include/distance.h @@ -15,17 +15,32 @@ enum Metric template class Distance { public: + Distance(diskann::Metric dist_metric) : _distance_metric(dist_metric) + { + } + // distance comparison function virtual float compare(const T *a, const T *b, uint32_t length) const = 0; + //Needed only for COSINE-BYTE and INNER_PRODUCT-BYTE + virtual float compare(const T *a, const T *b, const float normA, const float normB, uint32_t length) const + { + return std::numeric_limits::max(); + } + // For MIPS, normalization => a new dimension gets added to the vectors. // This function lets callers know if the normalization process // changes the dimension. - virtual uint32_t new_dimension(uint32_t orig_dimension) const + virtual uint32_t post_processed_dimension(uint32_t orig_dimension) const { return orig_dimension; } + virtual diskann::Metric get_metric() const + { + return _distance_metric; + } + // This is for efficiency. If no normalization is required, the callers //can simply ignore the normalize_data_for_build() function. virtual bool normalization_required() const @@ -38,20 +53,14 @@ template class Distance // // if (metric->normalization_required()){ // T* normalized_data_batch; - // if ( metric->new_dimension() == orig_dim ) { - // normalized_data_batch = nullptr; - // modify_data = true; - // } else { - // normalized_data_batch = new T[batch_size * metric->new_dimension()]; - // modify_data = false; - // } // Split data into batches of batch_size and for each, call: - // metric->normalize_data_for_build(data_batch, batch_size, orig_dim - // normalized_data_batch, modify_data); - //The default implementation is correct but very inefficient! - virtual void normalize_data_for_build(const T *original_data, - const uint32_t num_points, const uint32_t orig_dim, - T *normalized_data, bool modify_data = true) + // metric->normalize_data_for_build(data_batch, batch_size); + // + // TODO: This does not take into account the case for SSD inner product + // where the dimensions change after normalization. + // + virtual void normalize_data_for_build(T *original_data, + const uint32_t num_points) { } @@ -69,6 +78,9 @@ template class Distance virtual ~Distance() { } + + protected: + diskann::Metric _distance_metric; }; class DistanceCosineInt8 : public Distance diff --git a/include/in_mem_data_store.h b/include/in_mem_data_store.h index dd31d5afa..b0efc75ed 100644 --- a/include/in_mem_data_store.h +++ b/include/in_mem_data_store.h @@ -15,21 +15,20 @@ namespace diskann { -template class InMemDataStore : public AbstractDataStore +template class InMemDataStore : public AbstractDataStore { public: - InMemDataStore(const location_t max_pts, const location_t num_frozen_pts, const size_t dim); + InMemDataStore(const location_t max_pts, const size_t dim, std::shared_ptr> distance_metric); virtual ~InMemDataStore(); virtual void load(const std::string &filename); virtual void store(const std::string &filename); virtual data_t *get_vector(location_t i); - virtual data_t *get_vector_by_UID(id_t uid); - virtual void set_vector(const location_t i, const data_t *const vector); - virtual void get_distance(const T *query, const location_t *locations, const uint32_t location_count, - const shared_ptr> &metric, float *distances); + + virtual void get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, + const shared_ptr> &metric, float *distances); protected: virtual location_t load_data(const std::string &filename); @@ -55,6 +54,14 @@ template class InMemDataStore : public Abstract std::unique_ptr> _delete_set; std::shared_timed_mutex _lock; // Takes please of Index::_tag_lock + + // It may seem weird to put distance metric along with the data store class, but + // this gives us perf benefits as the datastore can do distance computations during + // search and compute norms of vectors internally without have to copy + // data back and forth. + shared_ptr> _distance_fn; + + shared_ptr _pre_computed_norms; }; } // namespace diskann \ No newline at end of file diff --git a/src/in_mem_data_store.cpp b/src/in_mem_data_store.cpp index 4e39297a9..aa90432bc 100644 --- a/src/in_mem_data_store.cpp +++ b/src/in_mem_data_store.cpp @@ -8,17 +8,17 @@ namespace diskann { -template -InMemDataStore::InMemDataStore(const location_t max_pts, const location_t num_frozen_pts, - const size_t dim) - : AbstractDataStore(max_pts, num_frozen_pts, dim), _aligned_dim(ROUND_UP(dim, 8)) +template +InMemDataStore::InMemDataStore(const location_t max_pts, + const size_t dim, std::shared_ptr> distance_metric) + : AbstractDataStore(max_pts, dim), _aligned_dim(ROUND_UP(dim, 8)), + _distance_metric(distance_metric) { - location_t total_internal_points = max_pts + num_frozen_pts; - alloc_aligned(((void **)&_data), total_internal_points * _aligned_dim * sizeof(data_t), 8 * sizeof(data_t)); - std::memset(_data, 0, total_internal_points * _aligned_dim * sizeof(data_t)); + alloc_aligned(((void **)&_data), max_pts * _aligned_dim * sizeof(data_t), 8 * sizeof(data_t)); + std::memset(_data, 0, max_pts * _aligned_dim * sizeof(data_t)); } -template InMemDataStore::~InMemDataStore() +template InMemDataStore::~InMemDataStore() { if (_data != nullptr) { @@ -26,16 +26,17 @@ template InMemDataStore::~InMemDa } } -template void InMemDataStore::load(const std::string &filename) +template void InMemDataStore::load(const std::string &filename) { + load_data(filename); } -template void InMemDataStore::store(const std::string &filename) +template void InMemDataStore::store(const std::string &filename) { } #ifdef EXEC_ENV_OLS -template location_t Index::load_data(AlignedFileReader &reader) +template location_t Index::load_data(AlignedFileReader &reader) { size_t file_dim, file_num_points; @@ -63,8 +64,8 @@ template location_t Index } #endif -template -location_t InMemDataStore::load_data(const std::string &filename) +template +location_t InMemDataStore::load_data(const std::string &filename) { size_t file_dim, file_num_points; if (!file_exists(filename)) @@ -90,9 +91,9 @@ location_t InMemDataStore::load_data(const std::string &filename) throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } - if (file_num_points > this->get_max_points() + this->get_num_frozen_pts()) + if (file_num_points > this->get_max_points()) { - resize(file_num_points - this->get_num_frozen_pts()); + resize(file_num_points); } copy_aligned_data_from_file(filename.c_str(), _data, file_num_points, file_dim, _aligned_dim); @@ -100,18 +101,13 @@ location_t InMemDataStore::load_data(const std::string &filename) return file_num_points; } -template data_t *InMemDataStore::get_vector(location_t i) +template data_t *InMemDataStore::get_vector(location_t i) { return _data + i * _aligned_dim; } -template data_t *InMemDataStore::get_vector_by_UID(id_t uid) -{ - return get_vector(_tag_to_location(uid)); -} - -template -void InMemDataStore::set_vector(const location_t loc, const data_t *const vector) +template +void InMemDataStore::set_vector(const location_t loc, const data_t *const vector) { memcpy(_data + loc * _aligned_dim, vector, this->_dim * sizeof(data_t)); } From 412ed828d7a5593de7c62deb7b2e1e17a4a2030a Mon Sep 17 00:00:00 2001 From: Gopal Srinivasa Date: Mon, 3 Apr 2023 18:38:33 +0530 Subject: [PATCH 06/69] Fixing compile errors in refactored code --- CMakeLists.txt | 1 + include/abstract_data_store.h | 2 +- include/distance.h | 50 +++++++++++++++++++++++++++++++++++ include/in_mem_data_store.h | 16 ++++++----- src/distance.cpp | 21 +++++++++++++-- 5 files changed, 80 insertions(+), 10 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index f29e13f3a..dfa1cb77e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -244,6 +244,7 @@ endif() add_subdirectory(src) add_subdirectory(tests) add_subdirectory(tests/utils) +add_subdirectory(tests/unittests) if (MSVC) message(STATUS "The ${PROJECT_NAME}.sln has been created, opened it from VisualStudio to build Release or Debug configurations.\n" diff --git a/include/abstract_data_store.h b/include/abstract_data_store.h index 9ee4e8108..9f43e176b 100644 --- a/include/abstract_data_store.h +++ b/include/abstract_data_store.h @@ -35,7 +35,7 @@ template class AbstractDataStore return _num_pts; } - virtual void get_distance(const T *query, const location_t *locations, const uint32_t location_count, + virtual void get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, float *distances) = 0; diff --git a/include/distance.h b/include/distance.h index ddf49e42b..f4f4b7794 100644 --- a/include/distance.h +++ b/include/distance.h @@ -86,12 +86,18 @@ template class Distance class DistanceCosineInt8 : public Distance { public: + DistanceCosineInt8() : Distance(diskann::Metric::COSINE) + { + } DISKANN_DLLEXPORT virtual float compare(const int8_t *a, const int8_t *b, uint32_t length) const; }; class DistanceL2Int8 : public Distance { public: + DistanceL2Int8() : Distance(diskann::Metric::L2) + { + } DISKANN_DLLEXPORT virtual float compare(const int8_t *a, const int8_t *b, uint32_t size) const; }; @@ -99,18 +105,28 @@ class DistanceL2Int8 : public Distance class AVXDistanceL2Int8 : public Distance { public: + AVXDistanceL2Int8() : Distance(diskann::Metric::L2) + { + } DISKANN_DLLEXPORT virtual float compare(const int8_t *a, const int8_t *b, uint32_t length) const; }; class DistanceCosineFloat : public Distance { public: + DistanceCosineFloat() : Distance(diskann::Metric::COSINE) + { + } DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t length) const; }; class DistanceL2Float : public Distance { public: + DistanceL2Float() : Distance(diskann::Metric::L2) + { + } + #ifdef _WINDOWS DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t size) const; #else @@ -121,24 +137,36 @@ class DistanceL2Float : public Distance class AVXDistanceL2Float : public Distance { public: + AVXDistanceL2Float() : Distance(diskann::Metric::L2) + { + } DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t length) const; }; class SlowDistanceL2Float : public Distance { public: + SlowDistanceL2Float() : Distance(diskann::Metric::L2) + { + } DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t length) const; }; class SlowDistanceCosineUInt8 : public Distance { public: + SlowDistanceCosineUInt8() : Distance(diskann::Metric::COSINE) + { + } DISKANN_DLLEXPORT virtual float compare(const uint8_t *a, const uint8_t *b, uint32_t length) const; }; class DistanceL2UInt8 : public Distance { public: + DistanceL2UInt8() : Distance(diskann::Metric::L2) + { + } DISKANN_DLLEXPORT virtual float compare(const uint8_t *a, const uint8_t *b, uint32_t size) const; }; @@ -146,6 +174,9 @@ class DistanceL2UInt8 : public Distance template class SlowDistanceL2Int : public Distance { public: + SlowDistanceL2Int() : Distance(diskann::Metric::L2) + { + } // Implementing here because this is a template function DISKANN_DLLEXPORT virtual float compare(const T *a, const T *b, uint32_t length) const { @@ -161,6 +192,13 @@ template class SlowDistanceL2Int : public Distance template class DistanceInnerProduct : public Distance { public: + DistanceInnerProduct() : Distance(diskann::Metric::INNER_PRODUCT) + { + } + + DistanceInnerProduct(diskann::Metric metric) : Distance(metric) + { + } inline float inner_product(const T *a, const T *b, unsigned size) const; inline float compare(const T *a, const T *b, unsigned size) const @@ -178,6 +216,9 @@ template class DistanceFastL2 : public DistanceInnerProduct // currently defined only for float. // templated for future use. public: + DistanceFastL2() : DistanceInnerProduct(diskann::Metric::FAST_L2) + { + } float norm(const T *a, unsigned size) const; float compare(const T *a, const T *b, float norm, unsigned size) const; }; @@ -185,6 +226,9 @@ template class DistanceFastL2 : public DistanceInnerProduct class AVXDistanceInnerProductFloat : public Distance { public: + AVXDistanceInnerProductFloat() : Distance(diskann::Metric::INNER_PRODUCT) + { + } DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t length) const; }; @@ -193,7 +237,13 @@ class AVXNormalizedCosineDistanceFloat : public Distance private: AVXDistanceInnerProductFloat _innerProduct; + protected: + void normalize_and_copy(const float *a, uint32_t length, float *a_norm ) const; + public: + AVXNormalizedCosineDistanceFloat() : Distance(diskann::Metric::COSINE) + { + } DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t length) const { // Inner product returns negative values to indicate distance. diff --git a/include/in_mem_data_store.h b/include/in_mem_data_store.h index b0efc75ed..9ec89bbf6 100644 --- a/include/in_mem_data_store.h +++ b/include/in_mem_data_store.h @@ -2,6 +2,7 @@ // Licensed under the MIT license. #include +#include #include "tsl/robin_map.h" #include "tsl/robin_set.h" @@ -10,12 +11,14 @@ #include "abstract_data_store.h" +#include "distance.h" #include "natural_number_map.h" #include "natural_number_set.h" namespace diskann { -template class InMemDataStore : public AbstractDataStore +template +class InMemDataStore : public AbstractDataStore { public: InMemDataStore(const location_t max_pts, const size_t dim, std::shared_ptr> distance_metric); @@ -27,8 +30,7 @@ template class InMemDataStore : public AbstractDataStore> &metric, float *distances); + virtual void get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, float *distances); protected: virtual location_t load_data(const std::string &filename); @@ -43,8 +45,8 @@ template class InMemDataStore : public AbstractDataStore _tag_to_location; - natural_number_map _location_to_tag; + //tsl::sparse_map _tag_to_location; + //natural_number_map _location_to_tag; // _empty_slots has unallocated slots and those freed by consolidate_delete. // _delete_set has locations marked deleted by lazy_delete. Will not be @@ -59,9 +61,9 @@ template class InMemDataStore : public AbstractDataStore> _distance_fn; + std::shared_ptr> _distance_fn; - shared_ptr _pre_computed_norms; + std::shared_ptr _pre_computed_norms; }; } // namespace diskann \ No newline at end of file diff --git a/src/distance.cpp b/src/distance.cpp index a988b0a56..0431ab443 100644 --- a/src/distance.cpp +++ b/src/distance.cpp @@ -543,10 +543,27 @@ void AVXNormalizedCosineDistanceFloat::normalize_data_for_build(const float *ori normalize(original_data + i * orig_dim, orig_dim); } } + void AVXNormalizedCosineDistanceFloat::normalize_vector_for_search(const float *query_vec, const uint32_t query_dim, - float *scratch_space) + float *query_scratch) +{ + normalize_and_copy(query_vec, query_dim, query_scratch); +} + +void AVXNormalizedCosineDistanceFloat::normalize_and_copy(const float *query_vec, const uint32_t query_dim, + float *query_target) const { - normalize(query_vec, query_dim); + float norm = 0.0f; + for (auto i = 0; i < query_dim; i++) + { + norm += query_vec[i]; + } + norm /= norm / query_dim; + + for (auto i = 0; i < query_dim; i++) + { + query_target[i] = query_vec[i] / norm; + } } From 2af4b1cc27a53bb17db99bca4374df6f605ec9bc Mon Sep 17 00:00:00 2001 From: Gopal Srinivasa Date: Mon, 3 Apr 2023 23:04:07 +0530 Subject: [PATCH 07/69] Fixing compile errors --- include/distance.h | 10 ++++------ include/index.h | 2 -- include/utils.h | 13 +++++++++---- src/distance.cpp | 22 ++++++---------------- 4 files changed, 19 insertions(+), 28 deletions(-) diff --git a/include/distance.h b/include/distance.h index f4f4b7794..e4ab0e51b 100644 --- a/include/distance.h +++ b/include/distance.h @@ -59,7 +59,7 @@ template class Distance // TODO: This does not take into account the case for SSD inner product // where the dimensions change after normalization. // - virtual void normalize_data_for_build(T *original_data, + virtual void normalize_data_for_build(T *original_data, const uint32_t orig_dim, const uint32_t num_points) { } @@ -250,16 +250,14 @@ class AVXNormalizedCosineDistanceFloat : public Distance // This will ensure that cosine is between -1 and 1. return 1.0f + _innerProduct.compare(a, b, length); } - DISKANN_DLLEXPORT virtual uint32_t new_dimension(uint32_t orig_dimension); + DISKANN_DLLEXPORT virtual uint32_t post_processed_dimension(uint32_t orig_dimension) const override; DISKANN_DLLEXPORT virtual bool normalization_required() const; - DISKANN_DLLEXPORT virtual void normalize_data_for_build(const float *original_data, const uint32_t num_points, - const uint32_t orig_dim, float *normalized_data, - bool modify_orig); + DISKANN_DLLEXPORT virtual void normalize_data_for_build(float *original_data, const uint32_t orig_dim, const uint32_t num_points) override; DISKANN_DLLEXPORT virtual void normalize_vector_for_search(const float *query_vec, const uint32_t query_dim, - float* scratch_query_vector); + float *scratch_query_vector) override; }; template Distance *get_distance_function(Metric m); diff --git a/include/index.h b/include/index.h index 1376d9987..f98da0c48 100644 --- a/include/index.h +++ b/include/index.h @@ -255,8 +255,6 @@ template clas // add reverse links from all the visited nodes to node n. void inter_insert(unsigned n, std::vector &pruned_list, const uint32_t range, InMemQueryScratch *scratch); - void inter_insert(uint32_t n, std::vector &pruned_list, const uint32_t range, - InMemQueryScratch *scratch); void inter_insert(uint32_t n, std::vector &pruned_list, InMemQueryScratch *scratch); diff --git a/include/utils.h b/include/utils.h index 018752792..77b116218 100644 --- a/include/utils.h +++ b/include/utils.h @@ -1004,18 +1004,23 @@ inline bool validate_index_file_size(std::ifstream &in) return true; } -// This function is valid only for float data type. -template inline void normalize(T *arr, size_t dim) +template inline float get_norm(T *arr, const size_t dim) { float sum = 0.0f; for (uint32_t i = 0; i < dim; i++) { sum += arr[i] * arr[i]; } - sum = sqrt(sum); + return sqrt(sum); +} + +// This function is valid only for float data type. +template inline void normalize(T *arr, const size_t dim) +{ + float norm = get_norm(arr, dim); for (uint32_t i = 0; i < dim; i++) { - arr[i] = (T)(arr[i] / sum); + arr[i] = (T)(arr[i] / norm); } } diff --git a/src/distance.cpp b/src/distance.cpp index 0431ab443..d0f446905 100644 --- a/src/distance.cpp +++ b/src/distance.cpp @@ -522,7 +522,7 @@ float AVXDistanceInnerProductFloat::compare(const float *a, const float *b, uint return -result; } -uint32_t AVXNormalizedCosineDistanceFloat::new_dimension(uint32_t orig_dimension) +uint32_t AVXNormalizedCosineDistanceFloat::post_processed_dimension(uint32_t orig_dimension) const { return orig_dimension; } @@ -530,17 +530,11 @@ bool AVXNormalizedCosineDistanceFloat::normalization_required() const { return true; } -void AVXNormalizedCosineDistanceFloat::normalize_data_for_build(const float *original_data, const uint32_t num_points, - const uint32_t orig_dim, float *normalized_data, - bool modify_orig) +void AVXNormalizedCosineDistanceFloat::normalize_data_for_build(float *original_data, const uint32_t orig_dim, const uint32_t num_points) { - if (!modify_orig) - { - throw diskann::ANNException("For cosine distance, modify_orig should be set to true for efficiency.", -1); - } for (auto i = 0; i < num_points; i++) { - normalize(original_data + i * orig_dim, orig_dim); + normalize((float *)(original_data + i * orig_dim), orig_dim); } } @@ -553,13 +547,9 @@ void AVXNormalizedCosineDistanceFloat::normalize_vector_for_search(const float * void AVXNormalizedCosineDistanceFloat::normalize_and_copy(const float *query_vec, const uint32_t query_dim, float *query_target) const { - float norm = 0.0f; - for (auto i = 0; i < query_dim; i++) - { - norm += query_vec[i]; - } - norm /= norm / query_dim; - + + float norm = get_norm(query_vec, query_dim); + for (auto i = 0; i < query_dim; i++) { query_target[i] = query_vec[i] / norm; From 8d6b832aaf580da0c975fce5ed00311176564054 Mon Sep 17 00:00:00 2001 From: Gopal Srinivasa Date: Tue, 4 Apr 2023 00:25:30 +0530 Subject: [PATCH 08/69] DiskANN Builds with initial refactoring changes --- include/abstract_graph_store.h | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/include/abstract_graph_store.h b/include/abstract_graph_store.h index fe088d51b..5fa944ab9 100644 --- a/include/abstract_graph_store.h +++ b/include/abstract_graph_store.h @@ -14,13 +14,18 @@ namespace diskann class AbstractGraphStore { public: - AbstractGraphStore(const size_t max_pts); + AbstractGraphStore(const size_t max_pts) + : _max_points(max_pts){} virtual int load(const std::string &index_path_prefix) = 0; virtual int store(const std::string &index_path_prefix) = 0; virtual void get_adj_list(const location_t i, std::vector &neighbors) = 0; virtual void set_adj_list(const location_t i, std::vector &neighbors) = 0; + + private: + size_t _max_points; + }; } // namespace diskann \ No newline at end of file From c417f9a150ad1d4819ca1435d5419d221d3fbf48 Mon Sep 17 00:00:00 2001 From: Gopal Srinivasa Date: Tue, 4 Apr 2023 16:34:59 +0530 Subject: [PATCH 09/69] Saving changes for Ravi --- include/abstract_data_store.h | 33 +-- include/in_mem_data_store.h | 18 +- include/index.h | 6 +- src/in_mem_data_store.cpp | 53 ++++- src/index.cpp | 417 ++++++++++++++++++---------------- 5 files changed, 306 insertions(+), 221 deletions(-) diff --git a/include/abstract_data_store.h b/include/abstract_data_store.h index 9f43e176b..b6b15e262 100644 --- a/include/abstract_data_store.h +++ b/include/abstract_data_store.h @@ -14,39 +14,46 @@ template class AbstractDataStore { public: AbstractDataStore(const location_t max_pts, const size_t dim) - : _max_pts(max_pts), _num_pts(0), _dim(dim) + : _max_points(max_pts), _num_points(0), _dim(dim) { } // Return number of points returned - virtual size_t load(const std::string &filename) = 0; + virtual location_t load(const std::string &filename) = 0; virtual void store(const std::string &filename) = 0; - virtual data_t *get_vector(location_t i) = 0; - virtual void set_vector(const location_t i, const data_t *const vector) = 0; + virtual void populate_data(const data_t * vectors, const location_t num_pts) = 0; + virtual void populate_data(const std::string &filename, const size_t offset) = 0; + + virtual void get_vector(const location_t i, data_t* dest) const = 0; + virtual void set_vector(const location_t i, const data_t *const vector) = 0; - location_t get_max_pts() + + virtual void reposition_points(const location_t start_loc, const location_t end_loc, + const location_t num_points) = 0; + + location_t get_max_points() const { - return _max_pts; + return _max_points; } - - location_t get_num_pts() + location_t get_num_points() const { - return _num_pts; + return _num_points; } virtual void get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, - float *distances) = 0; + float *distances) const = 0; + virtual float get_distance(const location_t loc1, const location_t loc2) const = 0; - size_t get_dims() + size_t get_dims() const { return _dim; } protected: - location_t _max_pts; - location_t _num_pts; + location_t _max_points; + location_t _num_points; const size_t _dim; }; diff --git a/include/in_mem_data_store.h b/include/in_mem_data_store.h index 9ec89bbf6..1c087f4ce 100644 --- a/include/in_mem_data_store.h +++ b/include/in_mem_data_store.h @@ -24,13 +24,21 @@ class InMemDataStore : public AbstractDataStore InMemDataStore(const location_t max_pts, const size_t dim, std::shared_ptr> distance_metric); virtual ~InMemDataStore(); - virtual void load(const std::string &filename); - virtual void store(const std::string &filename); + virtual location_t load(const std::string &filename) override; + virtual void store(const std::string &filename) override; - virtual data_t *get_vector(location_t i); - virtual void set_vector(const location_t i, const data_t *const vector); + //Populate internal data from unaligned data while doing alignment and any normalization that is required. + virtual void populate_data(const data_t *vectors, const location_t num_pts) override; + virtual void populate_data(const std::string &filename, const size_t offset) override; - virtual void get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, float *distances); + virtual void get_vector(const location_t i, data_t *target) const override; + virtual void set_vector(const location_t i, const data_t *const vector) override; + + virtual void reposition_points(const location_t start_loc, const location_t end_loc, + const location_t num_points) override; + + virtual void get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, float *distances) const override ; + virtual float get_distance(const location_t loc1, const location_t loc2) const override; protected: virtual location_t load_data(const std::string &filename); diff --git a/include/index.h b/include/index.h index f98da0c48..24a249531 100644 --- a/include/index.h +++ b/include/index.h @@ -18,6 +18,7 @@ #include "utils.h" #include "windows_customizations.h" #include "scratch.h" +#include "in_mem_data_store.h" #define OVERHEAD_FACTOR 1.1 #define EXPAND_IF_FULL 0 @@ -313,7 +314,8 @@ template clas Distance *_distance = nullptr; // Data - T *_data = nullptr; + //T *_data = nullptr; + std::shared_ptr> _data_store; char *_opt_graph = nullptr; // Graph related data structures @@ -321,7 +323,7 @@ template clas // Dimensions size_t _dim = 0; - size_t _aligned_dim = 0; + //size_t _aligned_dim = 0; 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 // Number of points which are used as initial candidates when iterating to diff --git a/src/in_mem_data_store.cpp b/src/in_mem_data_store.cpp index aa90432bc..1c09d7108 100644 --- a/src/in_mem_data_store.cpp +++ b/src/in_mem_data_store.cpp @@ -26,7 +26,7 @@ template InMemDataStore::~InMemDataStore() } } -template void InMemDataStore::load(const std::string &filename) +template location_t InMemDataStore::load(const std::string &filename) { load_data(filename); } @@ -35,6 +35,29 @@ template void InMemDataStore::store(const std::string { } +template void InMemDataStore::populate_data(const data_t *vectors, const location_t num_pts) +{ + for (auto i = 0; i < num_pts; i++) + { + memset(_data + i * _aligned_dim, 0, _aligned_dim * sizeof(data_t)); + std::memmove(_data + i * _aligned_dim, vectors + i * _dim, _dim * sizeof(data_t)); + } + + if (_distance_metric->normalization_required()) + { + _distance_metric->normalize(_data, num_pts); + } +} + +template void InMemDataStore::populate_data(const std::string &filename, const size_t offset) +{ + copy_aligned_data_from_file(filename.c_str(), _data, _num_points, _dim, _aligned_dim, offset); + if (_distance_metric->normalization_required()) + { + _distance_metric->normalize(_data, _num_points); + } +} + #ifdef EXEC_ENV_OLS template location_t Index::load_data(AlignedFileReader &reader) { @@ -101,15 +124,37 @@ location_t InMemDataStore::load_data(const std::string &filename) return file_num_points; } -template data_t *InMemDataStore::get_vector(location_t i) +template +void InMemDataStore::get_vector(const location_t i, data_t* dest) const { - return _data + i * _aligned_dim; + memcpy(dest, _data + i * _aligned_dim, this->_dim * sizeof(data_t)); } template void InMemDataStore::set_vector(const location_t loc, const data_t *const vector) { - memcpy(_data + loc * _aligned_dim, vector, this->_dim * sizeof(data_t)); + size_t offset_in_data = loc * _aligned_dim; + memset(_data + offset_in_data, 0, _aligned_dim * sizeof(data_t)); + memcpy(_data + offset_in_data, vector, this->_dim * sizeof(data_t)); + if (_distance_metric->normalization_required()) + { + _distance_metric->normalize(_data + offset_in_data, _aligned_dim); + } +} + +template +void InMemDataStore::get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, float *distances) const +{ + for (auto i = 0; i < location_count; i++) + { + distances[i] = _distance_metric->compare(query, _data + locations[i] * _aligned_dim, this->_aligned_dim); + } +} + +template +float InMemDataStore::get_distance(const location_t loc1, const location_t loc2) const +{ + return _distance_metric->compare(_data + loc1 * _aligned_dim, _data + loc2 * _aligned_dim, this->_aligned_dim); } } // namespace diskann \ No newline at end of file diff --git a/src/index.cpp b/src/index.cpp index 1aa2e11c1..814aed4c2 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -76,7 +76,8 @@ Index::Index(Metric m, const size_t dim, const size_t max_point } // data stored to _nd * aligned_dim matrix with necessary zero-padding - _aligned_dim = ROUND_UP(_dim, 8); + // REFACTOR + //_aligned_dim = ROUND_UP(_dim, 8); if (dynamic_index && _num_frozen_pts == 0) { @@ -97,13 +98,18 @@ Index::Index(Metric m, const size_t dim, const size_t max_point alloc_aligned(((void **)&_pq_data), total_internal_points * _num_pq_chunks * sizeof(char), 8 * sizeof(char)); std::memset(_pq_data, 0, total_internal_points * _num_pq_chunks * sizeof(char)); } - alloc_aligned(((void **)&_data), total_internal_points * _aligned_dim * sizeof(T), 8 * sizeof(T)); - std::memset(_data, 0, total_internal_points * _aligned_dim * sizeof(T)); + + //REFACTOR + //alloc_aligned(((void **)&_data), total_internal_points * _aligned_dim * sizeof(T), 8 * sizeof(T)); + //std::memset(_data, 0, total_internal_points * _aligned_dim * sizeof(T)); + //REFACTOR: This should move to a factory method. + _data_store = std::make_shared>(total_internal_points, _dim); _start = (uint32_t)_max_points; _final_graph.resize(total_internal_points); + //This should come from a factory. if (m == diskann::Metric::COSINE && std::is_floating_point::value) { // This is safe because T is float inside the if block. @@ -145,11 +151,12 @@ template Index::~I delete this->_distance; this->_distance = nullptr; } - if (this->_data != nullptr) - { - aligned_free(this->_data); - this->_data = nullptr; - } + //REFACTOR + //if (this->_data != nullptr) + //{ + // aligned_free(this->_data); + // this->_data = nullptr; + //} if (_opt_graph != nullptr) { delete[] _opt_graph; @@ -560,7 +567,8 @@ void Index::load(const char *filename, uint32_t num_threads, ui << 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; diskann::cerr << stream.str() << std::endl; - aligned_free(_data); + //REFACTOR + //aligned_free(_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } @@ -1505,7 +1513,7 @@ void Index::prune_all_nbrs(const Parameters ¶meters) { const uint32_t range = parameters.Get("R"); const uint32_t maxc = parameters.Get("C"); - const float alpha = parameters.Get("alpha"); + const float alpha = parameters.Get("alpha"); _filtered_index = true; diskann::Timer timer; @@ -1527,8 +1535,10 @@ void Index::prune_all_nbrs(const Parameters ¶meters) { 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, (uint32_t)_aligned_dim); + float dist = _data_store->compare(node, cur_nbr); + //REFACTOR + //float dist = _distance->compare(_data + _aligned_dim * (size_t)node, + // _data + _aligned_dim * (size_t)cur_nbr, (uint32_t)_aligned_dim); dummy_pool.emplace_back(Neighbor(cur_nbr, dist)); dummy_visited.insert(cur_nbr); } @@ -1566,49 +1576,50 @@ void Index::prune_all_nbrs(const Parameters ¶meters) } } -template -void Index::set_start_points(const T *data, size_t data_count) -{ - std::unique_lock ul(_update_lock); - std::unique_lock tl(_tag_lock); - if (_nd > 0) - throw ANNException("Can not set starting point for a non-empty index", -1, __FUNCSIG__, __FILE__, __LINE__); - - if (data_count != _num_frozen_pts * _aligned_dim) - throw ANNException("Invalid number of points", -1, __FUNCSIG__, __FILE__, __LINE__); - - memcpy(_data + _aligned_dim * _max_points, data, _aligned_dim * sizeof(T) * _num_frozen_pts); - _has_built = true; - diskann::cout << "Index start points set: #" << _num_frozen_pts << std::endl; -} - -template -void Index::set_start_points_at_random(T radius, uint32_t random_seed) -{ - std::mt19937 gen{random_seed}; - std::normal_distribution<> d{0.0, 1.0}; - - std::vector points_data; - points_data.reserve(_aligned_dim * _num_frozen_pts); - std::vector real_vec(_aligned_dim); - - for (size_t frozen_point = 0; frozen_point < _num_frozen_pts; frozen_point++) - { - double norm_sq = 0.0; - for (size_t i = 0; i < _dim; ++i) - { - auto r = d(gen); - real_vec[i] = r; - norm_sq += r * r; - } - - const double norm = std::sqrt(norm_sq); - for (auto iter : real_vec) - points_data.push_back(static_cast(iter * radius / norm)); - } - - set_start_points(points_data.data(), points_data.size()); -} +//REFACTOR +//template +//void Index::set_start_points(const T *data, size_t data_count) +//{ +// std::unique_lock ul(_update_lock); +// std::unique_lock tl(_tag_lock); +// if (_nd > 0) +// throw ANNException("Can not set starting point for a non-empty index", -1, __FUNCSIG__, __FILE__, __LINE__); +// +// if (data_count != _num_frozen_pts * _aligned_dim) +// throw ANNException("Invalid number of points", -1, __FUNCSIG__, __FILE__, __LINE__); +// +// memcpy(_data + _aligned_dim * _max_points, data, _aligned_dim * sizeof(T) * _num_frozen_pts); +// _has_built = true; +// diskann::cout << "Index start points set: #" << _num_frozen_pts << std::endl; +//} +// +//template +//void Index::set_start_points_at_random(T radius, uint32_t random_seed) +//{ +// std::mt19937 gen{random_seed}; +// std::normal_distribution<> d{0.0, 1.0}; +// +// std::vector points_data; +// points_data.reserve(_aligned_dim * _num_frozen_pts); +// std::vector real_vec(_aligned_dim); +// +// for (size_t frozen_point = 0; frozen_point < _num_frozen_pts; frozen_point++) +// { +// double norm_sq = 0.0; +// for (size_t i = 0; i < _dim; ++i) +// { +// auto r = d(gen); +// real_vec[i] = r; +// norm_sq += r * r; +// } +// +// const double norm = std::sqrt(norm_sq); +// for (auto iter : real_vec) +// points_data.push_back(static_cast(iter * radius / norm)); +// } +// +// set_start_points(points_data.data(), points_data.size()); +//} template void Index::build_with_data_populated(Parameters ¶meters, const std::vector &tags) @@ -1686,15 +1697,17 @@ void Index::build(const T *data, const size_t num_points_to_loa std::unique_lock tl(_tag_lock); _nd = num_points_to_load; - memcpy((char *)_data, (char *)data, _aligned_dim * _nd * sizeof(T)); + _data_store->populate_data(data, num_points_to_load); - if (_normalize_vecs) - { - for (size_t i = 0; i < num_points_to_load; i++) - { - normalize(_data + _aligned_dim * i, _aligned_dim); - } - } + //REFACTOR + //memcpy((char *)_data, (char *)data, _aligned_dim * _nd * sizeof(T)); + //if (_normalize_vecs) + //{ + // for (size_t i = 0; i < num_points_to_load; i++) + // { + // normalize(_data + _aligned_dim * i, _aligned_dim); + // } + //} } build_with_data_populated(parameters, tags); @@ -1732,8 +1745,9 @@ void Index::build(const char *filename, const size_t num_points if (_pq_dist) aligned_free(_pq_data); - else - aligned_free(_data); + //REFACTOR + //else + // aligned_free(_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } @@ -1745,8 +1759,9 @@ void Index::build(const char *filename, const size_t num_points if (_pq_dist) aligned_free(_pq_data); - else - aligned_free(_data); + //REFACTOR + //else + // aligned_free(_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } @@ -1759,8 +1774,9 @@ void Index::build(const char *filename, const size_t num_points if (_pq_dist) aligned_free(_pq_data); - else - aligned_free(_data); + //REFACTOR + //else + // aligned_free(_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } @@ -1786,14 +1802,16 @@ void Index::build(const char *filename, const size_t num_points #endif } - copy_aligned_data_from_file(filename, _data, file_num_points, file_dim, _aligned_dim); - if (_normalize_vecs) - { - for (size_t i = 0; i < file_num_points; i++) - { - normalize(_data + _aligned_dim * i, _aligned_dim); - } - } + _data_store->populate_data(filename, 0U); + //REFACTOR + //copy_aligned_data_from_file(filename, _data, file_num_points, file_dim, _aligned_dim); + //if (_normalize_vecs) + //{ + // for (size_t i = 0; i < file_num_points; i++) + // { + // normalize(_data + _aligned_dim * i, _aligned_dim); + // } + //} diskann::cout << "Using only first " << num_points_to_load << " from file.. " << std::endl; @@ -2192,7 +2210,9 @@ size_t Index::search_with_tags(const T *query, const uint64_t K if (res_vectors.size() > 0) { - memcpy(res_vectors[pos], _data + ((size_t)node.id) * _aligned_dim, _dim * sizeof(T)); + _data_store->get_vector(node.id, res_vectors[pos]); + //REFACTOR + //memcpy(res_vectors[pos], _data + ((size_t)node.id) * _aligned_dim, _dim * sizeof(T)); } if (distances != nullptr) @@ -2832,15 +2852,17 @@ int Index::insert_point(const T *point, const TagT tag) } tl.unlock(); + _data_store->set_vector(location, point); + //REFACTOR // Copy the vector in to the data array - auto offset_data = _data + (size_t)_aligned_dim * location; - memset((void *)offset_data, 0, sizeof(T) * _aligned_dim); - memcpy((void *)offset_data, point, sizeof(T) * _dim); + //auto offset_data = _data + (size_t)_aligned_dim * location; + //memset((void *)offset_data, 0, sizeof(T) * _aligned_dim); + //memcpy((void *)offset_data, point, sizeof(T) * _dim); - if (_normalize_vecs) - { - normalize((float *)offset_data, _dim); - } + //if (_normalize_vecs) + //{ + // normalize((float *)offset_data, _dim); + //} // Find and add appropriate graph edges ScratchStoreManager> manager(_query_scratch); @@ -3008,118 +3030,119 @@ template void Index void Index::optimize_index_layout() -{ // use after build or load - if (_dynamic_index) - { - throw diskann::ANNException("Optimize_index_layout not implemented for dyanmic indices", -1, __FUNCSIG__, - __FILE__, __LINE__); - } - - _data_len = (_aligned_dim + 1) * sizeof(float); - _neighbor_len = (_max_observed_degree + 1) * sizeof(uint32_t); - _node_size = _data_len + _neighbor_len; - _opt_graph = new char[_node_size * _nd]; - DistanceFastL2 *dist_fast = (DistanceFastL2 *)_distance; - for (uint32_t 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; - uint32_t k = _final_graph[i].size(); - std::memcpy(cur_node_offset, &k, sizeof(uint32_t)); - std::memcpy(cur_node_offset + sizeof(uint32_t), _final_graph[i].data(), k * sizeof(uint32_t)); - std::vector().swap(_final_graph[i]); - } - _final_graph.clear(); - _final_graph.shrink_to_fit(); -} - -template -void Index::search_with_optimized_layout(const T *query, size_t K, size_t L, uint32_t *indices) -{ - DistanceFastL2 *dist_fast = (DistanceFastL2 *)_distance; - - NeighborPriorityQueue retset(L); - std::vector init_ids(L); - - boost::dynamic_bitset<> flags{_nd, 0}; - uint32_t tmp_l = 0; - uint32_t *neighbors = (uint32_t *)(_opt_graph + _node_size * _start + _data_len); - uint32_t 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) - { - uint32_t id = rand() % _nd; - if (flags[id]) - continue; - flags[id] = true; - init_ids[tmp_l] = id; - tmp_l++; - } - - for (uint32_t i = 0; i < init_ids.size(); i++) - { - uint32_t id = init_ids[i]; - if (id >= _nd) - continue; - _mm_prefetch(_opt_graph + _node_size * id, _MM_HINT_T0); - } - L = 0; - for (uint32_t i = 0; i < init_ids.size(); i++) - { - uint32_t 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, (uint32_t)_aligned_dim); - retset.insert(Neighbor(id, dist)); - flags[id] = true; - L++; - } - - while (retset.has_unexpanded_node()) - { - auto nbr = retset.closest_unexpanded(); - auto n = nbr.id; - _mm_prefetch(_opt_graph + _node_size * n + _data_len, _MM_HINT_T0); - neighbors = (uint32_t *)(_opt_graph + _node_size * n + _data_len); - uint32_t MaxM = *neighbors; - neighbors++; - for (uint32_t m = 0; m < MaxM; ++m) - _mm_prefetch(_opt_graph + _node_size * neighbors[m], _MM_HINT_T0); - for (uint32_t m = 0; m < MaxM; ++m) - { - uint32_t id = neighbors[m]; - if (flags[id]) - continue; - flags[id] = 1; - T *data = (T *)(_opt_graph + _node_size * id); - float norm = *data; - data++; - float dist = dist_fast->compare(query, data, norm, (uint32_t)_aligned_dim); - Neighbor nn(id, dist); - retset.insert(nn); - } - } - - for (size_t i = 0; i < K; i++) - { - indices[i] = retset[i].id; - } -} +//REFACTOR: This should be an OptimizedDataStore class +//template void Index::optimize_index_layout() +//{ // use after build or load +// if (_dynamic_index) +// { +// throw diskann::ANNException("Optimize_index_layout not implemented for dyanmic indices", -1, __FUNCSIG__, +// __FILE__, __LINE__); +// } +// +// _data_len = (_aligned_dim + 1) * sizeof(float); +// _neighbor_len = (_max_observed_degree + 1) * sizeof(uint32_t); +// _node_size = _data_len + _neighbor_len; +// _opt_graph = new char[_node_size * _nd]; +// DistanceFastL2 *dist_fast = (DistanceFastL2 *)_distance; +// for (uint32_t 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; +// uint32_t k = _final_graph[i].size(); +// std::memcpy(cur_node_offset, &k, sizeof(uint32_t)); +// std::memcpy(cur_node_offset + sizeof(uint32_t), _final_graph[i].data(), k * sizeof(uint32_t)); +// std::vector().swap(_final_graph[i]); +// } +// _final_graph.clear(); +// _final_graph.shrink_to_fit(); +//} +// +//template +//void Index::search_with_optimized_layout(const T *query, size_t K, size_t L, uint32_t *indices) +//{ +// DistanceFastL2 *dist_fast = (DistanceFastL2 *)_distance; +// +// NeighborPriorityQueue retset(L); +// std::vector init_ids(L); +// +// boost::dynamic_bitset<> flags{_nd, 0}; +// uint32_t tmp_l = 0; +// uint32_t *neighbors = (uint32_t *)(_opt_graph + _node_size * _start + _data_len); +// uint32_t 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) +// { +// uint32_t id = rand() % _nd; +// if (flags[id]) +// continue; +// flags[id] = true; +// init_ids[tmp_l] = id; +// tmp_l++; +// } +// +// for (uint32_t i = 0; i < init_ids.size(); i++) +// { +// uint32_t id = init_ids[i]; +// if (id >= _nd) +// continue; +// _mm_prefetch(_opt_graph + _node_size * id, _MM_HINT_T0); +// } +// L = 0; +// for (uint32_t i = 0; i < init_ids.size(); i++) +// { +// uint32_t 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, (uint32_t)_aligned_dim); +// retset.insert(Neighbor(id, dist)); +// flags[id] = true; +// L++; +// } +// +// while (retset.has_unexpanded_node()) +// { +// auto nbr = retset.closest_unexpanded(); +// auto n = nbr.id; +// _mm_prefetch(_opt_graph + _node_size * n + _data_len, _MM_HINT_T0); +// neighbors = (uint32_t *)(_opt_graph + _node_size * n + _data_len); +// uint32_t MaxM = *neighbors; +// neighbors++; +// for (uint32_t m = 0; m < MaxM; ++m) +// _mm_prefetch(_opt_graph + _node_size * neighbors[m], _MM_HINT_T0); +// for (uint32_t m = 0; m < MaxM; ++m) +// { +// uint32_t id = neighbors[m]; +// if (flags[id]) +// continue; +// flags[id] = 1; +// T *data = (T *)(_opt_graph + _node_size * id); +// float norm = *data; +// data++; +// float dist = dist_fast->compare(query, data, norm, (uint32_t)_aligned_dim); +// Neighbor nn(id, dist); +// retset.insert(nn); +// } +// } +// +// for (size_t i = 0; i < K; i++) +// { +// indices[i] = retset[i].id; +// } +//} /* Internals of the library */ template const float Index::INDEX_GROWTH_FACTOR = 1.5f; From e82b784691fb323bb73007fc2f7fe6a4cfc182ab Mon Sep 17 00:00:00 2001 From: Gopal Srinivasa Date: Tue, 4 Apr 2023 23:54:38 +0530 Subject: [PATCH 10/69] More refactoring --- include/abstract_data_store.h | 23 ++++-- include/abstract_graph_store.h | 4 +- include/in_mem_data_store.h | 10 ++- src/in_mem_data_store.cpp | 115 ++++++++++++++++++++++++++++- src/index.cpp | 129 +++++++++++++++++++-------------- 5 files changed, 215 insertions(+), 66 deletions(-) diff --git a/include/abstract_data_store.h b/include/abstract_data_store.h index b6b15e262..380a1d77f 100644 --- a/include/abstract_data_store.h +++ b/include/abstract_data_store.h @@ -13,8 +13,8 @@ namespace diskann template class AbstractDataStore { public: - AbstractDataStore(const location_t max_pts, const size_t dim) - : _max_points(max_pts), _num_points(0), _dim(dim) + AbstractDataStore(const location_t capacity, const size_t dim) + : _capacity(capacity), _num_points(0), _dim(dim) { } @@ -31,28 +31,37 @@ template class AbstractDataStore virtual void reposition_points(const location_t start_loc, const location_t end_loc, const location_t num_points) = 0; + virtual void copy_points(const location_t from_loc, const location_t to_loc, const location_t num_points) = 0; + //Returns the point in the dataset that is closest to the mean of all points in the dataset + virtual location_t calculate_medoid() const = 0; - location_t get_max_points() const + virtual location_t capacity() const { - return _max_points; + return _capacity; } - location_t get_num_points() const + virtual location_t get_num_points() const { return _num_points; } + virtual float get_distance(const data_t* query, const location_t loc) const = 0; virtual void get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, float *distances) const = 0; virtual float get_distance(const location_t loc1, const location_t loc2) const = 0; - size_t get_dims() const + virtual size_t get_dims() const + { + return _dim; + } + + virtual size_t get_aligned_dim() const { return _dim; } protected: - location_t _max_points; + location_t _capacity; location_t _num_points; const size_t _dim; diff --git a/include/abstract_graph_store.h b/include/abstract_graph_store.h index 5fa944ab9..d12c43664 100644 --- a/include/abstract_graph_store.h +++ b/include/abstract_graph_store.h @@ -15,7 +15,7 @@ class AbstractGraphStore { public: AbstractGraphStore(const size_t max_pts) - : _max_points(max_pts){} + : _capacity(max_pts){} virtual int load(const std::string &index_path_prefix) = 0; virtual int store(const std::string &index_path_prefix) = 0; @@ -24,7 +24,7 @@ class AbstractGraphStore virtual void set_adj_list(const location_t i, std::vector &neighbors) = 0; private: - size_t _max_points; + size_t _capacity; }; diff --git a/include/in_mem_data_store.h b/include/in_mem_data_store.h index 1c087f4ce..05cb5bdde 100644 --- a/include/in_mem_data_store.h +++ b/include/in_mem_data_store.h @@ -34,12 +34,20 @@ class InMemDataStore : public AbstractDataStore virtual void get_vector(const location_t i, data_t *target) const override; virtual void set_vector(const location_t i, const data_t *const vector) override; - virtual void reposition_points(const location_t start_loc, const location_t end_loc, + virtual void reposition_points(const location_t old_location_start, const location_t new_location_start, const location_t num_points) override; + virtual void copy_points(const location_t from_loc, const location_t to_loc, const location_t num_points) override; + virtual location_t calculate_medoid() const override; + virtual float get_distance(const data_t *query, const location_t loc) const override; virtual void get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, float *distances) const override ; virtual float get_distance(const location_t loc1, const location_t loc2) const override; + virtual size_t get_aligned_dim() const override + { + return _aligned_dim; + } + protected: virtual location_t load_data(const std::string &filename); #ifdef EXEC_ENV_OLS diff --git a/src/in_mem_data_store.cpp b/src/in_mem_data_store.cpp index 1c09d7108..fcaaf3b96 100644 --- a/src/in_mem_data_store.cpp +++ b/src/in_mem_data_store.cpp @@ -114,7 +114,7 @@ location_t InMemDataStore::load_data(const std::string &filename) throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } - if (file_num_points > this->get_max_points()) + if (file_num_points > this->capacity()) { resize(file_num_points); } @@ -142,6 +142,11 @@ void InMemDataStore::set_vector(const location_t loc, const data_t *cons } } +template float InMemDataStore::get_distance(const data_t *query, const location_t loc) const +{ + return _distance_metric->compare(query, _data + _aligned_dim * loc, _aligned_dim); +} + template void InMemDataStore::get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, float *distances) const { @@ -157,4 +162,112 @@ float InMemDataStore::get_distance(const location_t loc1, const location return _distance_metric->compare(_data + loc1 * _aligned_dim, _data + loc2 * _aligned_dim, this->_aligned_dim); } +template +void InMemDataStore::reposition_points(const location_t old_location_start, const location_t new_location_start, const location_t num_locations) +{ + if (num_locations == 0 || old_location_start == new_location_start) + { + return; + } + + // Update pointers to the moved nodes. Note: the computation is correct even + // when new_location_start < old_location_start given the C++ uint32_t + // integer arithmetic rules. + const uint32_t location_delta = new_location_start - old_location_start; + + // 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; + + // Move the adjacency lists. Make sure that overlapping ranges are handled + // correctly. + 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_points(old_location_start, new_location_start, num_locations); + memset(_data + _aligned_dim * mem_clear_loc_start, 0, + sizeof(data_t) * _aligned_dim * (mem_clear_loc_end_limit - mem_clear_loc_start)); +} + +template +void InMemDataStore::copy_points(const location_t from_loc, const location_t to_loc, + const location_t num_points) +{ + assert(from_loc < _capacity); + assert(to_loc < _capacity); + assert(num_points < _capacity); + memmove(_data + _aligned_dim * to_loc, _data + _aligned_dim * from_loc, num_points * _aligned_dim * sizeof(data_t)); +} + +template location_t InMemDataStore::calculate_medoid() const +{ + // allocate and init centroid + float *center = new float[_aligned_dim]; + for (size_t j = 0; j < _aligned_dim; j++) + center[j] = 0; + + for (size_t i = 0; i < _num_points; i++) + for (size_t j = 0; j < _aligned_dim; j++) + center[j] += (float)_data[i * _aligned_dim + j]; + + for (size_t j = 0; j < _aligned_dim; j++) + center[j] /= (float)_num_points; + + // compute all to one distance + float *distances = new float[_num_points]; + +//TODO: REFACTOR. Removing pragma might make this slow. Must revisit. +// Problem is that we need to pass num_threads here, it is not clear +// if data store must be aware of threads! +//#pragma omp parallel for schedule(static, 65536) + for (int64_t i = 0; i < (int64_t)_num_points; i++) + { + // extract point and distance reference + float &dist = distances[i]; + const data_t *cur_vec = _data + (i * (size_t)_aligned_dim); + dist = 0; + float diff = 0; + for (size_t j = 0; j < _aligned_dim; j++) + { + diff = (center[j] - (float)cur_vec[j]) * (center[j] - (float)cur_vec[j]); + dist += diff; + } + } + // find imin + uint32_t min_idx = 0; + float min_dist = distances[0]; + for (uint32_t i = 1; i < _num_points; i++) + { + if (distances[i] < min_dist) + { + min_idx = i; + min_dist = distances[i]; + } + } + + delete[] distances; + delete[] center; + return min_idx; +} + } // namespace diskann \ No newline at end of file diff --git a/src/index.cpp b/src/index.cpp index 814aed4c2..8387d5cc4 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -811,8 +811,12 @@ template int Index return -1; } - size_t location = _tag_to_location[tag]; - memcpy((void *)vec, (void *)(_data + location * _aligned_dim), (size_t)_dim * sizeof(T)); + location_t location = _tag_to_location[tag]; + _data_store->get_vector(location, vec); + + //REFACTOR + //size_t location = _tag_to_location[tag]; + //memcpy((void *)vec, (void *)(_data + location * _aligned_dim), (size_t)_dim * sizeof(T)); return 0; } @@ -825,49 +829,54 @@ template uint32_t Indexcalculate_medoid(); + +// //REFACTOR +// // allocate and init centroid +// float *center = new float[_aligned_dim](); +// for (size_t j = 0; j < _aligned_dim; j++) +// center[j] = 0; +// +// for (size_t i = 0; i < _nd; i++) +// for (size_t j = 0; j < _aligned_dim; j++) +// center[j] += (float)_data[i * _aligned_dim + j]; +// +// for (size_t j = 0; j < _aligned_dim; j++) +// center[j] /= (float)_nd; +// +// // compute all to one distance +// float *distances = new float[_nd](); +//#pragma omp parallel for schedule(static, 65536) +// for (int64_t i = 0; i < (int64_t)_nd; i++) +// { +// // extract point and distance reference +// float &dist = distances[i]; +// const T *cur_vec = _data + (i * (size_t)_aligned_dim); +// dist = 0; +// float diff = 0; +// for (size_t j = 0; j < _aligned_dim; j++) +// { +// diff = (center[j] - (float)cur_vec[j]) * (center[j] - (float)cur_vec[j]); +// dist += diff; +// } +// } +// // find imin +// uint32_t min_idx = 0; +// float min_dist = distances[0]; +// for (uint32_t i = 1; i < _nd; i++) +// { +// if (distances[i] < min_dist) +// { +// min_idx = i; +// min_dist = distances[i]; +// } +// } +// +// delete[] distances; +// delete[] center; +// return min_idx; } template std::vector Index::get_init_ids() @@ -1008,9 +1017,16 @@ std::pair Index::iterate_to_fixed_point( float distance; if (_pq_dist) + { pq_dist_lookup(pq_coord_scratch, 1, this->_num_pq_chunks, pq_dists, &distance); + } else - distance = _distance->compare(_data + _aligned_dim * (size_t)id, aligned_query, (uint32_t)_aligned_dim); + { + _data_store->get_distance(aligned_query, id); + // REFACTOR + // distance = _distance->compare(_data + _aligned_dim * (size_t)id, aligned_query, + // (uint32_t)_aligned_dim); + } Neighbor nn = Neighbor(id, distance); best_L_nodes.insert(nn); } @@ -1535,7 +1551,7 @@ void Index::prune_all_nbrs(const Parameters ¶meters) { if (dummy_visited.find(cur_nbr) == dummy_visited.end() && cur_nbr != node) { - float dist = _data_store->compare(node, cur_nbr); + float dist = _data_store->get_distance(node, cur_nbr); //REFACTOR //float dist = _distance->compare(_data + _aligned_dim * (size_t)node, // _data + _aligned_dim * (size_t)cur_nbr, (uint32_t)_aligned_dim); @@ -1635,7 +1651,8 @@ void Index::build_with_data_populated(Parameters ¶meters, c stream << "ERROR: Driver requests loading " << _nd << " points from file," << "but tags vector is of size " << tags.size() << "." << std::endl; diskann::cerr << stream.str() << std::endl; - aligned_free(_data); + //REFACTOR + //aligned_free(_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } if (_enable_tags) @@ -1654,7 +1671,10 @@ void Index::build_with_data_populated(Parameters ¶meters, c if (_query_scratch.size() == 0) { - initialize_query_scratch(5 + num_threads_index, index_L, index_L, index_R, maxc, _aligned_dim); + //REFACTOR + //initialize_query_scratch(5 + num_threads_index, index_L, index_L, index_R, maxc, _aligned_dim); + //Comment: Why 5+? + initialize_query_scratch(5 + num_threads_index, index_L, index_L, index_R, maxc, _data_store->get_aligned_dim()); } generate_frozen_point(); @@ -2271,7 +2291,9 @@ template void Indexcopy_points(res, _max_points, 1); + //REFACTOR + //memcpy(_data + _max_points * _aligned_dim, _data + res * _aligned_dim, _aligned_dim * sizeof(T)); } } @@ -2722,12 +2744,8 @@ void Index::reposition_points(uint32_t old_location_start, uint mem_clear_loc_end_limit = new_location_start; } } + _data_store->reposition_points(old_location_start, new_location_start, num_locations); - // Use memmove to handle overlapping ranges. - memmove(_data + _aligned_dim * new_location_start, _data + _aligned_dim * old_location_start, - sizeof(T) * _aligned_dim * num_locations); - memset(_data + _aligned_dim * mem_clear_loc_start, 0, - sizeof(T) * _aligned_dim * (mem_clear_loc_end_limit - mem_clear_loc_start)); } template void Index::reposition_frozen_point_to_end() @@ -2741,6 +2759,7 @@ template void Index Date: Wed, 5 Apr 2023 15:30:21 +0530 Subject: [PATCH 11/69] Refactor --- CMakeLists.txt | 2 +- include/abstract_data_store.h | 5 ++- include/in_mem_data_store.h | 3 ++ src/in_mem_data_store.cpp | 20 +++++++++ src/index.cpp | 77 +++++++++++++++++++++++------------ 5 files changed, 79 insertions(+), 28 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index dfa1cb77e..825cb6df5 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -244,7 +244,7 @@ endif() add_subdirectory(src) add_subdirectory(tests) add_subdirectory(tests/utils) -add_subdirectory(tests/unittests) + if (MSVC) message(STATUS "The ${PROJECT_NAME}.sln has been created, opened it from VisualStudio to build Release or Debug configurations.\n" diff --git a/include/abstract_data_store.h b/include/abstract_data_store.h index 380a1d77f..1fbd30097 100644 --- a/include/abstract_data_store.h +++ b/include/abstract_data_store.h @@ -27,14 +27,17 @@ template class AbstractDataStore virtual void get_vector(const location_t i, data_t* dest) const = 0; virtual void set_vector(const location_t i, const data_t *const vector) = 0; + virtual void prefetch_vector(const location_t loc) = 0; + virtual void resize(const location_t new_size) = 0; + virtual void reposition_points(const location_t start_loc, const location_t end_loc, const location_t num_points) = 0; virtual void copy_points(const location_t from_loc, const location_t to_loc, const location_t num_points) = 0; //Returns the point in the dataset that is closest to the mean of all points in the dataset virtual location_t calculate_medoid() const = 0; - + virtual location_t capacity() const { return _capacity; diff --git a/include/in_mem_data_store.h b/include/in_mem_data_store.h index 05cb5bdde..e9d5ed666 100644 --- a/include/in_mem_data_store.h +++ b/include/in_mem_data_store.h @@ -33,6 +33,9 @@ class InMemDataStore : public AbstractDataStore virtual void get_vector(const location_t i, data_t *target) const override; virtual void set_vector(const location_t i, const data_t *const vector) override; + virtual void prefetch_vector(const location_t loc) override; + + virtual void resize(const location_t new_size) override; virtual void reposition_points(const location_t old_location_start, const location_t new_location_start, const location_t num_points) override; diff --git a/src/in_mem_data_store.cpp b/src/in_mem_data_store.cpp index fcaaf3b96..c13265d1a 100644 --- a/src/in_mem_data_store.cpp +++ b/src/in_mem_data_store.cpp @@ -142,6 +142,11 @@ void InMemDataStore::set_vector(const location_t loc, const data_t *cons } } +template void InMemDataStore::prefetch_vector(const location_t loc) +{ + diskann::prefetch_vector((const char *)_data + _aligned_dim * (size_t)loc, sizeof(T) * _aligned_dim); +} + template float InMemDataStore::get_distance(const data_t *query, const location_t loc) const { return _distance_metric->compare(query, _data + _aligned_dim * loc, _aligned_dim); @@ -162,6 +167,21 @@ float InMemDataStore::get_distance(const location_t loc1, const location return _distance_metric->compare(_data + loc1 * _aligned_dim, _data + loc2 * _aligned_dim, this->_aligned_dim); } +template void InMemDataStore::resize(const location_t new_size) +{ +#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, _data, (_max_points + _num_frozen_pts) * _aligned_dim * sizeof(data_t)); + aligned_free(_data); + _data = new_data; +#else + realloc_aligned((void **)&_data, new_size * _aligned_dim * sizeof(data_t), 8 * sizeof(data_t)); +#endif + _capacity = new_size; + +} + template void InMemDataStore::reposition_points(const location_t old_location_start, const location_t new_location_start, const location_t num_locations) { diff --git a/src/index.cpp b/src/index.cpp index 8387d5cc4..f7af37b4e 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -1126,12 +1126,16 @@ std::pair Index::iterate_to_fixed_point( if (m + 1 < id_scratch.size()) { auto nextn = id_scratch[m + 1]; - diskann::prefetch_vector((const char *)_data + _aligned_dim * (size_t)nextn, - sizeof(T) * _aligned_dim); + _data_store->prefetch_vector(nextn); + //REFACTOR + //diskann::prefetch_vector((const char *)_data + _aligned_dim * (size_t)nextn, + // sizeof(T) * _aligned_dim); } - dist_scratch.push_back( - _distance->compare(aligned_query, _data + _aligned_dim * (size_t)id, (uint32_t)_aligned_dim)); + dist_scratch.push_back(_data_store->get_distance(aligned_query, id)); + //REFACTOR + //dist_scratch.push_back( + // _distance->compare(aligned_query, _data + _aligned_dim * (size_t)id, (uint32_t)_aligned_dim)); } } cmps += id_scratch.size(); @@ -1156,16 +1160,25 @@ void Index::search_for_point_and_prune(int location, uint32_t L if (!use_filter) { - iterate_to_fixed_point(_data + _aligned_dim * location, Lindex, init_ids, scratch, false, unused_filter_label, - false); + _data_store->get_vector(location, scratch->aligned_query()); + iterate_to_fixed_point(scratch->aligned_query(), Lindex, init_ids, scratch, false, unused_filter_label, false); + + //REFACTOR + //iterate_to_fixed_point(_data + _aligned_dim * location, Lindex, init_ids, scratch, false, unused_filter_label, + // false); } else { std::vector filter_specific_start_nodes; for (auto &x : _pts_to_labels[location]) filter_specific_start_nodes.emplace_back(_label_to_medoid_id[x]); - iterate_to_fixed_point(_data + _aligned_dim * location, filteredLindex, filter_specific_start_nodes, scratch, + + _data_store->get_vector(location, scratch->aligned_query()); + iterate_to_fixed_point(scratch->aligned_query(), filteredLindex, filter_specific_start_nodes, scratch, true, _pts_to_labels[location], false); + //REFACTOR + //iterate_to_fixed_point(_data + _aligned_dim * location, filteredLindex, filter_specific_start_nodes, scratch, + // true, _pts_to_labels[location], false); } auto &pool = scratch->pool(); @@ -1261,8 +1274,10 @@ void Index::occlude_list(const uint32_t location, std::vectorcompare(_data + _aligned_dim * (size_t)iter2->id, - _data + _aligned_dim * (size_t)iter->id, (uint32_t)_aligned_dim); + float djk = _data_store->get_distance(iter2->id, iter->id); + //REFACTOR + //float djk = _distance->compare(_data + _aligned_dim * (size_t)iter2->id, + // _data + _aligned_dim * (size_t)iter->id, (uint32_t)_aligned_dim); if (_dist_metric == diskann::Metric::L2 || _dist_metric == diskann::Metric::COSINE) { occlude_factor[t] = (djk == 0) ? std::numeric_limits::max() @@ -1309,8 +1324,10 @@ void Index::prune_neighbors(const uint32_t location, std::vecto if (_pq_dist) { for (auto &ngh : pool) - ngh.distance = _distance->compare(_data + _aligned_dim * (size_t)ngh.id, - _data + _aligned_dim * (size_t)location, (uint32_t)_aligned_dim); + ngh.distance = _data_store->get_distance(ngh.id, location); + //REFACTOR + //ngh.distance = _distance->compare(_data + _aligned_dim * (size_t)ngh.id, + // _data + _aligned_dim * (size_t)location, (uint32_t)_aligned_dim); } // sort the pool based on distance to query and prune it with occlude_list @@ -1381,8 +1398,10 @@ void Index::inter_insert(uint32_t n, std::vector &pru { if (dummy_visited.find(cur_nbr) == dummy_visited.end() && cur_nbr != des) { - float dist = _distance->compare(_data + _aligned_dim * (size_t)des, - _data + _aligned_dim * (size_t)cur_nbr, (uint32_t)_aligned_dim); + float dist = _data_store->get_distance(des, cur_nbr); + //REFACTOR + //float dist = _distance->compare(_data + _aligned_dim * (size_t)des, + // _data + _aligned_dim * (size_t)cur_nbr, (uint32_t)_aligned_dim); dummy_pool.emplace_back(Neighbor(cur_nbr, dist)); dummy_visited.insert(cur_nbr); } @@ -1505,8 +1524,10 @@ template void Indexcompare(_data + _aligned_dim * (size_t)node, - _data + _aligned_dim * (size_t)cur_nbr, (uint32_t)_aligned_dim); + float dist = _data_store->get_distance(node, cur_nbr); + //REFACTOR + //float dist = _distance->compare(_data + _aligned_dim * (size_t)node, + // _data + _aligned_dim * (size_t)cur_nbr, (uint32_t)_aligned_dim); dummy_pool.emplace_back(Neighbor(cur_nbr, dist)); dummy_visited.insert(cur_nbr); } @@ -2583,8 +2604,10 @@ template void Indexcopy_points(old, new_location[old], 1); + //REFACTOR + //memcpy((void *)(_data + _aligned_dim * (size_t)new_location[old]), + // (void *)(_data + _aligned_dim * (size_t)old), _aligned_dim * sizeof(T)); } } else @@ -2770,15 +2793,17 @@ template void Indexresize(new_internal_points); + //REFACTOR +//#ifndef _WINDOWS +// T *new_data; +// alloc_aligned((void **)&new_data, new_internal_points * _aligned_dim * sizeof(T), 8 * sizeof(T)); +// memcpy(new_data, _data, (_max_points + _num_frozen_pts) * _aligned_dim * sizeof(T)); +// aligned_free(_data); +// _data = new_data; +//#else +// realloc_aligned((void **)&_data, new_internal_points * _aligned_dim * sizeof(T), 8 * sizeof(T)); +//#endif _final_graph.resize(new_internal_points); _locks = std::vector(new_internal_points); From 42becfb2ac152221b3214c2627770b454fb620cb Mon Sep 17 00:00:00 2001 From: Gopal Srinivasa Date: Wed, 5 Apr 2023 23:55:41 +0530 Subject: [PATCH 12/69] Fixed most of the bugs related to _data --- include/in_mem_data_store.h | 1 + src/index.cpp | 24 +++++++++++++++++------- 2 files changed, 18 insertions(+), 7 deletions(-) diff --git a/include/in_mem_data_store.h b/include/in_mem_data_store.h index e9d5ed666..95a7069ae 100644 --- a/include/in_mem_data_store.h +++ b/include/in_mem_data_store.h @@ -51,6 +51,7 @@ class InMemDataStore : public AbstractDataStore return _aligned_dim; } + protected: virtual location_t load_data(const std::string &filename); #ifdef EXEC_ENV_OLS diff --git a/src/index.cpp b/src/index.cpp index f7af37b4e..eca426719 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -449,7 +449,8 @@ size_t Index::load_data(std::string filename) std::stringstream stream; stream << "ERROR: data file " << filename << " does not exist." << std::endl; diskann::cerr << stream.str() << std::endl; - aligned_free(_data); + //REFACTOR + //aligned_free(_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } diskann::get_bin_metadata(filename, file_num_points, file_dim); @@ -464,7 +465,8 @@ size_t Index::load_data(std::string filename) stream << "ERROR: Driver requests loading " << _dim << " dimension," << "but file has " << file_dim << " dimension." << std::endl; diskann::cerr << stream.str() << std::endl; - aligned_free(_data); + //REFACTOR + //aligned_free(_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } @@ -475,9 +477,13 @@ size_t Index::load_data(std::string filename) } #ifdef EXEC_ENV_OLS + + //REFACTOR TODO: Must figure out how to support aligned reader in a clean manner. copy_aligned_data_from_file(reader, _data, file_num_points, file_dim, _aligned_dim); #else - copy_aligned_data_from_file(filename.c_str(), _data, file_num_points, file_dim, _aligned_dim); + _data_store->populate_data(filename, 0); //offset == 0. + //REFACTOR + //copy_aligned_data_from_file(filename.c_str(), _data, file_num_points, file_dim, _aligned_dim); #endif return file_num_points; } @@ -720,7 +726,8 @@ size_t Index::load_graph(std::string filename, size_t expected_ << std::endl; } diskann::cerr << stream.str() << std::endl; - aligned_free(_data); + //REFACTOR + //aligned_free(_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } @@ -2398,9 +2405,12 @@ inline void Index::process_delete(const tsl::robin_setcompare(_data + _aligned_dim * loc, _data + _aligned_dim * ngh, (uint32_t)_aligned_dim)); + expanded_nghrs_vec.emplace_back(ngh, _data_store->get_distance(loc, ngh)); + + //REFACTOR + //expanded_nghrs_vec.emplace_back( + // ngh, + // _distance->compare(_data + _aligned_dim * loc, _data + _aligned_dim * ngh, (uint32_t)_aligned_dim)); } std::sort(expanded_nghrs_vec.begin(), expanded_nghrs_vec.end()); std::vector &occlude_list_output = scratch->occlude_list_output(); From f62259cf508135d8d476f5ebb33009844ccd1995 Mon Sep 17 00:00:00 2001 From: harsha vardhan simhadri Date: Fri, 31 Mar 2023 16:00:54 -0700 Subject: [PATCH 13/69] add seed files --- include/abstract_data_store.h | 58 ++++++++++++++++ include/abstract_graph_store.h | 26 +++++++ include/in_mem_data_store.h | 58 ++++++++++++++++ include/in_mem_graph_store.h | 23 +++++++ src/CMakeLists.txt | 1 + src/dll/CMakeLists.txt | 3 +- src/in_mem_data_store.cpp | 119 +++++++++++++++++++++++++++++++++ src/in_mem_graph_store.cpp | 28 ++++++++ 8 files changed, 315 insertions(+), 1 deletion(-) create mode 100644 include/abstract_data_store.h create mode 100644 include/abstract_graph_store.h create mode 100644 include/in_mem_data_store.h create mode 100644 include/in_mem_graph_store.h create mode 100644 src/in_mem_data_store.cpp create mode 100644 src/in_mem_graph_store.cpp diff --git a/include/abstract_data_store.h b/include/abstract_data_store.h new file mode 100644 index 000000000..82d7d646a --- /dev/null +++ b/include/abstract_data_store.h @@ -0,0 +1,58 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#pragma once + +#include + +#include "types.h" + +namespace diskann +{ + +template class AbstractDataStore +{ + public: + AbstractDataStore(const location_t max_pts, const location_t num_frozen_pts, const size_t dim) + : _max_pts(max_pts), _num_pts(0), _num_frozen_pts(num_frozen_pts), _dim(dim) + { + } + + // Return number of points returned + virtual size_t load(const std::string &filename) = 0; + virtual void store(const std::string &filename) = 0; + + virtual data_t *get_vector(location_t i) = 0; + virtual data_t *get_vector_by_UID(id_t uid) = 0; + + virtual void set_vector(const location_t i, const data_t *const vector) = 0; + + location_t get_max_pts() + { + return _max_pts; + } + + location_t get_num_pts() + { + return _num_pts; + } + + location_t get_num_frozen_pts() + { + return _num_frozen_pts; + } + + size_t get_dims() + { + return _dim; + } + + protected: + location_t _max_pts; + location_t _num_pts; + location_t _num_frozen_pts; + + const size_t _dim; +}; + +} // namespace diskann \ No newline at end of file diff --git a/include/abstract_graph_store.h b/include/abstract_graph_store.h new file mode 100644 index 000000000..fe088d51b --- /dev/null +++ b/include/abstract_graph_store.h @@ -0,0 +1,26 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#pragma once + +#include +#include + +#include "types.h" + +namespace diskann +{ + +class AbstractGraphStore +{ + public: + AbstractGraphStore(const size_t max_pts); + + virtual int load(const std::string &index_path_prefix) = 0; + virtual int store(const std::string &index_path_prefix) = 0; + + virtual void get_adj_list(const location_t i, std::vector &neighbors) = 0; + virtual void set_adj_list(const location_t i, std::vector &neighbors) = 0; +}; + +} // namespace diskann \ No newline at end of file diff --git a/include/in_mem_data_store.h b/include/in_mem_data_store.h new file mode 100644 index 000000000..d948f74ee --- /dev/null +++ b/include/in_mem_data_store.h @@ -0,0 +1,58 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#include + +#include "tsl/robin_map.h" +#include "tsl/robin_set.h" +#include "tsl/sparse_map.h" +#include "boost/dynamic_bitset.hpp" + +#include "abstract_data_store.h" + +#include "natural_number_map.h" +#include "natural_number_set.h" + +namespace diskann +{ +template class InMemDataStore : public AbstractDataStore +{ + public: + InMemDataStore(const location_t max_pts, const location_t num_frozen_pts, const size_t dim); + ~InMemDataStore(); + + void load(const std::string &filename); + void store(const std::string &filename); + + data_t *get_vector(location_t i); + data_t *get_vector_by_UID(id_t uid); + + void set_vector(const location_t i, const data_t *const vector); + + protected: + location_t load_data(const std::string &filename); +#ifdef EXEC_ENV_OLS + location_t load_data(AlignedFileReader &reader); +#endif + + private: + data_t *_data = nullptr; + + const size_t _aligned_dim; + + // lazy_delete removes entry from _location_to_tag and _tag_to_location. If + // _location_to_tag does not resolve a location, infer that it was deleted. + tsl::sparse_map _tag_to_location; + natural_number_map _location_to_tag; + + // _empty_slots has unallocated slots and those freed by consolidate_delete. + // _delete_set has locations marked deleted by lazy_delete. Will not be + // immediately available for insert. consolidate_delete will release these + // slots to _empty_slots. + natural_number_set _empty_slots; + std::unique_ptr> _delete_set; + + std::shared_timed_mutex _lock; // Takes please of Index::_tag_lock +}; + +} // namespace diskann \ No newline at end of file diff --git a/include/in_mem_graph_store.h b/include/in_mem_graph_store.h new file mode 100644 index 000000000..31b502fd7 --- /dev/null +++ b/include/in_mem_graph_store.h @@ -0,0 +1,23 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#pragma once + +#include "abstract_graph_store.h" + +namespace diskann +{ + +class InMemGraphStore : public AbstractGraphStore +{ + public: + InMemGraphStore(const size_t max_pts); + + int load(const std::string &index_path_prefix); + int store(const std::string &index_path_prefix); + + void get_adj_list(const location_t i, std::vector &neighbors); + void set_adj_list(const location_t i, std::vector &neighbors); +}; + +} // namespace diskann \ No newline at end of file diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index f8dda23c7..b4a6a4e89 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -9,6 +9,7 @@ else() #file(GLOB CPP_SOURCES *.cpp) set(CPP_SOURCES ann_exception.cpp disk_utils.cpp distance.cpp index.cpp 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) if (RESTAPI) diff --git a/src/dll/CMakeLists.txt b/src/dll/CMakeLists.txt index 1423d6c2c..0f511a8ae 100644 --- a/src/dll/CMakeLists.txt +++ b/src/dll/CMakeLists.txt @@ -2,7 +2,8 @@ #Licensed under the MIT license. add_library(${PROJECT_NAME} SHARED dllmain.cpp ../partition.cpp ../pq.cpp ../pq_flash_index.cpp ../logger.cpp ../utils.cpp - ../windows_aligned_file_reader.cpp ../distance.cpp ../memory_mapper.cpp ../index.cpp ../math_utils.cpp ../disk_utils.cpp ../filter_utils.cpp + ../windows_aligned_file_reader.cpp ../distance.cpp ../memory_mapper.cpp ../index.cpp + ../in_mem_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) set(TARGET_DIR "$<$:${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}>$<$:${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}>") diff --git a/src/in_mem_data_store.cpp b/src/in_mem_data_store.cpp new file mode 100644 index 000000000..4e39297a9 --- /dev/null +++ b/src/in_mem_data_store.cpp @@ -0,0 +1,119 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#include "in_mem_data_store.h" + +#include "utils.h" + +namespace diskann +{ + +template +InMemDataStore::InMemDataStore(const location_t max_pts, const location_t num_frozen_pts, + const size_t dim) + : AbstractDataStore(max_pts, num_frozen_pts, dim), _aligned_dim(ROUND_UP(dim, 8)) +{ + location_t total_internal_points = max_pts + num_frozen_pts; + alloc_aligned(((void **)&_data), total_internal_points * _aligned_dim * sizeof(data_t), 8 * sizeof(data_t)); + std::memset(_data, 0, total_internal_points * _aligned_dim * sizeof(data_t)); +} + +template InMemDataStore::~InMemDataStore() +{ + if (_data != nullptr) + { + aligned_free(this->_data); + } +} + +template void InMemDataStore::load(const std::string &filename) +{ +} + +template void InMemDataStore::store(const std::string &filename) +{ +} + +#ifdef EXEC_ENV_OLS +template location_t Index::load_data(AlignedFileReader &reader) +{ + size_t file_dim, file_num_points; + + diskann::get_bin_metadata(reader, file_num_points, file_dim); + + // since we are loading a new dataset, _empty_slots must be cleared + _empty_slots.clear(); + + if (file_dim != _dim) + { + std::stringstream stream; + stream << "ERROR: Driver requests loading " << _dim << " dimension," + << "but file has " << file_dim << " dimension." << std::endl; + diskann::cerr << stream.str() << std::endl; + aligned_free(_data); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } + + if (file_num_points > _max_points + _num_frozen_pts) + { + resize(file_num_points - _num_frozen_pts); + } + + return file_num_points; +} +#endif + +template +location_t InMemDataStore::load_data(const std::string &filename) +{ + size_t file_dim, file_num_points; + if (!file_exists(filename)) + { + std::stringstream stream; + stream << "ERROR: data file " << filename << " does not exist." << std::endl; + diskann::cerr << stream.str() << std::endl; + aligned_free(_data); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } + diskann::get_bin_metadata(filename, file_num_points, file_dim); + + // since we are loading a new dataset, _empty_slots must be cleared + _empty_slots.clear(); + + if (file_dim != this->_dim) + { + std::stringstream stream; + stream << "ERROR: Driver requests loading " << this->_dim << " dimension," + << "but file has " << file_dim << " dimension." << std::endl; + diskann::cerr << stream.str() << std::endl; + aligned_free(_data); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } + + if (file_num_points > this->get_max_points() + this->get_num_frozen_pts()) + { + resize(file_num_points - this->get_num_frozen_pts()); + } + + copy_aligned_data_from_file(filename.c_str(), _data, file_num_points, file_dim, _aligned_dim); + + return file_num_points; +} + +template data_t *InMemDataStore::get_vector(location_t i) +{ + return _data + i * _aligned_dim; +} + +template data_t *InMemDataStore::get_vector_by_UID(id_t uid) +{ + return get_vector(_tag_to_location(uid)); +} + +template +void InMemDataStore::set_vector(const location_t loc, const data_t *const vector) +{ + memcpy(_data + loc * _aligned_dim, vector, this->_dim * sizeof(data_t)); +} + +} // namespace diskann \ No newline at end of file diff --git a/src/in_mem_graph_store.cpp b/src/in_mem_graph_store.cpp new file mode 100644 index 000000000..df9933080 --- /dev/null +++ b/src/in_mem_graph_store.cpp @@ -0,0 +1,28 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#include "in_mem_graph_store.h" + +namespace diskann +{ + +InMemGraphStore::InMemGraphStore(const size_t max_pts) : AbstractGraphStore(max_pts) +{ +} + +int InMemGraphStore::load(const std::string &index_path_prefix) +{ +} +int InMemGraphStore::store(const std::string &index_path_prefix) +{ +} + +void InMemGraphStore::get_adj_list(const location_t i, std::vector &neighbors) +{ +} + +void InMemGraphStore::set_adj_list(const location_t i, std::vector &neighbors) +{ +} + +} // namespace diskann \ No newline at end of file From d18804c16e38e4435b3389cb7a6e032b2b998089 Mon Sep 17 00:00:00 2001 From: harsha vardhan simhadri Date: Tue, 28 Mar 2023 01:19:47 -0700 Subject: [PATCH 14/69] gi# This is a combination of 2 commits. remove _u, _s typedefs --- include/utils.h | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/include/utils.h b/include/utils.h index f34cb6989..018752792 100644 --- a/include/utils.h +++ b/include/utils.h @@ -878,13 +878,13 @@ template void save_Tvecs(const char *filename, T *data, size_t npts // create cached ofstream with 64MB cache cached_ofstream writer(fname, 64 * 1048576); - unsigned dims_u32 = (unsigned)ndims; + unsigned dimsuint32_t = (unsigned)ndims; // start writing for (size_t i = 0; i < npts; i++) { // write dims in u32 - writer.write((char *)&dims_u32, sizeof(unsigned)); + writer.write((char *)&dimsuint32_t, sizeof(unsigned)); // get cur point in data T *cur_pt = data + i * ndims; From 4de26eb6b6dc7978a5e38642d474aa0725a5551b Mon Sep 17 00:00:00 2001 From: harsha vardhan simhadri Date: Tue, 28 Mar 2023 01:45:55 -0700 Subject: [PATCH 15/69] added some seed files --- src/CMakeLists.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index b4a6a4e89..bac0ee6f7 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -8,6 +8,7 @@ if(MSVC) else() #file(GLOB CPP_SOURCES *.cpp) set(CPP_SOURCES ann_exception.cpp disk_utils.cpp distance.cpp index.cpp + in_mem_graph_store.cpp in_mem_data_store.cpp 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 From a1c07a17381967508dc7ef69667f215602083986 Mon Sep 17 00:00:00 2001 From: Gopal Srinivasa Date: Sun, 2 Apr 2023 14:08:48 +0530 Subject: [PATCH 16/69] New distance metric hierarchy --- include/distance.h | 63 +++++++++++++++++++++++++++++++++++++++++++++- src/distance.cpp | 28 +++++++++++++++++++++ 2 files changed, 90 insertions(+), 1 deletion(-) diff --git a/include/distance.h b/include/distance.h index e04be7ee2..39db1a0a7 100644 --- a/include/distance.h +++ b/include/distance.h @@ -1,6 +1,7 @@ #pragma once #include "windows_customizations.h" + namespace diskann { enum Metric @@ -14,8 +15,58 @@ enum Metric template class Distance { public: + // distance comparison function virtual float compare(const T *a, const T *b, uint32_t length) const = 0; - virtual ~Distance() + + // For MIPS, normalization => a new dimension gets added to the vectors. + // This function lets callers know if the normalization process + // changes the dimension. + virtual uint32_t new_dimension(uint32_t orig_dimension) const + { + return orig_dimension; + } + + // This is for efficiency. If no normalization is required, the callers + //can simply ignore the normalize_data_for_build() function. + virtual bool normalization_required() const + { + return false; + } + + //Check the normalization_required() function before calling this. + //Clients can call the function like this: + // + // if (metric->normalization_required()){ + // T* normalized_data_batch; + // if ( metric->new_dimension() == orig_dim ) { + // normalized_data_batch = nullptr; + // modify_data = true; + // } else { + // normalized_data_batch = new T[batch_size * metric->new_dimension()]; + // modify_data = false; + // } + // Split data into batches of batch_size and for each, call: + // metric->normalize_data_for_build(data_batch, batch_size, orig_dim + // normalized_data_batch, modify_data); + //The default implementation is correct but very inefficient! + virtual void normalize_data_for_build(const T *original_data, + const uint32_t num_points, const uint32_t orig_dim, + T *normalized_data, bool modify_data = true) + { + } + + //Invokes normalization for a single vector during search. The scratch space + //has to be created by the caller keeping track of the fact that normalization + //might change the dimension of the query vector. + virtual void normalize_vector_for_search(const T *query_vec, const uint32_t query_dim, + T* scratch_query) + { + memcpy(scratch_query, query_vec, query_dim * sizeof(T)); + } + + //Providing a default implementation for the virtual destructor because we don't + //expect most metric implementations to need it. + virtual ~Distance() { } }; @@ -137,6 +188,16 @@ class AVXNormalizedCosineDistanceFloat : public Distance // This will ensure that cosine is between -1 and 1. return 1.0f + _innerProduct.compare(a, b, length); } + DISKANN_DLLEXPORT virtual uint32_t new_dimension(uint32_t orig_dimension); + + DISKANN_DLLEXPORT virtual bool normalization_required() const; + + DISKANN_DLLEXPORT virtual void normalize_data_for_build(const float *original_data, const uint32_t num_points, + const uint32_t orig_dim, float *normalized_data, + bool modify_orig); + + DISKANN_DLLEXPORT virtual void normalize_vector_for_search(const float *query_vec, const uint32_t query_dim, + float* scratch_query_vector); }; template Distance *get_distance_function(Metric m); diff --git a/src/distance.cpp b/src/distance.cpp index 99676bb6f..a988b0a56 100644 --- a/src/distance.cpp +++ b/src/distance.cpp @@ -522,6 +522,34 @@ float AVXDistanceInnerProductFloat::compare(const float *a, const float *b, uint return -result; } +uint32_t AVXNormalizedCosineDistanceFloat::new_dimension(uint32_t orig_dimension) +{ + return orig_dimension; +} +bool AVXNormalizedCosineDistanceFloat::normalization_required() const +{ + return true; +} +void AVXNormalizedCosineDistanceFloat::normalize_data_for_build(const float *original_data, const uint32_t num_points, + const uint32_t orig_dim, float *normalized_data, + bool modify_orig) +{ + if (!modify_orig) + { + throw diskann::ANNException("For cosine distance, modify_orig should be set to true for efficiency.", -1); + } + for (auto i = 0; i < num_points; i++) + { + normalize(original_data + i * orig_dim, orig_dim); + } +} +void AVXNormalizedCosineDistanceFloat::normalize_vector_for_search(const float *query_vec, const uint32_t query_dim, + float *scratch_space) +{ + normalize(query_vec, query_dim); +} + + // Get the right distance function for the given metric. template <> diskann::Distance *get_distance_function(diskann::Metric m) { From e9c6697a18fe772d8c6011737e3d38cf0c14e22f Mon Sep 17 00:00:00 2001 From: Gopal Srinivasa Date: Mon, 3 Apr 2023 14:51:34 +0530 Subject: [PATCH 17/69] Refactoring changes --- include/abstract_data_store.h | 16 +++++-------- include/distance.h | 40 +++++++++++++++++++++------------ include/in_mem_data_store.h | 21 +++++++++++++----- src/in_mem_data_store.cpp | 42 ++++++++++++++++------------------- 4 files changed, 66 insertions(+), 53 deletions(-) diff --git a/include/abstract_data_store.h b/include/abstract_data_store.h index 82d7d646a..9ee4e8108 100644 --- a/include/abstract_data_store.h +++ b/include/abstract_data_store.h @@ -10,11 +10,11 @@ namespace diskann { -template class AbstractDataStore +template class AbstractDataStore { public: - AbstractDataStore(const location_t max_pts, const location_t num_frozen_pts, const size_t dim) - : _max_pts(max_pts), _num_pts(0), _num_frozen_pts(num_frozen_pts), _dim(dim) + AbstractDataStore(const location_t max_pts, const size_t dim) + : _max_pts(max_pts), _num_pts(0), _dim(dim) { } @@ -23,8 +23,6 @@ template class AbstractDataStore virtual void store(const std::string &filename) = 0; virtual data_t *get_vector(location_t i) = 0; - virtual data_t *get_vector_by_UID(id_t uid) = 0; - virtual void set_vector(const location_t i, const data_t *const vector) = 0; location_t get_max_pts() @@ -37,10 +35,9 @@ template class AbstractDataStore return _num_pts; } - location_t get_num_frozen_pts() - { - return _num_frozen_pts; - } + virtual void get_distance(const T *query, const location_t *locations, const uint32_t location_count, + float *distances) = 0; + size_t get_dims() { @@ -50,7 +47,6 @@ template class AbstractDataStore protected: location_t _max_pts; location_t _num_pts; - location_t _num_frozen_pts; const size_t _dim; }; diff --git a/include/distance.h b/include/distance.h index 39db1a0a7..ddf49e42b 100644 --- a/include/distance.h +++ b/include/distance.h @@ -15,17 +15,32 @@ enum Metric template class Distance { public: + Distance(diskann::Metric dist_metric) : _distance_metric(dist_metric) + { + } + // distance comparison function virtual float compare(const T *a, const T *b, uint32_t length) const = 0; + //Needed only for COSINE-BYTE and INNER_PRODUCT-BYTE + virtual float compare(const T *a, const T *b, const float normA, const float normB, uint32_t length) const + { + return std::numeric_limits::max(); + } + // For MIPS, normalization => a new dimension gets added to the vectors. // This function lets callers know if the normalization process // changes the dimension. - virtual uint32_t new_dimension(uint32_t orig_dimension) const + virtual uint32_t post_processed_dimension(uint32_t orig_dimension) const { return orig_dimension; } + virtual diskann::Metric get_metric() const + { + return _distance_metric; + } + // This is for efficiency. If no normalization is required, the callers //can simply ignore the normalize_data_for_build() function. virtual bool normalization_required() const @@ -38,20 +53,14 @@ template class Distance // // if (metric->normalization_required()){ // T* normalized_data_batch; - // if ( metric->new_dimension() == orig_dim ) { - // normalized_data_batch = nullptr; - // modify_data = true; - // } else { - // normalized_data_batch = new T[batch_size * metric->new_dimension()]; - // modify_data = false; - // } // Split data into batches of batch_size and for each, call: - // metric->normalize_data_for_build(data_batch, batch_size, orig_dim - // normalized_data_batch, modify_data); - //The default implementation is correct but very inefficient! - virtual void normalize_data_for_build(const T *original_data, - const uint32_t num_points, const uint32_t orig_dim, - T *normalized_data, bool modify_data = true) + // metric->normalize_data_for_build(data_batch, batch_size); + // + // TODO: This does not take into account the case for SSD inner product + // where the dimensions change after normalization. + // + virtual void normalize_data_for_build(T *original_data, + const uint32_t num_points) { } @@ -69,6 +78,9 @@ template class Distance virtual ~Distance() { } + + protected: + diskann::Metric _distance_metric; }; class DistanceCosineInt8 : public Distance diff --git a/include/in_mem_data_store.h b/include/in_mem_data_store.h index d948f74ee..815a3c8d3 100644 --- a/include/in_mem_data_store.h +++ b/include/in_mem_data_store.h @@ -15,19 +15,20 @@ namespace diskann { -template class InMemDataStore : public AbstractDataStore +template class InMemDataStore : public AbstractDataStore { public: - InMemDataStore(const location_t max_pts, const location_t num_frozen_pts, const size_t dim); - ~InMemDataStore(); + InMemDataStore(const location_t max_pts, const size_t dim, std::shared_ptr> distance_metric); + virtual ~InMemDataStore(); void load(const std::string &filename); void store(const std::string &filename); - data_t *get_vector(location_t i); - data_t *get_vector_by_UID(id_t uid); + virtual data_t *get_vector(location_t i); + virtual void set_vector(const location_t i, const data_t *const vector); - void set_vector(const location_t i, const data_t *const vector); + virtual void get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, + const shared_ptr> &metric, float *distances); protected: location_t load_data(const std::string &filename); @@ -53,6 +54,14 @@ template class InMemDataStore : public Abstract std::unique_ptr> _delete_set; std::shared_timed_mutex _lock; // Takes please of Index::_tag_lock + + // It may seem weird to put distance metric along with the data store class, but + // this gives us perf benefits as the datastore can do distance computations during + // search and compute norms of vectors internally without have to copy + // data back and forth. + shared_ptr> _distance_fn; + + shared_ptr _pre_computed_norms; }; } // namespace diskann \ No newline at end of file diff --git a/src/in_mem_data_store.cpp b/src/in_mem_data_store.cpp index 4e39297a9..aa90432bc 100644 --- a/src/in_mem_data_store.cpp +++ b/src/in_mem_data_store.cpp @@ -8,17 +8,17 @@ namespace diskann { -template -InMemDataStore::InMemDataStore(const location_t max_pts, const location_t num_frozen_pts, - const size_t dim) - : AbstractDataStore(max_pts, num_frozen_pts, dim), _aligned_dim(ROUND_UP(dim, 8)) +template +InMemDataStore::InMemDataStore(const location_t max_pts, + const size_t dim, std::shared_ptr> distance_metric) + : AbstractDataStore(max_pts, dim), _aligned_dim(ROUND_UP(dim, 8)), + _distance_metric(distance_metric) { - location_t total_internal_points = max_pts + num_frozen_pts; - alloc_aligned(((void **)&_data), total_internal_points * _aligned_dim * sizeof(data_t), 8 * sizeof(data_t)); - std::memset(_data, 0, total_internal_points * _aligned_dim * sizeof(data_t)); + alloc_aligned(((void **)&_data), max_pts * _aligned_dim * sizeof(data_t), 8 * sizeof(data_t)); + std::memset(_data, 0, max_pts * _aligned_dim * sizeof(data_t)); } -template InMemDataStore::~InMemDataStore() +template InMemDataStore::~InMemDataStore() { if (_data != nullptr) { @@ -26,16 +26,17 @@ template InMemDataStore::~InMemDa } } -template void InMemDataStore::load(const std::string &filename) +template void InMemDataStore::load(const std::string &filename) { + load_data(filename); } -template void InMemDataStore::store(const std::string &filename) +template void InMemDataStore::store(const std::string &filename) { } #ifdef EXEC_ENV_OLS -template location_t Index::load_data(AlignedFileReader &reader) +template location_t Index::load_data(AlignedFileReader &reader) { size_t file_dim, file_num_points; @@ -63,8 +64,8 @@ template location_t Index } #endif -template -location_t InMemDataStore::load_data(const std::string &filename) +template +location_t InMemDataStore::load_data(const std::string &filename) { size_t file_dim, file_num_points; if (!file_exists(filename)) @@ -90,9 +91,9 @@ location_t InMemDataStore::load_data(const std::string &filename) throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } - if (file_num_points > this->get_max_points() + this->get_num_frozen_pts()) + if (file_num_points > this->get_max_points()) { - resize(file_num_points - this->get_num_frozen_pts()); + resize(file_num_points); } copy_aligned_data_from_file(filename.c_str(), _data, file_num_points, file_dim, _aligned_dim); @@ -100,18 +101,13 @@ location_t InMemDataStore::load_data(const std::string &filename) return file_num_points; } -template data_t *InMemDataStore::get_vector(location_t i) +template data_t *InMemDataStore::get_vector(location_t i) { return _data + i * _aligned_dim; } -template data_t *InMemDataStore::get_vector_by_UID(id_t uid) -{ - return get_vector(_tag_to_location(uid)); -} - -template -void InMemDataStore::set_vector(const location_t loc, const data_t *const vector) +template +void InMemDataStore::set_vector(const location_t loc, const data_t *const vector) { memcpy(_data + loc * _aligned_dim, vector, this->_dim * sizeof(data_t)); } From 7684cf58ede37df016809de386aae84853719062 Mon Sep 17 00:00:00 2001 From: Gopal Srinivasa Date: Mon, 3 Apr 2023 18:38:33 +0530 Subject: [PATCH 18/69] Fixing compile errors in refactored code --- CMakeLists.txt | 1 + include/abstract_data_store.h | 2 +- include/distance.h | 50 +++++++++++++++++++++++++++++++++++ include/in_mem_data_store.h | 16 ++++++----- src/distance.cpp | 21 +++++++++++++-- 5 files changed, 80 insertions(+), 10 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index f29e13f3a..dfa1cb77e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -244,6 +244,7 @@ endif() add_subdirectory(src) add_subdirectory(tests) add_subdirectory(tests/utils) +add_subdirectory(tests/unittests) if (MSVC) message(STATUS "The ${PROJECT_NAME}.sln has been created, opened it from VisualStudio to build Release or Debug configurations.\n" diff --git a/include/abstract_data_store.h b/include/abstract_data_store.h index 9ee4e8108..9f43e176b 100644 --- a/include/abstract_data_store.h +++ b/include/abstract_data_store.h @@ -35,7 +35,7 @@ template class AbstractDataStore return _num_pts; } - virtual void get_distance(const T *query, const location_t *locations, const uint32_t location_count, + virtual void get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, float *distances) = 0; diff --git a/include/distance.h b/include/distance.h index ddf49e42b..f4f4b7794 100644 --- a/include/distance.h +++ b/include/distance.h @@ -86,12 +86,18 @@ template class Distance class DistanceCosineInt8 : public Distance { public: + DistanceCosineInt8() : Distance(diskann::Metric::COSINE) + { + } DISKANN_DLLEXPORT virtual float compare(const int8_t *a, const int8_t *b, uint32_t length) const; }; class DistanceL2Int8 : public Distance { public: + DistanceL2Int8() : Distance(diskann::Metric::L2) + { + } DISKANN_DLLEXPORT virtual float compare(const int8_t *a, const int8_t *b, uint32_t size) const; }; @@ -99,18 +105,28 @@ class DistanceL2Int8 : public Distance class AVXDistanceL2Int8 : public Distance { public: + AVXDistanceL2Int8() : Distance(diskann::Metric::L2) + { + } DISKANN_DLLEXPORT virtual float compare(const int8_t *a, const int8_t *b, uint32_t length) const; }; class DistanceCosineFloat : public Distance { public: + DistanceCosineFloat() : Distance(diskann::Metric::COSINE) + { + } DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t length) const; }; class DistanceL2Float : public Distance { public: + DistanceL2Float() : Distance(diskann::Metric::L2) + { + } + #ifdef _WINDOWS DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t size) const; #else @@ -121,24 +137,36 @@ class DistanceL2Float : public Distance class AVXDistanceL2Float : public Distance { public: + AVXDistanceL2Float() : Distance(diskann::Metric::L2) + { + } DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t length) const; }; class SlowDistanceL2Float : public Distance { public: + SlowDistanceL2Float() : Distance(diskann::Metric::L2) + { + } DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t length) const; }; class SlowDistanceCosineUInt8 : public Distance { public: + SlowDistanceCosineUInt8() : Distance(diskann::Metric::COSINE) + { + } DISKANN_DLLEXPORT virtual float compare(const uint8_t *a, const uint8_t *b, uint32_t length) const; }; class DistanceL2UInt8 : public Distance { public: + DistanceL2UInt8() : Distance(diskann::Metric::L2) + { + } DISKANN_DLLEXPORT virtual float compare(const uint8_t *a, const uint8_t *b, uint32_t size) const; }; @@ -146,6 +174,9 @@ class DistanceL2UInt8 : public Distance template class SlowDistanceL2Int : public Distance { public: + SlowDistanceL2Int() : Distance(diskann::Metric::L2) + { + } // Implementing here because this is a template function DISKANN_DLLEXPORT virtual float compare(const T *a, const T *b, uint32_t length) const { @@ -161,6 +192,13 @@ template class SlowDistanceL2Int : public Distance template class DistanceInnerProduct : public Distance { public: + DistanceInnerProduct() : Distance(diskann::Metric::INNER_PRODUCT) + { + } + + DistanceInnerProduct(diskann::Metric metric) : Distance(metric) + { + } inline float inner_product(const T *a, const T *b, unsigned size) const; inline float compare(const T *a, const T *b, unsigned size) const @@ -178,6 +216,9 @@ template class DistanceFastL2 : public DistanceInnerProduct // currently defined only for float. // templated for future use. public: + DistanceFastL2() : DistanceInnerProduct(diskann::Metric::FAST_L2) + { + } float norm(const T *a, unsigned size) const; float compare(const T *a, const T *b, float norm, unsigned size) const; }; @@ -185,6 +226,9 @@ template class DistanceFastL2 : public DistanceInnerProduct class AVXDistanceInnerProductFloat : public Distance { public: + AVXDistanceInnerProductFloat() : Distance(diskann::Metric::INNER_PRODUCT) + { + } DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t length) const; }; @@ -193,7 +237,13 @@ class AVXNormalizedCosineDistanceFloat : public Distance private: AVXDistanceInnerProductFloat _innerProduct; + protected: + void normalize_and_copy(const float *a, uint32_t length, float *a_norm ) const; + public: + AVXNormalizedCosineDistanceFloat() : Distance(diskann::Metric::COSINE) + { + } DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t length) const { // Inner product returns negative values to indicate distance. diff --git a/include/in_mem_data_store.h b/include/in_mem_data_store.h index 815a3c8d3..981b4b095 100644 --- a/include/in_mem_data_store.h +++ b/include/in_mem_data_store.h @@ -2,6 +2,7 @@ // Licensed under the MIT license. #include +#include #include "tsl/robin_map.h" #include "tsl/robin_set.h" @@ -10,12 +11,14 @@ #include "abstract_data_store.h" +#include "distance.h" #include "natural_number_map.h" #include "natural_number_set.h" namespace diskann { -template class InMemDataStore : public AbstractDataStore +template +class InMemDataStore : public AbstractDataStore { public: InMemDataStore(const location_t max_pts, const size_t dim, std::shared_ptr> distance_metric); @@ -27,8 +30,7 @@ template class InMemDataStore : public AbstractDataStore> &metric, float *distances); + virtual void get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, float *distances); protected: location_t load_data(const std::string &filename); @@ -43,8 +45,8 @@ template class InMemDataStore : public AbstractDataStore _tag_to_location; - natural_number_map _location_to_tag; + //tsl::sparse_map _tag_to_location; + //natural_number_map _location_to_tag; // _empty_slots has unallocated slots and those freed by consolidate_delete. // _delete_set has locations marked deleted by lazy_delete. Will not be @@ -59,9 +61,9 @@ template class InMemDataStore : public AbstractDataStore> _distance_fn; + std::shared_ptr> _distance_fn; - shared_ptr _pre_computed_norms; + std::shared_ptr _pre_computed_norms; }; } // namespace diskann \ No newline at end of file diff --git a/src/distance.cpp b/src/distance.cpp index a988b0a56..0431ab443 100644 --- a/src/distance.cpp +++ b/src/distance.cpp @@ -543,10 +543,27 @@ void AVXNormalizedCosineDistanceFloat::normalize_data_for_build(const float *ori normalize(original_data + i * orig_dim, orig_dim); } } + void AVXNormalizedCosineDistanceFloat::normalize_vector_for_search(const float *query_vec, const uint32_t query_dim, - float *scratch_space) + float *query_scratch) +{ + normalize_and_copy(query_vec, query_dim, query_scratch); +} + +void AVXNormalizedCosineDistanceFloat::normalize_and_copy(const float *query_vec, const uint32_t query_dim, + float *query_target) const { - normalize(query_vec, query_dim); + float norm = 0.0f; + for (auto i = 0; i < query_dim; i++) + { + norm += query_vec[i]; + } + norm /= norm / query_dim; + + for (auto i = 0; i < query_dim; i++) + { + query_target[i] = query_vec[i] / norm; + } } From 713e4ff91638eb7176f091fcfebe67628a0b9ec7 Mon Sep 17 00:00:00 2001 From: Gopal Srinivasa Date: Mon, 3 Apr 2023 23:04:07 +0530 Subject: [PATCH 19/69] Fixing compile errors --- include/distance.h | 10 ++++------ include/index.h | 1 + include/utils.h | 13 +++++++++---- src/distance.cpp | 22 ++++++---------------- 4 files changed, 20 insertions(+), 26 deletions(-) diff --git a/include/distance.h b/include/distance.h index f4f4b7794..e4ab0e51b 100644 --- a/include/distance.h +++ b/include/distance.h @@ -59,7 +59,7 @@ template class Distance // TODO: This does not take into account the case for SSD inner product // where the dimensions change after normalization. // - virtual void normalize_data_for_build(T *original_data, + virtual void normalize_data_for_build(T *original_data, const uint32_t orig_dim, const uint32_t num_points) { } @@ -250,16 +250,14 @@ class AVXNormalizedCosineDistanceFloat : public Distance // This will ensure that cosine is between -1 and 1. return 1.0f + _innerProduct.compare(a, b, length); } - DISKANN_DLLEXPORT virtual uint32_t new_dimension(uint32_t orig_dimension); + DISKANN_DLLEXPORT virtual uint32_t post_processed_dimension(uint32_t orig_dimension) const override; DISKANN_DLLEXPORT virtual bool normalization_required() const; - DISKANN_DLLEXPORT virtual void normalize_data_for_build(const float *original_data, const uint32_t num_points, - const uint32_t orig_dim, float *normalized_data, - bool modify_orig); + DISKANN_DLLEXPORT virtual void normalize_data_for_build(float *original_data, const uint32_t orig_dim, const uint32_t num_points) override; DISKANN_DLLEXPORT virtual void normalize_vector_for_search(const float *query_vec, const uint32_t query_dim, - float* scratch_query_vector); + float *scratch_query_vector) override; }; template Distance *get_distance_function(Metric m); diff --git a/include/index.h b/include/index.h index b49059192..6fb00b382 100644 --- a/include/index.h +++ b/include/index.h @@ -254,6 +254,7 @@ template clas const tsl::robin_set *const delete_set_ptr = nullptr); // add reverse links from all the visited nodes to node n. + void inter_insert(uint32_t n, std::vector &pruned_list, const uint32_t range, InMemQueryScratch *scratch); diff --git a/include/utils.h b/include/utils.h index 018752792..77b116218 100644 --- a/include/utils.h +++ b/include/utils.h @@ -1004,18 +1004,23 @@ inline bool validate_index_file_size(std::ifstream &in) return true; } -// This function is valid only for float data type. -template inline void normalize(T *arr, size_t dim) +template inline float get_norm(T *arr, const size_t dim) { float sum = 0.0f; for (uint32_t i = 0; i < dim; i++) { sum += arr[i] * arr[i]; } - sum = sqrt(sum); + return sqrt(sum); +} + +// This function is valid only for float data type. +template inline void normalize(T *arr, const size_t dim) +{ + float norm = get_norm(arr, dim); for (uint32_t i = 0; i < dim; i++) { - arr[i] = (T)(arr[i] / sum); + arr[i] = (T)(arr[i] / norm); } } diff --git a/src/distance.cpp b/src/distance.cpp index 0431ab443..d0f446905 100644 --- a/src/distance.cpp +++ b/src/distance.cpp @@ -522,7 +522,7 @@ float AVXDistanceInnerProductFloat::compare(const float *a, const float *b, uint return -result; } -uint32_t AVXNormalizedCosineDistanceFloat::new_dimension(uint32_t orig_dimension) +uint32_t AVXNormalizedCosineDistanceFloat::post_processed_dimension(uint32_t orig_dimension) const { return orig_dimension; } @@ -530,17 +530,11 @@ bool AVXNormalizedCosineDistanceFloat::normalization_required() const { return true; } -void AVXNormalizedCosineDistanceFloat::normalize_data_for_build(const float *original_data, const uint32_t num_points, - const uint32_t orig_dim, float *normalized_data, - bool modify_orig) +void AVXNormalizedCosineDistanceFloat::normalize_data_for_build(float *original_data, const uint32_t orig_dim, const uint32_t num_points) { - if (!modify_orig) - { - throw diskann::ANNException("For cosine distance, modify_orig should be set to true for efficiency.", -1); - } for (auto i = 0; i < num_points; i++) { - normalize(original_data + i * orig_dim, orig_dim); + normalize((float *)(original_data + i * orig_dim), orig_dim); } } @@ -553,13 +547,9 @@ void AVXNormalizedCosineDistanceFloat::normalize_vector_for_search(const float * void AVXNormalizedCosineDistanceFloat::normalize_and_copy(const float *query_vec, const uint32_t query_dim, float *query_target) const { - float norm = 0.0f; - for (auto i = 0; i < query_dim; i++) - { - norm += query_vec[i]; - } - norm /= norm / query_dim; - + + float norm = get_norm(query_vec, query_dim); + for (auto i = 0; i < query_dim; i++) { query_target[i] = query_vec[i] / norm; From 260766eb16366682fafac527ec7ab614b9e160a9 Mon Sep 17 00:00:00 2001 From: Gopal Srinivasa Date: Tue, 4 Apr 2023 00:25:30 +0530 Subject: [PATCH 20/69] DiskANN Builds with initial refactoring changes --- include/abstract_graph_store.h | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/include/abstract_graph_store.h b/include/abstract_graph_store.h index fe088d51b..5fa944ab9 100644 --- a/include/abstract_graph_store.h +++ b/include/abstract_graph_store.h @@ -14,13 +14,18 @@ namespace diskann class AbstractGraphStore { public: - AbstractGraphStore(const size_t max_pts); + AbstractGraphStore(const size_t max_pts) + : _max_points(max_pts){} virtual int load(const std::string &index_path_prefix) = 0; virtual int store(const std::string &index_path_prefix) = 0; virtual void get_adj_list(const location_t i, std::vector &neighbors) = 0; virtual void set_adj_list(const location_t i, std::vector &neighbors) = 0; + + private: + size_t _max_points; + }; } // namespace diskann \ No newline at end of file From 516fb993c6b621fac567a1be762b1553a0b8b2c9 Mon Sep 17 00:00:00 2001 From: Gopal Srinivasa Date: Tue, 4 Apr 2023 16:34:59 +0530 Subject: [PATCH 21/69] Saving changes for Ravi --- include/abstract_data_store.h | 33 +-- include/in_mem_data_store.h | 18 +- include/index.h | 6 +- src/in_mem_data_store.cpp | 53 ++++- src/index.cpp | 420 ++++++++++++++++++---------------- 5 files changed, 310 insertions(+), 220 deletions(-) diff --git a/include/abstract_data_store.h b/include/abstract_data_store.h index 9f43e176b..b6b15e262 100644 --- a/include/abstract_data_store.h +++ b/include/abstract_data_store.h @@ -14,39 +14,46 @@ template class AbstractDataStore { public: AbstractDataStore(const location_t max_pts, const size_t dim) - : _max_pts(max_pts), _num_pts(0), _dim(dim) + : _max_points(max_pts), _num_points(0), _dim(dim) { } // Return number of points returned - virtual size_t load(const std::string &filename) = 0; + virtual location_t load(const std::string &filename) = 0; virtual void store(const std::string &filename) = 0; - virtual data_t *get_vector(location_t i) = 0; - virtual void set_vector(const location_t i, const data_t *const vector) = 0; + virtual void populate_data(const data_t * vectors, const location_t num_pts) = 0; + virtual void populate_data(const std::string &filename, const size_t offset) = 0; + + virtual void get_vector(const location_t i, data_t* dest) const = 0; + virtual void set_vector(const location_t i, const data_t *const vector) = 0; - location_t get_max_pts() + + virtual void reposition_points(const location_t start_loc, const location_t end_loc, + const location_t num_points) = 0; + + location_t get_max_points() const { - return _max_pts; + return _max_points; } - - location_t get_num_pts() + location_t get_num_points() const { - return _num_pts; + return _num_points; } virtual void get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, - float *distances) = 0; + float *distances) const = 0; + virtual float get_distance(const location_t loc1, const location_t loc2) const = 0; - size_t get_dims() + size_t get_dims() const { return _dim; } protected: - location_t _max_pts; - location_t _num_pts; + location_t _max_points; + location_t _num_points; const size_t _dim; }; diff --git a/include/in_mem_data_store.h b/include/in_mem_data_store.h index 981b4b095..bc1b1d401 100644 --- a/include/in_mem_data_store.h +++ b/include/in_mem_data_store.h @@ -24,13 +24,21 @@ class InMemDataStore : public AbstractDataStore InMemDataStore(const location_t max_pts, const size_t dim, std::shared_ptr> distance_metric); virtual ~InMemDataStore(); - void load(const std::string &filename); - void store(const std::string &filename); + virtual location_t load(const std::string &filename) override; + virtual void store(const std::string &filename) override; - virtual data_t *get_vector(location_t i); - virtual void set_vector(const location_t i, const data_t *const vector); + //Populate internal data from unaligned data while doing alignment and any normalization that is required. + virtual void populate_data(const data_t *vectors, const location_t num_pts) override; + virtual void populate_data(const std::string &filename, const size_t offset) override; - virtual void get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, float *distances); + virtual void get_vector(const location_t i, data_t *target) const override; + virtual void set_vector(const location_t i, const data_t *const vector) override; + + virtual void reposition_points(const location_t start_loc, const location_t end_loc, + const location_t num_points) override; + + virtual void get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, float *distances) const override ; + virtual float get_distance(const location_t loc1, const location_t loc2) const override; protected: location_t load_data(const std::string &filename); diff --git a/include/index.h b/include/index.h index 6fb00b382..45baa6f2a 100644 --- a/include/index.h +++ b/include/index.h @@ -18,6 +18,7 @@ #include "utils.h" #include "windows_customizations.h" #include "scratch.h" +#include "in_mem_data_store.h" #define OVERHEAD_FACTOR 1.1 #define EXPAND_IF_FULL 0 @@ -315,7 +316,8 @@ template clas Distance *_distance = nullptr; // Data - T *_data = nullptr; + //T *_data = nullptr; + std::shared_ptr> _data_store; char *_opt_graph = nullptr; // Graph related data structures @@ -323,7 +325,7 @@ template clas // Dimensions size_t _dim = 0; - size_t _aligned_dim = 0; + //size_t _aligned_dim = 0; 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 // Number of points which are used as initial candidates when iterating to diff --git a/src/in_mem_data_store.cpp b/src/in_mem_data_store.cpp index aa90432bc..1c09d7108 100644 --- a/src/in_mem_data_store.cpp +++ b/src/in_mem_data_store.cpp @@ -26,7 +26,7 @@ template InMemDataStore::~InMemDataStore() } } -template void InMemDataStore::load(const std::string &filename) +template location_t InMemDataStore::load(const std::string &filename) { load_data(filename); } @@ -35,6 +35,29 @@ template void InMemDataStore::store(const std::string { } +template void InMemDataStore::populate_data(const data_t *vectors, const location_t num_pts) +{ + for (auto i = 0; i < num_pts; i++) + { + memset(_data + i * _aligned_dim, 0, _aligned_dim * sizeof(data_t)); + std::memmove(_data + i * _aligned_dim, vectors + i * _dim, _dim * sizeof(data_t)); + } + + if (_distance_metric->normalization_required()) + { + _distance_metric->normalize(_data, num_pts); + } +} + +template void InMemDataStore::populate_data(const std::string &filename, const size_t offset) +{ + copy_aligned_data_from_file(filename.c_str(), _data, _num_points, _dim, _aligned_dim, offset); + if (_distance_metric->normalization_required()) + { + _distance_metric->normalize(_data, _num_points); + } +} + #ifdef EXEC_ENV_OLS template location_t Index::load_data(AlignedFileReader &reader) { @@ -101,15 +124,37 @@ location_t InMemDataStore::load_data(const std::string &filename) return file_num_points; } -template data_t *InMemDataStore::get_vector(location_t i) +template +void InMemDataStore::get_vector(const location_t i, data_t* dest) const { - return _data + i * _aligned_dim; + memcpy(dest, _data + i * _aligned_dim, this->_dim * sizeof(data_t)); } template void InMemDataStore::set_vector(const location_t loc, const data_t *const vector) { - memcpy(_data + loc * _aligned_dim, vector, this->_dim * sizeof(data_t)); + size_t offset_in_data = loc * _aligned_dim; + memset(_data + offset_in_data, 0, _aligned_dim * sizeof(data_t)); + memcpy(_data + offset_in_data, vector, this->_dim * sizeof(data_t)); + if (_distance_metric->normalization_required()) + { + _distance_metric->normalize(_data + offset_in_data, _aligned_dim); + } +} + +template +void InMemDataStore::get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, float *distances) const +{ + for (auto i = 0; i < location_count; i++) + { + distances[i] = _distance_metric->compare(query, _data + locations[i] * _aligned_dim, this->_aligned_dim); + } +} + +template +float InMemDataStore::get_distance(const location_t loc1, const location_t loc2) const +{ + return _distance_metric->compare(_data + loc1 * _aligned_dim, _data + loc2 * _aligned_dim, this->_aligned_dim); } } // namespace diskann \ No newline at end of file diff --git a/src/index.cpp b/src/index.cpp index 81e3464b7..74b0a5084 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -75,7 +75,8 @@ Index::Index(Metric m, const size_t dim, const size_t max_point } // data stored to _nd * aligned_dim matrix with necessary zero-padding - _aligned_dim = ROUND_UP(_dim, 8); + // REFACTOR + //_aligned_dim = ROUND_UP(_dim, 8); if (dynamic_index && _num_frozen_pts == 0) { @@ -96,13 +97,18 @@ Index::Index(Metric m, const size_t dim, const size_t max_point alloc_aligned(((void **)&_pq_data), total_internal_points * _num_pq_chunks * sizeof(char), 8 * sizeof(char)); std::memset(_pq_data, 0, total_internal_points * _num_pq_chunks * sizeof(char)); } - alloc_aligned(((void **)&_data), total_internal_points * _aligned_dim * sizeof(T), 8 * sizeof(T)); - std::memset(_data, 0, total_internal_points * _aligned_dim * sizeof(T)); + + //REFACTOR + //alloc_aligned(((void **)&_data), total_internal_points * _aligned_dim * sizeof(T), 8 * sizeof(T)); + //std::memset(_data, 0, total_internal_points * _aligned_dim * sizeof(T)); + //REFACTOR: This should move to a factory method. + _data_store = std::make_shared>(total_internal_points, _dim); _start = (uint32_t)_max_points; _final_graph.resize(total_internal_points); + //This should come from a factory. if (m == diskann::Metric::COSINE && std::is_floating_point::value) { // This is safe because T is float inside the if block. @@ -144,11 +150,12 @@ template Index::~I delete this->_distance; this->_distance = nullptr; } - if (this->_data != nullptr) - { - aligned_free(this->_data); - this->_data = nullptr; - } + //REFACTOR + //if (this->_data != nullptr) + //{ + // aligned_free(this->_data); + // this->_data = nullptr; + //} if (_opt_graph != nullptr) { delete[] _opt_graph; @@ -559,7 +566,8 @@ void Index::load(const char *filename, uint32_t num_threads, ui << 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; diskann::cerr << stream.str() << std::endl; - aligned_free(_data); + //REFACTOR + //aligned_free(_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } @@ -1506,6 +1514,11 @@ void Index::prune_all_neighbors(const uint32_t max_degree, cons { const uint32_t range = max_degree; const uint32_t maxc = max_occlusion_size; + +// const uint32_t range = parameters.Get("R"); +// const uint32_t maxc = parameters.Get("C"); +// const float alpha = parameters.Get("alpha"); + _filtered_index = true; diskann::Timer timer; @@ -1527,8 +1540,10 @@ void Index::prune_all_neighbors(const uint32_t max_degree, cons { 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, (uint32_t)_aligned_dim); + float dist = _data_store->compare(node, cur_nbr); + //REFACTOR + //float dist = _distance->compare(_data + _aligned_dim * (size_t)node, + // _data + _aligned_dim * (size_t)cur_nbr, (uint32_t)_aligned_dim); dummy_pool.emplace_back(Neighbor(cur_nbr, dist)); dummy_visited.insert(cur_nbr); } @@ -1566,49 +1581,50 @@ void Index::prune_all_neighbors(const uint32_t max_degree, cons } } -template -void Index::set_start_points(const T *data, size_t data_count) -{ - std::unique_lock ul(_update_lock); - std::unique_lock tl(_tag_lock); - if (_nd > 0) - throw ANNException("Can not set starting point for a non-empty index", -1, __FUNCSIG__, __FILE__, __LINE__); - - if (data_count != _num_frozen_pts * _aligned_dim) - throw ANNException("Invalid number of points", -1, __FUNCSIG__, __FILE__, __LINE__); - - memcpy(_data + _aligned_dim * _max_points, data, _aligned_dim * sizeof(T) * _num_frozen_pts); - _has_built = true; - diskann::cout << "Index start points set: #" << _num_frozen_pts << std::endl; -} - -template -void Index::set_start_points_at_random(T radius, uint32_t random_seed) -{ - std::mt19937 gen{random_seed}; - std::normal_distribution<> d{0.0, 1.0}; - - std::vector points_data; - points_data.reserve(_aligned_dim * _num_frozen_pts); - std::vector real_vec(_aligned_dim); - - for (size_t frozen_point = 0; frozen_point < _num_frozen_pts; frozen_point++) - { - double norm_sq = 0.0; - for (size_t i = 0; i < _dim; ++i) - { - auto r = d(gen); - real_vec[i] = r; - norm_sq += r * r; - } - - const double norm = std::sqrt(norm_sq); - for (auto iter : real_vec) - points_data.push_back(static_cast(iter * radius / norm)); - } - - set_start_points(points_data.data(), points_data.size()); -} +//REFACTOR +//template +//void Index::set_start_points(const T *data, size_t data_count) +//{ +// std::unique_lock ul(_update_lock); +// std::unique_lock tl(_tag_lock); +// if (_nd > 0) +// throw ANNException("Can not set starting point for a non-empty index", -1, __FUNCSIG__, __FILE__, __LINE__); +// +// if (data_count != _num_frozen_pts * _aligned_dim) +// throw ANNException("Invalid number of points", -1, __FUNCSIG__, __FILE__, __LINE__); +// +// memcpy(_data + _aligned_dim * _max_points, data, _aligned_dim * sizeof(T) * _num_frozen_pts); +// _has_built = true; +// diskann::cout << "Index start points set: #" << _num_frozen_pts << std::endl; +//} +// +//template +//void Index::set_start_points_at_random(T radius, uint32_t random_seed) +//{ +// std::mt19937 gen{random_seed}; +// std::normal_distribution<> d{0.0, 1.0}; +// +// std::vector points_data; +// points_data.reserve(_aligned_dim * _num_frozen_pts); +// std::vector real_vec(_aligned_dim); +// +// for (size_t frozen_point = 0; frozen_point < _num_frozen_pts; frozen_point++) +// { +// double norm_sq = 0.0; +// for (size_t i = 0; i < _dim; ++i) +// { +// auto r = d(gen); +// real_vec[i] = r; +// norm_sq += r * r; +// } +// +// const double norm = std::sqrt(norm_sq); +// for (auto iter : real_vec) +// points_data.push_back(static_cast(iter * radius / norm)); +// } +// +// set_start_points(points_data.data(), points_data.size()); +//} template void Index::build_with_data_populated(IndexWriteParameters ¶meters, const std::vector &tags) @@ -1686,15 +1702,17 @@ void Index::build(const T *data, const size_t num_points_to_loa std::unique_lock tl(_tag_lock); _nd = num_points_to_load; - memcpy((char *)_data, (char *)data, _aligned_dim * _nd * sizeof(T)); + _data_store->populate_data(data, num_points_to_load); - if (_normalize_vecs) - { - for (size_t i = 0; i < num_points_to_load; i++) - { - normalize(_data + _aligned_dim * i, _aligned_dim); - } - } + //REFACTOR + //memcpy((char *)_data, (char *)data, _aligned_dim * _nd * sizeof(T)); + //if (_normalize_vecs) + //{ + // for (size_t i = 0; i < num_points_to_load; i++) + // { + // normalize(_data + _aligned_dim * i, _aligned_dim); + // } + //} } build_with_data_populated(parameters, tags); @@ -1732,8 +1750,9 @@ void Index::build(const char *filename, const size_t num_points if (_pq_dist) aligned_free(_pq_data); - else - aligned_free(_data); + //REFACTOR + //else + // aligned_free(_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } @@ -1745,8 +1764,9 @@ void Index::build(const char *filename, const size_t num_points if (_pq_dist) aligned_free(_pq_data); - else - aligned_free(_data); + //REFACTOR + //else + // aligned_free(_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } @@ -1759,8 +1779,9 @@ void Index::build(const char *filename, const size_t num_points if (_pq_dist) aligned_free(_pq_data); - else - aligned_free(_data); + //REFACTOR + //else + // aligned_free(_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } @@ -1786,14 +1807,16 @@ void Index::build(const char *filename, const size_t num_points #endif } - copy_aligned_data_from_file(filename, _data, file_num_points, file_dim, _aligned_dim); - if (_normalize_vecs) - { - for (size_t i = 0; i < file_num_points; i++) - { - normalize(_data + _aligned_dim * i, _aligned_dim); - } - } + _data_store->populate_data(filename, 0U); + //REFACTOR + //copy_aligned_data_from_file(filename, _data, file_num_points, file_dim, _aligned_dim); + //if (_normalize_vecs) + //{ + // for (size_t i = 0; i < file_num_points; i++) + // { + // normalize(_data + _aligned_dim * i, _aligned_dim); + // } + //} diskann::cout << "Using only first " << num_points_to_load << " from file.. " << std::endl; @@ -2192,7 +2215,9 @@ size_t Index::search_with_tags(const T *query, const uint64_t K if (res_vectors.size() > 0) { - memcpy(res_vectors[pos], _data + ((size_t)node.id) * _aligned_dim, _dim * sizeof(T)); + _data_store->get_vector(node.id, res_vectors[pos]); + //REFACTOR + //memcpy(res_vectors[pos], _data + ((size_t)node.id) * _aligned_dim, _dim * sizeof(T)); } if (distances != nullptr) @@ -2831,15 +2856,17 @@ int Index::insert_point(const T *point, const TagT tag) } tl.unlock(); + _data_store->set_vector(location, point); + //REFACTOR // Copy the vector in to the data array - auto offset_data = _data + (size_t)_aligned_dim * location; - memset((void *)offset_data, 0, sizeof(T) * _aligned_dim); - memcpy((void *)offset_data, point, sizeof(T) * _dim); + //auto offset_data = _data + (size_t)_aligned_dim * location; + //memset((void *)offset_data, 0, sizeof(T) * _aligned_dim); + //memcpy((void *)offset_data, point, sizeof(T) * _dim); - if (_normalize_vecs) - { - normalize((float *)offset_data, _dim); - } + //if (_normalize_vecs) + //{ + // normalize((float *)offset_data, _dim); + //} // Find and add appropriate graph edges ScratchStoreManager> manager(_query_scratch); @@ -3007,118 +3034,119 @@ template void Index void Index::optimize_index_layout() -{ // use after build or load - if (_dynamic_index) - { - throw diskann::ANNException("Optimize_index_layout not implemented for dyanmic indices", -1, __FUNCSIG__, - __FILE__, __LINE__); - } - - _data_len = (_aligned_dim + 1) * sizeof(float); - _neighbor_len = (_max_observed_degree + 1) * sizeof(uint32_t); - _node_size = _data_len + _neighbor_len; - _opt_graph = new char[_node_size * _nd]; - DistanceFastL2 *dist_fast = (DistanceFastL2 *)_distance; - for (uint32_t 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; - uint32_t k = _final_graph[i].size(); - std::memcpy(cur_node_offset, &k, sizeof(uint32_t)); - std::memcpy(cur_node_offset + sizeof(uint32_t), _final_graph[i].data(), k * sizeof(uint32_t)); - std::vector().swap(_final_graph[i]); - } - _final_graph.clear(); - _final_graph.shrink_to_fit(); -} - -template -void Index::search_with_optimized_layout(const T *query, size_t K, size_t L, uint32_t *indices) -{ - DistanceFastL2 *dist_fast = (DistanceFastL2 *)_distance; - - NeighborPriorityQueue retset(L); - std::vector init_ids(L); - - boost::dynamic_bitset<> flags{_nd, 0}; - uint32_t tmp_l = 0; - uint32_t *neighbors = (uint32_t *)(_opt_graph + _node_size * _start + _data_len); - uint32_t 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) - { - uint32_t id = rand() % _nd; - if (flags[id]) - continue; - flags[id] = true; - init_ids[tmp_l] = id; - tmp_l++; - } - - for (uint32_t i = 0; i < init_ids.size(); i++) - { - uint32_t id = init_ids[i]; - if (id >= _nd) - continue; - _mm_prefetch(_opt_graph + _node_size * id, _MM_HINT_T0); - } - L = 0; - for (uint32_t i = 0; i < init_ids.size(); i++) - { - uint32_t 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, (uint32_t)_aligned_dim); - retset.insert(Neighbor(id, dist)); - flags[id] = true; - L++; - } - - while (retset.has_unexpanded_node()) - { - auto nbr = retset.closest_unexpanded(); - auto n = nbr.id; - _mm_prefetch(_opt_graph + _node_size * n + _data_len, _MM_HINT_T0); - neighbors = (uint32_t *)(_opt_graph + _node_size * n + _data_len); - uint32_t MaxM = *neighbors; - neighbors++; - for (uint32_t m = 0; m < MaxM; ++m) - _mm_prefetch(_opt_graph + _node_size * neighbors[m], _MM_HINT_T0); - for (uint32_t m = 0; m < MaxM; ++m) - { - uint32_t id = neighbors[m]; - if (flags[id]) - continue; - flags[id] = 1; - T *data = (T *)(_opt_graph + _node_size * id); - float norm = *data; - data++; - float dist = dist_fast->compare(query, data, norm, (uint32_t)_aligned_dim); - Neighbor nn(id, dist); - retset.insert(nn); - } - } - - for (size_t i = 0; i < K; i++) - { - indices[i] = retset[i].id; - } -} +//REFACTOR: This should be an OptimizedDataStore class +//template void Index::optimize_index_layout() +//{ // use after build or load +// if (_dynamic_index) +// { +// throw diskann::ANNException("Optimize_index_layout not implemented for dyanmic indices", -1, __FUNCSIG__, +// __FILE__, __LINE__); +// } +// +// _data_len = (_aligned_dim + 1) * sizeof(float); +// _neighbor_len = (_max_observed_degree + 1) * sizeof(uint32_t); +// _node_size = _data_len + _neighbor_len; +// _opt_graph = new char[_node_size * _nd]; +// DistanceFastL2 *dist_fast = (DistanceFastL2 *)_distance; +// for (uint32_t 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; +// uint32_t k = _final_graph[i].size(); +// std::memcpy(cur_node_offset, &k, sizeof(uint32_t)); +// std::memcpy(cur_node_offset + sizeof(uint32_t), _final_graph[i].data(), k * sizeof(uint32_t)); +// std::vector().swap(_final_graph[i]); +// } +// _final_graph.clear(); +// _final_graph.shrink_to_fit(); +//} +// +//template +//void Index::search_with_optimized_layout(const T *query, size_t K, size_t L, uint32_t *indices) +//{ +// DistanceFastL2 *dist_fast = (DistanceFastL2 *)_distance; +// +// NeighborPriorityQueue retset(L); +// std::vector init_ids(L); +// +// boost::dynamic_bitset<> flags{_nd, 0}; +// uint32_t tmp_l = 0; +// uint32_t *neighbors = (uint32_t *)(_opt_graph + _node_size * _start + _data_len); +// uint32_t 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) +// { +// uint32_t id = rand() % _nd; +// if (flags[id]) +// continue; +// flags[id] = true; +// init_ids[tmp_l] = id; +// tmp_l++; +// } +// +// for (uint32_t i = 0; i < init_ids.size(); i++) +// { +// uint32_t id = init_ids[i]; +// if (id >= _nd) +// continue; +// _mm_prefetch(_opt_graph + _node_size * id, _MM_HINT_T0); +// } +// L = 0; +// for (uint32_t i = 0; i < init_ids.size(); i++) +// { +// uint32_t 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, (uint32_t)_aligned_dim); +// retset.insert(Neighbor(id, dist)); +// flags[id] = true; +// L++; +// } +// +// while (retset.has_unexpanded_node()) +// { +// auto nbr = retset.closest_unexpanded(); +// auto n = nbr.id; +// _mm_prefetch(_opt_graph + _node_size * n + _data_len, _MM_HINT_T0); +// neighbors = (uint32_t *)(_opt_graph + _node_size * n + _data_len); +// uint32_t MaxM = *neighbors; +// neighbors++; +// for (uint32_t m = 0; m < MaxM; ++m) +// _mm_prefetch(_opt_graph + _node_size * neighbors[m], _MM_HINT_T0); +// for (uint32_t m = 0; m < MaxM; ++m) +// { +// uint32_t id = neighbors[m]; +// if (flags[id]) +// continue; +// flags[id] = 1; +// T *data = (T *)(_opt_graph + _node_size * id); +// float norm = *data; +// data++; +// float dist = dist_fast->compare(query, data, norm, (uint32_t)_aligned_dim); +// Neighbor nn(id, dist); +// retset.insert(nn); +// } +// } +// +// for (size_t i = 0; i < K; i++) +// { +// indices[i] = retset[i].id; +// } +//} /* Internals of the library */ template const float Index::INDEX_GROWTH_FACTOR = 1.5f; From f66f975714b0e46c1a43ff4b982923691da18980 Mon Sep 17 00:00:00 2001 From: Gopal Srinivasa Date: Tue, 4 Apr 2023 23:54:38 +0530 Subject: [PATCH 22/69] More refactoring --- include/abstract_data_store.h | 23 ++++-- include/abstract_graph_store.h | 4 +- include/in_mem_data_store.h | 10 ++- src/in_mem_data_store.cpp | 115 ++++++++++++++++++++++++++++- src/index.cpp | 129 +++++++++++++++++++-------------- 5 files changed, 215 insertions(+), 66 deletions(-) diff --git a/include/abstract_data_store.h b/include/abstract_data_store.h index b6b15e262..380a1d77f 100644 --- a/include/abstract_data_store.h +++ b/include/abstract_data_store.h @@ -13,8 +13,8 @@ namespace diskann template class AbstractDataStore { public: - AbstractDataStore(const location_t max_pts, const size_t dim) - : _max_points(max_pts), _num_points(0), _dim(dim) + AbstractDataStore(const location_t capacity, const size_t dim) + : _capacity(capacity), _num_points(0), _dim(dim) { } @@ -31,28 +31,37 @@ template class AbstractDataStore virtual void reposition_points(const location_t start_loc, const location_t end_loc, const location_t num_points) = 0; + virtual void copy_points(const location_t from_loc, const location_t to_loc, const location_t num_points) = 0; + //Returns the point in the dataset that is closest to the mean of all points in the dataset + virtual location_t calculate_medoid() const = 0; - location_t get_max_points() const + virtual location_t capacity() const { - return _max_points; + return _capacity; } - location_t get_num_points() const + virtual location_t get_num_points() const { return _num_points; } + virtual float get_distance(const data_t* query, const location_t loc) const = 0; virtual void get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, float *distances) const = 0; virtual float get_distance(const location_t loc1, const location_t loc2) const = 0; - size_t get_dims() const + virtual size_t get_dims() const + { + return _dim; + } + + virtual size_t get_aligned_dim() const { return _dim; } protected: - location_t _max_points; + location_t _capacity; location_t _num_points; const size_t _dim; diff --git a/include/abstract_graph_store.h b/include/abstract_graph_store.h index 5fa944ab9..d12c43664 100644 --- a/include/abstract_graph_store.h +++ b/include/abstract_graph_store.h @@ -15,7 +15,7 @@ class AbstractGraphStore { public: AbstractGraphStore(const size_t max_pts) - : _max_points(max_pts){} + : _capacity(max_pts){} virtual int load(const std::string &index_path_prefix) = 0; virtual int store(const std::string &index_path_prefix) = 0; @@ -24,7 +24,7 @@ class AbstractGraphStore virtual void set_adj_list(const location_t i, std::vector &neighbors) = 0; private: - size_t _max_points; + size_t _capacity; }; diff --git a/include/in_mem_data_store.h b/include/in_mem_data_store.h index bc1b1d401..ae3473f72 100644 --- a/include/in_mem_data_store.h +++ b/include/in_mem_data_store.h @@ -34,12 +34,20 @@ class InMemDataStore : public AbstractDataStore virtual void get_vector(const location_t i, data_t *target) const override; virtual void set_vector(const location_t i, const data_t *const vector) override; - virtual void reposition_points(const location_t start_loc, const location_t end_loc, + virtual void reposition_points(const location_t old_location_start, const location_t new_location_start, const location_t num_points) override; + virtual void copy_points(const location_t from_loc, const location_t to_loc, const location_t num_points) override; + virtual location_t calculate_medoid() const override; + virtual float get_distance(const data_t *query, const location_t loc) const override; virtual void get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, float *distances) const override ; virtual float get_distance(const location_t loc1, const location_t loc2) const override; + virtual size_t get_aligned_dim() const override + { + return _aligned_dim; + } + protected: location_t load_data(const std::string &filename); #ifdef EXEC_ENV_OLS diff --git a/src/in_mem_data_store.cpp b/src/in_mem_data_store.cpp index 1c09d7108..fcaaf3b96 100644 --- a/src/in_mem_data_store.cpp +++ b/src/in_mem_data_store.cpp @@ -114,7 +114,7 @@ location_t InMemDataStore::load_data(const std::string &filename) throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } - if (file_num_points > this->get_max_points()) + if (file_num_points > this->capacity()) { resize(file_num_points); } @@ -142,6 +142,11 @@ void InMemDataStore::set_vector(const location_t loc, const data_t *cons } } +template float InMemDataStore::get_distance(const data_t *query, const location_t loc) const +{ + return _distance_metric->compare(query, _data + _aligned_dim * loc, _aligned_dim); +} + template void InMemDataStore::get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, float *distances) const { @@ -157,4 +162,112 @@ float InMemDataStore::get_distance(const location_t loc1, const location return _distance_metric->compare(_data + loc1 * _aligned_dim, _data + loc2 * _aligned_dim, this->_aligned_dim); } +template +void InMemDataStore::reposition_points(const location_t old_location_start, const location_t new_location_start, const location_t num_locations) +{ + if (num_locations == 0 || old_location_start == new_location_start) + { + return; + } + + // Update pointers to the moved nodes. Note: the computation is correct even + // when new_location_start < old_location_start given the C++ uint32_t + // integer arithmetic rules. + const uint32_t location_delta = new_location_start - old_location_start; + + // 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; + + // Move the adjacency lists. Make sure that overlapping ranges are handled + // correctly. + 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_points(old_location_start, new_location_start, num_locations); + memset(_data + _aligned_dim * mem_clear_loc_start, 0, + sizeof(data_t) * _aligned_dim * (mem_clear_loc_end_limit - mem_clear_loc_start)); +} + +template +void InMemDataStore::copy_points(const location_t from_loc, const location_t to_loc, + const location_t num_points) +{ + assert(from_loc < _capacity); + assert(to_loc < _capacity); + assert(num_points < _capacity); + memmove(_data + _aligned_dim * to_loc, _data + _aligned_dim * from_loc, num_points * _aligned_dim * sizeof(data_t)); +} + +template location_t InMemDataStore::calculate_medoid() const +{ + // allocate and init centroid + float *center = new float[_aligned_dim]; + for (size_t j = 0; j < _aligned_dim; j++) + center[j] = 0; + + for (size_t i = 0; i < _num_points; i++) + for (size_t j = 0; j < _aligned_dim; j++) + center[j] += (float)_data[i * _aligned_dim + j]; + + for (size_t j = 0; j < _aligned_dim; j++) + center[j] /= (float)_num_points; + + // compute all to one distance + float *distances = new float[_num_points]; + +//TODO: REFACTOR. Removing pragma might make this slow. Must revisit. +// Problem is that we need to pass num_threads here, it is not clear +// if data store must be aware of threads! +//#pragma omp parallel for schedule(static, 65536) + for (int64_t i = 0; i < (int64_t)_num_points; i++) + { + // extract point and distance reference + float &dist = distances[i]; + const data_t *cur_vec = _data + (i * (size_t)_aligned_dim); + dist = 0; + float diff = 0; + for (size_t j = 0; j < _aligned_dim; j++) + { + diff = (center[j] - (float)cur_vec[j]) * (center[j] - (float)cur_vec[j]); + dist += diff; + } + } + // find imin + uint32_t min_idx = 0; + float min_dist = distances[0]; + for (uint32_t i = 1; i < _num_points; i++) + { + if (distances[i] < min_dist) + { + min_idx = i; + min_dist = distances[i]; + } + } + + delete[] distances; + delete[] center; + return min_idx; +} + } // namespace diskann \ No newline at end of file diff --git a/src/index.cpp b/src/index.cpp index 74b0a5084..3d55677de 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -810,8 +810,12 @@ template int Index return -1; } - size_t location = _tag_to_location[tag]; - memcpy((void *)vec, (void *)(_data + location * _aligned_dim), (size_t)_dim * sizeof(T)); + location_t location = _tag_to_location[tag]; + _data_store->get_vector(location, vec); + + //REFACTOR + //size_t location = _tag_to_location[tag]; + //memcpy((void *)vec, (void *)(_data + location * _aligned_dim), (size_t)_dim * sizeof(T)); return 0; } @@ -824,49 +828,54 @@ template uint32_t Indexcalculate_medoid(); + +// //REFACTOR +// // allocate and init centroid +// float *center = new float[_aligned_dim](); +// for (size_t j = 0; j < _aligned_dim; j++) +// center[j] = 0; +// +// for (size_t i = 0; i < _nd; i++) +// for (size_t j = 0; j < _aligned_dim; j++) +// center[j] += (float)_data[i * _aligned_dim + j]; +// +// for (size_t j = 0; j < _aligned_dim; j++) +// center[j] /= (float)_nd; +// +// // compute all to one distance +// float *distances = new float[_nd](); +//#pragma omp parallel for schedule(static, 65536) +// for (int64_t i = 0; i < (int64_t)_nd; i++) +// { +// // extract point and distance reference +// float &dist = distances[i]; +// const T *cur_vec = _data + (i * (size_t)_aligned_dim); +// dist = 0; +// float diff = 0; +// for (size_t j = 0; j < _aligned_dim; j++) +// { +// diff = (center[j] - (float)cur_vec[j]) * (center[j] - (float)cur_vec[j]); +// dist += diff; +// } +// } +// // find imin +// uint32_t min_idx = 0; +// float min_dist = distances[0]; +// for (uint32_t i = 1; i < _nd; i++) +// { +// if (distances[i] < min_dist) +// { +// min_idx = i; +// min_dist = distances[i]; +// } +// } +// +// delete[] distances; +// delete[] center; +// return min_idx; } template std::vector Index::get_init_ids() @@ -1007,9 +1016,16 @@ std::pair Index::iterate_to_fixed_point( float distance; if (_pq_dist) + { pq_dist_lookup(pq_coord_scratch, 1, this->_num_pq_chunks, pq_dists, &distance); + } else - distance = _distance->compare(_data + _aligned_dim * (size_t)id, aligned_query, (uint32_t)_aligned_dim); + { + _data_store->get_distance(aligned_query, id); + // REFACTOR + // distance = _distance->compare(_data + _aligned_dim * (size_t)id, aligned_query, + // (uint32_t)_aligned_dim); + } Neighbor nn = Neighbor(id, distance); best_L_nodes.insert(nn); } @@ -1540,7 +1556,7 @@ void Index::prune_all_neighbors(const uint32_t max_degree, cons { if (dummy_visited.find(cur_nbr) == dummy_visited.end() && cur_nbr != node) { - float dist = _data_store->compare(node, cur_nbr); + float dist = _data_store->get_distance(node, cur_nbr); //REFACTOR //float dist = _distance->compare(_data + _aligned_dim * (size_t)node, // _data + _aligned_dim * (size_t)cur_nbr, (uint32_t)_aligned_dim); @@ -1640,7 +1656,8 @@ void Index::build_with_data_populated(IndexWriteParameters &par stream << "ERROR: Driver requests loading " << _nd << " points from file," << "but tags vector is of size " << tags.size() << "." << std::endl; diskann::cerr << stream.str() << std::endl; - aligned_free(_data); + //REFACTOR + //aligned_free(_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } if (_enable_tags) @@ -1659,7 +1676,10 @@ void Index::build_with_data_populated(IndexWriteParameters &par if (_query_scratch.size() == 0) { - initialize_query_scratch(5 + num_threads_index, index_L, index_L, index_R, maxc, _aligned_dim); + //REFACTOR + //initialize_query_scratch(5 + num_threads_index, index_L, index_L, index_R, maxc, _aligned_dim); + //Comment: Why 5+? + initialize_query_scratch(5 + num_threads_index, index_L, index_L, index_R, maxc, _data_store->get_aligned_dim()); } generate_frozen_point(); @@ -2276,7 +2296,9 @@ template void Indexcopy_points(res, _max_points, 1); + //REFACTOR + //memcpy(_data + _max_points * _aligned_dim, _data + res * _aligned_dim, _aligned_dim * sizeof(T)); } } @@ -2726,12 +2748,8 @@ void Index::reposition_points(uint32_t old_location_start, uint mem_clear_loc_end_limit = new_location_start; } } + _data_store->reposition_points(old_location_start, new_location_start, num_locations); - // Use memmove to handle overlapping ranges. - memmove(_data + _aligned_dim * new_location_start, _data + _aligned_dim * old_location_start, - sizeof(T) * _aligned_dim * num_locations); - memset(_data + _aligned_dim * mem_clear_loc_start, 0, - sizeof(T) * _aligned_dim * (mem_clear_loc_end_limit - mem_clear_loc_start)); } template void Index::reposition_frozen_point_to_end() @@ -2745,6 +2763,7 @@ template void Index Date: Wed, 5 Apr 2023 15:30:21 +0530 Subject: [PATCH 23/69] Refactor --- CMakeLists.txt | 2 +- include/abstract_data_store.h | 5 ++- include/in_mem_data_store.h | 3 ++ src/in_mem_data_store.cpp | 20 +++++++++ src/index.cpp | 77 +++++++++++++++++++++++------------ 5 files changed, 79 insertions(+), 28 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index dfa1cb77e..825cb6df5 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -244,7 +244,7 @@ endif() add_subdirectory(src) add_subdirectory(tests) add_subdirectory(tests/utils) -add_subdirectory(tests/unittests) + if (MSVC) message(STATUS "The ${PROJECT_NAME}.sln has been created, opened it from VisualStudio to build Release or Debug configurations.\n" diff --git a/include/abstract_data_store.h b/include/abstract_data_store.h index 380a1d77f..1fbd30097 100644 --- a/include/abstract_data_store.h +++ b/include/abstract_data_store.h @@ -27,14 +27,17 @@ template class AbstractDataStore virtual void get_vector(const location_t i, data_t* dest) const = 0; virtual void set_vector(const location_t i, const data_t *const vector) = 0; + virtual void prefetch_vector(const location_t loc) = 0; + virtual void resize(const location_t new_size) = 0; + virtual void reposition_points(const location_t start_loc, const location_t end_loc, const location_t num_points) = 0; virtual void copy_points(const location_t from_loc, const location_t to_loc, const location_t num_points) = 0; //Returns the point in the dataset that is closest to the mean of all points in the dataset virtual location_t calculate_medoid() const = 0; - + virtual location_t capacity() const { return _capacity; diff --git a/include/in_mem_data_store.h b/include/in_mem_data_store.h index ae3473f72..d2e6ae935 100644 --- a/include/in_mem_data_store.h +++ b/include/in_mem_data_store.h @@ -33,6 +33,9 @@ class InMemDataStore : public AbstractDataStore virtual void get_vector(const location_t i, data_t *target) const override; virtual void set_vector(const location_t i, const data_t *const vector) override; + virtual void prefetch_vector(const location_t loc) override; + + virtual void resize(const location_t new_size) override; virtual void reposition_points(const location_t old_location_start, const location_t new_location_start, const location_t num_points) override; diff --git a/src/in_mem_data_store.cpp b/src/in_mem_data_store.cpp index fcaaf3b96..c13265d1a 100644 --- a/src/in_mem_data_store.cpp +++ b/src/in_mem_data_store.cpp @@ -142,6 +142,11 @@ void InMemDataStore::set_vector(const location_t loc, const data_t *cons } } +template void InMemDataStore::prefetch_vector(const location_t loc) +{ + diskann::prefetch_vector((const char *)_data + _aligned_dim * (size_t)loc, sizeof(T) * _aligned_dim); +} + template float InMemDataStore::get_distance(const data_t *query, const location_t loc) const { return _distance_metric->compare(query, _data + _aligned_dim * loc, _aligned_dim); @@ -162,6 +167,21 @@ float InMemDataStore::get_distance(const location_t loc1, const location return _distance_metric->compare(_data + loc1 * _aligned_dim, _data + loc2 * _aligned_dim, this->_aligned_dim); } +template void InMemDataStore::resize(const location_t new_size) +{ +#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, _data, (_max_points + _num_frozen_pts) * _aligned_dim * sizeof(data_t)); + aligned_free(_data); + _data = new_data; +#else + realloc_aligned((void **)&_data, new_size * _aligned_dim * sizeof(data_t), 8 * sizeof(data_t)); +#endif + _capacity = new_size; + +} + template void InMemDataStore::reposition_points(const location_t old_location_start, const location_t new_location_start, const location_t num_locations) { diff --git a/src/index.cpp b/src/index.cpp index 3d55677de..97696fec2 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -1125,12 +1125,16 @@ std::pair Index::iterate_to_fixed_point( if (m + 1 < id_scratch.size()) { auto nextn = id_scratch[m + 1]; - diskann::prefetch_vector((const char *)_data + _aligned_dim * (size_t)nextn, - sizeof(T) * _aligned_dim); + _data_store->prefetch_vector(nextn); + //REFACTOR + //diskann::prefetch_vector((const char *)_data + _aligned_dim * (size_t)nextn, + // sizeof(T) * _aligned_dim); } - dist_scratch.push_back( - _distance->compare(aligned_query, _data + _aligned_dim * (size_t)id, (uint32_t)_aligned_dim)); + dist_scratch.push_back(_data_store->get_distance(aligned_query, id)); + //REFACTOR + //dist_scratch.push_back( + // _distance->compare(aligned_query, _data + _aligned_dim * (size_t)id, (uint32_t)_aligned_dim)); } } cmps += id_scratch.size(); @@ -1155,16 +1159,25 @@ void Index::search_for_point_and_prune(int location, uint32_t L if (!use_filter) { - iterate_to_fixed_point(_data + _aligned_dim * location, Lindex, init_ids, scratch, false, unused_filter_label, - false); + _data_store->get_vector(location, scratch->aligned_query()); + iterate_to_fixed_point(scratch->aligned_query(), Lindex, init_ids, scratch, false, unused_filter_label, false); + + //REFACTOR + //iterate_to_fixed_point(_data + _aligned_dim * location, Lindex, init_ids, scratch, false, unused_filter_label, + // false); } else { std::vector filter_specific_start_nodes; for (auto &x : _pts_to_labels[location]) filter_specific_start_nodes.emplace_back(_label_to_medoid_id[x]); - iterate_to_fixed_point(_data + _aligned_dim * location, filteredLindex, filter_specific_start_nodes, scratch, + + _data_store->get_vector(location, scratch->aligned_query()); + iterate_to_fixed_point(scratch->aligned_query(), filteredLindex, filter_specific_start_nodes, scratch, true, _pts_to_labels[location], false); + //REFACTOR + //iterate_to_fixed_point(_data + _aligned_dim * location, filteredLindex, filter_specific_start_nodes, scratch, + // true, _pts_to_labels[location], false); } auto &pool = scratch->pool(); @@ -1260,8 +1273,10 @@ void Index::occlude_list(const uint32_t location, std::vectorcompare(_data + _aligned_dim * (size_t)iter2->id, - _data + _aligned_dim * (size_t)iter->id, (uint32_t)_aligned_dim); + float djk = _data_store->get_distance(iter2->id, iter->id); + //REFACTOR + //float djk = _distance->compare(_data + _aligned_dim * (size_t)iter2->id, + // _data + _aligned_dim * (size_t)iter->id, (uint32_t)_aligned_dim); if (_dist_metric == diskann::Metric::L2 || _dist_metric == diskann::Metric::COSINE) { occlude_factor[t] = (djk == 0) ? std::numeric_limits::max() @@ -1308,8 +1323,10 @@ void Index::prune_neighbors(const uint32_t location, std::vecto if (_pq_dist) { for (auto &ngh : pool) - ngh.distance = _distance->compare(_data + _aligned_dim * (size_t)ngh.id, - _data + _aligned_dim * (size_t)location, (uint32_t)_aligned_dim); + ngh.distance = _data_store->get_distance(ngh.id, location); + //REFACTOR + //ngh.distance = _distance->compare(_data + _aligned_dim * (size_t)ngh.id, + // _data + _aligned_dim * (size_t)location, (uint32_t)_aligned_dim); } // sort the pool based on distance to query and prune it with occlude_list @@ -1380,8 +1397,10 @@ void Index::inter_insert(uint32_t n, std::vector &pru { if (dummy_visited.find(cur_nbr) == dummy_visited.end() && cur_nbr != des) { - float dist = _distance->compare(_data + _aligned_dim * (size_t)des, - _data + _aligned_dim * (size_t)cur_nbr, (uint32_t)_aligned_dim); + float dist = _data_store->get_distance(des, cur_nbr); + //REFACTOR + //float dist = _distance->compare(_data + _aligned_dim * (size_t)des, + // _data + _aligned_dim * (size_t)cur_nbr, (uint32_t)_aligned_dim); dummy_pool.emplace_back(Neighbor(cur_nbr, dist)); dummy_visited.insert(cur_nbr); } @@ -1505,8 +1524,10 @@ void Index::link(IndexWriteParameters ¶meters) { 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, (uint32_t)_aligned_dim); + float dist = _data_store->get_distance(node, cur_nbr); + //REFACTOR + //float dist = _distance->compare(_data + _aligned_dim * (size_t)node, + // _data + _aligned_dim * (size_t)cur_nbr, (uint32_t)_aligned_dim); dummy_pool.emplace_back(Neighbor(cur_nbr, dist)); dummy_visited.insert(cur_nbr); } @@ -2587,8 +2608,10 @@ template void Indexcopy_points(old, new_location[old], 1); + //REFACTOR + //memcpy((void *)(_data + _aligned_dim * (size_t)new_location[old]), + // (void *)(_data + _aligned_dim * (size_t)old), _aligned_dim * sizeof(T)); } } else @@ -2774,15 +2797,17 @@ template void Indexresize(new_internal_points); + //REFACTOR +//#ifndef _WINDOWS +// T *new_data; +// alloc_aligned((void **)&new_data, new_internal_points * _aligned_dim * sizeof(T), 8 * sizeof(T)); +// memcpy(new_data, _data, (_max_points + _num_frozen_pts) * _aligned_dim * sizeof(T)); +// aligned_free(_data); +// _data = new_data; +//#else +// realloc_aligned((void **)&_data, new_internal_points * _aligned_dim * sizeof(T), 8 * sizeof(T)); +//#endif _final_graph.resize(new_internal_points); _locks = std::vector(new_internal_points); From b348a07dc2f4606e15d79d0769b5d0fc7fd26f92 Mon Sep 17 00:00:00 2001 From: Gopal Srinivasa Date: Wed, 5 Apr 2023 23:55:41 +0530 Subject: [PATCH 24/69] Fixed most of the bugs related to _data --- include/in_mem_data_store.h | 1 + src/index.cpp | 24 +++++++++++++++++------- 2 files changed, 18 insertions(+), 7 deletions(-) diff --git a/include/in_mem_data_store.h b/include/in_mem_data_store.h index d2e6ae935..d456ae28d 100644 --- a/include/in_mem_data_store.h +++ b/include/in_mem_data_store.h @@ -51,6 +51,7 @@ class InMemDataStore : public AbstractDataStore return _aligned_dim; } + protected: location_t load_data(const std::string &filename); #ifdef EXEC_ENV_OLS diff --git a/src/index.cpp b/src/index.cpp index 97696fec2..53579cf63 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -448,7 +448,8 @@ size_t Index::load_data(std::string filename) std::stringstream stream; stream << "ERROR: data file " << filename << " does not exist." << std::endl; diskann::cerr << stream.str() << std::endl; - aligned_free(_data); + //REFACTOR + //aligned_free(_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } diskann::get_bin_metadata(filename, file_num_points, file_dim); @@ -463,7 +464,8 @@ size_t Index::load_data(std::string filename) stream << "ERROR: Driver requests loading " << _dim << " dimension," << "but file has " << file_dim << " dimension." << std::endl; diskann::cerr << stream.str() << std::endl; - aligned_free(_data); + //REFACTOR + //aligned_free(_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } @@ -474,9 +476,13 @@ size_t Index::load_data(std::string filename) } #ifdef EXEC_ENV_OLS + + //REFACTOR TODO: Must figure out how to support aligned reader in a clean manner. copy_aligned_data_from_file(reader, _data, file_num_points, file_dim, _aligned_dim); #else - copy_aligned_data_from_file(filename.c_str(), _data, file_num_points, file_dim, _aligned_dim); + _data_store->populate_data(filename, 0); //offset == 0. + //REFACTOR + //copy_aligned_data_from_file(filename.c_str(), _data, file_num_points, file_dim, _aligned_dim); #endif return file_num_points; } @@ -719,7 +725,8 @@ size_t Index::load_graph(std::string filename, size_t expected_ << std::endl; } diskann::cerr << stream.str() << std::endl; - aligned_free(_data); + //REFACTOR + //aligned_free(_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } @@ -2403,9 +2410,12 @@ inline void Index::process_delete(const tsl::robin_setcompare(_data + _aligned_dim * loc, _data + _aligned_dim * ngh, (uint32_t)_aligned_dim)); + expanded_nghrs_vec.emplace_back(ngh, _data_store->get_distance(loc, ngh)); + + //REFACTOR + //expanded_nghrs_vec.emplace_back( + // ngh, + // _distance->compare(_data + _aligned_dim * loc, _data + _aligned_dim * ngh, (uint32_t)_aligned_dim)); } std::sort(expanded_nghrs_vec.begin(), expanded_nghrs_vec.end()); std::vector &occlude_list_output = scratch->occlude_list_output(); From aadafaddf12440112643d3f693a6cee5ad6514a0 Mon Sep 17 00:00:00 2001 From: Gopal Srinivasa Date: Thu, 6 Apr 2023 15:14:36 +0530 Subject: [PATCH 25/69] Post merge with main --- include/abstract_data_store.h | 29 ++++++++++++++++++++++------- include/in_mem_data_store.h | 8 +++++--- src/in_mem_data_store.cpp | 24 +++++++++++++++++++----- src/index.cpp | 3 +-- 4 files changed, 47 insertions(+), 17 deletions(-) diff --git a/include/abstract_data_store.h b/include/abstract_data_store.h index 1fbd30097..f9689d5ac 100644 --- a/include/abstract_data_store.h +++ b/include/abstract_data_store.h @@ -30,7 +30,21 @@ template class AbstractDataStore virtual void prefetch_vector(const location_t loc) = 0; - virtual void resize(const location_t new_size) = 0; + virtual void resize(const location_t num_points) + { + if (num_points > _capacity) + { + expand(num_points); + } + else if (num_point < _capacity) + { + shrink(num_points); + } + else + { + //ignore. + } + } virtual void reposition_points(const location_t start_loc, const location_t end_loc, const location_t num_points) = 0; @@ -42,10 +56,6 @@ template class AbstractDataStore { return _capacity; } - virtual location_t get_num_points() const - { - return _num_points; - } virtual float get_distance(const data_t* query, const location_t loc) const = 0; virtual void get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, @@ -64,9 +74,14 @@ template class AbstractDataStore } protected: - location_t _capacity; - location_t _num_points; + // Expand the datastore to new_num_points. + virtual void expand(const location_t new_num_points) = 0; + // Shrink the datastore to new_num_points. This function should be called after compaction to free unused memory. + virtual void shrink(const location_t new_num_points) = 0; + + + location_t _capacity; const size_t _dim; }; diff --git a/include/in_mem_data_store.h b/include/in_mem_data_store.h index 95a7069ae..bf450a589 100644 --- a/include/in_mem_data_store.h +++ b/include/in_mem_data_store.h @@ -7,7 +7,7 @@ #include "tsl/robin_map.h" #include "tsl/robin_set.h" #include "tsl/sparse_map.h" -#include "boost/dynamic_bitset.hpp" +//#include "boost/dynamic_bitset.hpp" #include "abstract_data_store.h" @@ -35,8 +35,6 @@ class InMemDataStore : public AbstractDataStore virtual void set_vector(const location_t i, const data_t *const vector) override; virtual void prefetch_vector(const location_t loc) override; - virtual void resize(const location_t new_size) override; - virtual void reposition_points(const location_t old_location_start, const location_t new_location_start, const location_t num_points) override; virtual void copy_points(const location_t from_loc, const location_t to_loc, const location_t num_points) override; @@ -53,6 +51,10 @@ class InMemDataStore : public AbstractDataStore protected: + virtual void expand(const location_t new_size) override; + virtual void shrink(const location_t new_size) override; + + virtual location_t load_data(const std::string &filename); #ifdef EXEC_ENV_OLS virtual location_t load_data(AlignedFileReader &reader); diff --git a/src/in_mem_data_store.cpp b/src/in_mem_data_store.cpp index c13265d1a..1c0aefa2e 100644 --- a/src/in_mem_data_store.cpp +++ b/src/in_mem_data_store.cpp @@ -9,13 +9,13 @@ namespace diskann { template -InMemDataStore::InMemDataStore(const location_t max_pts, +InMemDataStore::InMemDataStore(const location_t capacity, const size_t dim, std::shared_ptr> distance_metric) - : AbstractDataStore(max_pts, dim), _aligned_dim(ROUND_UP(dim, 8)), + : AbstractDataStore(capacity, dim), _aligned_dim(ROUND_UP(dim, 8)), _distance_metric(distance_metric) { - alloc_aligned(((void **)&_data), max_pts * _aligned_dim * sizeof(data_t), 8 * sizeof(data_t)); - std::memset(_data, 0, max_pts * _aligned_dim * sizeof(data_t)); + alloc_aligned(((void **)&_data), capacity * _aligned_dim * sizeof(data_t), 8 * sizeof(data_t)); + std::memset(_data, 0, capacity * _aligned_dim * sizeof(data_t)); } template InMemDataStore::~InMemDataStore() @@ -33,6 +33,7 @@ template location_t InMemDataStore::load(const std::st template void InMemDataStore::store(const std::string &filename) { + } template void InMemDataStore::populate_data(const data_t *vectors, const location_t num_pts) @@ -167,7 +168,7 @@ float InMemDataStore::get_distance(const location_t loc1, const location return _distance_metric->compare(_data + loc1 * _aligned_dim, _data + loc2 * _aligned_dim, this->_aligned_dim); } -template void InMemDataStore::resize(const location_t new_size) +template void InMemDataStore::expand(const location_t new_size) { #ifndef _WINDOWS data_t *new_data; @@ -182,6 +183,19 @@ template void InMemDataStore::resize(const location_t } +template void InMemDataStore::shrink(const location_t new_size) +{ +#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, _data, new_size * _aligned_dim * sizeof(data_t)); + aligned_free(_data); + _data = new_data; +#else + realloc_aligned((void **)&_data, new_size * _aligned_dim * sizeof(data_t), 8 * sizeof(data_t)); +#endif +} + template void InMemDataStore::reposition_points(const location_t old_location_start, const location_t new_location_start, const location_t num_locations) { diff --git a/src/index.cpp b/src/index.cpp index 53579cf63..dee1ce5c8 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -480,7 +480,7 @@ size_t Index::load_data(std::string filename) //REFACTOR TODO: Must figure out how to support aligned reader in a clean manner. copy_aligned_data_from_file(reader, _data, file_num_points, file_dim, _aligned_dim); #else - _data_store->populate_data(filename, 0); //offset == 0. + _data_store->populate_data(filename, 0U); //offset == 0. //REFACTOR //copy_aligned_data_from_file(filename.c_str(), _data, file_num_points, file_dim, _aligned_dim); #endif @@ -2652,7 +2652,6 @@ template void Index Date: Thu, 6 Apr 2023 17:31:38 +0530 Subject: [PATCH 26/69] Refactored version which compiles on Windows --- include/abstract_data_store.h | 15 +++-- include/in_mem_data_store.h | 5 +- include/index.h | 2 +- include/natural_number_map.h | 3 +- src/in_mem_data_store.cpp | 108 ++++++++++++++++++++++++---------- src/index.cpp | 24 ++++---- 6 files changed, 108 insertions(+), 49 deletions(-) diff --git a/include/abstract_data_store.h b/include/abstract_data_store.h index f9689d5ac..2ecc9d5f6 100644 --- a/include/abstract_data_store.h +++ b/include/abstract_data_store.h @@ -14,13 +14,18 @@ template class AbstractDataStore { public: AbstractDataStore(const location_t capacity, const size_t dim) - : _capacity(capacity), _num_points(0), _dim(dim) + : _capacity(capacity), _dim(dim) { } // Return number of points returned virtual location_t load(const std::string &filename) = 0; - virtual void store(const std::string &filename) = 0; + + //Why does store take num_pts? Since store only has capacity, but we allow resizing + //we can end up in a situation where the store has spare capacity. To optimize disk + //utilization, we pass the number of points that are "true" points, so that the store + //can discard the empty locations before saving. + virtual size_t save(const std::string &filename, const location_t num_pts) = 0; virtual void populate_data(const data_t * vectors, const location_t num_pts) = 0; virtual void populate_data(const std::string &filename, const size_t offset) = 0; @@ -32,11 +37,11 @@ template class AbstractDataStore virtual void resize(const location_t num_points) { - if (num_points > _capacity) + if (num_points > _capacity) { expand(num_points); } - else if (num_point < _capacity) + else if (num_points < _capacity) { shrink(num_points); } @@ -82,7 +87,7 @@ template class AbstractDataStore location_t _capacity; - const size_t _dim; + size_t _dim; }; } // namespace diskann \ No newline at end of file diff --git a/include/in_mem_data_store.h b/include/in_mem_data_store.h index bf450a589..84bdf4147 100644 --- a/include/in_mem_data_store.h +++ b/include/in_mem_data_store.h @@ -25,7 +25,7 @@ class InMemDataStore : public AbstractDataStore virtual ~InMemDataStore(); virtual location_t load(const std::string &filename) override; - virtual void store(const std::string &filename) override; + virtual size_t save(const std::string &filename, const location_t num_points) override; //Populate internal data from unaligned data while doing alignment and any normalization that is required. virtual void populate_data(const data_t *vectors, const location_t num_pts) override; @@ -41,8 +41,9 @@ class InMemDataStore : public AbstractDataStore virtual location_t calculate_medoid() const override; virtual float get_distance(const data_t *query, const location_t loc) const override; - virtual void get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, float *distances) const override ; virtual float get_distance(const location_t loc1, const location_t loc2) const override; + virtual void get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, float *distances) const override ; + virtual size_t get_aligned_dim() const override { diff --git a/include/index.h b/include/index.h index 45baa6f2a..22ebe6cea 100644 --- a/include/index.h +++ b/include/index.h @@ -313,7 +313,7 @@ template clas private: // Distance functions Metric _dist_metric = diskann::L2; - Distance *_distance = nullptr; + std::shared_ptr> _distance; // Data //T *_data = nullptr; diff --git a/include/natural_number_map.h b/include/natural_number_map.h index 3bbdcb640..7678c2aa0 100644 --- a/include/natural_number_map.h +++ b/include/natural_number_map.h @@ -7,7 +7,8 @@ #include #include -#include "boost_dynamic_bitset_fwd.h" +#include + //#include "boost_dynamic_bitset_fwd.h" namespace diskann { diff --git a/src/in_mem_data_store.cpp b/src/in_mem_data_store.cpp index 1c0aefa2e..8a9804b93 100644 --- a/src/in_mem_data_store.cpp +++ b/src/in_mem_data_store.cpp @@ -1,6 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT license. +#include #include "in_mem_data_store.h" #include "utils.h" @@ -9,13 +10,13 @@ namespace diskann { template -InMemDataStore::InMemDataStore(const location_t capacity, +InMemDataStore::InMemDataStore(const location_t num_points, const size_t dim, std::shared_ptr> distance_metric) - : AbstractDataStore(capacity, dim), _aligned_dim(ROUND_UP(dim, 8)), - _distance_metric(distance_metric) + : AbstractDataStore(num_points, dim), _aligned_dim(ROUND_UP(dim, 8)), + _distance_fn(distance_metric) { - alloc_aligned(((void **)&_data), capacity * _aligned_dim * sizeof(data_t), 8 * sizeof(data_t)); - std::memset(_data, 0, capacity * _aligned_dim * sizeof(data_t)); + alloc_aligned(((void **)&_data), this->_capacity * _aligned_dim * sizeof(data_t), 8 * sizeof(data_t)); + std::memset(_data, 0, this->_capacity * _aligned_dim * sizeof(data_t)); } template InMemDataStore::~InMemDataStore() @@ -28,12 +29,13 @@ template InMemDataStore::~InMemDataStore() template location_t InMemDataStore::load(const std::string &filename) { - load_data(filename); + return load_data(filename); } -template void InMemDataStore::store(const std::string &filename) +template size_t InMemDataStore::save(const std::string &filename, const location_t num_points) { - + shrink(num_points); + return save_data_in_base_dimensions(filename, _data, num_points, this->get_dims(), this->get_aligned_dim(), 0U); } template void InMemDataStore::populate_data(const data_t *vectors, const location_t num_pts) @@ -41,21 +43,40 @@ template void InMemDataStore::populate_data(const data for (auto i = 0; i < num_pts; i++) { memset(_data + i * _aligned_dim, 0, _aligned_dim * sizeof(data_t)); - std::memmove(_data + i * _aligned_dim, vectors + i * _dim, _dim * sizeof(data_t)); + std::memmove(_data + i * _aligned_dim, vectors + i * this->_dim, this->_dim * sizeof(data_t)); } - if (_distance_metric->normalization_required()) + if (_distance_fn->normalization_required()) { - _distance_metric->normalize(_data, num_pts); + _distance_fn->normalize_data_for_build(_data, this->_dim, num_pts); } } template void InMemDataStore::populate_data(const std::string &filename, const size_t offset) { - copy_aligned_data_from_file(filename.c_str(), _data, _num_points, _dim, _aligned_dim, offset); - if (_distance_metric->normalization_required()) + size_t npts, ndim; + copy_aligned_data_from_file(filename.c_str(), _data, npts, ndim, _aligned_dim, offset); + + if ((location_t)npts > this->capacity()) + { + std::stringstream ss; + ss << "Number of points in the file: " << filename + << " is greater than the capacity of data store: " << this->capacity() + << ". Must invoke resize before calling populate_data()" << std::endl; + throw diskann::ANNException(ss.str(), -1); + } + + if ((location_t)ndim != this->get_dims()) { - _distance_metric->normalize(_data, _num_points); + std::stringstream ss; + ss << "Number of dimensions of a point in the file: " << filename + << " is not equal to dimensions of data store: " << this->capacity() << "." << std::endl; + throw diskann::ANNException(ss.str(), -1); + } + + if (_distance_fn->normalization_required()) + { + _distance_fn->normalize_data_for_build(_data, this->_dim, this->capacity()); } } @@ -117,7 +138,7 @@ location_t InMemDataStore::load_data(const std::string &filename) if (file_num_points > this->capacity()) { - resize(file_num_points); + this->resize(file_num_points); } copy_aligned_data_from_file(filename.c_str(), _data, file_num_points, file_dim, _aligned_dim); @@ -137,20 +158,20 @@ void InMemDataStore::set_vector(const location_t loc, const data_t *cons size_t offset_in_data = loc * _aligned_dim; memset(_data + offset_in_data, 0, _aligned_dim * sizeof(data_t)); memcpy(_data + offset_in_data, vector, this->_dim * sizeof(data_t)); - if (_distance_metric->normalization_required()) + if (_distance_fn->normalization_required()) { - _distance_metric->normalize(_data + offset_in_data, _aligned_dim); + _distance_fn->normalize_data_for_build(_data + offset_in_data, _aligned_dim, 1); } } template void InMemDataStore::prefetch_vector(const location_t loc) { - diskann::prefetch_vector((const char *)_data + _aligned_dim * (size_t)loc, sizeof(T) * _aligned_dim); + diskann::prefetch_vector((const char *)_data + _aligned_dim * (size_t)loc, sizeof(data_t) * _aligned_dim); } template float InMemDataStore::get_distance(const data_t *query, const location_t loc) const { - return _distance_metric->compare(query, _data + _aligned_dim * loc, _aligned_dim); + return _distance_fn->compare(query, _data + _aligned_dim * loc, _aligned_dim); } template @@ -158,18 +179,29 @@ void InMemDataStore::get_distance(const data_t *query, const location_t { for (auto i = 0; i < location_count; i++) { - distances[i] = _distance_metric->compare(query, _data + locations[i] * _aligned_dim, this->_aligned_dim); + distances[i] = _distance_fn->compare(query, _data + locations[i] * _aligned_dim, this->_aligned_dim); } } template float InMemDataStore::get_distance(const location_t loc1, const location_t loc2) const { - return _distance_metric->compare(_data + loc1 * _aligned_dim, _data + loc2 * _aligned_dim, this->_aligned_dim); + return _distance_fn->compare(_data + loc1 * _aligned_dim, _data + loc2 * _aligned_dim, this->_aligned_dim); } template void InMemDataStore::expand(const location_t new_size) { + if (new_size == this->capacity()) + { + return; + } + 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)); @@ -179,12 +211,23 @@ template void InMemDataStore::expand(const location_t #else realloc_aligned((void **)&_data, new_size * _aligned_dim * sizeof(data_t), 8 * sizeof(data_t)); #endif - _capacity = new_size; + this->_capacity = new_size; } template void InMemDataStore::shrink(const location_t new_size) { + if (new_size == this->capacity()) + { + return; + } + 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)); @@ -247,9 +290,9 @@ template void InMemDataStore::copy_points(const location_t from_loc, const location_t to_loc, const location_t num_points) { - assert(from_loc < _capacity); - assert(to_loc < _capacity); - assert(num_points < _capacity); + assert(from_loc < this->_capacity); + assert(to_loc < this->_capacity); + assert(num_points < this->_capacity); memmove(_data + _aligned_dim * to_loc, _data + _aligned_dim * from_loc, num_points * _aligned_dim * sizeof(data_t)); } @@ -260,21 +303,21 @@ template location_t InMemDataStore::calculate_medoid() for (size_t j = 0; j < _aligned_dim; j++) center[j] = 0; - for (size_t i = 0; i < _num_points; i++) + for (size_t i = 0; i < this->capacity(); i++) for (size_t j = 0; j < _aligned_dim; j++) center[j] += (float)_data[i * _aligned_dim + j]; for (size_t j = 0; j < _aligned_dim; j++) - center[j] /= (float)_num_points; + center[j] /= (float)this->capacity(); // compute all to one distance - float *distances = new float[_num_points]; + float *distances = new float[this->capacity()]; //TODO: REFACTOR. Removing pragma might make this slow. Must revisit. // Problem is that we need to pass num_threads here, it is not clear // if data store must be aware of threads! //#pragma omp parallel for schedule(static, 65536) - for (int64_t i = 0; i < (int64_t)_num_points; i++) + for (int64_t i = 0; i < (int64_t)this->capacity(); i++) { // extract point and distance reference float &dist = distances[i]; @@ -290,7 +333,7 @@ template location_t InMemDataStore::calculate_medoid() // find imin uint32_t min_idx = 0; float min_dist = distances[0]; - for (uint32_t i = 1; i < _num_points; i++) + for (uint32_t i = 1; i < this->capacity(); i++) { if (distances[i] < min_dist) { @@ -304,4 +347,9 @@ template location_t InMemDataStore::calculate_medoid() return min_idx; } +template DISKANN_DLLEXPORT class InMemDataStore; +template DISKANN_DLLEXPORT class InMemDataStore; +template DISKANN_DLLEXPORT class InMemDataStore; + + } // namespace diskann \ No newline at end of file diff --git a/src/index.cpp b/src/index.cpp index dee1ce5c8..3a59a2cd4 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -101,8 +101,6 @@ Index::Index(Metric m, const size_t dim, const size_t max_point //REFACTOR //alloc_aligned(((void **)&_data), total_internal_points * _aligned_dim * sizeof(T), 8 * sizeof(T)); //std::memset(_data, 0, total_internal_points * _aligned_dim * sizeof(T)); - //REFACTOR: This should move to a factory method. - _data_store = std::make_shared>(total_internal_points, _dim); _start = (uint32_t)_max_points; @@ -112,7 +110,7 @@ Index::Index(Metric m, const size_t dim, const size_t max_point if (m == diskann::Metric::COSINE && std::is_floating_point::value) { // This is safe because T is float inside the if block. - this->_distance = (Distance *)new AVXNormalizedCosineDistanceFloat(); + this->_distance.reset((Distance*)new AVXNormalizedCosineDistanceFloat()); this->_normalize_vecs = true; diskann::cout << "Normalizing vectors and using L2 for cosine " "AVXNormalizedCosineDistanceFloat()." @@ -120,8 +118,12 @@ Index::Index(Metric m, const size_t dim, const size_t max_point } else { - this->_distance = get_distance_function(m); + this->_distance.reset((Distance *)get_distance_function(m)); } + // REFACTOR: TODO This should move to a factory method. + + _data_store = std::make_shared>((location_t)total_internal_points, _dim, this->_distance); + _locks = std::vector(total_internal_points); @@ -145,11 +147,11 @@ template Index::~I LockGuard lg(lock); } - if (this->_distance != nullptr) - { - delete this->_distance; - this->_distance = nullptr; - } + //if (this->_distance != nullptr) + //{ + // delete this->_distance; + // this->_distance = nullptr; + //} //REFACTOR //if (this->_data != nullptr) //{ @@ -222,7 +224,9 @@ template size_t Indexsave(data_file, _nd + _num_frozen_pts); + // REFACTOR + //return save_data_in_base_dimensions(data_file, _data, _nd + _num_frozen_pts, _dim, _aligned_dim); } // save the graph index on a file as an adjacency list. For each point, From 469c6a5f009aa8e8e647ad284a623b5158e9ef67 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Thu, 6 Apr 2023 17:41:38 +0000 Subject: [PATCH 27/69] now compiles on linux --- include/distance.h | 4 ++-- src/in_mem_data_store.cpp | 2 +- src/index.cpp | 20 +++++++++++++++++++- 3 files changed, 22 insertions(+), 4 deletions(-) diff --git a/include/distance.h b/include/distance.h index e4ab0e51b..7a92a03c1 100644 --- a/include/distance.h +++ b/include/distance.h @@ -1,6 +1,6 @@ #pragma once #include "windows_customizations.h" - +#include namespace diskann { @@ -70,7 +70,7 @@ template class Distance virtual void normalize_vector_for_search(const T *query_vec, const uint32_t query_dim, T* scratch_query) { - memcpy(scratch_query, query_vec, query_dim * sizeof(T)); + std::memcpy(scratch_query, query_vec, query_dim * sizeof(T)); } //Providing a default implementation for the virtual destructor because we don't diff --git a/src/in_mem_data_store.cpp b/src/in_mem_data_store.cpp index 8a9804b93..6a06bad2c 100644 --- a/src/in_mem_data_store.cpp +++ b/src/in_mem_data_store.cpp @@ -205,7 +205,7 @@ template void InMemDataStore::expand(const location_t #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, _data, (_max_points + _num_frozen_pts) * _aligned_dim * sizeof(data_t)); + memcpy(new_data, _data, new_size * _aligned_dim * sizeof(data_t)); aligned_free(_data); _data = new_data; #else diff --git a/src/index.cpp b/src/index.cpp index 3a59a2cd4..2b02ef8ef 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -1646,6 +1646,14 @@ void Index::prune_all_neighbors(const uint32_t max_degree, cons // diskann::cout << "Index start points set: #" << _num_frozen_pts << std::endl; //} // + +// REFACTOR: added dummy implementation for now. +template +void Index::set_start_points_at_random(T radius, uint32_t random_seed) { + +} + + //template //void Index::set_start_points_at_random(T radius, uint32_t random_seed) //{ @@ -3091,6 +3099,11 @@ template void Index void Index::optimize_index_layout() +{ // use after build or load +} + //REFACTOR: This should be an OptimizedDataStore class //template void Index::optimize_index_layout() //{ // use after build or load @@ -3121,7 +3134,12 @@ template void Index +void Index::search_with_optimized_layout(const T *query, size_t K, size_t L, uint32_t *indices) +{ +} + //template //void Index::search_with_optimized_layout(const T *query, size_t K, size_t L, uint32_t *indices) //{ From 874c53c1348f8059151009cf3425be57c86dd779 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Sun, 9 Apr 2023 17:49:27 +0000 Subject: [PATCH 28/69] minor clean-up --- include/abstract_data_store.h | 50 +++++++++------- include/in_mem_data_store.h | 34 ++++------- src/in_mem_data_store.cpp | 108 ++++++++++++++++++---------------- 3 files changed, 98 insertions(+), 94 deletions(-) diff --git a/include/abstract_data_store.h b/include/abstract_data_store.h index 2ecc9d5f6..282f91696 100644 --- a/include/abstract_data_store.h +++ b/include/abstract_data_store.h @@ -27,13 +27,28 @@ template class AbstractDataStore //can discard the empty locations before saving. virtual size_t save(const std::string &filename, const location_t num_pts) = 0; + virtual location_t capacity() const + { + return _capacity; + } + + virtual size_t get_dims() const + { + return _dim; + } + +// by default, aligned dim = dim, some stores can align the data differently, so we may have different values + virtual size_t get_aligned_dim() const + { + return _dim; + } + +// populate the store with bulk vectors (either as pointer or bin file), potentially after normalizing the vectors if the metric deems so virtual void populate_data(const data_t * vectors, const location_t num_pts) = 0; virtual void populate_data(const std::string &filename, const size_t offset) = 0; - - virtual void get_vector(const location_t i, data_t* dest) const = 0; - virtual void set_vector(const location_t i, const data_t *const vector) = 0; - virtual void prefetch_vector(const location_t loc) = 0; +// reverse of populate, save the first num_pts many points back to bin file + virtual void save_data_to_bin(const std::string &filename, const location_t num_pts); virtual void resize(const location_t num_points) { @@ -51,32 +66,26 @@ template class AbstractDataStore } } +// operations on vectors + virtual void get_vector(const location_t i, data_t* dest) const = 0; + virtual void set_vector(const location_t i, const data_t *const vector) = 0; + virtual void prefetch_vector(const location_t loc) = 0; + +// internal shuffle operations to move around vectors virtual void reposition_points(const location_t start_loc, const location_t end_loc, const location_t num_points) = 0; virtual void copy_points(const location_t from_loc, const location_t to_loc, const location_t num_points) = 0; - //Returns the point in the dataset that is closest to the mean of all points in the dataset - virtual location_t calculate_medoid() const = 0; - virtual location_t capacity() const - { - return _capacity; - } +// metric specific operations virtual float get_distance(const data_t* query, const location_t loc) const = 0; virtual void get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, float *distances) const = 0; virtual float get_distance(const location_t loc1, const location_t loc2) const = 0; - - virtual size_t get_dims() const - { - return _dim; - } - - virtual size_t get_aligned_dim() const - { - return _dim; - } +// stats of the data stored in store + //Returns the point in the dataset that is closest to the mean of all points in the dataset + virtual location_t calculate_medoid() const = 0; protected: // Expand the datastore to new_num_points. @@ -85,7 +94,6 @@ template class AbstractDataStore virtual void shrink(const location_t new_num_points) = 0; - location_t _capacity; size_t _dim; }; diff --git a/include/in_mem_data_store.h b/include/in_mem_data_store.h index 84bdf4147..2a66c479a 100644 --- a/include/in_mem_data_store.h +++ b/include/in_mem_data_store.h @@ -21,16 +21,24 @@ template class InMemDataStore : public AbstractDataStore { public: - InMemDataStore(const location_t max_pts, const size_t dim, std::shared_ptr> distance_metric); + InMemDataStore(const location_t capacity, const size_t dim, std::shared_ptr> distance_metric); virtual ~InMemDataStore(); virtual location_t load(const std::string &filename) override; virtual size_t save(const std::string &filename, const location_t num_points) override; + virtual size_t get_aligned_dim() const override + { + return _aligned_dim; + } + //Populate internal data from unaligned data while doing alignment and any normalization that is required. virtual void populate_data(const data_t *vectors, const location_t num_pts) override; virtual void populate_data(const std::string &filename, const size_t offset) override; + virtual void save_data_to_bin(const std::string &filename, const location_t num_pts) override; + + virtual void get_vector(const location_t i, data_t *target) const override; virtual void set_vector(const location_t i, const data_t *const vector) override; virtual void prefetch_vector(const location_t loc) override; @@ -38,17 +46,14 @@ class InMemDataStore : public AbstractDataStore virtual void reposition_points(const location_t old_location_start, const location_t new_location_start, const location_t num_points) override; virtual void copy_points(const location_t from_loc, const location_t to_loc, const location_t num_points) override; - virtual location_t calculate_medoid() const override; + virtual float get_distance(const data_t *query, const location_t loc) const override; virtual float get_distance(const location_t loc1, const location_t loc2) const override; virtual void get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, float *distances) const override ; - - virtual size_t get_aligned_dim() const override - { - return _aligned_dim; - } + virtual location_t calculate_medoid() const override; + protected: @@ -66,26 +71,13 @@ class InMemDataStore : public AbstractDataStore const size_t _aligned_dim; - // lazy_delete removes entry from _location_to_tag and _tag_to_location. If - // _location_to_tag does not resolve a location, infer that it was deleted. - //tsl::sparse_map _tag_to_location; - //natural_number_map _location_to_tag; - - // _empty_slots has unallocated slots and those freed by consolidate_delete. - // _delete_set has locations marked deleted by lazy_delete. Will not be - // immediately available for insert. consolidate_delete will release these - // slots to _empty_slots. - natural_number_set _empty_slots; - std::unique_ptr> _delete_set; - - std::shared_timed_mutex _lock; // Takes please of Index::_tag_lock - // It may seem weird to put distance metric along with the data store class, but // this gives us perf benefits as the datastore can do distance computations during // search and compute norms of vectors internally without have to copy // data back and forth. std::shared_ptr> _distance_fn; +// in case we need to save vector norms for optimization std::shared_ptr _pre_computed_norms; }; diff --git a/src/in_mem_data_store.cpp b/src/in_mem_data_store.cpp index 6a06bad2c..a66d642f6 100644 --- a/src/in_mem_data_store.cpp +++ b/src/in_mem_data_store.cpp @@ -32,54 +32,6 @@ template location_t InMemDataStore::load(const std::st return load_data(filename); } -template size_t InMemDataStore::save(const std::string &filename, const location_t num_points) -{ - shrink(num_points); - return save_data_in_base_dimensions(filename, _data, num_points, this->get_dims(), this->get_aligned_dim(), 0U); -} - -template void InMemDataStore::populate_data(const data_t *vectors, const location_t num_pts) -{ - for (auto i = 0; i < num_pts; i++) - { - memset(_data + i * _aligned_dim, 0, _aligned_dim * sizeof(data_t)); - std::memmove(_data + i * _aligned_dim, vectors + i * this->_dim, this->_dim * sizeof(data_t)); - } - - if (_distance_fn->normalization_required()) - { - _distance_fn->normalize_data_for_build(_data, this->_dim, num_pts); - } -} - -template void InMemDataStore::populate_data(const std::string &filename, const size_t offset) -{ - size_t npts, ndim; - copy_aligned_data_from_file(filename.c_str(), _data, npts, ndim, _aligned_dim, offset); - - if ((location_t)npts > this->capacity()) - { - std::stringstream ss; - ss << "Number of points in the file: " << filename - << " is greater than the capacity of data store: " << this->capacity() - << ". Must invoke resize before calling populate_data()" << std::endl; - throw diskann::ANNException(ss.str(), -1); - } - - if ((location_t)ndim != this->get_dims()) - { - std::stringstream ss; - ss << "Number of dimensions of a point in the file: " << filename - << " is not equal to dimensions of data store: " << this->capacity() << "." << std::endl; - throw diskann::ANNException(ss.str(), -1); - } - - if (_distance_fn->normalization_required()) - { - _distance_fn->normalize_data_for_build(_data, this->_dim, this->capacity()); - } -} - #ifdef EXEC_ENV_OLS template location_t Index::load_data(AlignedFileReader &reader) { @@ -87,9 +39,6 @@ template location_t Index::load_data(AlignedFileReader diskann::get_bin_metadata(reader, file_num_points, file_dim); - // since we are loading a new dataset, _empty_slots must be cleared - _empty_slots.clear(); - if (file_dim != _dim) { std::stringstream stream; @@ -124,7 +73,7 @@ location_t InMemDataStore::load_data(const std::string &filename) diskann::get_bin_metadata(filename, file_num_points, file_dim); // since we are loading a new dataset, _empty_slots must be cleared - _empty_slots.clear(); +// _empty_slots.clear(); if (file_dim != this->_dim) { @@ -146,6 +95,61 @@ location_t InMemDataStore::load_data(const std::string &filename) return file_num_points; } + +template size_t InMemDataStore::save(const std::string &filename, const location_t num_points) +{ + return save_data_in_base_dimensions(filename, _data, num_points, this->get_dims(), this->get_aligned_dim(), 0U); +} + +template void InMemDataStore::populate_data(const data_t *vectors, const location_t num_pts) +{ + for (auto i = 0; i < num_pts; i++) + { + memset(_data + i * _aligned_dim, 0, _aligned_dim * sizeof(data_t)); + std::memmove(_data + i * _aligned_dim, vectors + i * this->_dim, this->_dim * sizeof(data_t)); + } + + if (_distance_fn->normalization_required()) + { + _distance_fn->normalize_data_for_build(_data, this->_dim, num_pts); + } +} + +template void InMemDataStore::populate_data(const std::string &filename, const size_t offset) +{ + size_t npts, ndim; + copy_aligned_data_from_file(filename.c_str(), _data, npts, ndim, _aligned_dim, offset); + + if ((location_t)npts > this->capacity()) + { + std::stringstream ss; + ss << "Number of points in the file: " << filename + << " is greater than the capacity of data store: " << this->capacity() + << ". Must invoke resize before calling populate_data()" << std::endl; + throw diskann::ANNException(ss.str(), -1); + } + + if ((location_t)ndim != this->get_dims()) + { + std::stringstream ss; + ss << "Number of dimensions of a point in the file: " << filename + << " is not equal to dimensions of data store: " << this->capacity() << "." << std::endl; + throw diskann::ANNException(ss.str(), -1); + } + + if (_distance_fn->normalization_required()) + { + _distance_fn->normalize_data_for_build(_data, this->_dim, this->capacity()); + } +} + +template void InMemDataStore::save_data_to_bin(const std::string &filename, const location_t num_points) +{ + save_data_in_base_dimensions(filename, _data, num_points, this->get_dims(), this->get_aligned_dim(), 0U); +} + + + template void InMemDataStore::get_vector(const location_t i, data_t* dest) const { From c50214790433626bb5ad97ae663048899187a2b0 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Sun, 9 Apr 2023 18:34:43 +0000 Subject: [PATCH 29/69] minor bug fix --- src/index.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/index.cpp b/src/index.cpp index 2b02ef8ef..533f16587 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -1032,7 +1032,7 @@ std::pair Index::iterate_to_fixed_point( } else { - _data_store->get_distance(aligned_query, id); + distance = _data_store->get_distance(aligned_query, id); // REFACTOR // distance = _distance->compare(_data + _aligned_dim * (size_t)id, aligned_query, // (uint32_t)_aligned_dim); From 310f37317715ad6cbf7057047557972722cbfcdc Mon Sep 17 00:00:00 2001 From: ravishankar Date: Mon, 10 Apr 2023 16:35:35 +0000 Subject: [PATCH 30/69] minor bug --- include/in_mem_data_store.h | 4 ++-- src/in_mem_data_store.cpp | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/include/in_mem_data_store.h b/include/in_mem_data_store.h index 2a66c479a..5dcd1ce40 100644 --- a/include/in_mem_data_store.h +++ b/include/in_mem_data_store.h @@ -61,9 +61,9 @@ class InMemDataStore : public AbstractDataStore virtual void shrink(const location_t new_size) override; - virtual location_t load_data(const std::string &filename); + virtual location_t load_impl(const std::string &filename); #ifdef EXEC_ENV_OLS - virtual location_t load_data(AlignedFileReader &reader); + virtual location_t load_impl(AlignedFileReader &reader); #endif private: diff --git a/src/in_mem_data_store.cpp b/src/in_mem_data_store.cpp index a66d642f6..fe8a57a50 100644 --- a/src/in_mem_data_store.cpp +++ b/src/in_mem_data_store.cpp @@ -29,11 +29,11 @@ template InMemDataStore::~InMemDataStore() template location_t InMemDataStore::load(const std::string &filename) { - return load_data(filename); + return load_impl(filename); } #ifdef EXEC_ENV_OLS -template location_t Index::load_data(AlignedFileReader &reader) +template location_t Index::load_impl(AlignedFileReader &reader) { size_t file_dim, file_num_points; @@ -59,7 +59,7 @@ template location_t Index::load_data(AlignedFileReader #endif template -location_t InMemDataStore::load_data(const std::string &filename) +location_t InMemDataStore::load_impl(const std::string &filename) { size_t file_dim, file_num_points; if (!file_exists(filename)) @@ -209,7 +209,7 @@ template void InMemDataStore::expand(const location_t #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, _data, new_size * _aligned_dim * sizeof(data_t)); + memcpy(new_data, _data, this->capacity() * _aligned_dim * sizeof(data_t)); aligned_free(_data); _data = new_data; #else From bf11919a095ad21695a82a37e1570206641fcd23 Mon Sep 17 00:00:00 2001 From: yashpatel007 Date: Mon, 10 Apr 2023 15:31:11 -0400 Subject: [PATCH 31/69] clang format fix + build error fix --- include/abstract_data_store.h | 53 ++- include/abstract_graph_store.h | 6 +- include/in_mem_data_store.h | 19 +- include/natural_number_map.h | 2 +- src/distance.cpp | 8 +- src/in_mem_data_store.cpp | 62 ++-- src/in_mem_graph_store.cpp | 1 + src/index.cpp | 629 +++++++++++++++++---------------- 8 files changed, 385 insertions(+), 395 deletions(-) diff --git a/include/abstract_data_store.h b/include/abstract_data_store.h index 282f91696..7d60dd34e 100644 --- a/include/abstract_data_store.h +++ b/include/abstract_data_store.h @@ -13,46 +13,46 @@ namespace diskann template class AbstractDataStore { public: - AbstractDataStore(const location_t capacity, const size_t dim) - : _capacity(capacity), _dim(dim) + AbstractDataStore(const location_t capacity, const size_t dim) : _capacity(capacity), _dim(dim) { } // Return number of points returned virtual location_t load(const std::string &filename) = 0; - //Why does store take num_pts? Since store only has capacity, but we allow resizing - //we can end up in a situation where the store has spare capacity. To optimize disk - //utilization, we pass the number of points that are "true" points, so that the store - //can discard the empty locations before saving. + // Why does store take num_pts? Since store only has capacity, but we allow resizing + // we can end up in a situation where the store has spare capacity. To optimize disk + // utilization, we pass the number of points that are "true" points, so that the store + // can discard the empty locations before saving. virtual size_t save(const std::string &filename, const location_t num_pts) = 0; - virtual location_t capacity() const + virtual location_t capacity() const { return _capacity; } - virtual size_t get_dims() const + virtual size_t get_dims() const { return _dim; } -// by default, aligned dim = dim, some stores can align the data differently, so we may have different values + // by default, aligned dim = dim, some stores can align the data differently, so we may have different values virtual size_t get_aligned_dim() const { return _dim; } -// populate the store with bulk vectors (either as pointer or bin file), potentially after normalizing the vectors if the metric deems so - virtual void populate_data(const data_t * vectors, const location_t num_pts) = 0; + // populate the store with bulk vectors (either as pointer or bin file), potentially after normalizing the vectors + // if the metric deems so + virtual void populate_data(const data_t *vectors, const location_t num_pts) = 0; virtual void populate_data(const std::string &filename, const size_t offset) = 0; -// reverse of populate, save the first num_pts many points back to bin file - virtual void save_data_to_bin(const std::string &filename, const location_t num_pts); + // reverse of populate, save the first num_pts many points back to bin file + virtual void save_data_to_bin(const std::string &filename, const location_t num_pts) = 0; - virtual void resize(const location_t num_points) + virtual void resize(const location_t num_points) { - if (num_points > _capacity) + if (num_points > _capacity) { expand(num_points); } @@ -62,29 +62,29 @@ template class AbstractDataStore } else { - //ignore. + // ignore. } } -// operations on vectors - virtual void get_vector(const location_t i, data_t* dest) const = 0; - virtual void set_vector(const location_t i, const data_t *const vector) = 0; + // operations on vectors + virtual void get_vector(const location_t i, data_t *dest) const = 0; + virtual void set_vector(const location_t i, const data_t *const vector) = 0; virtual void prefetch_vector(const location_t loc) = 0; -// internal shuffle operations to move around vectors + // internal shuffle operations to move around vectors virtual void reposition_points(const location_t start_loc, const location_t end_loc, const location_t num_points) = 0; virtual void copy_points(const location_t from_loc, const location_t to_loc, const location_t num_points) = 0; - -// metric specific operations - virtual float get_distance(const data_t* query, const location_t loc) const = 0; + // metric specific operations + + virtual float get_distance(const data_t *query, const location_t loc) const = 0; virtual void get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, - float *distances) const = 0; + float *distances) const = 0; virtual float get_distance(const location_t loc1, const location_t loc2) const = 0; -// stats of the data stored in store - //Returns the point in the dataset that is closest to the mean of all points in the dataset + // stats of the data stored in store + // Returns the point in the dataset that is closest to the mean of all points in the dataset virtual location_t calculate_medoid() const = 0; protected: @@ -93,7 +93,6 @@ template class AbstractDataStore // Shrink the datastore to new_num_points. This function should be called after compaction to free unused memory. virtual void shrink(const location_t new_num_points) = 0; - location_t _capacity; size_t _dim; }; diff --git a/include/abstract_graph_store.h b/include/abstract_graph_store.h index d12c43664..692149783 100644 --- a/include/abstract_graph_store.h +++ b/include/abstract_graph_store.h @@ -14,8 +14,9 @@ namespace diskann class AbstractGraphStore { public: - AbstractGraphStore(const size_t max_pts) - : _capacity(max_pts){} + AbstractGraphStore(const size_t max_pts) : _capacity(max_pts) + { + } virtual int load(const std::string &index_path_prefix) = 0; virtual int store(const std::string &index_path_prefix) = 0; @@ -25,7 +26,6 @@ class AbstractGraphStore private: size_t _capacity; - }; } // namespace diskann \ No newline at end of file diff --git a/include/in_mem_data_store.h b/include/in_mem_data_store.h index 5dcd1ce40..e637422ed 100644 --- a/include/in_mem_data_store.h +++ b/include/in_mem_data_store.h @@ -7,7 +7,7 @@ #include "tsl/robin_map.h" #include "tsl/robin_set.h" #include "tsl/sparse_map.h" -//#include "boost/dynamic_bitset.hpp" +// #include "boost/dynamic_bitset.hpp" #include "abstract_data_store.h" @@ -17,8 +17,7 @@ namespace diskann { -template -class InMemDataStore : public AbstractDataStore +template class InMemDataStore : public AbstractDataStore { public: InMemDataStore(const location_t capacity, const size_t dim, std::shared_ptr> distance_metric); @@ -32,13 +31,12 @@ class InMemDataStore : public AbstractDataStore return _aligned_dim; } - //Populate internal data from unaligned data while doing alignment and any normalization that is required. + // Populate internal data from unaligned data while doing alignment and any normalization that is required. virtual void populate_data(const data_t *vectors, const location_t num_pts) override; virtual void populate_data(const std::string &filename, const size_t offset) override; virtual void save_data_to_bin(const std::string &filename, const location_t num_pts) override; - virtual void get_vector(const location_t i, data_t *target) const override; virtual void set_vector(const location_t i, const data_t *const vector) override; virtual void prefetch_vector(const location_t loc) override; @@ -47,20 +45,17 @@ class InMemDataStore : public AbstractDataStore const location_t num_points) override; virtual void copy_points(const location_t from_loc, const location_t to_loc, const location_t num_points) override; - virtual float get_distance(const data_t *query, const location_t loc) const override; virtual float get_distance(const location_t loc1, const location_t loc2) const override; - virtual void get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, float *distances) const override ; + virtual void get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, + float *distances) const override; virtual location_t calculate_medoid() const override; - - protected: virtual void expand(const location_t new_size) override; virtual void shrink(const location_t new_size) override; - virtual location_t load_impl(const std::string &filename); #ifdef EXEC_ENV_OLS virtual location_t load_impl(AlignedFileReader &reader); @@ -74,10 +69,10 @@ class InMemDataStore : public AbstractDataStore // It may seem weird to put distance metric along with the data store class, but // this gives us perf benefits as the datastore can do distance computations during // search and compute norms of vectors internally without have to copy - // data back and forth. + // data back and forth. std::shared_ptr> _distance_fn; -// in case we need to save vector norms for optimization + // in case we need to save vector norms for optimization std::shared_ptr _pre_computed_norms; }; diff --git a/include/natural_number_map.h b/include/natural_number_map.h index 7678c2aa0..67986e24b 100644 --- a/include/natural_number_map.h +++ b/include/natural_number_map.h @@ -8,7 +8,7 @@ #include #include - //#include "boost_dynamic_bitset_fwd.h" +// #include "boost_dynamic_bitset_fwd.h" namespace diskann { diff --git a/src/distance.cpp b/src/distance.cpp index d0f446905..5872856b9 100644 --- a/src/distance.cpp +++ b/src/distance.cpp @@ -530,7 +530,8 @@ bool AVXNormalizedCosineDistanceFloat::normalization_required() const { return true; } -void AVXNormalizedCosineDistanceFloat::normalize_data_for_build(float *original_data, const uint32_t orig_dim, const uint32_t num_points) +void AVXNormalizedCosineDistanceFloat::normalize_data_for_build(float *original_data, const uint32_t orig_dim, + const uint32_t num_points) { for (auto i = 0; i < num_points; i++) { @@ -547,16 +548,15 @@ void AVXNormalizedCosineDistanceFloat::normalize_vector_for_search(const float * void AVXNormalizedCosineDistanceFloat::normalize_and_copy(const float *query_vec, const uint32_t query_dim, float *query_target) const { - + float norm = get_norm(query_vec, query_dim); - + for (auto i = 0; i < query_dim; i++) { query_target[i] = query_vec[i] / norm; } } - // Get the right distance function for the given metric. template <> diskann::Distance *get_distance_function(diskann::Metric m) { diff --git a/src/in_mem_data_store.cpp b/src/in_mem_data_store.cpp index fe8a57a50..972e72cfa 100644 --- a/src/in_mem_data_store.cpp +++ b/src/in_mem_data_store.cpp @@ -10,10 +10,9 @@ namespace diskann { template -InMemDataStore::InMemDataStore(const location_t num_points, - const size_t dim, std::shared_ptr> distance_metric) - : AbstractDataStore(num_points, dim), _aligned_dim(ROUND_UP(dim, 8)), - _distance_fn(distance_metric) +InMemDataStore::InMemDataStore(const location_t num_points, const size_t dim, + std::shared_ptr> distance_metric) + : AbstractDataStore(num_points, dim), _aligned_dim(ROUND_UP(dim, 8)), _distance_fn(distance_metric) { alloc_aligned(((void **)&_data), this->_capacity * _aligned_dim * sizeof(data_t), 8 * sizeof(data_t)); std::memset(_data, 0, this->_capacity * _aligned_dim * sizeof(data_t)); @@ -58,8 +57,7 @@ template location_t Index::load_impl(AlignedFileReader } #endif -template -location_t InMemDataStore::load_impl(const std::string &filename) +template location_t InMemDataStore::load_impl(const std::string &filename) { size_t file_dim, file_num_points; if (!file_exists(filename)) @@ -73,7 +71,7 @@ location_t InMemDataStore::load_impl(const std::string &filename) diskann::get_bin_metadata(filename, file_num_points, file_dim); // since we are loading a new dataset, _empty_slots must be cleared -// _empty_slots.clear(); + // _empty_slots.clear(); if (file_dim != this->_dim) { @@ -95,7 +93,6 @@ location_t InMemDataStore::load_impl(const std::string &filename) return file_num_points; } - template size_t InMemDataStore::save(const std::string &filename, const location_t num_points) { return save_data_in_base_dimensions(filename, _data, num_points, this->get_dims(), this->get_aligned_dim(), 0U); @@ -115,7 +112,7 @@ template void InMemDataStore::populate_data(const data } } -template void InMemDataStore::populate_data(const std::string &filename, const size_t offset) +template void InMemDataStore::populate_data(const std::string &filename, const size_t offset) { size_t npts, ndim; copy_aligned_data_from_file(filename.c_str(), _data, npts, ndim, _aligned_dim, offset); @@ -143,21 +140,18 @@ template void InMemDataStore::populate_data(const std:: } } -template void InMemDataStore::save_data_to_bin(const std::string &filename, const location_t num_points) +template +void InMemDataStore::save_data_to_bin(const std::string &filename, const location_t num_points) { save_data_in_base_dimensions(filename, _data, num_points, this->get_dims(), this->get_aligned_dim(), 0U); } - - -template -void InMemDataStore::get_vector(const location_t i, data_t* dest) const +template void InMemDataStore::get_vector(const location_t i, data_t *dest) const { memcpy(dest, _data + i * _aligned_dim, this->_dim * sizeof(data_t)); } -template -void InMemDataStore::set_vector(const location_t loc, const data_t *const vector) +template void InMemDataStore::set_vector(const location_t loc, const data_t *const vector) { size_t offset_in_data = loc * _aligned_dim; memset(_data + offset_in_data, 0, _aligned_dim * sizeof(data_t)); @@ -168,18 +162,19 @@ void InMemDataStore::set_vector(const location_t loc, const data_t *cons } } -template void InMemDataStore::prefetch_vector(const location_t loc) +template void InMemDataStore::prefetch_vector(const location_t loc) { diskann::prefetch_vector((const char *)_data + _aligned_dim * (size_t)loc, sizeof(data_t) * _aligned_dim); } -template float InMemDataStore::get_distance(const data_t *query, const location_t loc) const +template float InMemDataStore::get_distance(const data_t *query, const location_t loc) const { return _distance_fn->compare(query, _data + _aligned_dim * loc, _aligned_dim); } -template -void InMemDataStore::get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, float *distances) const +template +void InMemDataStore::get_distance(const data_t *query, const location_t *locations, + const uint32_t location_count, float *distances) const { for (auto i = 0; i < location_count; i++) { @@ -187,7 +182,7 @@ void InMemDataStore::get_distance(const data_t *query, const location_t } } -template +template float InMemDataStore::get_distance(const location_t loc1, const location_t loc2) const { return _distance_fn->compare(_data + loc1 * _aligned_dim, _data + loc2 * _aligned_dim, this->_aligned_dim); @@ -202,8 +197,8 @@ template void InMemDataStore::expand(const location_t else if (new_size < this->capacity()) { std::stringstream ss; - ss << "Cannot 'expand' datastore when new capacity (" << new_size << ") < existing capacity(" << this->capacity() - << ")" << std::endl; + ss << "Cannot 'expand' datastore when new capacity (" << new_size << ") < existing capacity(" + << this->capacity() << ")" << std::endl; throw diskann::ANNException(ss.str(), -1); } #ifndef _WINDOWS @@ -216,7 +211,6 @@ template void InMemDataStore::expand(const location_t realloc_aligned((void **)&_data, new_size * _aligned_dim * sizeof(data_t), 8 * sizeof(data_t)); #endif this->_capacity = new_size; - } template void InMemDataStore::shrink(const location_t new_size) @@ -228,8 +222,8 @@ template void InMemDataStore::shrink(const location_t else if (new_size > this->capacity()) { std::stringstream ss; - ss << "Cannot 'shrink' datastore when new capacity (" << new_size << ") > existing capacity(" << this->capacity() - << ")" << std::endl; + ss << "Cannot 'shrink' datastore when new capacity (" << new_size << ") > existing capacity(" + << this->capacity() << ")" << std::endl; throw diskann::ANNException(ss.str(), -1); } #ifndef _WINDOWS @@ -243,8 +237,9 @@ template void InMemDataStore::shrink(const location_t #endif } -template -void InMemDataStore::reposition_points(const location_t old_location_start, const location_t new_location_start, const location_t num_locations) +template +void InMemDataStore::reposition_points(const location_t old_location_start, const location_t new_location_start, + const location_t num_locations) { if (num_locations == 0 || old_location_start == new_location_start) { @@ -290,7 +285,7 @@ void InMemDataStore::reposition_points(const location_t old_location_sta sizeof(data_t) * _aligned_dim * (mem_clear_loc_end_limit - mem_clear_loc_start)); } -template +template void InMemDataStore::copy_points(const location_t from_loc, const location_t to_loc, const location_t num_points) { @@ -317,10 +312,10 @@ template location_t InMemDataStore::calculate_medoid() // compute all to one distance float *distances = new float[this->capacity()]; -//TODO: REFACTOR. Removing pragma might make this slow. Must revisit. -// Problem is that we need to pass num_threads here, it is not clear -// if data store must be aware of threads! -//#pragma omp parallel for schedule(static, 65536) + // TODO: REFACTOR. Removing pragma might make this slow. Must revisit. + // Problem is that we need to pass num_threads here, it is not clear + // if data store must be aware of threads! + // #pragma omp parallel for schedule(static, 65536) for (int64_t i = 0; i < (int64_t)this->capacity(); i++) { // extract point and distance reference @@ -355,5 +350,4 @@ template DISKANN_DLLEXPORT class InMemDataStore; template DISKANN_DLLEXPORT class InMemDataStore; template DISKANN_DLLEXPORT class InMemDataStore; - } // namespace diskann \ No newline at end of file diff --git a/src/in_mem_graph_store.cpp b/src/in_mem_graph_store.cpp index df9933080..cb83481d3 100644 --- a/src/in_mem_graph_store.cpp +++ b/src/in_mem_graph_store.cpp @@ -2,6 +2,7 @@ // Licensed under the MIT license. #include "in_mem_graph_store.h" +#include "utils.h" namespace diskann { diff --git a/src/index.cpp b/src/index.cpp index 533f16587..dba931e46 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -98,19 +98,19 @@ Index::Index(Metric m, const size_t dim, const size_t max_point std::memset(_pq_data, 0, total_internal_points * _num_pq_chunks * sizeof(char)); } - //REFACTOR - //alloc_aligned(((void **)&_data), total_internal_points * _aligned_dim * sizeof(T), 8 * sizeof(T)); - //std::memset(_data, 0, total_internal_points * _aligned_dim * sizeof(T)); + // REFACTOR + // alloc_aligned(((void **)&_data), total_internal_points * _aligned_dim * sizeof(T), 8 * sizeof(T)); + // std::memset(_data, 0, total_internal_points * _aligned_dim * sizeof(T)); _start = (uint32_t)_max_points; _final_graph.resize(total_internal_points); - //This should come from a factory. + // This should come from a factory. if (m == diskann::Metric::COSINE && std::is_floating_point::value) { // This is safe because T is float inside the if block. - this->_distance.reset((Distance*)new AVXNormalizedCosineDistanceFloat()); + this->_distance.reset((Distance *)new AVXNormalizedCosineDistanceFloat()); this->_normalize_vecs = true; diskann::cout << "Normalizing vectors and using L2 for cosine " "AVXNormalizedCosineDistanceFloat()." @@ -122,8 +122,8 @@ Index::Index(Metric m, const size_t dim, const size_t max_point } // REFACTOR: TODO This should move to a factory method. - _data_store = std::make_shared>((location_t)total_internal_points, _dim, this->_distance); - + _data_store = + std::make_shared>((location_t)total_internal_points, _dim, this->_distance); _locks = std::vector(total_internal_points); @@ -147,17 +147,17 @@ template Index::~I LockGuard lg(lock); } - //if (this->_distance != nullptr) + // if (this->_distance != nullptr) //{ - // delete this->_distance; - // this->_distance = nullptr; - //} - //REFACTOR - //if (this->_data != nullptr) + // delete this->_distance; + // this->_distance = nullptr; + // } + // REFACTOR + // if (this->_data != nullptr) //{ - // aligned_free(this->_data); - // this->_data = nullptr; - //} + // aligned_free(this->_data); + // this->_data = nullptr; + // } if (_opt_graph != nullptr) { delete[] _opt_graph; @@ -226,7 +226,7 @@ template size_t Indexsave(data_file, _nd + _num_frozen_pts); // REFACTOR - //return save_data_in_base_dimensions(data_file, _data, _nd + _num_frozen_pts, _dim, _aligned_dim); + // return save_data_in_base_dimensions(data_file, _data, _nd + _num_frozen_pts, _dim, _aligned_dim); } // save the graph index on a file as an adjacency list. For each point, @@ -452,8 +452,8 @@ size_t Index::load_data(std::string filename) std::stringstream stream; stream << "ERROR: data file " << filename << " does not exist." << std::endl; diskann::cerr << stream.str() << std::endl; - //REFACTOR - //aligned_free(_data); + // REFACTOR + // aligned_free(_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } diskann::get_bin_metadata(filename, file_num_points, file_dim); @@ -468,8 +468,8 @@ size_t Index::load_data(std::string filename) stream << "ERROR: Driver requests loading " << _dim << " dimension," << "but file has " << file_dim << " dimension." << std::endl; diskann::cerr << stream.str() << std::endl; - //REFACTOR - //aligned_free(_data); + // REFACTOR + // aligned_free(_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } @@ -481,12 +481,12 @@ size_t Index::load_data(std::string filename) #ifdef EXEC_ENV_OLS - //REFACTOR TODO: Must figure out how to support aligned reader in a clean manner. + // REFACTOR TODO: Must figure out how to support aligned reader in a clean manner. copy_aligned_data_from_file(reader, _data, file_num_points, file_dim, _aligned_dim); #else - _data_store->populate_data(filename, 0U); //offset == 0. - //REFACTOR - //copy_aligned_data_from_file(filename.c_str(), _data, file_num_points, file_dim, _aligned_dim); + _data_store->populate_data(filename, 0U); // offset == 0. + // REFACTOR + // copy_aligned_data_from_file(filename.c_str(), _data, file_num_points, file_dim, _aligned_dim); #endif return file_num_points; } @@ -576,8 +576,8 @@ void Index::load(const char *filename, uint32_t num_threads, ui << 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; diskann::cerr << stream.str() << std::endl; - //REFACTOR - //aligned_free(_data); + // REFACTOR + // aligned_free(_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } @@ -729,8 +729,8 @@ size_t Index::load_graph(std::string filename, size_t expected_ << std::endl; } diskann::cerr << stream.str() << std::endl; - //REFACTOR - //aligned_free(_data); + // REFACTOR + // aligned_free(_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } @@ -824,9 +824,9 @@ template int Index location_t location = _tag_to_location[tag]; _data_store->get_vector(location, vec); - //REFACTOR - //size_t location = _tag_to_location[tag]; - //memcpy((void *)vec, (void *)(_data + location * _aligned_dim), (size_t)_dim * sizeof(T)); + // REFACTOR + // size_t location = _tag_to_location[tag]; + // memcpy((void *)vec, (void *)(_data + location * _aligned_dim), (size_t)_dim * sizeof(T)); return 0; } @@ -839,54 +839,54 @@ template uint32_t Indexcalculate_medoid(); - -// //REFACTOR -// // allocate and init centroid -// float *center = new float[_aligned_dim](); -// for (size_t j = 0; j < _aligned_dim; j++) -// center[j] = 0; -// -// for (size_t i = 0; i < _nd; i++) -// for (size_t j = 0; j < _aligned_dim; j++) -// center[j] += (float)_data[i * _aligned_dim + j]; -// -// for (size_t j = 0; j < _aligned_dim; j++) -// center[j] /= (float)_nd; -// -// // compute all to one distance -// float *distances = new float[_nd](); -//#pragma omp parallel for schedule(static, 65536) -// for (int64_t i = 0; i < (int64_t)_nd; i++) -// { -// // extract point and distance reference -// float &dist = distances[i]; -// const T *cur_vec = _data + (i * (size_t)_aligned_dim); -// dist = 0; -// float diff = 0; -// for (size_t j = 0; j < _aligned_dim; j++) -// { -// diff = (center[j] - (float)cur_vec[j]) * (center[j] - (float)cur_vec[j]); -// dist += diff; -// } -// } -// // find imin -// uint32_t min_idx = 0; -// float min_dist = distances[0]; -// for (uint32_t i = 1; i < _nd; i++) -// { -// if (distances[i] < min_dist) -// { -// min_idx = i; -// min_dist = distances[i]; -// } -// } -// -// delete[] distances; -// delete[] center; -// return min_idx; + + // //REFACTOR + // // allocate and init centroid + // float *center = new float[_aligned_dim](); + // for (size_t j = 0; j < _aligned_dim; j++) + // center[j] = 0; + // + // for (size_t i = 0; i < _nd; i++) + // for (size_t j = 0; j < _aligned_dim; j++) + // center[j] += (float)_data[i * _aligned_dim + j]; + // + // for (size_t j = 0; j < _aligned_dim; j++) + // center[j] /= (float)_nd; + // + // // compute all to one distance + // float *distances = new float[_nd](); + // #pragma omp parallel for schedule(static, 65536) + // for (int64_t i = 0; i < (int64_t)_nd; i++) + // { + // // extract point and distance reference + // float &dist = distances[i]; + // const T *cur_vec = _data + (i * (size_t)_aligned_dim); + // dist = 0; + // float diff = 0; + // for (size_t j = 0; j < _aligned_dim; j++) + // { + // diff = (center[j] - (float)cur_vec[j]) * (center[j] - (float)cur_vec[j]); + // dist += diff; + // } + // } + // // find imin + // uint32_t min_idx = 0; + // float min_dist = distances[0]; + // for (uint32_t i = 1; i < _nd; i++) + // { + // if (distances[i] < min_dist) + // { + // min_idx = i; + // min_dist = distances[i]; + // } + // } + // + // delete[] distances; + // delete[] center; + // return min_idx; } template std::vector Index::get_init_ids() @@ -1137,15 +1137,15 @@ std::pair Index::iterate_to_fixed_point( { auto nextn = id_scratch[m + 1]; _data_store->prefetch_vector(nextn); - //REFACTOR - //diskann::prefetch_vector((const char *)_data + _aligned_dim * (size_t)nextn, - // sizeof(T) * _aligned_dim); + // REFACTOR + // diskann::prefetch_vector((const char *)_data + _aligned_dim * (size_t)nextn, + // sizeof(T) * _aligned_dim); } dist_scratch.push_back(_data_store->get_distance(aligned_query, id)); - //REFACTOR - //dist_scratch.push_back( - // _distance->compare(aligned_query, _data + _aligned_dim * (size_t)id, (uint32_t)_aligned_dim)); + // REFACTOR + // dist_scratch.push_back( + // _distance->compare(aligned_query, _data + _aligned_dim * (size_t)id, (uint32_t)_aligned_dim)); } } cmps += id_scratch.size(); @@ -1173,9 +1173,10 @@ void Index::search_for_point_and_prune(int location, uint32_t L _data_store->get_vector(location, scratch->aligned_query()); iterate_to_fixed_point(scratch->aligned_query(), Lindex, init_ids, scratch, false, unused_filter_label, false); - //REFACTOR - //iterate_to_fixed_point(_data + _aligned_dim * location, Lindex, init_ids, scratch, false, unused_filter_label, - // false); + // REFACTOR + // iterate_to_fixed_point(_data + _aligned_dim * location, Lindex, init_ids, scratch, false, + // unused_filter_label, + // false); } else { @@ -1184,11 +1185,11 @@ void Index::search_for_point_and_prune(int location, uint32_t L filter_specific_start_nodes.emplace_back(_label_to_medoid_id[x]); _data_store->get_vector(location, scratch->aligned_query()); - iterate_to_fixed_point(scratch->aligned_query(), filteredLindex, filter_specific_start_nodes, scratch, - true, _pts_to_labels[location], false); - //REFACTOR - //iterate_to_fixed_point(_data + _aligned_dim * location, filteredLindex, filter_specific_start_nodes, scratch, - // true, _pts_to_labels[location], false); + iterate_to_fixed_point(scratch->aligned_query(), filteredLindex, filter_specific_start_nodes, scratch, true, + _pts_to_labels[location], false); + // REFACTOR + // iterate_to_fixed_point(_data + _aligned_dim * location, filteredLindex, filter_specific_start_nodes, scratch, + // true, _pts_to_labels[location], false); } auto &pool = scratch->pool(); @@ -1285,9 +1286,9 @@ void Index::occlude_list(const uint32_t location, std::vectorget_distance(iter2->id, iter->id); - //REFACTOR - //float djk = _distance->compare(_data + _aligned_dim * (size_t)iter2->id, - // _data + _aligned_dim * (size_t)iter->id, (uint32_t)_aligned_dim); + // REFACTOR + // float djk = _distance->compare(_data + _aligned_dim * (size_t)iter2->id, + // _data + _aligned_dim * (size_t)iter->id, (uint32_t)_aligned_dim); if (_dist_metric == diskann::Metric::L2 || _dist_metric == diskann::Metric::COSINE) { occlude_factor[t] = (djk == 0) ? std::numeric_limits::max() @@ -1335,9 +1336,9 @@ void Index::prune_neighbors(const uint32_t location, std::vecto { for (auto &ngh : pool) ngh.distance = _data_store->get_distance(ngh.id, location); - //REFACTOR - //ngh.distance = _distance->compare(_data + _aligned_dim * (size_t)ngh.id, - // _data + _aligned_dim * (size_t)location, (uint32_t)_aligned_dim); + // REFACTOR + // ngh.distance = _distance->compare(_data + _aligned_dim * (size_t)ngh.id, + // _data + _aligned_dim * (size_t)location, (uint32_t)_aligned_dim); } // sort the pool based on distance to query and prune it with occlude_list @@ -1409,9 +1410,9 @@ void Index::inter_insert(uint32_t n, std::vector &pru if (dummy_visited.find(cur_nbr) == dummy_visited.end() && cur_nbr != des) { float dist = _data_store->get_distance(des, cur_nbr); - //REFACTOR - //float dist = _distance->compare(_data + _aligned_dim * (size_t)des, - // _data + _aligned_dim * (size_t)cur_nbr, (uint32_t)_aligned_dim); + // REFACTOR + // float dist = _distance->compare(_data + _aligned_dim * (size_t)des, + // _data + _aligned_dim * (size_t)cur_nbr, (uint32_t)_aligned_dim); dummy_pool.emplace_back(Neighbor(cur_nbr, dist)); dummy_visited.insert(cur_nbr); } @@ -1536,9 +1537,9 @@ void Index::link(IndexWriteParameters ¶meters) if (dummy_visited.find(cur_nbr) == dummy_visited.end() && cur_nbr != node) { float dist = _data_store->get_distance(node, cur_nbr); - //REFACTOR - //float dist = _distance->compare(_data + _aligned_dim * (size_t)node, - // _data + _aligned_dim * (size_t)cur_nbr, (uint32_t)_aligned_dim); + // REFACTOR + // float dist = _distance->compare(_data + _aligned_dim * (size_t)node, + // _data + _aligned_dim * (size_t)cur_nbr, (uint32_t)_aligned_dim); dummy_pool.emplace_back(Neighbor(cur_nbr, dist)); dummy_visited.insert(cur_nbr); } @@ -1563,9 +1564,9 @@ void Index::prune_all_neighbors(const uint32_t max_degree, cons const uint32_t range = max_degree; const uint32_t maxc = max_occlusion_size; -// const uint32_t range = parameters.Get("R"); -// const uint32_t maxc = parameters.Get("C"); -// const float alpha = parameters.Get("alpha"); + // const uint32_t range = parameters.Get("R"); + // const uint32_t maxc = parameters.Get("C"); + // const float alpha = parameters.Get("alpha"); _filtered_index = true; @@ -1589,9 +1590,10 @@ void Index::prune_all_neighbors(const uint32_t max_degree, cons if (dummy_visited.find(cur_nbr) == dummy_visited.end() && cur_nbr != node) { float dist = _data_store->get_distance(node, cur_nbr); - //REFACTOR - //float dist = _distance->compare(_data + _aligned_dim * (size_t)node, - // _data + _aligned_dim * (size_t)cur_nbr, (uint32_t)_aligned_dim); + // REFACTOR + // float dist = _distance->compare(_data + _aligned_dim * (size_t)node, + // _data + _aligned_dim * (size_t)cur_nbr, + // (uint32_t)_aligned_dim); dummy_pool.emplace_back(Neighbor(cur_nbr, dist)); dummy_visited.insert(cur_nbr); } @@ -1629,58 +1631,57 @@ void Index::prune_all_neighbors(const uint32_t max_degree, cons } } -//REFACTOR -//template -//void Index::set_start_points(const T *data, size_t data_count) +// REFACTOR +// template +// void Index::set_start_points(const T *data, size_t data_count) //{ -// std::unique_lock ul(_update_lock); -// std::unique_lock tl(_tag_lock); -// if (_nd > 0) -// throw ANNException("Can not set starting point for a non-empty index", -1, __FUNCSIG__, __FILE__, __LINE__); +// std::unique_lock ul(_update_lock); +// std::unique_lock tl(_tag_lock); +// if (_nd > 0) +// throw ANNException("Can not set starting point for a non-empty index", -1, __FUNCSIG__, __FILE__, __LINE__); // -// if (data_count != _num_frozen_pts * _aligned_dim) -// throw ANNException("Invalid number of points", -1, __FUNCSIG__, __FILE__, __LINE__); +// if (data_count != _num_frozen_pts * _aligned_dim) +// throw ANNException("Invalid number of points", -1, __FUNCSIG__, __FILE__, __LINE__); // -// memcpy(_data + _aligned_dim * _max_points, data, _aligned_dim * sizeof(T) * _num_frozen_pts); -// _has_built = true; -// diskann::cout << "Index start points set: #" << _num_frozen_pts << std::endl; -//} +// memcpy(_data + _aligned_dim * _max_points, data, _aligned_dim * sizeof(T) * _num_frozen_pts); +// _has_built = true; +// diskann::cout << "Index start points set: #" << _num_frozen_pts << std::endl; +// } // // REFACTOR: added dummy implementation for now. template -void Index::set_start_points_at_random(T radius, uint32_t random_seed) { - +void Index::set_start_points_at_random(T radius, uint32_t random_seed) +{ } - -//template -//void Index::set_start_points_at_random(T radius, uint32_t random_seed) +// template +// void Index::set_start_points_at_random(T radius, uint32_t random_seed) //{ -// std::mt19937 gen{random_seed}; -// std::normal_distribution<> d{0.0, 1.0}; +// std::mt19937 gen{random_seed}; +// std::normal_distribution<> d{0.0, 1.0}; // -// std::vector points_data; -// points_data.reserve(_aligned_dim * _num_frozen_pts); -// std::vector real_vec(_aligned_dim); +// std::vector points_data; +// points_data.reserve(_aligned_dim * _num_frozen_pts); +// std::vector real_vec(_aligned_dim); // -// for (size_t frozen_point = 0; frozen_point < _num_frozen_pts; frozen_point++) -// { -// double norm_sq = 0.0; -// for (size_t i = 0; i < _dim; ++i) -// { -// auto r = d(gen); -// real_vec[i] = r; -// norm_sq += r * r; -// } +// for (size_t frozen_point = 0; frozen_point < _num_frozen_pts; frozen_point++) +// { +// double norm_sq = 0.0; +// for (size_t i = 0; i < _dim; ++i) +// { +// auto r = d(gen); +// real_vec[i] = r; +// norm_sq += r * r; +// } // -// const double norm = std::sqrt(norm_sq); -// for (auto iter : real_vec) -// points_data.push_back(static_cast(iter * radius / norm)); -// } +// const double norm = std::sqrt(norm_sq); +// for (auto iter : real_vec) +// points_data.push_back(static_cast(iter * radius / norm)); +// } // -// set_start_points(points_data.data(), points_data.size()); -//} +// set_start_points(points_data.data(), points_data.size()); +// } template void Index::build_with_data_populated(IndexWriteParameters ¶meters, const std::vector &tags) @@ -1696,8 +1697,8 @@ void Index::build_with_data_populated(IndexWriteParameters &par stream << "ERROR: Driver requests loading " << _nd << " points from file," << "but tags vector is of size " << tags.size() << "." << std::endl; diskann::cerr << stream.str() << std::endl; - //REFACTOR - //aligned_free(_data); + // REFACTOR + // aligned_free(_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } if (_enable_tags) @@ -1716,10 +1717,11 @@ void Index::build_with_data_populated(IndexWriteParameters &par if (_query_scratch.size() == 0) { - //REFACTOR - //initialize_query_scratch(5 + num_threads_index, index_L, index_L, index_R, maxc, _aligned_dim); - //Comment: Why 5+? - initialize_query_scratch(5 + num_threads_index, index_L, index_L, index_R, maxc, _data_store->get_aligned_dim()); + // REFACTOR + // initialize_query_scratch(5 + num_threads_index, index_L, index_L, index_R, maxc, _aligned_dim); + // Comment: Why 5+? + initialize_query_scratch(5 + num_threads_index, index_L, index_L, index_R, maxc, + _data_store->get_aligned_dim()); } generate_frozen_point(); @@ -1764,15 +1766,15 @@ void Index::build(const T *data, const size_t num_points_to_loa _data_store->populate_data(data, num_points_to_load); - //REFACTOR - //memcpy((char *)_data, (char *)data, _aligned_dim * _nd * sizeof(T)); - //if (_normalize_vecs) + // REFACTOR + // memcpy((char *)_data, (char *)data, _aligned_dim * _nd * sizeof(T)); + // if (_normalize_vecs) //{ - // for (size_t i = 0; i < num_points_to_load; i++) - // { - // normalize(_data + _aligned_dim * i, _aligned_dim); - // } - //} + // for (size_t i = 0; i < num_points_to_load; i++) + // { + // normalize(_data + _aligned_dim * i, _aligned_dim); + // } + // } } build_with_data_populated(parameters, tags); @@ -1810,9 +1812,9 @@ void Index::build(const char *filename, const size_t num_points if (_pq_dist) aligned_free(_pq_data); - //REFACTOR - //else - // aligned_free(_data); + // REFACTOR + // else + // aligned_free(_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } @@ -1824,9 +1826,9 @@ void Index::build(const char *filename, const size_t num_points if (_pq_dist) aligned_free(_pq_data); - //REFACTOR - //else - // aligned_free(_data); + // REFACTOR + // else + // aligned_free(_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } @@ -1839,9 +1841,9 @@ void Index::build(const char *filename, const size_t num_points if (_pq_dist) aligned_free(_pq_data); - //REFACTOR - //else - // aligned_free(_data); + // REFACTOR + // else + // aligned_free(_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } @@ -1868,15 +1870,15 @@ void Index::build(const char *filename, const size_t num_points } _data_store->populate_data(filename, 0U); - //REFACTOR - //copy_aligned_data_from_file(filename, _data, file_num_points, file_dim, _aligned_dim); - //if (_normalize_vecs) + // REFACTOR + // copy_aligned_data_from_file(filename, _data, file_num_points, file_dim, _aligned_dim); + // if (_normalize_vecs) //{ - // for (size_t i = 0; i < file_num_points; i++) - // { - // normalize(_data + _aligned_dim * i, _aligned_dim); - // } - //} + // for (size_t i = 0; i < file_num_points; i++) + // { + // normalize(_data + _aligned_dim * i, _aligned_dim); + // } + // } diskann::cout << "Using only first " << num_points_to_load << " from file.. " << std::endl; @@ -2276,8 +2278,8 @@ 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]); - //REFACTOR - //memcpy(res_vectors[pos], _data + ((size_t)node.id) * _aligned_dim, _dim * sizeof(T)); + // REFACTOR + // memcpy(res_vectors[pos], _data + ((size_t)node.id) * _aligned_dim, _dim * sizeof(T)); } if (distances != nullptr) @@ -2337,8 +2339,8 @@ template void Indexcopy_points(res, _max_points, 1); - //REFACTOR - //memcpy(_data + _max_points * _aligned_dim, _data + res * _aligned_dim, _aligned_dim * sizeof(T)); + // REFACTOR + // memcpy(_data + _max_points * _aligned_dim, _data + res * _aligned_dim, _aligned_dim * sizeof(T)); } } @@ -2424,10 +2426,11 @@ inline void Index::process_delete(const tsl::robin_setget_distance(loc, ngh)); - //REFACTOR - //expanded_nghrs_vec.emplace_back( - // ngh, - // _distance->compare(_data + _aligned_dim * loc, _data + _aligned_dim * ngh, (uint32_t)_aligned_dim)); + // REFACTOR + // expanded_nghrs_vec.emplace_back( + // ngh, + // _distance->compare(_data + _aligned_dim * loc, _data + _aligned_dim * ngh, + // (uint32_t)_aligned_dim)); } std::sort(expanded_nghrs_vec.begin(), expanded_nghrs_vec.end()); std::vector &occlude_list_output = scratch->occlude_list_output(); @@ -2631,9 +2634,9 @@ template void Indexcopy_points(old, new_location[old], 1); - //REFACTOR - //memcpy((void *)(_data + _aligned_dim * (size_t)new_location[old]), - // (void *)(_data + _aligned_dim * (size_t)old), _aligned_dim * sizeof(T)); + // REFACTOR + // memcpy((void *)(_data + _aligned_dim * (size_t)new_location[old]), + // (void *)(_data + _aligned_dim * (size_t)old), _aligned_dim * sizeof(T)); } } else @@ -2793,7 +2796,6 @@ void Index::reposition_points(uint32_t old_location_start, uint } } _data_store->reposition_points(old_location_start, new_location_start, num_locations); - } template void Index::reposition_frozen_point_to_end() @@ -2807,7 +2809,6 @@ template void Index void Indexresize(new_internal_points); - //REFACTOR -//#ifndef _WINDOWS -// T *new_data; -// alloc_aligned((void **)&new_data, new_internal_points * _aligned_dim * sizeof(T), 8 * sizeof(T)); -// memcpy(new_data, _data, (_max_points + _num_frozen_pts) * _aligned_dim * sizeof(T)); -// aligned_free(_data); -// _data = new_data; -//#else -// realloc_aligned((void **)&_data, new_internal_points * _aligned_dim * sizeof(T), 8 * sizeof(T)); -//#endif + // REFACTOR + // #ifndef _WINDOWS + // T *new_data; + // alloc_aligned((void **)&new_data, new_internal_points * _aligned_dim * sizeof(T), 8 * sizeof(T)); + // memcpy(new_data, _data, (_max_points + _num_frozen_pts) * _aligned_dim * sizeof(T)); + // aligned_free(_data); + // _data = new_data; + // #else + // realloc_aligned((void **)&_data, new_internal_points * _aligned_dim * sizeof(T), 8 * sizeof(T)); + // #endif _final_graph.resize(new_internal_points); _locks = std::vector(new_internal_points); @@ -2922,16 +2923,16 @@ int Index::insert_point(const T *point, const TagT tag) tl.unlock(); _data_store->set_vector(location, point); - //REFACTOR - // Copy the vector in to the data array - //auto offset_data = _data + (size_t)_aligned_dim * location; - //memset((void *)offset_data, 0, sizeof(T) * _aligned_dim); - //memcpy((void *)offset_data, point, sizeof(T) * _dim); + // REFACTOR + // Copy the vector in to the data array + // auto offset_data = _data + (size_t)_aligned_dim * location; + // memset((void *)offset_data, 0, sizeof(T) * _aligned_dim); + // memcpy((void *)offset_data, point, sizeof(T) * _dim); - //if (_normalize_vecs) + // if (_normalize_vecs) //{ - // normalize((float *)offset_data, _dim); - //} + // normalize((float *)offset_data, _dim); + // } // Find and add appropriate graph edges ScratchStoreManager> manager(_query_scratch); @@ -3099,129 +3100,129 @@ template void Index void Index::optimize_index_layout() { // use after build or load } -//REFACTOR: This should be an OptimizedDataStore class -//template void Index::optimize_index_layout() +// REFACTOR: This should be an OptimizedDataStore class +// template void Index::optimize_index_layout() //{ // use after build or load -// if (_dynamic_index) -// { -// throw diskann::ANNException("Optimize_index_layout not implemented for dyanmic indices", -1, __FUNCSIG__, -// __FILE__, __LINE__); -// } +// if (_dynamic_index) +// { +// throw diskann::ANNException("Optimize_index_layout not implemented for dyanmic indices", -1, __FUNCSIG__, +// __FILE__, __LINE__); +// } // -// _data_len = (_aligned_dim + 1) * sizeof(float); -// _neighbor_len = (_max_observed_degree + 1) * sizeof(uint32_t); -// _node_size = _data_len + _neighbor_len; -// _opt_graph = new char[_node_size * _nd]; -// DistanceFastL2 *dist_fast = (DistanceFastL2 *)_distance; -// for (uint32_t 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)); +// _data_len = (_aligned_dim + 1) * sizeof(float); +// _neighbor_len = (_max_observed_degree + 1) * sizeof(uint32_t); +// _node_size = _data_len + _neighbor_len; +// _opt_graph = new char[_node_size * _nd]; +// DistanceFastL2 *dist_fast = (DistanceFastL2 *)_distance; +// for (uint32_t 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; -// uint32_t k = _final_graph[i].size(); -// std::memcpy(cur_node_offset, &k, sizeof(uint32_t)); -// std::memcpy(cur_node_offset + sizeof(uint32_t), _final_graph[i].data(), k * sizeof(uint32_t)); -// std::vector().swap(_final_graph[i]); -// } -// _final_graph.clear(); -// _final_graph.shrink_to_fit(); -//} -// REFACTOR: once optimized layout becomes its own Data+Graph store, we should just invoke regular search +// cur_node_offset += _data_len; +// uint32_t k = _final_graph[i].size(); +// std::memcpy(cur_node_offset, &k, sizeof(uint32_t)); +// std::memcpy(cur_node_offset + sizeof(uint32_t), _final_graph[i].data(), k * sizeof(uint32_t)); +// std::vector().swap(_final_graph[i]); +// } +// _final_graph.clear(); +// _final_graph.shrink_to_fit(); +// } +// REFACTOR: once optimized layout becomes its own Data+Graph store, we should just invoke regular search template void Index::search_with_optimized_layout(const T *query, size_t K, size_t L, uint32_t *indices) { } -//template -//void Index::search_with_optimized_layout(const T *query, size_t K, size_t L, uint32_t *indices) +// template +// void Index::search_with_optimized_layout(const T *query, size_t K, size_t L, uint32_t *indices) //{ -// DistanceFastL2 *dist_fast = (DistanceFastL2 *)_distance; +// DistanceFastL2 *dist_fast = (DistanceFastL2 *)_distance; // -// NeighborPriorityQueue retset(L); -// std::vector init_ids(L); +// NeighborPriorityQueue retset(L); +// std::vector init_ids(L); // -// boost::dynamic_bitset<> flags{_nd, 0}; -// uint32_t tmp_l = 0; -// uint32_t *neighbors = (uint32_t *)(_opt_graph + _node_size * _start + _data_len); -// uint32_t MaxM_ep = *neighbors; -// neighbors++; +// boost::dynamic_bitset<> flags{_nd, 0}; +// uint32_t tmp_l = 0; +// uint32_t *neighbors = (uint32_t *)(_opt_graph + _node_size * _start + _data_len); +// uint32_t 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; -// } +// 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) -// { -// uint32_t id = rand() % _nd; -// if (flags[id]) -// continue; -// flags[id] = true; -// init_ids[tmp_l] = id; -// tmp_l++; -// } +// while (tmp_l < L) +// { +// uint32_t id = rand() % _nd; +// if (flags[id]) +// continue; +// flags[id] = true; +// init_ids[tmp_l] = id; +// tmp_l++; +// } // -// for (uint32_t i = 0; i < init_ids.size(); i++) -// { -// uint32_t id = init_ids[i]; -// if (id >= _nd) -// continue; -// _mm_prefetch(_opt_graph + _node_size * id, _MM_HINT_T0); -// } -// L = 0; -// for (uint32_t i = 0; i < init_ids.size(); i++) -// { -// uint32_t 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, (uint32_t)_aligned_dim); -// retset.insert(Neighbor(id, dist)); -// flags[id] = true; -// L++; -// } +// for (uint32_t i = 0; i < init_ids.size(); i++) +// { +// uint32_t id = init_ids[i]; +// if (id >= _nd) +// continue; +// _mm_prefetch(_opt_graph + _node_size * id, _MM_HINT_T0); +// } +// L = 0; +// for (uint32_t i = 0; i < init_ids.size(); i++) +// { +// uint32_t 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, (uint32_t)_aligned_dim); +// retset.insert(Neighbor(id, dist)); +// flags[id] = true; +// L++; +// } // -// while (retset.has_unexpanded_node()) -// { -// auto nbr = retset.closest_unexpanded(); -// auto n = nbr.id; -// _mm_prefetch(_opt_graph + _node_size * n + _data_len, _MM_HINT_T0); -// neighbors = (uint32_t *)(_opt_graph + _node_size * n + _data_len); -// uint32_t MaxM = *neighbors; -// neighbors++; -// for (uint32_t m = 0; m < MaxM; ++m) -// _mm_prefetch(_opt_graph + _node_size * neighbors[m], _MM_HINT_T0); -// for (uint32_t m = 0; m < MaxM; ++m) -// { -// uint32_t id = neighbors[m]; -// if (flags[id]) -// continue; -// flags[id] = 1; -// T *data = (T *)(_opt_graph + _node_size * id); -// float norm = *data; -// data++; -// float dist = dist_fast->compare(query, data, norm, (uint32_t)_aligned_dim); -// Neighbor nn(id, dist); -// retset.insert(nn); -// } -// } +// while (retset.has_unexpanded_node()) +// { +// auto nbr = retset.closest_unexpanded(); +// auto n = nbr.id; +// _mm_prefetch(_opt_graph + _node_size * n + _data_len, _MM_HINT_T0); +// neighbors = (uint32_t *)(_opt_graph + _node_size * n + _data_len); +// uint32_t MaxM = *neighbors; +// neighbors++; +// for (uint32_t m = 0; m < MaxM; ++m) +// _mm_prefetch(_opt_graph + _node_size * neighbors[m], _MM_HINT_T0); +// for (uint32_t m = 0; m < MaxM; ++m) +// { +// uint32_t id = neighbors[m]; +// if (flags[id]) +// continue; +// flags[id] = 1; +// T *data = (T *)(_opt_graph + _node_size * id); +// float norm = *data; +// data++; +// float dist = dist_fast->compare(query, data, norm, (uint32_t)_aligned_dim); +// Neighbor nn(id, dist); +// retset.insert(nn); +// } +// } // -// for (size_t i = 0; i < K; i++) -// { -// indices[i] = retset[i].id; -// } -//} +// for (size_t i = 0; i < K; i++) +// { +// indices[i] = retset[i].id; +// } +// } /* Internals of the library */ template const float Index::INDEX_GROWTH_FACTOR = 1.5f; From f7e3424dcf65bb631d91562aa7274e03aa6f172c Mon Sep 17 00:00:00 2001 From: yashpatel007 Date: Mon, 10 Apr 2023 15:34:50 -0400 Subject: [PATCH 32/69] clang format fix --- include/distance.h | 53 +++++++++++++++++++++++----------------------- include/index.h | 4 ++-- 2 files changed, 28 insertions(+), 29 deletions(-) diff --git a/include/distance.h b/include/distance.h index 7a92a03c1..069721a04 100644 --- a/include/distance.h +++ b/include/distance.h @@ -15,14 +15,14 @@ enum Metric template class Distance { public: - Distance(diskann::Metric dist_metric) : _distance_metric(dist_metric) + Distance(diskann::Metric dist_metric) : _distance_metric(dist_metric) { } // distance comparison function virtual float compare(const T *a, const T *b, uint32_t length) const = 0; - //Needed only for COSINE-BYTE and INNER_PRODUCT-BYTE + // Needed only for COSINE-BYTE and INNER_PRODUCT-BYTE virtual float compare(const T *a, const T *b, const float normA, const float normB, uint32_t length) const { return std::numeric_limits::max(); @@ -42,40 +42,38 @@ template class Distance } // This is for efficiency. If no normalization is required, the callers - //can simply ignore the normalize_data_for_build() function. + // can simply ignore the normalize_data_for_build() function. virtual bool normalization_required() const { return false; } - - //Check the normalization_required() function before calling this. - //Clients can call the function like this: - // - // if (metric->normalization_required()){ - // T* normalized_data_batch; - // Split data into batches of batch_size and for each, call: - // metric->normalize_data_for_build(data_batch, batch_size); + + // Check the normalization_required() function before calling this. + // Clients can call the function like this: + // + // if (metric->normalization_required()){ + // T* normalized_data_batch; + // Split data into batches of batch_size and for each, call: + // metric->normalize_data_for_build(data_batch, batch_size); // - // TODO: This does not take into account the case for SSD inner product - // where the dimensions change after normalization. + // TODO: This does not take into account the case for SSD inner product + // where the dimensions change after normalization. // - virtual void normalize_data_for_build(T *original_data, const uint32_t orig_dim, - const uint32_t num_points) + virtual void normalize_data_for_build(T *original_data, const uint32_t orig_dim, const uint32_t num_points) { } - //Invokes normalization for a single vector during search. The scratch space - //has to be created by the caller keeping track of the fact that normalization - //might change the dimension of the query vector. - virtual void normalize_vector_for_search(const T *query_vec, const uint32_t query_dim, - T* scratch_query) + // Invokes normalization for a single vector during search. The scratch space + // has to be created by the caller keeping track of the fact that normalization + // might change the dimension of the query vector. + virtual void normalize_vector_for_search(const T *query_vec, const uint32_t query_dim, T *scratch_query) { std::memcpy(scratch_query, query_vec, query_dim * sizeof(T)); } - //Providing a default implementation for the virtual destructor because we don't - //expect most metric implementations to need it. - virtual ~Distance() + // Providing a default implementation for the virtual destructor because we don't + // expect most metric implementations to need it. + virtual ~Distance() { } @@ -126,7 +124,7 @@ class DistanceL2Float : public Distance DistanceL2Float() : Distance(diskann::Metric::L2) { } - + #ifdef _WINDOWS DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t size) const; #else @@ -195,7 +193,7 @@ template class DistanceInnerProduct : public Distance DistanceInnerProduct() : Distance(diskann::Metric::INNER_PRODUCT) { } - + DistanceInnerProduct(diskann::Metric metric) : Distance(metric) { } @@ -238,7 +236,7 @@ class AVXNormalizedCosineDistanceFloat : public Distance AVXDistanceInnerProductFloat _innerProduct; protected: - void normalize_and_copy(const float *a, uint32_t length, float *a_norm ) const; + void normalize_and_copy(const float *a, uint32_t length, float *a_norm) const; public: AVXNormalizedCosineDistanceFloat() : Distance(diskann::Metric::COSINE) @@ -254,7 +252,8 @@ class AVXNormalizedCosineDistanceFloat : public Distance DISKANN_DLLEXPORT virtual bool normalization_required() const; - DISKANN_DLLEXPORT virtual void normalize_data_for_build(float *original_data, const uint32_t orig_dim, const uint32_t num_points) override; + DISKANN_DLLEXPORT virtual void normalize_data_for_build(float *original_data, const uint32_t orig_dim, + const uint32_t num_points) override; DISKANN_DLLEXPORT virtual void normalize_vector_for_search(const float *query_vec, const uint32_t query_dim, float *scratch_query_vector) override; diff --git a/include/index.h b/include/index.h index 22ebe6cea..0ef97825f 100644 --- a/include/index.h +++ b/include/index.h @@ -316,7 +316,7 @@ template clas std::shared_ptr> _distance; // Data - //T *_data = nullptr; + // T *_data = nullptr; std::shared_ptr> _data_store; char *_opt_graph = nullptr; @@ -325,7 +325,7 @@ template clas // Dimensions size_t _dim = 0; - //size_t _aligned_dim = 0; + // size_t _aligned_dim = 0; 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 // Number of points which are used as initial candidates when iterating to From f890a969184dea0c11c33a708ef5e46118a02c50 Mon Sep 17 00:00:00 2001 From: harsha vardhan simhadri Date: Mon, 10 Apr 2023 17:03:41 -0700 Subject: [PATCH 33/69] minor changes --- CMakeLists.txt | 1 - include/abstract_data_store.h | 6 +- include/abstract_graph_store.h | 2 +- include/distance.h | 6 +- include/in_mem_graph_store.h | 2 +- include/index.h | 12 +-- include/natural_number_map.h | 1 - include/utils.h | 4 +- src/distance.cpp | 2 +- src/in_mem_data_store.cpp | 3 - src/index.cpp | 173 +-------------------------------- 11 files changed, 19 insertions(+), 193 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 825cb6df5..f29e13f3a 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -245,7 +245,6 @@ add_subdirectory(src) add_subdirectory(tests) add_subdirectory(tests/utils) - if (MSVC) message(STATUS "The ${PROJECT_NAME}.sln has been created, opened it from VisualStudio to build Release or Debug configurations.\n" "Alternatively, use MSBuild to build:\n\n" diff --git a/include/abstract_data_store.h b/include/abstract_data_store.h index 7d60dd34e..6ad3c3eba 100644 --- a/include/abstract_data_store.h +++ b/include/abstract_data_store.h @@ -42,8 +42,8 @@ template class AbstractDataStore return _dim; } - // populate the store with bulk vectors (either as pointer or bin file), potentially after normalizing the vectors - // if the metric deems so + // populate the store with vectors (either from a pointer or bin file), + // potentially after normalizing the vectors if the metric deems so virtual void populate_data(const data_t *vectors, const location_t num_pts) = 0; virtual void populate_data(const std::string &filename, const size_t offset) = 0; @@ -97,4 +97,4 @@ template class AbstractDataStore size_t _dim; }; -} // namespace diskann \ No newline at end of file +} // namespace diskann diff --git a/include/abstract_graph_store.h b/include/abstract_graph_store.h index 692149783..f7735b79a 100644 --- a/include/abstract_graph_store.h +++ b/include/abstract_graph_store.h @@ -28,4 +28,4 @@ class AbstractGraphStore size_t _capacity; }; -} // namespace diskann \ No newline at end of file +} // namespace diskann diff --git a/include/distance.h b/include/distance.h index 069721a04..4da3d2d3c 100644 --- a/include/distance.h +++ b/include/distance.h @@ -28,10 +28,10 @@ template class Distance return std::numeric_limits::max(); } - // For MIPS, normalization => a new dimension gets added to the vectors. + // For MIPS, normalization adds an extra dimension to the vectors. // This function lets callers know if the normalization process // changes the dimension. - virtual uint32_t post_processed_dimension(uint32_t orig_dimension) const + virtual uint32_t post_normalization_dimension(uint32_t orig_dimension) const { return orig_dimension; } @@ -248,7 +248,7 @@ class AVXNormalizedCosineDistanceFloat : public Distance // This will ensure that cosine is between -1 and 1. return 1.0f + _innerProduct.compare(a, b, length); } - DISKANN_DLLEXPORT virtual uint32_t post_processed_dimension(uint32_t orig_dimension) const override; + DISKANN_DLLEXPORT virtual uint32_t post_normalization_dimension(uint32_t orig_dimension) const override; DISKANN_DLLEXPORT virtual bool normalization_required() const; diff --git a/include/in_mem_graph_store.h b/include/in_mem_graph_store.h index 31b502fd7..98a9e4dc5 100644 --- a/include/in_mem_graph_store.h +++ b/include/in_mem_graph_store.h @@ -20,4 +20,4 @@ class InMemGraphStore : public AbstractGraphStore void set_adj_list(const location_t i, std::vector &neighbors); }; -} // namespace diskann \ No newline at end of file +} // namespace diskann diff --git a/include/index.h b/include/index.h index 0ef97825f..056341320 100644 --- a/include/index.h +++ b/include/index.h @@ -255,7 +255,6 @@ template clas const tsl::robin_set *const delete_set_ptr = nullptr); // add reverse links from all the visited nodes to node n. - void inter_insert(uint32_t n, std::vector &pruned_list, const uint32_t range, InMemQueryScratch *scratch); @@ -316,7 +315,6 @@ template clas std::shared_ptr> _distance; // Data - // T *_data = nullptr; std::shared_ptr> _data_store; char *_opt_graph = nullptr; @@ -325,12 +323,12 @@ template clas // Dimensions size_t _dim = 0; - // size_t _aligned_dim = 0; 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 - // Number of points which are used as initial candidates when iterating to - // closest point(s). These are not visible externally and won't be returned - // by search. DiskANN forces at least 1 frozen point for dynamic index. + + // _num_frozen_pts is the number of points which are used as initial candidates + // when iterating to closest point(s). These are not visible externally and won't + // be returned by search. At least 1 frozen point is needed for a dynamic index. // The frozen points have consecutive locations. See also _start below. size_t _num_frozen_pts = 0; size_t _max_range_of_loaded_graph = 0; @@ -398,7 +396,7 @@ template clas std::unique_ptr> _delete_set; bool _data_compacted = true; // true if data has been compacted - bool _is_saved = false; // Gopal. Checking if the index is already saved. + bool _is_saved = false; // Checking if the index is already saved. bool _conc_consolidate = false; // use _lock while searching // Acquire locks in the order below when acquiring multiple locks diff --git a/include/natural_number_map.h b/include/natural_number_map.h index 67986e24b..820ac3fdf 100644 --- a/include/natural_number_map.h +++ b/include/natural_number_map.h @@ -8,7 +8,6 @@ #include #include -// #include "boost_dynamic_bitset_fwd.h" namespace diskann { diff --git a/include/utils.h b/include/utils.h index 77b116218..5c6a2287a 100644 --- a/include/utils.h +++ b/include/utils.h @@ -878,13 +878,13 @@ template void save_Tvecs(const char *filename, T *data, size_t npts // create cached ofstream with 64MB cache cached_ofstream writer(fname, 64 * 1048576); - unsigned dimsuint32_t = (unsigned)ndims; + unsigned dims_u32 = (unsigned)ndims; // start writing for (size_t i = 0; i < npts; i++) { // write dims in u32 - writer.write((char *)&dimsuint32_t, sizeof(unsigned)); + writer.write((char *)&dims_u32, sizeof(unsigned)); // get cur point in data T *cur_pt = data + i * ndims; diff --git a/src/distance.cpp b/src/distance.cpp index 5872856b9..421b2b21b 100644 --- a/src/distance.cpp +++ b/src/distance.cpp @@ -522,7 +522,7 @@ float AVXDistanceInnerProductFloat::compare(const float *a, const float *b, uint return -result; } -uint32_t AVXNormalizedCosineDistanceFloat::post_processed_dimension(uint32_t orig_dimension) const +uint32_t AVXNormalizedCosineDistanceFloat::post_normalization_dimension(uint32_t orig_dimension) const { return orig_dimension; } diff --git a/src/in_mem_data_store.cpp b/src/in_mem_data_store.cpp index 972e72cfa..0b4f9f9e0 100644 --- a/src/in_mem_data_store.cpp +++ b/src/in_mem_data_store.cpp @@ -70,9 +70,6 @@ template location_t InMemDataStore::load_impl(const st } diskann::get_bin_metadata(filename, file_num_points, file_dim); - // since we are loading a new dataset, _empty_slots must be cleared - // _empty_slots.clear(); - if (file_dim != this->_dim) { std::stringstream stream; diff --git a/src/index.cpp b/src/index.cpp index dba931e46..dc6d5bb57 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -74,10 +74,6 @@ Index::Index(Metric m, const size_t dim, const size_t max_point -1, __FUNCSIG__, __FILE__, __LINE__); } - // data stored to _nd * aligned_dim matrix with necessary zero-padding - // REFACTOR - //_aligned_dim = ROUND_UP(_dim, 8); - if (dynamic_index && _num_frozen_pts == 0) { _num_frozen_pts = 1; @@ -98,11 +94,7 @@ Index::Index(Metric m, const size_t dim, const size_t max_point std::memset(_pq_data, 0, total_internal_points * _num_pq_chunks * sizeof(char)); } - // REFACTOR - // alloc_aligned(((void **)&_data), total_internal_points * _aligned_dim * sizeof(T), 8 * sizeof(T)); - // std::memset(_data, 0, total_internal_points * _aligned_dim * sizeof(T)); - - _start = (uint32_t)_max_points; + _start = (uint32_t)_max_points; _final_graph.resize(total_internal_points); @@ -153,11 +145,7 @@ template Index::~I // this->_distance = nullptr; // } // REFACTOR - // if (this->_data != nullptr) - //{ - // aligned_free(this->_data); - // this->_data = nullptr; - // } + if (_opt_graph != nullptr) { delete[] _opt_graph; @@ -225,8 +213,6 @@ template size_t Indexsave(data_file, _nd + _num_frozen_pts); - // REFACTOR - // return save_data_in_base_dimensions(data_file, _data, _nd + _num_frozen_pts, _dim, _aligned_dim); } // save the graph index on a file as an adjacency list. For each point, @@ -452,8 +438,6 @@ size_t Index::load_data(std::string filename) std::stringstream stream; stream << "ERROR: data file " << filename << " does not exist." << std::endl; diskann::cerr << stream.str() << std::endl; - // REFACTOR - // aligned_free(_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } diskann::get_bin_metadata(filename, file_num_points, file_dim); @@ -468,8 +452,6 @@ size_t Index::load_data(std::string filename) stream << "ERROR: Driver requests loading " << _dim << " dimension," << "but file has " << file_dim << " dimension." << std::endl; diskann::cerr << stream.str() << std::endl; - // REFACTOR - // aligned_free(_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } @@ -485,8 +467,6 @@ size_t Index::load_data(std::string filename) copy_aligned_data_from_file(reader, _data, file_num_points, file_dim, _aligned_dim); #else _data_store->populate_data(filename, 0U); // offset == 0. - // REFACTOR - // copy_aligned_data_from_file(filename.c_str(), _data, file_num_points, file_dim, _aligned_dim); #endif return file_num_points; } @@ -576,8 +556,6 @@ void Index::load(const char *filename, uint32_t num_threads, ui << 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; diskann::cerr << stream.str() << std::endl; - // REFACTOR - // aligned_free(_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } @@ -729,8 +707,6 @@ size_t Index::load_graph(std::string filename, size_t expected_ << std::endl; } diskann::cerr << stream.str() << std::endl; - // REFACTOR - // aligned_free(_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } @@ -824,9 +800,6 @@ template int Index location_t location = _tag_to_location[tag]; _data_store->get_vector(location, vec); - // REFACTOR - // size_t location = _tag_to_location[tag]; - // memcpy((void *)vec, (void *)(_data + location * _aligned_dim), (size_t)_dim * sizeof(T)); return 0; } @@ -842,51 +815,6 @@ template uint32_t Indexcalculate_medoid(); - - // //REFACTOR - // // allocate and init centroid - // float *center = new float[_aligned_dim](); - // for (size_t j = 0; j < _aligned_dim; j++) - // center[j] = 0; - // - // for (size_t i = 0; i < _nd; i++) - // for (size_t j = 0; j < _aligned_dim; j++) - // center[j] += (float)_data[i * _aligned_dim + j]; - // - // for (size_t j = 0; j < _aligned_dim; j++) - // center[j] /= (float)_nd; - // - // // compute all to one distance - // float *distances = new float[_nd](); - // #pragma omp parallel for schedule(static, 65536) - // for (int64_t i = 0; i < (int64_t)_nd; i++) - // { - // // extract point and distance reference - // float &dist = distances[i]; - // const T *cur_vec = _data + (i * (size_t)_aligned_dim); - // dist = 0; - // float diff = 0; - // for (size_t j = 0; j < _aligned_dim; j++) - // { - // diff = (center[j] - (float)cur_vec[j]) * (center[j] - (float)cur_vec[j]); - // dist += diff; - // } - // } - // // find imin - // uint32_t min_idx = 0; - // float min_dist = distances[0]; - // for (uint32_t i = 1; i < _nd; i++) - // { - // if (distances[i] < min_dist) - // { - // min_idx = i; - // min_dist = distances[i]; - // } - // } - // - // delete[] distances; - // delete[] center; - // return min_idx; } template std::vector Index::get_init_ids() @@ -1033,9 +961,6 @@ std::pair Index::iterate_to_fixed_point( else { distance = _data_store->get_distance(aligned_query, id); - // REFACTOR - // distance = _distance->compare(_data + _aligned_dim * (size_t)id, aligned_query, - // (uint32_t)_aligned_dim); } Neighbor nn = Neighbor(id, distance); best_L_nodes.insert(nn); @@ -1137,15 +1062,9 @@ std::pair Index::iterate_to_fixed_point( { auto nextn = id_scratch[m + 1]; _data_store->prefetch_vector(nextn); - // REFACTOR - // diskann::prefetch_vector((const char *)_data + _aligned_dim * (size_t)nextn, - // sizeof(T) * _aligned_dim); } dist_scratch.push_back(_data_store->get_distance(aligned_query, id)); - // REFACTOR - // dist_scratch.push_back( - // _distance->compare(aligned_query, _data + _aligned_dim * (size_t)id, (uint32_t)_aligned_dim)); } } cmps += id_scratch.size(); @@ -1172,11 +1091,6 @@ void Index::search_for_point_and_prune(int location, uint32_t L { _data_store->get_vector(location, scratch->aligned_query()); iterate_to_fixed_point(scratch->aligned_query(), Lindex, init_ids, scratch, false, unused_filter_label, false); - - // REFACTOR - // iterate_to_fixed_point(_data + _aligned_dim * location, Lindex, init_ids, scratch, false, - // unused_filter_label, - // false); } else { @@ -1187,9 +1101,6 @@ void Index::search_for_point_and_prune(int location, uint32_t L _data_store->get_vector(location, scratch->aligned_query()); iterate_to_fixed_point(scratch->aligned_query(), filteredLindex, filter_specific_start_nodes, scratch, true, _pts_to_labels[location], false); - // REFACTOR - // iterate_to_fixed_point(_data + _aligned_dim * location, filteredLindex, filter_specific_start_nodes, scratch, - // true, _pts_to_labels[location], false); } auto &pool = scratch->pool(); @@ -1286,9 +1197,6 @@ void Index::occlude_list(const uint32_t location, std::vectorget_distance(iter2->id, iter->id); - // REFACTOR - // float djk = _distance->compare(_data + _aligned_dim * (size_t)iter2->id, - // _data + _aligned_dim * (size_t)iter->id, (uint32_t)_aligned_dim); if (_dist_metric == diskann::Metric::L2 || _dist_metric == diskann::Metric::COSINE) { occlude_factor[t] = (djk == 0) ? std::numeric_limits::max() @@ -1336,9 +1244,6 @@ void Index::prune_neighbors(const uint32_t location, std::vecto { for (auto &ngh : pool) ngh.distance = _data_store->get_distance(ngh.id, location); - // REFACTOR - // ngh.distance = _distance->compare(_data + _aligned_dim * (size_t)ngh.id, - // _data + _aligned_dim * (size_t)location, (uint32_t)_aligned_dim); } // sort the pool based on distance to query and prune it with occlude_list @@ -1410,9 +1315,6 @@ void Index::inter_insert(uint32_t n, std::vector &pru if (dummy_visited.find(cur_nbr) == dummy_visited.end() && cur_nbr != des) { float dist = _data_store->get_distance(des, cur_nbr); - // REFACTOR - // float dist = _distance->compare(_data + _aligned_dim * (size_t)des, - // _data + _aligned_dim * (size_t)cur_nbr, (uint32_t)_aligned_dim); dummy_pool.emplace_back(Neighbor(cur_nbr, dist)); dummy_visited.insert(cur_nbr); } @@ -1537,9 +1439,6 @@ void Index::link(IndexWriteParameters ¶meters) if (dummy_visited.find(cur_nbr) == dummy_visited.end() && cur_nbr != node) { float dist = _data_store->get_distance(node, cur_nbr); - // REFACTOR - // float dist = _distance->compare(_data + _aligned_dim * (size_t)node, - // _data + _aligned_dim * (size_t)cur_nbr, (uint32_t)_aligned_dim); dummy_pool.emplace_back(Neighbor(cur_nbr, dist)); dummy_visited.insert(cur_nbr); } @@ -1564,10 +1463,6 @@ void Index::prune_all_neighbors(const uint32_t max_degree, cons const uint32_t range = max_degree; const uint32_t maxc = max_occlusion_size; - // const uint32_t range = parameters.Get("R"); - // const uint32_t maxc = parameters.Get("C"); - // const float alpha = parameters.Get("alpha"); - _filtered_index = true; diskann::Timer timer; @@ -1590,10 +1485,6 @@ void Index::prune_all_neighbors(const uint32_t max_degree, cons if (dummy_visited.find(cur_nbr) == dummy_visited.end() && cur_nbr != node) { float dist = _data_store->get_distance(node, cur_nbr); - // REFACTOR - // float dist = _distance->compare(_data + _aligned_dim * (size_t)node, - // _data + _aligned_dim * (size_t)cur_nbr, - // (uint32_t)_aligned_dim); dummy_pool.emplace_back(Neighbor(cur_nbr, dist)); dummy_visited.insert(cur_nbr); } @@ -1697,8 +1588,6 @@ void Index::build_with_data_populated(IndexWriteParameters &par stream << "ERROR: Driver requests loading " << _nd << " points from file," << "but tags vector is of size " << tags.size() << "." << std::endl; diskann::cerr << stream.str() << std::endl; - // REFACTOR - // aligned_free(_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } if (_enable_tags) @@ -1717,9 +1606,6 @@ void Index::build_with_data_populated(IndexWriteParameters &par if (_query_scratch.size() == 0) { - // REFACTOR - // initialize_query_scratch(5 + num_threads_index, index_L, index_L, index_R, maxc, _aligned_dim); - // Comment: Why 5+? initialize_query_scratch(5 + num_threads_index, index_L, index_L, index_R, maxc, _data_store->get_aligned_dim()); } @@ -1812,9 +1698,6 @@ void Index::build(const char *filename, const size_t num_points if (_pq_dist) aligned_free(_pq_data); - // REFACTOR - // else - // aligned_free(_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } @@ -1826,9 +1709,6 @@ void Index::build(const char *filename, const size_t num_points if (_pq_dist) aligned_free(_pq_data); - // REFACTOR - // else - // aligned_free(_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } @@ -1841,9 +1721,6 @@ void Index::build(const char *filename, const size_t num_points if (_pq_dist) aligned_free(_pq_data); - // REFACTOR - // else - // aligned_free(_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } @@ -1870,16 +1747,6 @@ void Index::build(const char *filename, const size_t num_points } _data_store->populate_data(filename, 0U); - // REFACTOR - // copy_aligned_data_from_file(filename, _data, file_num_points, file_dim, _aligned_dim); - // if (_normalize_vecs) - //{ - // for (size_t i = 0; i < file_num_points; i++) - // { - // normalize(_data + _aligned_dim * i, _aligned_dim); - // } - // } - diskann::cout << "Using only first " << num_points_to_load << " from file.. " << std::endl; { @@ -2035,8 +1902,7 @@ void Index::build_filtered_index(const char *filename, const st _label_to_medoid_id.clear(); size_t num_points_labels = 0; parse_label_file(label_file, - num_points_labels); // determines medoid for each label and - // identifies the points to label mapping + num_points_labels); // determines medoid for each label and identifies the points to label mapping std::unordered_map> label_to_points; @@ -2278,8 +2144,6 @@ 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]); - // REFACTOR - // memcpy(res_vectors[pos], _data + ((size_t)node.id) * _aligned_dim, _dim * sizeof(T)); } if (distances != nullptr) @@ -2339,8 +2203,6 @@ template void Indexcopy_points(res, _max_points, 1); - // REFACTOR - // memcpy(_data + _max_points * _aligned_dim, _data + res * _aligned_dim, _aligned_dim * sizeof(T)); } } @@ -2425,12 +2287,6 @@ inline void Index::process_delete(const tsl::robin_setget_distance(loc, ngh)); - - // REFACTOR - // expanded_nghrs_vec.emplace_back( - // ngh, - // _distance->compare(_data + _aligned_dim * loc, _data + _aligned_dim * ngh, - // (uint32_t)_aligned_dim)); } std::sort(expanded_nghrs_vec.begin(), expanded_nghrs_vec.end()); std::vector &occlude_list_output = scratch->occlude_list_output(); @@ -2634,9 +2490,6 @@ template void Indexcopy_points(old, new_location[old], 1); - // REFACTOR - // memcpy((void *)(_data + _aligned_dim * (size_t)new_location[old]), - // (void *)(_data + _aligned_dim * (size_t)old), _aligned_dim * sizeof(T)); } } else @@ -2820,16 +2673,6 @@ template void Indexresize(new_internal_points); - // REFACTOR - // #ifndef _WINDOWS - // T *new_data; - // alloc_aligned((void **)&new_data, new_internal_points * _aligned_dim * sizeof(T), 8 * sizeof(T)); - // memcpy(new_data, _data, (_max_points + _num_frozen_pts) * _aligned_dim * sizeof(T)); - // aligned_free(_data); - // _data = new_data; - // #else - // realloc_aligned((void **)&_data, new_internal_points * _aligned_dim * sizeof(T), 8 * sizeof(T)); - // #endif _final_graph.resize(new_internal_points); _locks = std::vector(new_internal_points); @@ -2923,16 +2766,6 @@ int Index::insert_point(const T *point, const TagT tag) tl.unlock(); _data_store->set_vector(location, point); - // REFACTOR - // Copy the vector in to the data array - // auto offset_data = _data + (size_t)_aligned_dim * location; - // memset((void *)offset_data, 0, sizeof(T) * _aligned_dim); - // memcpy((void *)offset_data, point, sizeof(T) * _dim); - - // if (_normalize_vecs) - //{ - // normalize((float *)offset_data, _dim); - // } // Find and add appropriate graph edges ScratchStoreManager> manager(_query_scratch); From 05280460b598307ac2dad06ed7620ea3fd83e3be Mon Sep 17 00:00:00 2001 From: ravishankar Date: Tue, 11 Apr 2023 17:17:03 +0000 Subject: [PATCH 34/69] added back the fast_l2 feature --- include/in_mem_data_store.h | 2 + src/in_mem_data_store.cpp | 5 + src/index.cpp | 241 ++++++++++++++++++------------------ 3 files changed, 130 insertions(+), 118 deletions(-) diff --git a/include/in_mem_data_store.h b/include/in_mem_data_store.h index e637422ed..0b7b2ee61 100644 --- a/include/in_mem_data_store.h +++ b/include/in_mem_data_store.h @@ -52,6 +52,8 @@ template class InMemDataStore : public AbstractDataStore * get_dist_fn(); + protected: virtual void expand(const location_t new_size) override; virtual void shrink(const location_t new_size) override; diff --git a/src/in_mem_data_store.cpp b/src/in_mem_data_store.cpp index 0b4f9f9e0..5a4319cf3 100644 --- a/src/in_mem_data_store.cpp +++ b/src/in_mem_data_store.cpp @@ -343,6 +343,11 @@ template location_t InMemDataStore::calculate_medoid() return min_idx; } +template +Distance* InMemDataStore::get_dist_fn() { + return this->_distance_fn.get(); +} + template DISKANN_DLLEXPORT class InMemDataStore; template DISKANN_DLLEXPORT class InMemDataStore; template DISKANN_DLLEXPORT class InMemDataStore; diff --git a/src/index.cpp b/src/index.cpp index dc6d5bb57..197e9b3ea 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -2934,128 +2934,133 @@ template void Index void Index::optimize_index_layout() -{ // use after build or load -} +//template void Index::optimize_index_layout() +//{ // use after build or load +//} // REFACTOR: This should be an OptimizedDataStore class -// template void Index::optimize_index_layout() -//{ // use after build or load -// if (_dynamic_index) -// { -// throw diskann::ANNException("Optimize_index_layout not implemented for dyanmic indices", -1, __FUNCSIG__, -// __FILE__, __LINE__); -// } -// -// _data_len = (_aligned_dim + 1) * sizeof(float); -// _neighbor_len = (_max_observed_degree + 1) * sizeof(uint32_t); -// _node_size = _data_len + _neighbor_len; -// _opt_graph = new char[_node_size * _nd]; -// DistanceFastL2 *dist_fast = (DistanceFastL2 *)_distance; -// for (uint32_t 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; -// uint32_t k = _final_graph[i].size(); -// std::memcpy(cur_node_offset, &k, sizeof(uint32_t)); -// std::memcpy(cur_node_offset + sizeof(uint32_t), _final_graph[i].data(), k * sizeof(uint32_t)); -// std::vector().swap(_final_graph[i]); -// } -// _final_graph.clear(); -// _final_graph.shrink_to_fit(); -// } -// REFACTOR: once optimized layout becomes its own Data+Graph store, we should just invoke regular search -template -void Index::search_with_optimized_layout(const T *query, size_t K, size_t L, uint32_t *indices) -{ -} + template void Index::optimize_index_layout() +{ // use after build or load + if (_dynamic_index) + { + throw diskann::ANNException("Optimize_index_layout not implemented for dyanmic indices", -1, __FUNCSIG__, + __FILE__, __LINE__); + } + + float* cur_vec = new float[_data_store->get_aligned_dim()]; + std::memset(cur_vec, 0, _data_store->get_aligned_dim()*sizeof(float)); + _data_len = (_data_store->get_aligned_dim() + 1) * sizeof(float); + _neighbor_len = (_max_observed_degree + 1) * sizeof(uint32_t); + _node_size = _data_len + _neighbor_len; + _opt_graph = new char[_node_size * _nd]; + DistanceFastL2 *dist_fast = (DistanceFastL2 *)_data_store->get_dist_fn(); + for (uint32_t i = 0; i < _nd; i++) + { + char *cur_node_offset = _opt_graph + i * _node_size; + _data_store->get_vector(i, (T*) cur_vec); + float cur_norm = dist_fast->norm((T*) cur_vec, _data_store->get_aligned_dim()); + std::memcpy(cur_node_offset, &cur_norm, sizeof(float)); + std::memcpy(cur_node_offset + sizeof(float), cur_vec, _data_len - sizeof(float)); + + cur_node_offset += _data_len; + uint32_t k = _final_graph[i].size(); + std::memcpy(cur_node_offset, &k, sizeof(uint32_t)); + std::memcpy(cur_node_offset + sizeof(uint32_t), _final_graph[i].data(), k * sizeof(uint32_t)); + std::vector().swap(_final_graph[i]); + } + _final_graph.clear(); + _final_graph.shrink_to_fit(); + delete[] cur_vec; + } -// template -// void Index::search_with_optimized_layout(const T *query, size_t K, size_t L, uint32_t *indices) +// REFACTOR: once optimized layout becomes its own Data+Graph store, we should just invoke regular search +//template +//void Index::search_with_optimized_layout(const T *query, size_t K, size_t L, uint32_t *indices) //{ -// DistanceFastL2 *dist_fast = (DistanceFastL2 *)_distance; -// -// NeighborPriorityQueue retset(L); -// std::vector init_ids(L); -// -// boost::dynamic_bitset<> flags{_nd, 0}; -// uint32_t tmp_l = 0; -// uint32_t *neighbors = (uint32_t *)(_opt_graph + _node_size * _start + _data_len); -// uint32_t 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) -// { -// uint32_t id = rand() % _nd; -// if (flags[id]) -// continue; -// flags[id] = true; -// init_ids[tmp_l] = id; -// tmp_l++; -// } -// -// for (uint32_t i = 0; i < init_ids.size(); i++) -// { -// uint32_t id = init_ids[i]; -// if (id >= _nd) -// continue; -// _mm_prefetch(_opt_graph + _node_size * id, _MM_HINT_T0); -// } -// L = 0; -// for (uint32_t i = 0; i < init_ids.size(); i++) -// { -// uint32_t 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, (uint32_t)_aligned_dim); -// retset.insert(Neighbor(id, dist)); -// flags[id] = true; -// L++; -// } -// -// while (retset.has_unexpanded_node()) -// { -// auto nbr = retset.closest_unexpanded(); -// auto n = nbr.id; -// _mm_prefetch(_opt_graph + _node_size * n + _data_len, _MM_HINT_T0); -// neighbors = (uint32_t *)(_opt_graph + _node_size * n + _data_len); -// uint32_t MaxM = *neighbors; -// neighbors++; -// for (uint32_t m = 0; m < MaxM; ++m) -// _mm_prefetch(_opt_graph + _node_size * neighbors[m], _MM_HINT_T0); -// for (uint32_t m = 0; m < MaxM; ++m) -// { -// uint32_t id = neighbors[m]; -// if (flags[id]) -// continue; -// flags[id] = 1; -// T *data = (T *)(_opt_graph + _node_size * id); -// float norm = *data; -// data++; -// float dist = dist_fast->compare(query, data, norm, (uint32_t)_aligned_dim); -// Neighbor nn(id, dist); -// retset.insert(nn); -// } -// } -// -// for (size_t i = 0; i < K; i++) -// { -// indices[i] = retset[i].id; -// } -// } +//} + + template + void Index::search_with_optimized_layout(const T *query, size_t K, size_t L, uint32_t *indices) +{ + DistanceFastL2 *dist_fast = (DistanceFastL2 *) _data_store->get_dist_fn(); + + NeighborPriorityQueue retset(L); + std::vector init_ids(L); + + boost::dynamic_bitset<> flags{_nd, 0}; + uint32_t tmp_l = 0; + uint32_t *neighbors = (uint32_t *)(_opt_graph + _node_size * _start + _data_len); + uint32_t 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) + { + uint32_t id = rand() % _nd; + if (flags[id]) + continue; + flags[id] = true; + init_ids[tmp_l] = id; + tmp_l++; + } + + for (uint32_t i = 0; i < init_ids.size(); i++) + { + uint32_t id = init_ids[i]; + if (id >= _nd) + continue; + _mm_prefetch(_opt_graph + _node_size * id, _MM_HINT_T0); + } + L = 0; + for (uint32_t i = 0; i < init_ids.size(); i++) + { + uint32_t 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, (uint32_t)_data_store->get_aligned_dim()); + retset.insert(Neighbor(id, dist)); + flags[id] = true; + L++; + } + + while (retset.has_unexpanded_node()) + { + auto nbr = retset.closest_unexpanded(); + auto n = nbr.id; + _mm_prefetch(_opt_graph + _node_size * n + _data_len, _MM_HINT_T0); + neighbors = (uint32_t *)(_opt_graph + _node_size * n + _data_len); + uint32_t MaxM = *neighbors; + neighbors++; + for (uint32_t m = 0; m < MaxM; ++m) + _mm_prefetch(_opt_graph + _node_size * neighbors[m], _MM_HINT_T0); + for (uint32_t m = 0; m < MaxM; ++m) + { + uint32_t id = neighbors[m]; + if (flags[id]) + continue; + flags[id] = 1; + T *data = (T *)(_opt_graph + _node_size * id); + float norm = *data; + data++; + float dist = dist_fast->compare(query, data, norm, (uint32_t)_data_store->get_aligned_dim()); + Neighbor nn(id, dist); + retset.insert(nn); + } + } + + for (size_t i = 0; i < K; i++) + { + indices[i] = retset[i].id; + } + } /* Internals of the library */ template const float Index::INDEX_GROWTH_FACTOR = 1.5f; From 85b44e1103b8bb6dc919018984a894756ca11bcf Mon Sep 17 00:00:00 2001 From: ravishankar Date: Tue, 11 Apr 2023 17:42:06 +0000 Subject: [PATCH 35/69] added back set_start_points in index.cpp --- src/index.cpp | 91 +++++++++++++++++++++++++-------------------------- 1 file changed, 44 insertions(+), 47 deletions(-) diff --git a/src/index.cpp b/src/index.cpp index 197e9b3ea..cc2ac3450 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -1522,57 +1522,54 @@ void Index::prune_all_neighbors(const uint32_t max_degree, cons } } -// REFACTOR -// template -// void Index::set_start_points(const T *data, size_t data_count) -//{ -// std::unique_lock ul(_update_lock); -// std::unique_lock tl(_tag_lock); -// if (_nd > 0) -// throw ANNException("Can not set starting point for a non-empty index", -1, __FUNCSIG__, __FILE__, __LINE__); -// -// if (data_count != _num_frozen_pts * _aligned_dim) -// throw ANNException("Invalid number of points", -1, __FUNCSIG__, __FILE__, __LINE__); -// + //REFACTOR + template + void Index::set_start_points(const T *data, size_t data_count) +{ + std::unique_lock ul(_update_lock); + std::unique_lock tl(_tag_lock); + if (_nd > 0) + throw ANNException("Can not set starting point for a non-empty index", -1, __FUNCSIG__, __FILE__, __LINE__); + + if (data_count != _num_frozen_pts * _dim) + throw ANNException("Invalid number of points", -1, __FUNCSIG__, __FILE__, __LINE__); + // memcpy(_data + _aligned_dim * _max_points, data, _aligned_dim * sizeof(T) * _num_frozen_pts); -// _has_built = true; -// diskann::cout << "Index start points set: #" << _num_frozen_pts << std::endl; -// } -// + for (location_t i = _max_points; i < _max_points + _num_frozen_pts; i++) { + _data_store->set_vector(i, data + i*_dim); + } + _has_built = true; + diskann::cout << "Index start points set: #" << _num_frozen_pts << std::endl; + } -// REFACTOR: added dummy implementation for now. -template -void Index::set_start_points_at_random(T radius, uint32_t random_seed) + + template + void Index::set_start_points_at_random(T radius, uint32_t random_seed) { -} + std::mt19937 gen{random_seed}; + std::normal_distribution<> d{0.0, 1.0}; -// template -// void Index::set_start_points_at_random(T radius, uint32_t random_seed) -//{ -// std::mt19937 gen{random_seed}; -// std::normal_distribution<> d{0.0, 1.0}; -// -// std::vector points_data; -// points_data.reserve(_aligned_dim * _num_frozen_pts); -// std::vector real_vec(_aligned_dim); -// -// for (size_t frozen_point = 0; frozen_point < _num_frozen_pts; frozen_point++) -// { -// double norm_sq = 0.0; -// for (size_t i = 0; i < _dim; ++i) -// { -// auto r = d(gen); -// real_vec[i] = r; -// norm_sq += r * r; -// } -// -// const double norm = std::sqrt(norm_sq); -// for (auto iter : real_vec) -// points_data.push_back(static_cast(iter * radius / norm)); -// } -// -// set_start_points(points_data.data(), points_data.size()); -// } + std::vector points_data; + points_data.reserve(_dim * _num_frozen_pts); + std::vector real_vec(_dim); + + for (size_t frozen_point = 0; frozen_point < _num_frozen_pts; frozen_point++) + { + double norm_sq = 0.0; + for (size_t i = 0; i < _dim; ++i) + { + auto r = d(gen); + real_vec[i] = r; + norm_sq += r * r; + } + + const double norm = std::sqrt(norm_sq); + for (auto iter : real_vec) + points_data.push_back(static_cast(iter * radius / norm)); + } + + set_start_points(points_data.data(), points_data.size()); + } template void Index::build_with_data_populated(IndexWriteParameters ¶meters, const std::vector &tags) From f7fec4584b185dec897ed7ec99dd4d4ca5389fb6 Mon Sep 17 00:00:00 2001 From: Gopal Srinivasa Date: Wed, 12 Apr 2023 00:03:20 +0530 Subject: [PATCH 36/69] Version for review --- tests/build_memory_index.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/build_memory_index.cpp b/tests/build_memory_index.cpp index 8d68feb0e..ba63bdeb7 100644 --- a/tests/build_memory_index.cpp +++ b/tests/build_memory_index.cpp @@ -76,7 +76,7 @@ int main(int argc, char **argv) { desc.add_options()("help,h", "Print information on arguments"); desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); - desc.add_options()("dist_fn", po::value(&dist_fn)->required(), "distance function "); + desc.add_options()("dist_fn", po::value(&dist_fn)->required(), "distance function "); desc.add_options()("data_path", po::value(&data_path)->required(), "Input data file in bin format"); desc.add_options()("index_path_prefix", po::value(&index_path_prefix)->required(), From 921cec536baa8c5f47ba1bb3ed53193f85a76525 Mon Sep 17 00:00:00 2001 From: Gopal Srinivasa Date: Wed, 12 Apr 2023 16:10:15 +0530 Subject: [PATCH 37/69] Incorporating Harsha's comments - 2 --- include/abstract_data_store.h | 56 ++++++++++++++++------------- include/distance.h | 68 +++++++++-------------------------- include/in_mem_data_store.h | 4 +-- include/index.h | 2 +- src/distance.cpp | 56 ++++++++++++++++++++++++++--- src/in_mem_data_store.cpp | 13 +++---- src/index.cpp | 2 +- 7 files changed, 111 insertions(+), 90 deletions(-) diff --git a/include/abstract_data_store.h b/include/abstract_data_store.h index 6ad3c3eba..1dcaed820 100644 --- a/include/abstract_data_store.h +++ b/include/abstract_data_store.h @@ -26,21 +26,19 @@ template class AbstractDataStore // can discard the empty locations before saving. virtual size_t save(const std::string &filename, const location_t num_pts) = 0; - virtual location_t capacity() const - { + virtual location_t capacity() const{ return _capacity; } - virtual size_t get_dims() const - { + virtual size_t get_dims() const { return _dim; } - // by default, aligned dim = dim, some stores can align the data differently, so we may have different values - virtual size_t get_aligned_dim() const - { - return _dim; - } + // Implementers can choose to return _dim if they are not + // concerned about memory alignment. + // Returns _dim aligned to a 8-byte value. Used for allocating + // aligned memory for efficiency/simplicity of code. + virtual size_t get_aligned_dim() const = 0; // populate the store with vectors (either from a pointer or bin file), // potentially after normalizing the vectors if the metric deems so @@ -50,19 +48,22 @@ template class AbstractDataStore // reverse of populate, save the first num_pts many points back to bin file virtual void save_data_to_bin(const std::string &filename, const location_t num_pts) = 0; - virtual void resize(const location_t num_points) + //Returns the updated capacity of the datastore. Clients should check + //if resize actually changed the capacity to new_num_points before + //proceeding with operations. See the code below: + // auto new_capcity = data_store->resize(new_num_points); + // if ( new_capacity >= new_num_points) { + // //PROCEED + // else + // //ERROR. + virtual location_t resize(const location_t new_num_points) { - if (num_points > _capacity) - { - expand(num_points); - } - else if (num_points < _capacity) - { - shrink(num_points); - } - else - { - // ignore. + if (new_num_points > _capacity) { + return expand(new_num_points); + } else if (new_num_points < _capacity) { + return shrink(new_num_points); + } else { + return _capacity; } } @@ -88,10 +89,15 @@ template class AbstractDataStore virtual location_t calculate_medoid() const = 0; protected: - // Expand the datastore to new_num_points. - virtual void expand(const location_t new_num_points) = 0; - // Shrink the datastore to new_num_points. This function should be called after compaction to free unused memory. - virtual void shrink(const location_t new_num_points) = 0; + // Expand the datastore to new_num_points. Returns the new capacity created, which should be == new_num_points + // in the normal case. Implementers can also return _capacity to indicate that there are not implementing this + // method. + virtual location_t expand(const location_t new_num_points) = 0; + + // Shrink the datastore to new_num_points. It is NOT an error if shrink doesn't reduce the capacity + // so callers need to check this correctly. See also for "default" implementation + virtual location_t shrink(const location_t new_num_points) = 0; + location_t _capacity; size_t _dim; diff --git a/include/distance.h b/include/distance.h index 4da3d2d3c..22a82b7d0 100644 --- a/include/distance.h +++ b/include/distance.h @@ -15,38 +15,28 @@ enum Metric template class Distance { public: - Distance(diskann::Metric dist_metric) : _distance_metric(dist_metric) + DISKANN_DLLEXPORT Distance(diskann::Metric dist_metric) : _distance_metric(dist_metric) { } // distance comparison function - virtual float compare(const T *a, const T *b, uint32_t length) const = 0; + DISKANN_DLLEXPORT virtual float compare(const T *a, const T *b, uint32_t length) const = 0; // Needed only for COSINE-BYTE and INNER_PRODUCT-BYTE - virtual float compare(const T *a, const T *b, const float normA, const float normB, uint32_t length) const - { - return std::numeric_limits::max(); - } + DISKANN_DLLEXPORT virtual float compare(const T *a, const T *b, const float normA, const float normB, + uint32_t length) const = 0; // For MIPS, normalization adds an extra dimension to the vectors. // This function lets callers know if the normalization process // changes the dimension. - virtual uint32_t post_normalization_dimension(uint32_t orig_dimension) const - { - return orig_dimension; - } + DISKANN_DLLEXPORT virtual uint32_t post_normalization_dimension(uint32_t orig_dimension) const; - virtual diskann::Metric get_metric() const - { - return _distance_metric; - } + DISKANN_DLLEXPORT virtual diskann::Metric get_metric() const; + // This is for efficiency. If no normalization is required, the callers // can simply ignore the normalize_data_for_build() function. - virtual bool normalization_required() const - { - return false; - } + DISKANN_DLLEXPORT virtual bool normalization_required() const; // Check the normalization_required() function before calling this. // Clients can call the function like this: @@ -58,24 +48,18 @@ template class Distance // // TODO: This does not take into account the case for SSD inner product // where the dimensions change after normalization. - // - virtual void normalize_data_for_build(T *original_data, const uint32_t orig_dim, const uint32_t num_points) - { - } + DISKANN_DLLEXPORT virtual void normalize_data_for_build(T *original_data, const uint32_t orig_dim, + const uint32_t num_points); // Invokes normalization for a single vector during search. The scratch space // has to be created by the caller keeping track of the fact that normalization // might change the dimension of the query vector. - virtual void normalize_vector_for_search(const T *query_vec, const uint32_t query_dim, T *scratch_query) - { - std::memcpy(scratch_query, query_vec, query_dim * sizeof(T)); - } + DISKANN_DLLEXPORT virtual void normalize_vector_for_search(const T *query_vec, const uint32_t query_dim, + T *scratch_query); // Providing a default implementation for the virtual destructor because we don't // expect most metric implementations to need it. - virtual ~Distance() - { - } + DISKANN_DLLEXPORT virtual ~Distance(); protected: diskann::Metric _distance_metric; @@ -141,13 +125,14 @@ class AVXDistanceL2Float : public Distance DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t length) const; }; -class SlowDistanceL2Float : public Distance +template +class SlowDistanceL2 : public Distance { public: - SlowDistanceL2Float() : Distance(diskann::Metric::L2) + SlowDistanceL2() : Distance(diskann::Metric::L2) { } - DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t length) const; + DISKANN_DLLEXPORT virtual float compare(const T *a, const T *b, uint32_t length) const; }; class SlowDistanceCosineUInt8 : public Distance @@ -168,25 +153,6 @@ class DistanceL2UInt8 : public Distance DISKANN_DLLEXPORT virtual float compare(const uint8_t *a, const uint8_t *b, uint32_t size) const; }; -// Simple implementations for non-AVX machines. Compiler can optimize. -template class SlowDistanceL2Int : public Distance -{ - public: - SlowDistanceL2Int() : Distance(diskann::Metric::L2) - { - } - // Implementing here because this is a template function - DISKANN_DLLEXPORT virtual float compare(const T *a, const T *b, uint32_t length) const - { - uint32_t result = 0; - 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])); - } - return (float)result; - } -}; - template class DistanceInnerProduct : public Distance { public: diff --git a/include/in_mem_data_store.h b/include/in_mem_data_store.h index 0b7b2ee61..2dd261c2a 100644 --- a/include/in_mem_data_store.h +++ b/include/in_mem_data_store.h @@ -55,8 +55,8 @@ template class InMemDataStore : public AbstractDataStore * get_dist_fn(); protected: - virtual void expand(const location_t new_size) override; - virtual void shrink(const location_t new_size) override; + virtual location_t expand(const location_t new_size) override; + virtual location_t shrink(const location_t new_size) override; virtual location_t load_impl(const std::string &filename); #ifdef EXEC_ENV_OLS diff --git a/include/index.h b/include/index.h index 056341320..83d58f48a 100644 --- a/include/index.h +++ b/include/index.h @@ -315,7 +315,7 @@ template clas std::shared_ptr> _distance; // Data - std::shared_ptr> _data_store; + std::unique_ptr> _data_store; char *_opt_graph = nullptr; // Graph related data structures diff --git a/src/distance.cpp b/src/distance.cpp index 421b2b21b..5ed336afc 100644 --- a/src/distance.cpp +++ b/src/distance.cpp @@ -22,6 +22,46 @@ namespace diskann { +// +// Base Class Implementatons +// +template +float Distance::compare(const T *a, const T *b, const float normA, const float normB, uint32_t length) const +{ + throw std::logic_error("This function is not implemented."); +} + +template uint32_t Distance::post_normalization_dimension(uint32_t orig_dimension) const +{ + return orig_dimension; +} + +template diskann::Metric Distance::get_metric() const +{ + return _distance_metric; +} + +template bool Distance::normalization_required() const +{ + return false; +} + +template +void Distance::normalize_data_for_build(T *original_data, const uint32_t orig_dim, const uint32_t num_points) +{ +} + +template +void Distance::normalize_vector_for_search(const T *query_vec, const uint32_t query_dim, T *scratch_query) +{ + std::memcpy(scratch_query, query_vec, query_dim * sizeof(T)); +} + +template Distance::~Distance() +{ +} + + // // Cosine distance functions. // @@ -181,12 +221,13 @@ float DistanceL2Float::compare(const float *a, const float *b, uint32_t size) co return result; } -float SlowDistanceL2Float::compare(const float *a, const float *b, uint32_t length) const +template +float SlowDistanceL2::compare(const T *a, const T *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]); + result += ((float)(a[i] - b[i])) * (a[i] - b[i]); } return result; } @@ -557,6 +598,9 @@ void AVXNormalizedCosineDistanceFloat::normalize_and_copy(const float *query_vec } } + + + // Get the right distance function for the given metric. template <> diskann::Distance *get_distance_function(diskann::Metric m) { @@ -575,7 +619,7 @@ template <> diskann::Distance *get_distance_function(diskann::Metric m) else { diskann::cout << "L2: Older CPU. Using slow distance computation" << std::endl; - return new diskann::SlowDistanceL2Float(); + return new diskann::SlowDistanceL2(); } } else if (m == diskann::Metric::COSINE) @@ -627,7 +671,7 @@ template <> diskann::Distance *get_distance_function(diskann::Metric m) diskann::cout << "Older CPU. Using slow distance computation " "SlowDistanceL2Int." << std::endl; - return new diskann::SlowDistanceL2Int(); + return new diskann::SlowDistanceL2(); } } else if (m == diskann::Metric::COSINE) @@ -684,4 +728,8 @@ template DISKANN_DLLEXPORT class DistanceFastL2; template DISKANN_DLLEXPORT class DistanceFastL2; template DISKANN_DLLEXPORT class DistanceFastL2; +template DISKANN_DLLEXPORT class SlowDistanceL2; +template DISKANN_DLLEXPORT class SlowDistanceL2; +template DISKANN_DLLEXPORT class SlowDistanceL2; + } // namespace diskann diff --git a/src/in_mem_data_store.cpp b/src/in_mem_data_store.cpp index 5a4319cf3..84bf0ddfe 100644 --- a/src/in_mem_data_store.cpp +++ b/src/in_mem_data_store.cpp @@ -185,11 +185,11 @@ float InMemDataStore::get_distance(const location_t loc1, const location return _distance_fn->compare(_data + loc1 * _aligned_dim, _data + loc2 * _aligned_dim, this->_aligned_dim); } -template void InMemDataStore::expand(const location_t new_size) +template location_t InMemDataStore::expand(const location_t new_size) { if (new_size == this->capacity()) { - return; + return this->capacity(); } else if (new_size < this->capacity()) { @@ -208,13 +208,14 @@ template void InMemDataStore::expand(const location_t realloc_aligned((void **)&_data, new_size * _aligned_dim * sizeof(data_t), 8 * sizeof(data_t)); #endif this->_capacity = new_size; + return this->_capacity; } -template void InMemDataStore::shrink(const location_t new_size) +template location_t InMemDataStore::shrink(const location_t new_size) { if (new_size == this->capacity()) { - return; + return this->capacity(); } else if (new_size > this->capacity()) { @@ -232,6 +233,8 @@ template void InMemDataStore::shrink(const location_t #else realloc_aligned((void **)&_data, new_size * _aligned_dim * sizeof(data_t), 8 * sizeof(data_t)); #endif + this->_capacity = new_size; + return this->_capacity; } template @@ -253,8 +256,6 @@ void InMemDataStore::reposition_points(const location_t old_location_sta uint32_t mem_clear_loc_start = old_location_start; uint32_t mem_clear_loc_end_limit = old_location_start + num_locations; - // Move the adjacency lists. Make sure that overlapping ranges are handled - // correctly. if (new_location_start < old_location_start) { // If ranges are overlapping, make sure not to clear the newly copied diff --git a/src/index.cpp b/src/index.cpp index cc2ac3450..fc56b9aad 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -115,7 +115,7 @@ Index::Index(Metric m, const size_t dim, const size_t max_point // REFACTOR: TODO This should move to a factory method. _data_store = - std::make_shared>((location_t)total_internal_points, _dim, this->_distance); + std::make_unique>((location_t)total_internal_points, _dim, this->_distance); _locks = std::vector(total_internal_points); From a310e44778e3b8390534950117c90a69d8498e44 Mon Sep 17 00:00:00 2001 From: harsha vardhan simhadri Date: Wed, 12 Apr 2023 13:00:06 -0700 Subject: [PATCH 38/69] move implementation of abstract data store methods to a cpp file --- include/abstract_data_store.h | 24 ++++--------------- src/CMakeLists.txt | 4 ++-- src/abstract_data_store.cpp | 44 +++++++++++++++++++++++++++++++++++ 3 files changed, 51 insertions(+), 21 deletions(-) create mode 100644 src/abstract_data_store.cpp diff --git a/include/abstract_data_store.h b/include/abstract_data_store.h index 1dcaed820..08808adab 100644 --- a/include/abstract_data_store.h +++ b/include/abstract_data_store.h @@ -4,6 +4,7 @@ #pragma once #include +#include #include "types.h" @@ -13,9 +14,7 @@ namespace diskann template class AbstractDataStore { public: - AbstractDataStore(const location_t capacity, const size_t dim) : _capacity(capacity), _dim(dim) - { - } + AbstractDataStore(const location_t capacity, const size_t dim); // Return number of points returned virtual location_t load(const std::string &filename) = 0; @@ -26,13 +25,9 @@ template class AbstractDataStore // can discard the empty locations before saving. virtual size_t save(const std::string &filename, const location_t num_pts) = 0; - virtual location_t capacity() const{ - return _capacity; - } + virtual location_t capacity() const; - virtual size_t get_dims() const { - return _dim; - } + virtual size_t get_dims() const; // Implementers can choose to return _dim if they are not // concerned about memory alignment. @@ -56,16 +51,7 @@ template class AbstractDataStore // //PROCEED // else // //ERROR. - virtual location_t resize(const location_t new_num_points) - { - if (new_num_points > _capacity) { - return expand(new_num_points); - } else if (new_num_points < _capacity) { - return shrink(new_num_points); - } else { - return _capacity; - } - } + virtual location_t resize(const location_t new_num_points); // operations on vectors virtual void get_vector(const location_t i, data_t *dest) const = 0; diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index bac0ee6f7..2cb835c2b 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -7,8 +7,8 @@ if(MSVC) add_subdirectory(dll) else() #file(GLOB CPP_SOURCES *.cpp) - set(CPP_SOURCES ann_exception.cpp disk_utils.cpp distance.cpp index.cpp - in_mem_graph_store.cpp in_mem_data_store.cpp + set(CPP_SOURCES abstract_data_store.cpp ann_exception.cpp disk_utils.cpp + distance.cpp index.cpp in_mem_graph_store.cpp in_mem_data_store.cpp 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 diff --git a/src/abstract_data_store.cpp b/src/abstract_data_store.cpp new file mode 100644 index 000000000..cb291e6d7 --- /dev/null +++ b/src/abstract_data_store.cpp @@ -0,0 +1,44 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#pragma once + +#include + +#include "abstract_data_store.h" + +namespace diskann +{ + +template +AbstractDataStore::AbstractDataStore(const location_t capacity, const size_t dim) : _capacity(capacity), _dim(dim) +{ +} + +template location_t AbstractDataStore::capacity() const +{ + return _capacity; +} + +template size_t AbstractDataStore::get_dims() const +{ + return _dim; +} + +template location_t AbstractDataStore::resize(const location_t new_num_points) +{ + if (new_num_points > _capacity) + { + return expand(new_num_points); + } + else if (new_num_points < _capacity) + { + return shrink(new_num_points); + } + else + { + return _capacity; + } +} + +} // namespace diskann From d751fa4327fa4252f1598635f60a96e1307439d7 Mon Sep 17 00:00:00 2001 From: harsha vardhan simhadri Date: Wed, 12 Apr 2023 13:13:08 -0700 Subject: [PATCH 39/69] clang format --- include/abstract_data_store.h | 21 ++- include/distance.h | 8 +- include/in_mem_data_store.h | 2 +- include/index.h | 2 +- src/abstract_data_store.cpp | 5 +- src/distance.cpp | 7 +- src/in_mem_data_store.cpp | 4 +- src/index.cpp | 312 +++++++++++++++++----------------- 8 files changed, 177 insertions(+), 184 deletions(-) diff --git a/include/abstract_data_store.h b/include/abstract_data_store.h index 08808adab..ee1ad4d33 100644 --- a/include/abstract_data_store.h +++ b/include/abstract_data_store.h @@ -29,7 +29,7 @@ template class AbstractDataStore virtual size_t get_dims() const; - // Implementers can choose to return _dim if they are not + // Implementers can choose to return _dim if they are not // concerned about memory alignment. // Returns _dim aligned to a 8-byte value. Used for allocating // aligned memory for efficiency/simplicity of code. @@ -43,14 +43,14 @@ template class AbstractDataStore // reverse of populate, save the first num_pts many points back to bin file virtual void save_data_to_bin(const std::string &filename, const location_t num_pts) = 0; - //Returns the updated capacity of the datastore. Clients should check - //if resize actually changed the capacity to new_num_points before - //proceeding with operations. See the code below: - // auto new_capcity = data_store->resize(new_num_points); - // if ( new_capacity >= new_num_points) { - // //PROCEED - // else - // //ERROR. + // Returns the updated capacity of the datastore. Clients should check + // if resize actually changed the capacity to new_num_points before + // proceeding with operations. See the code below: + // auto new_capcity = data_store->resize(new_num_points); + // if ( new_capacity >= new_num_points) { + // //PROCEED + // else + // //ERROR. virtual location_t resize(const location_t new_num_points); // operations on vectors @@ -77,14 +77,13 @@ template class AbstractDataStore protected: // Expand the datastore to new_num_points. Returns the new capacity created, which should be == new_num_points // in the normal case. Implementers can also return _capacity to indicate that there are not implementing this - // method. + // method. virtual location_t expand(const location_t new_num_points) = 0; // Shrink the datastore to new_num_points. It is NOT an error if shrink doesn't reduce the capacity // so callers need to check this correctly. See also for "default" implementation virtual location_t shrink(const location_t new_num_points) = 0; - location_t _capacity; size_t _dim; }; diff --git a/include/distance.h b/include/distance.h index 22a82b7d0..1ba4f32fb 100644 --- a/include/distance.h +++ b/include/distance.h @@ -32,7 +32,6 @@ template class Distance DISKANN_DLLEXPORT virtual uint32_t post_normalization_dimension(uint32_t orig_dimension) const; DISKANN_DLLEXPORT virtual diskann::Metric get_metric() const; - // This is for efficiency. If no normalization is required, the callers // can simply ignore the normalize_data_for_build() function. @@ -49,13 +48,13 @@ template class Distance // TODO: This does not take into account the case for SSD inner product // where the dimensions change after normalization. DISKANN_DLLEXPORT virtual void normalize_data_for_build(T *original_data, const uint32_t orig_dim, - const uint32_t num_points); + const uint32_t num_points); // Invokes normalization for a single vector during search. The scratch space // has to be created by the caller keeping track of the fact that normalization // might change the dimension of the query vector. DISKANN_DLLEXPORT virtual void normalize_vector_for_search(const T *query_vec, const uint32_t query_dim, - T *scratch_query); + T *scratch_query); // Providing a default implementation for the virtual destructor because we don't // expect most metric implementations to need it. @@ -125,8 +124,7 @@ class AVXDistanceL2Float : public Distance DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t length) const; }; -template -class SlowDistanceL2 : public Distance +template class SlowDistanceL2 : public Distance { public: SlowDistanceL2() : Distance(diskann::Metric::L2) diff --git a/include/in_mem_data_store.h b/include/in_mem_data_store.h index 2dd261c2a..fb098c1aa 100644 --- a/include/in_mem_data_store.h +++ b/include/in_mem_data_store.h @@ -52,7 +52,7 @@ template class InMemDataStore : public AbstractDataStore * get_dist_fn(); + virtual Distance *get_dist_fn(); protected: virtual location_t expand(const location_t new_size) override; diff --git a/include/index.h b/include/index.h index 83d58f48a..d32314383 100644 --- a/include/index.h +++ b/include/index.h @@ -326,7 +326,7 @@ template clas 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 - // _num_frozen_pts is the number of points which are used as initial candidates + // _num_frozen_pts is the number of points which are used as initial candidates // when iterating to closest point(s). These are not visible externally and won't // be returned by search. At least 1 frozen point is needed for a dynamic index. // The frozen points have consecutive locations. See also _start below. diff --git a/src/abstract_data_store.cpp b/src/abstract_data_store.cpp index cb291e6d7..b895f1499 100644 --- a/src/abstract_data_store.cpp +++ b/src/abstract_data_store.cpp @@ -10,8 +10,9 @@ namespace diskann { -template -AbstractDataStore::AbstractDataStore(const location_t capacity, const size_t dim) : _capacity(capacity), _dim(dim) +template +AbstractDataStore::AbstractDataStore(const location_t capacity, const size_t dim) + : _capacity(capacity), _dim(dim) { } diff --git a/src/distance.cpp b/src/distance.cpp index 5ed336afc..88c314fd9 100644 --- a/src/distance.cpp +++ b/src/distance.cpp @@ -61,7 +61,6 @@ template Distance::~Distance() { } - // // Cosine distance functions. // @@ -221,8 +220,7 @@ float DistanceL2Float::compare(const float *a, const float *b, uint32_t size) co return result; } -template -float SlowDistanceL2::compare(const T *a, const T *b, uint32_t length) const +template float SlowDistanceL2::compare(const T *a, const T *b, uint32_t length) const { float result = 0.0f; for (uint32_t i = 0; i < length; i++) @@ -598,9 +596,6 @@ void AVXNormalizedCosineDistanceFloat::normalize_and_copy(const float *query_vec } } - - - // Get the right distance function for the given metric. template <> diskann::Distance *get_distance_function(diskann::Metric m) { diff --git a/src/in_mem_data_store.cpp b/src/in_mem_data_store.cpp index 84bf0ddfe..8ff07b854 100644 --- a/src/in_mem_data_store.cpp +++ b/src/in_mem_data_store.cpp @@ -344,8 +344,8 @@ template location_t InMemDataStore::calculate_medoid() return min_idx; } -template -Distance* InMemDataStore::get_dist_fn() { +template Distance *InMemDataStore::get_dist_fn() +{ return this->_distance_fn.get(); } diff --git a/src/index.cpp b/src/index.cpp index fc56b9aad..e8195cdad 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -94,7 +94,7 @@ Index::Index(Metric m, const size_t dim, const size_t max_point std::memset(_pq_data, 0, total_internal_points * _num_pq_chunks * sizeof(char)); } - _start = (uint32_t)_max_points; + _start = (uint32_t)_max_points; _final_graph.resize(total_internal_points); @@ -1522,54 +1522,54 @@ void Index::prune_all_neighbors(const uint32_t max_degree, cons } } - //REFACTOR - template - void Index::set_start_points(const T *data, size_t data_count) +// REFACTOR +template +void Index::set_start_points(const T *data, size_t data_count) { - std::unique_lock ul(_update_lock); - std::unique_lock tl(_tag_lock); - if (_nd > 0) - throw ANNException("Can not set starting point for a non-empty index", -1, __FUNCSIG__, __FILE__, __LINE__); + std::unique_lock ul(_update_lock); + std::unique_lock tl(_tag_lock); + if (_nd > 0) + throw ANNException("Can not set starting point for a non-empty index", -1, __FUNCSIG__, __FILE__, __LINE__); - if (data_count != _num_frozen_pts * _dim) - throw ANNException("Invalid number of points", -1, __FUNCSIG__, __FILE__, __LINE__); + if (data_count != _num_frozen_pts * _dim) + throw ANNException("Invalid number of points", -1, __FUNCSIG__, __FILE__, __LINE__); -// memcpy(_data + _aligned_dim * _max_points, data, _aligned_dim * sizeof(T) * _num_frozen_pts); - for (location_t i = _max_points; i < _max_points + _num_frozen_pts; i++) { - _data_store->set_vector(i, data + i*_dim); + // memcpy(_data + _aligned_dim * _max_points, data, _aligned_dim * sizeof(T) * _num_frozen_pts); + for (location_t i = _max_points; i < _max_points + _num_frozen_pts; i++) + { + _data_store->set_vector(i, data + i * _dim); } - _has_built = true; - diskann::cout << "Index start points set: #" << _num_frozen_pts << std::endl; - } - + _has_built = true; + diskann::cout << "Index start points set: #" << _num_frozen_pts << std::endl; +} - template - void Index::set_start_points_at_random(T radius, uint32_t random_seed) +template +void Index::set_start_points_at_random(T radius, uint32_t random_seed) { - std::mt19937 gen{random_seed}; - std::normal_distribution<> d{0.0, 1.0}; + std::mt19937 gen{random_seed}; + std::normal_distribution<> d{0.0, 1.0}; - std::vector points_data; - points_data.reserve(_dim * _num_frozen_pts); - std::vector real_vec(_dim); + std::vector points_data; + points_data.reserve(_dim * _num_frozen_pts); + std::vector real_vec(_dim); - for (size_t frozen_point = 0; frozen_point < _num_frozen_pts; frozen_point++) - { - double norm_sq = 0.0; - for (size_t i = 0; i < _dim; ++i) - { - auto r = d(gen); - real_vec[i] = r; - norm_sq += r * r; - } + for (size_t frozen_point = 0; frozen_point < _num_frozen_pts; frozen_point++) + { + double norm_sq = 0.0; + for (size_t i = 0; i < _dim; ++i) + { + auto r = d(gen); + real_vec[i] = r; + norm_sq += r * r; + } - const double norm = std::sqrt(norm_sq); - for (auto iter : real_vec) - points_data.push_back(static_cast(iter * radius / norm)); - } + const double norm = std::sqrt(norm_sq); + for (auto iter : real_vec) + points_data.push_back(static_cast(iter * radius / norm)); + } - set_start_points(points_data.data(), points_data.size()); - } + set_start_points(points_data.data(), points_data.size()); +} template void Index::build_with_data_populated(IndexWriteParameters ¶meters, const std::vector &tags) @@ -2931,133 +2931,133 @@ template void Index void Index::optimize_index_layout() +// template void Index::optimize_index_layout() //{ // use after build or load //} // REFACTOR: This should be an OptimizedDataStore class - template void Index::optimize_index_layout() +template void Index::optimize_index_layout() { // use after build or load - if (_dynamic_index) - { - throw diskann::ANNException("Optimize_index_layout not implemented for dyanmic indices", -1, __FUNCSIG__, - __FILE__, __LINE__); - } - - float* cur_vec = new float[_data_store->get_aligned_dim()]; - std::memset(cur_vec, 0, _data_store->get_aligned_dim()*sizeof(float)); - _data_len = (_data_store->get_aligned_dim() + 1) * sizeof(float); - _neighbor_len = (_max_observed_degree + 1) * sizeof(uint32_t); - _node_size = _data_len + _neighbor_len; - _opt_graph = new char[_node_size * _nd]; - DistanceFastL2 *dist_fast = (DistanceFastL2 *)_data_store->get_dist_fn(); - for (uint32_t i = 0; i < _nd; i++) - { - char *cur_node_offset = _opt_graph + i * _node_size; - _data_store->get_vector(i, (T*) cur_vec); - float cur_norm = dist_fast->norm((T*) cur_vec, _data_store->get_aligned_dim()); - std::memcpy(cur_node_offset, &cur_norm, sizeof(float)); - std::memcpy(cur_node_offset + sizeof(float), cur_vec, _data_len - sizeof(float)); - - cur_node_offset += _data_len; - uint32_t k = _final_graph[i].size(); - std::memcpy(cur_node_offset, &k, sizeof(uint32_t)); - std::memcpy(cur_node_offset + sizeof(uint32_t), _final_graph[i].data(), k * sizeof(uint32_t)); - std::vector().swap(_final_graph[i]); - } - _final_graph.clear(); - _final_graph.shrink_to_fit(); - delete[] cur_vec; - } + if (_dynamic_index) + { + throw diskann::ANNException("Optimize_index_layout not implemented for dyanmic indices", -1, __FUNCSIG__, + __FILE__, __LINE__); + } + + float *cur_vec = new float[_data_store->get_aligned_dim()]; + std::memset(cur_vec, 0, _data_store->get_aligned_dim() * sizeof(float)); + _data_len = (_data_store->get_aligned_dim() + 1) * sizeof(float); + _neighbor_len = (_max_observed_degree + 1) * sizeof(uint32_t); + _node_size = _data_len + _neighbor_len; + _opt_graph = new char[_node_size * _nd]; + DistanceFastL2 *dist_fast = (DistanceFastL2 *)_data_store->get_dist_fn(); + for (uint32_t i = 0; i < _nd; i++) + { + char *cur_node_offset = _opt_graph + i * _node_size; + _data_store->get_vector(i, (T *)cur_vec); + float cur_norm = dist_fast->norm((T *)cur_vec, _data_store->get_aligned_dim()); + std::memcpy(cur_node_offset, &cur_norm, sizeof(float)); + std::memcpy(cur_node_offset + sizeof(float), cur_vec, _data_len - sizeof(float)); + + cur_node_offset += _data_len; + uint32_t k = _final_graph[i].size(); + std::memcpy(cur_node_offset, &k, sizeof(uint32_t)); + std::memcpy(cur_node_offset + sizeof(uint32_t), _final_graph[i].data(), k * sizeof(uint32_t)); + std::vector().swap(_final_graph[i]); + } + _final_graph.clear(); + _final_graph.shrink_to_fit(); + delete[] cur_vec; +} // REFACTOR: once optimized layout becomes its own Data+Graph store, we should just invoke regular search -//template -//void Index::search_with_optimized_layout(const T *query, size_t K, size_t L, uint32_t *indices) +// template +// void Index::search_with_optimized_layout(const T *query, size_t K, size_t L, uint32_t *indices) //{ //} - template - void Index::search_with_optimized_layout(const T *query, size_t K, size_t L, uint32_t *indices) -{ - DistanceFastL2 *dist_fast = (DistanceFastL2 *) _data_store->get_dist_fn(); - - NeighborPriorityQueue retset(L); - std::vector init_ids(L); - - boost::dynamic_bitset<> flags{_nd, 0}; - uint32_t tmp_l = 0; - uint32_t *neighbors = (uint32_t *)(_opt_graph + _node_size * _start + _data_len); - uint32_t 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) - { - uint32_t id = rand() % _nd; - if (flags[id]) - continue; - flags[id] = true; - init_ids[tmp_l] = id; - tmp_l++; - } - - for (uint32_t i = 0; i < init_ids.size(); i++) - { - uint32_t id = init_ids[i]; - if (id >= _nd) - continue; - _mm_prefetch(_opt_graph + _node_size * id, _MM_HINT_T0); - } - L = 0; - for (uint32_t i = 0; i < init_ids.size(); i++) - { - uint32_t 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, (uint32_t)_data_store->get_aligned_dim()); - retset.insert(Neighbor(id, dist)); - flags[id] = true; - L++; - } - - while (retset.has_unexpanded_node()) - { - auto nbr = retset.closest_unexpanded(); - auto n = nbr.id; - _mm_prefetch(_opt_graph + _node_size * n + _data_len, _MM_HINT_T0); - neighbors = (uint32_t *)(_opt_graph + _node_size * n + _data_len); - uint32_t MaxM = *neighbors; - neighbors++; - for (uint32_t m = 0; m < MaxM; ++m) - _mm_prefetch(_opt_graph + _node_size * neighbors[m], _MM_HINT_T0); - for (uint32_t m = 0; m < MaxM; ++m) - { - uint32_t id = neighbors[m]; - if (flags[id]) - continue; - flags[id] = 1; - T *data = (T *)(_opt_graph + _node_size * id); - float norm = *data; - data++; - float dist = dist_fast->compare(query, data, norm, (uint32_t)_data_store->get_aligned_dim()); - Neighbor nn(id, dist); - retset.insert(nn); - } - } - - for (size_t i = 0; i < K; i++) - { - indices[i] = retset[i].id; - } - } +template +void Index::search_with_optimized_layout(const T *query, size_t K, size_t L, uint32_t *indices) +{ + DistanceFastL2 *dist_fast = (DistanceFastL2 *)_data_store->get_dist_fn(); + + NeighborPriorityQueue retset(L); + std::vector init_ids(L); + + boost::dynamic_bitset<> flags{_nd, 0}; + uint32_t tmp_l = 0; + uint32_t *neighbors = (uint32_t *)(_opt_graph + _node_size * _start + _data_len); + uint32_t 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) + { + uint32_t id = rand() % _nd; + if (flags[id]) + continue; + flags[id] = true; + init_ids[tmp_l] = id; + tmp_l++; + } + + for (uint32_t i = 0; i < init_ids.size(); i++) + { + uint32_t id = init_ids[i]; + if (id >= _nd) + continue; + _mm_prefetch(_opt_graph + _node_size * id, _MM_HINT_T0); + } + L = 0; + for (uint32_t i = 0; i < init_ids.size(); i++) + { + uint32_t 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, (uint32_t)_data_store->get_aligned_dim()); + retset.insert(Neighbor(id, dist)); + flags[id] = true; + L++; + } + + while (retset.has_unexpanded_node()) + { + auto nbr = retset.closest_unexpanded(); + auto n = nbr.id; + _mm_prefetch(_opt_graph + _node_size * n + _data_len, _MM_HINT_T0); + neighbors = (uint32_t *)(_opt_graph + _node_size * n + _data_len); + uint32_t MaxM = *neighbors; + neighbors++; + for (uint32_t m = 0; m < MaxM; ++m) + _mm_prefetch(_opt_graph + _node_size * neighbors[m], _MM_HINT_T0); + for (uint32_t m = 0; m < MaxM; ++m) + { + uint32_t id = neighbors[m]; + if (flags[id]) + continue; + flags[id] = 1; + T *data = (T *)(_opt_graph + _node_size * id); + float norm = *data; + data++; + float dist = dist_fast->compare(query, data, norm, (uint32_t)_data_store->get_aligned_dim()); + Neighbor nn(id, dist); + retset.insert(nn); + } + } + + for (size_t i = 0; i < K; i++) + { + indices[i] = retset[i].id; + } +} /* Internals of the library */ template const float Index::INDEX_GROWTH_FACTOR = 1.5f; From e7aa49b347280f939ed16a9ef7e9f1eee755696c Mon Sep 17 00:00:00 2001 From: harsha vardhan simhadri Date: Wed, 12 Apr 2023 13:19:08 -0700 Subject: [PATCH 40/69] clang format --- tests/build_memory_index.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/build_memory_index.cpp b/tests/build_memory_index.cpp index ba63bdeb7..3712350c3 100644 --- a/tests/build_memory_index.cpp +++ b/tests/build_memory_index.cpp @@ -76,7 +76,8 @@ int main(int argc, char **argv) { desc.add_options()("help,h", "Print information on arguments"); desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); - desc.add_options()("dist_fn", po::value(&dist_fn)->required(), "distance function "); + desc.add_options()("dist_fn", po::value(&dist_fn)->required(), + "distance function "); desc.add_options()("data_path", po::value(&data_path)->required(), "Input data file in bin format"); desc.add_options()("index_path_prefix", po::value(&index_path_prefix)->required(), From e6e5fc935a089276e7d132b79fcf144c61f9e9bb Mon Sep 17 00:00:00 2001 From: Gopal Srinivasa Date: Thu, 13 Apr 2023 23:38:20 +0530 Subject: [PATCH 41/69] Added slot manager file (empty) and fixed compile errors --- include/abstract_data_store.h | 8 +++++--- include/distance.h | 2 +- include/slot_manager.h | 1 + src/abstract_data_store.cpp | 6 ++++++ src/dll/CMakeLists.txt | 4 ++-- src/slot_manager.cpp | 1 + 6 files changed, 16 insertions(+), 6 deletions(-) create mode 100644 include/slot_manager.h create mode 100644 src/slot_manager.cpp diff --git a/include/abstract_data_store.h b/include/abstract_data_store.h index ee1ad4d33..53baa28be 100644 --- a/include/abstract_data_store.h +++ b/include/abstract_data_store.h @@ -7,11 +7,13 @@ #include #include "types.h" +#include "windows_customizations.h" namespace diskann { -template class AbstractDataStore +template +class AbstractDataStore { public: AbstractDataStore(const location_t capacity, const size_t dim); @@ -25,9 +27,9 @@ template class AbstractDataStore // can discard the empty locations before saving. virtual size_t save(const std::string &filename, const location_t num_pts) = 0; - virtual location_t capacity() const; + DISKANN_DLLEXPORT virtual location_t capacity() const; - virtual size_t get_dims() const; + DISKANN_DLLEXPORT virtual size_t get_dims() const; // Implementers can choose to return _dim if they are not // concerned about memory alignment. diff --git a/include/distance.h b/include/distance.h index 1ba4f32fb..d178846b8 100644 --- a/include/distance.h +++ b/include/distance.h @@ -24,7 +24,7 @@ template class Distance // Needed only for COSINE-BYTE and INNER_PRODUCT-BYTE DISKANN_DLLEXPORT virtual float compare(const T *a, const T *b, const float normA, const float normB, - uint32_t length) const = 0; + uint32_t length) const; // For MIPS, normalization adds an extra dimension to the vectors. // This function lets callers know if the normalization process diff --git a/include/slot_manager.h b/include/slot_manager.h new file mode 100644 index 000000000..7b9637ef9 --- /dev/null +++ b/include/slot_manager.h @@ -0,0 +1 @@ +#pragma once \ No newline at end of file diff --git a/src/abstract_data_store.cpp b/src/abstract_data_store.cpp index b895f1499..faee00ac3 100644 --- a/src/abstract_data_store.cpp +++ b/src/abstract_data_store.cpp @@ -42,4 +42,10 @@ template location_t AbstractDataStore::resize(const lo } } +#ifdef _WINDOWS +template DISKANN_DLLEXPORT class AbstractDataStore; +template DISKANN_DLLEXPORT class AbstractDataStore; +template DISKANN_DLLEXPORT class AbstractDataStore; +#endif + } // namespace diskann diff --git a/src/dll/CMakeLists.txt b/src/dll/CMakeLists.txt index 0f511a8ae..9e97a35b2 100644 --- a/src/dll/CMakeLists.txt +++ b/src/dll/CMakeLists.txt @@ -1,10 +1,10 @@ #Copyright(c) Microsoft Corporation.All rights reserved. #Licensed under the MIT license. -add_library(${PROJECT_NAME} SHARED dllmain.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 ../pq.cpp ../pq_flash_index.cpp ../logger.cpp ../utils.cpp ../windows_aligned_file_reader.cpp ../distance.cpp ../memory_mapper.cpp ../index.cpp ../in_mem_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) + ../ann_exception.cpp ../natural_number_set.cpp ../natural_number_map.cpp ../scratch.cpp ../slot_manager.cpp) set(TARGET_DIR "$<$:${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}>$<$:${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}>") set(DISKANN_DLL_IMPLIB "${TARGET_DIR}/${PROJECT_NAME}.lib") diff --git a/src/slot_manager.cpp b/src/slot_manager.cpp new file mode 100644 index 000000000..e4dc4a48c --- /dev/null +++ b/src/slot_manager.cpp @@ -0,0 +1 @@ +#include "slot_manager.h" \ No newline at end of file From b7809058108b99346f5ab21ab510e36bdee8785c Mon Sep 17 00:00:00 2001 From: ravishankar Date: Fri, 14 Apr 2023 00:09:36 +0000 Subject: [PATCH 42/69] fixed a linux compile error --- src/abstract_data_store.cpp | 4 ---- src/distance.cpp | 4 ++-- src/in_mem_data_store.cpp | 8 ++++---- src/index.cpp | 10 +++++----- 4 files changed, 11 insertions(+), 15 deletions(-) diff --git a/src/abstract_data_store.cpp b/src/abstract_data_store.cpp index faee00ac3..c8864b2aa 100644 --- a/src/abstract_data_store.cpp +++ b/src/abstract_data_store.cpp @@ -1,8 +1,6 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT license. -#pragma once - #include #include "abstract_data_store.h" @@ -42,10 +40,8 @@ template location_t AbstractDataStore::resize(const lo } } -#ifdef _WINDOWS template DISKANN_DLLEXPORT class AbstractDataStore; template DISKANN_DLLEXPORT class AbstractDataStore; template DISKANN_DLLEXPORT class AbstractDataStore; -#endif } // namespace diskann diff --git a/src/distance.cpp b/src/distance.cpp index 88c314fd9..2990bb128 100644 --- a/src/distance.cpp +++ b/src/distance.cpp @@ -572,7 +572,7 @@ bool AVXNormalizedCosineDistanceFloat::normalization_required() const void AVXNormalizedCosineDistanceFloat::normalize_data_for_build(float *original_data, const uint32_t orig_dim, const uint32_t num_points) { - for (auto i = 0; i < num_points; i++) + for (uint32_t i = 0; i < num_points; i++) { normalize((float *)(original_data + i * orig_dim), orig_dim); } @@ -590,7 +590,7 @@ void AVXNormalizedCosineDistanceFloat::normalize_and_copy(const float *query_vec float norm = get_norm(query_vec, query_dim); - for (auto i = 0; i < query_dim; i++) + for (uint32_t i = 0; i < query_dim; i++) { query_target[i] = query_vec[i] / norm; } diff --git a/src/in_mem_data_store.cpp b/src/in_mem_data_store.cpp index 8ff07b854..7e773ec44 100644 --- a/src/in_mem_data_store.cpp +++ b/src/in_mem_data_store.cpp @@ -97,7 +97,7 @@ template size_t InMemDataStore::save(const std::string template void InMemDataStore::populate_data(const data_t *vectors, const location_t num_pts) { - for (auto i = 0; i < num_pts; i++) + for (location_t i = 0; i < num_pts; i++) { memset(_data + i * _aligned_dim, 0, _aligned_dim * sizeof(data_t)); std::memmove(_data + i * _aligned_dim, vectors + i * this->_dim, this->_dim * sizeof(data_t)); @@ -173,7 +173,7 @@ template void InMemDataStore::get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, float *distances) const { - for (auto i = 0; i < location_count; i++) + for (location_t i = 0; i < location_count; i++) { distances[i] = _distance_fn->compare(query, _data + locations[i] * _aligned_dim, this->_aligned_dim); } @@ -246,11 +246,11 @@ void InMemDataStore::reposition_points(const location_t old_location_sta return; } - // Update pointers to the moved nodes. Note: the computation is correct even +/* // Update pointers to the moved nodes. Note: the computation is correct even // when new_location_start < old_location_start given the C++ uint32_t // integer arithmetic rules. const uint32_t location_delta = new_location_start - old_location_start; - +*/ // The [start, end) interval which will contain obsolete points to be // cleared. uint32_t mem_clear_loc_start = old_location_start; diff --git a/src/index.cpp b/src/index.cpp index e8195cdad..1e7a83fdc 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -51,10 +51,10 @@ template Index::Index(Metric m, const size_t dim, const size_t max_points, 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 size_t num_frozen_pts) - : _dist_metric(m), _dim(dim), _num_frozen_pts(num_frozen_pts), _max_points(max_points), - _dynamic_index(dynamic_index), _enable_tags(enable_tags), _indexingMaxC(DEFAULT_MAXC), _query_scratch(nullptr), - _conc_consolidate(concurrent_consolidate), _delete_set(new tsl::robin_set), _pq_dist(pq_dist_build), - _use_opq(use_opq), _num_pq_chunks(num_pq_chunks) + : _dist_metric(m), _dim(dim), _max_points(max_points), _num_frozen_pts(num_frozen_pts), + _dynamic_index(dynamic_index), _enable_tags(enable_tags), _indexingMaxC(DEFAULT_MAXC), _query_scratch(nullptr), _pq_dist(pq_dist_build), + _use_opq(use_opq), _num_pq_chunks(num_pq_chunks), + _delete_set(new tsl::robin_set), _conc_consolidate(concurrent_consolidate) { if (dynamic_index && !enable_tags) { @@ -1903,7 +1903,7 @@ void Index::build_filtered_index(const char *filename, const st std::unordered_map> label_to_points; - for (int lbl = 0; lbl < _labels.size(); lbl++) + for (uint32_t lbl = 0; lbl < _labels.size(); lbl++) { auto itr = _labels.begin(); std::advance(itr, lbl); From 23d5a10218b69ad82d195d08628030c90cdf9300 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Fri, 14 Apr 2023 00:14:11 +0000 Subject: [PATCH 43/69] clang --- include/abstract_data_store.h | 3 +-- src/in_mem_data_store.cpp | 10 +++++----- src/index.cpp | 6 +++--- 3 files changed, 9 insertions(+), 10 deletions(-) diff --git a/include/abstract_data_store.h b/include/abstract_data_store.h index 53baa28be..860e86f91 100644 --- a/include/abstract_data_store.h +++ b/include/abstract_data_store.h @@ -12,8 +12,7 @@ namespace diskann { -template -class AbstractDataStore +template class AbstractDataStore { public: AbstractDataStore(const location_t capacity, const size_t dim); diff --git a/src/in_mem_data_store.cpp b/src/in_mem_data_store.cpp index 7e773ec44..c3b2f5369 100644 --- a/src/in_mem_data_store.cpp +++ b/src/in_mem_data_store.cpp @@ -246,11 +246,11 @@ void InMemDataStore::reposition_points(const location_t old_location_sta return; } -/* // Update pointers to the moved nodes. Note: the computation is correct even - // when new_location_start < old_location_start given the C++ uint32_t - // integer arithmetic rules. - const uint32_t location_delta = new_location_start - old_location_start; -*/ + /* // Update pointers to the moved nodes. Note: the computation is correct even + // when new_location_start < old_location_start given the C++ uint32_t + // integer arithmetic rules. + const uint32_t location_delta = new_location_start - old_location_start; + */ // The [start, end) interval which will contain obsolete points to be // cleared. uint32_t mem_clear_loc_start = old_location_start; diff --git a/src/index.cpp b/src/index.cpp index 1e7a83fdc..fa2c0ea2a 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -52,9 +52,9 @@ Index::Index(Metric m, const size_t dim, const size_t max_point const bool enable_tags, const bool concurrent_consolidate, const bool pq_dist_build, const size_t num_pq_chunks, const bool use_opq, const size_t num_frozen_pts) : _dist_metric(m), _dim(dim), _max_points(max_points), _num_frozen_pts(num_frozen_pts), - _dynamic_index(dynamic_index), _enable_tags(enable_tags), _indexingMaxC(DEFAULT_MAXC), _query_scratch(nullptr), _pq_dist(pq_dist_build), - _use_opq(use_opq), _num_pq_chunks(num_pq_chunks), - _delete_set(new tsl::robin_set), _conc_consolidate(concurrent_consolidate) + _dynamic_index(dynamic_index), _enable_tags(enable_tags), _indexingMaxC(DEFAULT_MAXC), _query_scratch(nullptr), + _pq_dist(pq_dist_build), _use_opq(use_opq), _num_pq_chunks(num_pq_chunks), + _delete_set(new tsl::robin_set), _conc_consolidate(concurrent_consolidate) { if (dynamic_index && !enable_tags) { From 362e896a8970492185a60470dc82dd3f445aa310 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Sun, 16 Apr 2023 18:56:44 +0000 Subject: [PATCH 44/69] debugging workflow failure --- tests/test_streaming_scenario.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_streaming_scenario.cpp b/tests/test_streaming_scenario.cpp index 2e8d6363c..a34c7e8ff 100644 --- a/tests/test_streaming_scenario.cpp +++ b/tests/test_streaming_scenario.cpp @@ -202,7 +202,7 @@ void build_incremental_index(const std::string &data_path, const uint32_t L, con } if (num_points < max_points_to_insert) - throw diskann::ANNException("num_points < max_points_to_insert", -1, __FUNCSIG__, __FILE__, __LINE__); + throw diskann::ANNException("num_points " + std::to_string(num_points) + " < max_points_to_insert " + std::to_string(max_points_to_insert), -1, __FUNCSIG__, __FILE__, __LINE__); if (max_points_to_insert < active_window + consolidate_interval) throw diskann::ANNException("ERROR: max_points_to_insert < " From 85e008decb25ee0191019aee3fc682cc623fb5a8 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Sun, 16 Apr 2023 18:59:02 +0000 Subject: [PATCH 45/69] clang --- tests/test_streaming_scenario.cpp | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/test_streaming_scenario.cpp b/tests/test_streaming_scenario.cpp index a34c7e8ff..f82e10f49 100644 --- a/tests/test_streaming_scenario.cpp +++ b/tests/test_streaming_scenario.cpp @@ -202,7 +202,9 @@ void build_incremental_index(const std::string &data_path, const uint32_t L, con } if (num_points < max_points_to_insert) - throw diskann::ANNException("num_points " + std::to_string(num_points) + " < max_points_to_insert " + std::to_string(max_points_to_insert), -1, __FUNCSIG__, __FILE__, __LINE__); + throw diskann::ANNException("num_points " + std::to_string(num_points) + " < max_points_to_insert " + + std::to_string(max_points_to_insert), + -1, __FUNCSIG__, __FILE__, __LINE__); if (max_points_to_insert < active_window + consolidate_interval) throw diskann::ANNException("ERROR: max_points_to_insert < " From 0b8fc55475bc00a0b4dba65e504cab36fff4b11f Mon Sep 17 00:00:00 2001 From: ravishankar Date: Sun, 16 Apr 2023 19:14:32 +0000 Subject: [PATCH 46/69] more debug --- tests/test_streaming_scenario.cpp | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_streaming_scenario.cpp b/tests/test_streaming_scenario.cpp index f82e10f49..09204cbd7 100644 --- a/tests/test_streaming_scenario.cpp +++ b/tests/test_streaming_scenario.cpp @@ -194,6 +194,7 @@ void build_incremental_index(const std::string &data_path, const uint32_t L, con size_t num_points; diskann::get_bin_metadata(data_path, num_points, dim); + diskann::cout<<"metadata: file " << data_path << " has " << num_points << " points in " << dim << " dims" << std::endl; aligned_dim = ROUND_UP(dim, 8); if (max_points_to_insert == 0) From dc183e91bb4a9c10e2865a49a1dbcd01d27d2a76 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Sun, 16 Apr 2023 19:16:02 +0000 Subject: [PATCH 47/69] more debug --- tests/test_streaming_scenario.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_streaming_scenario.cpp b/tests/test_streaming_scenario.cpp index 09204cbd7..389712afa 100644 --- a/tests/test_streaming_scenario.cpp +++ b/tests/test_streaming_scenario.cpp @@ -194,7 +194,8 @@ void build_incremental_index(const std::string &data_path, const uint32_t L, con size_t num_points; diskann::get_bin_metadata(data_path, num_points, dim); - diskann::cout<<"metadata: file " << data_path << " has " << num_points << " points in " << dim << " dims" << std::endl; + diskann::cout << "metadata: file " << data_path << " has " << num_points << " points in " << dim << " dims" + << std::endl; aligned_dim = ROUND_UP(dim, 8); if (max_points_to_insert == 0) From a191354ab81d1b0a8080c25c9fe7332a513a2560 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Sun, 16 Apr 2023 19:27:36 +0000 Subject: [PATCH 48/69] debug for workflow --- .github/workflows/pr-test.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index f64a201e1..8cb04bb76 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -172,7 +172,7 @@ jobs: - name: test a streaming index if: success() || failure() run: | - ${{ env.diskann_built_tests }}/test_streaming_scenario --data_type int8 --dist_fn l2 --data_path rand_int8_10D_10K_norm50.0.bin --index_path_prefix index_stream -R 64 -L 600 --alpha 1.2 --insert_threads 4 --consolidate_threads 4 --max_points_to_insert 10000 --active_window 4000 --consolidate_interval 2000 --start_point_norm 200 + ${{ env.diskann_built_tests }}/test_streaming_scenario --data_type int8 --dist_fn l2 --data_path ./rand_int8_10D_10K_norm50.0.bin --index_path_prefix index_stream -R 64 -L 600 --alpha 1.2 --insert_threads 4 --consolidate_threads 4 --max_points_to_insert 10000 --active_window 4000 --consolidate_interval 2000 --start_point_norm 200 ${{ env.diskann_built_utils }}/compute_groundtruth --data_type int8 --dist_fn l2 --base_file index_stream.after-streaming-act4000-cons2000-max10000.data --query_file rand_int8_10D_1K_norm50.0.bin --K 100 --gt_file gt100_base-act4000-cons2000-max10000 --tags_file index_stream.after-streaming-act4000-cons2000-max10000.tags ${{ env.diskann_built_tests }}/search_memory_index --data_type int8 --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix index_stream.after-streaming-act4000-cons2000-max10000 --result_path res_stream --query_file ./rand_int8_10D_1K_norm50.0.bin --gt_file gt100_base-act4000-cons2000-max10000 -K 10 -L 20 40 60 80 100 -T 64 --dynamic true --tags 1 From 6ea93ea6796805ab23a275e07e7b35dfb117d2db Mon Sep 17 00:00:00 2001 From: harsha vardhan simhadri Date: Mon, 17 Apr 2023 01:11:34 -0700 Subject: [PATCH 49/69] remove slot manager --- include/slot_manager.h | 1 - src/slot_manager.cpp | 1 - 2 files changed, 2 deletions(-) delete mode 100644 include/slot_manager.h delete mode 100644 src/slot_manager.cpp diff --git a/include/slot_manager.h b/include/slot_manager.h deleted file mode 100644 index 7b9637ef9..000000000 --- a/include/slot_manager.h +++ /dev/null @@ -1 +0,0 @@ -#pragma once \ No newline at end of file diff --git a/src/slot_manager.cpp b/src/slot_manager.cpp deleted file mode 100644 index e4dc4a48c..000000000 --- a/src/slot_manager.cpp +++ /dev/null @@ -1 +0,0 @@ -#include "slot_manager.h" \ No newline at end of file From fc2e5ad2522e047575298bc628dcc4e75cd98d4d Mon Sep 17 00:00:00 2001 From: Gopal Srinivasa Date: Tue, 18 Apr 2023 11:12:58 +0530 Subject: [PATCH 50/69] Removed the #ifdef WINDOWS directive from class definitions --- src/abstract_data_store.cpp | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/abstract_data_store.cpp b/src/abstract_data_store.cpp index faee00ac3..46674149c 100644 --- a/src/abstract_data_store.cpp +++ b/src/abstract_data_store.cpp @@ -42,10 +42,9 @@ template location_t AbstractDataStore::resize(const lo } } -#ifdef _WINDOWS + template DISKANN_DLLEXPORT class AbstractDataStore; template DISKANN_DLLEXPORT class AbstractDataStore; template DISKANN_DLLEXPORT class AbstractDataStore; -#endif } // namespace diskann From e301aebd2564d42592ce3e4daa2084a268c3b87b Mon Sep 17 00:00:00 2001 From: Gopal Srinivasa Date: Tue, 18 Apr 2023 17:44:29 +0530 Subject: [PATCH 51/69] Refactoring alignment factor into distance hierarchy --- include/abstract_data_store.h | 3 +++ include/distance.h | 7 +++++++ include/in_mem_data_store.h | 10 +++++----- include/scratch.h | 5 ++++- include/slot_manager.h | 3 +++ src/distance.cpp | 5 +++++ src/in_mem_data_store.cpp | 13 ++++++++++++- src/index.cpp | 4 +++- src/scratch.cpp | 9 +++++++-- src/slot_manager.cpp | 1 + 10 files changed, 50 insertions(+), 10 deletions(-) create mode 100644 include/slot_manager.h create mode 100644 src/slot_manager.cpp diff --git a/include/abstract_data_store.h b/include/abstract_data_store.h index 860e86f91..4c59efd2f 100644 --- a/include/abstract_data_store.h +++ b/include/abstract_data_store.h @@ -75,6 +75,9 @@ template class AbstractDataStore // Returns the point in the dataset that is closest to the mean of all points in the dataset virtual location_t calculate_medoid() const = 0; + //search helpers + virtual size_t get_alignment_factor() const = 0; + protected: // Expand the datastore to new_num_points. Returns the new capacity created, which should be == new_num_points // in the normal case. Implementers can also return _capacity to indicate that there are not implementing this diff --git a/include/distance.h b/include/distance.h index d178846b8..d4fd08f51 100644 --- a/include/distance.h +++ b/include/distance.h @@ -56,12 +56,19 @@ template class Distance DISKANN_DLLEXPORT virtual void normalize_vector_for_search(const T *query_vec, const uint32_t query_dim, T *scratch_query); + //If an algorithm has a requirement that some data be aligned to a certain boundary + //it can use this function to indicate that requirement. Currently, we are setting it to 8 + //because that works well for AVX2. If we have AVX512 implementations of distance algos, + //they might have to set this to 16 (depending on how they are implemented) + DISKANN_DLLEXPORT virtual size_t get_required_alignment() const; + // Providing a default implementation for the virtual destructor because we don't // expect most metric implementations to need it. DISKANN_DLLEXPORT virtual ~Distance(); protected: diskann::Metric _distance_metric; + size_t _alignment_factor = 8; }; class DistanceCosineInt8 : public Distance diff --git a/include/in_mem_data_store.h b/include/in_mem_data_store.h index fb098c1aa..60dd78313 100644 --- a/include/in_mem_data_store.h +++ b/include/in_mem_data_store.h @@ -26,10 +26,8 @@ template class InMemDataStore : public AbstractDataStore class InMemDataStore : public AbstractDataStore *get_dist_fn(); + virtual size_t get_alignment_factor() const override; + protected: virtual location_t expand(const location_t new_size) override; virtual location_t shrink(const location_t new_size) override; @@ -66,7 +66,7 @@ template class InMemDataStore : public AbstractDataStore class InMemQueryScratch { public: ~InMemQueryScratch(); - InMemQueryScratch(uint32_t search_l, uint32_t indexing_l, uint32_t r, uint32_t maxc, size_t dim, + //REFACTOR TODO: move all parameters to a new class. + InMemQueryScratch(uint32_t search_l, uint32_t indexing_l, uint32_t r, uint32_t maxc, size_t dim, size_t aligned_dim, size_t alignment_factor, bool init_pq_scratch = false); void resize_for_new_L(uint32_t new_search_l); void clear(); diff --git a/include/slot_manager.h b/include/slot_manager.h new file mode 100644 index 000000000..831382462 --- /dev/null +++ b/include/slot_manager.h @@ -0,0 +1,3 @@ +#pragma once + +#include \ No newline at end of file diff --git a/src/distance.cpp b/src/distance.cpp index 2990bb128..8e68a7b09 100644 --- a/src/distance.cpp +++ b/src/distance.cpp @@ -57,6 +57,11 @@ void Distance::normalize_vector_for_search(const T *query_vec, const uint32_t std::memcpy(scratch_query, query_vec, query_dim * sizeof(T)); } +template size_t Distance::get_required_alignment() const +{ + return _alignment_factor; +} + template Distance::~Distance() { } diff --git a/src/in_mem_data_store.cpp b/src/in_mem_data_store.cpp index c3b2f5369..0e7831b6c 100644 --- a/src/in_mem_data_store.cpp +++ b/src/in_mem_data_store.cpp @@ -12,8 +12,9 @@ namespace diskann template InMemDataStore::InMemDataStore(const location_t num_points, const size_t dim, std::shared_ptr> distance_metric) - : AbstractDataStore(num_points, dim), _aligned_dim(ROUND_UP(dim, 8)), _distance_fn(distance_metric) + : AbstractDataStore(num_points, dim), _distance_fn(distance_metric) { + _aligned_dim = ROUND_UP(dim, _distance_fn->get_required_alignment()); alloc_aligned(((void **)&_data), this->_capacity * _aligned_dim * sizeof(data_t), 8 * sizeof(data_t)); std::memset(_data, 0, this->_capacity * _aligned_dim * sizeof(data_t)); } @@ -26,6 +27,16 @@ template InMemDataStore::~InMemDataStore() } } +template size_t InMemDataStore::get_aligned_dim() const +{ + return _aligned_dim; +} + +template size_t InMemDataStore::get_alignment_factor() const +{ + return _distance_fn->get_required_alignment(); +} + template location_t InMemDataStore::load(const std::string &filename) { return load_impl(filename); diff --git a/src/index.cpp b/src/index.cpp index 04dadd8f8..92eeff9bf 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -164,7 +164,7 @@ void Index::initialize_query_scratch(uint32_t num_threads, uint { for (uint32_t i = 0; i < num_threads; i++) { - auto scratch = new InMemQueryScratch(search_l, indexing_l, r, maxc, dim, _pq_dist); + auto scratch = new InMemQueryScratch(search_l, indexing_l, r, maxc, dim, _data_store->get_aligned_dim(), _data_store->get_alignment_factor(), _pq_dist); _query_scratch.push(scratch); } } @@ -855,6 +855,8 @@ std::pair Index::iterate_to_fixed_point( normalize((float *)aligned_query, _dim); } + + float *query_float = nullptr; float *query_rotated = nullptr; float *pq_dists = nullptr; diff --git a/src/scratch.cpp b/src/scratch.cpp index 7b14b14fa..cc0e74e19 100644 --- a/src/scratch.cpp +++ b/src/scratch.cpp @@ -13,6 +13,7 @@ namespace diskann // template InMemQueryScratch::InMemQueryScratch(uint32_t search_l, uint32_t indexing_l, uint32_t r, uint32_t maxc, size_t dim, + size_t aligned_dim, size_t alignment_factor, bool init_pq_scratch) : _L(0), _R(r), _maxc(maxc) { @@ -24,8 +25,12 @@ InMemQueryScratch::InMemQueryScratch(uint32_t search_l, uint32_t indexing_l, throw diskann::ANNException(ss.str(), -1); } - auto aligned_dim = ROUND_UP(dim, 8); - alloc_aligned(((void **)&_aligned_query), aligned_dim * sizeof(T), 8 * sizeof(T)); + //REFACTOR + //auto aligned_dim = ROUND_UP(dim, 8); + //alloc_aligned(((void **)&_aligned_query), aligned_dim * sizeof(T), 8 * sizeof(T)); + //memset(_aligned_query, 0, aligned_dim * sizeof(T)); + + alloc_aligned(((void **)&_aligned_query), aligned_dim * sizeof(T), alignment_factor * sizeof(T)); memset(_aligned_query, 0, aligned_dim * sizeof(T)); if (init_pq_scratch) diff --git a/src/slot_manager.cpp b/src/slot_manager.cpp new file mode 100644 index 000000000..e4dc4a48c --- /dev/null +++ b/src/slot_manager.cpp @@ -0,0 +1 @@ +#include "slot_manager.h" \ No newline at end of file From 879a37e97081e512668bd49a5f98ab94e0057a11 Mon Sep 17 00:00:00 2001 From: Gopal Srinivasa Date: Tue, 18 Apr 2023 22:12:10 +0530 Subject: [PATCH 52/69] Fixing cosine distance --- include/distance.h | 22 +++++++++++----------- src/distance.cpp | 14 +++++++------- src/in_mem_data_store.cpp | 12 ++++++------ src/index.cpp | 19 ++++++++++++------- 4 files changed, 36 insertions(+), 31 deletions(-) diff --git a/include/distance.h b/include/distance.h index d4fd08f51..da820a4f1 100644 --- a/include/distance.h +++ b/include/distance.h @@ -35,25 +35,25 @@ template class Distance // This is for efficiency. If no normalization is required, the callers // can simply ignore the normalize_data_for_build() function. - DISKANN_DLLEXPORT virtual bool normalization_required() const; + DISKANN_DLLEXPORT virtual bool preprocessing_required() const; - // Check the normalization_required() function before calling this. + // Check the preprocessing_required() function before calling this. // Clients can call the function like this: // - // if (metric->normalization_required()){ + // if (metric->preprocessing_required()){ // T* normalized_data_batch; // Split data into batches of batch_size and for each, call: - // metric->normalize_data_for_build(data_batch, batch_size); + // metric->preprocess_base_points(data_batch, batch_size); // // TODO: This does not take into account the case for SSD inner product // where the dimensions change after normalization. - DISKANN_DLLEXPORT virtual void normalize_data_for_build(T *original_data, const uint32_t orig_dim, - const uint32_t num_points); + DISKANN_DLLEXPORT virtual void preprocess_base_points(T *original_data, const size_t orig_dim, + const size_t num_points); // Invokes normalization for a single vector during search. The scratch space // has to be created by the caller keeping track of the fact that normalization // might change the dimension of the query vector. - DISKANN_DLLEXPORT virtual void normalize_vector_for_search(const T *query_vec, const uint32_t query_dim, + DISKANN_DLLEXPORT virtual void preprocess_query(const T *query_vec, const size_t query_dim, T *scratch_query); //If an algorithm has a requirement that some data be aligned to a certain boundary @@ -221,12 +221,12 @@ class AVXNormalizedCosineDistanceFloat : public Distance } DISKANN_DLLEXPORT virtual uint32_t post_normalization_dimension(uint32_t orig_dimension) const override; - DISKANN_DLLEXPORT virtual bool normalization_required() const; + DISKANN_DLLEXPORT virtual bool preprocessing_required() const; - DISKANN_DLLEXPORT virtual void normalize_data_for_build(float *original_data, const uint32_t orig_dim, - const uint32_t num_points) override; + DISKANN_DLLEXPORT virtual void preprocess_base_points(float *original_data, const size_t orig_dim, + const size_t num_points) override; - DISKANN_DLLEXPORT virtual void normalize_vector_for_search(const float *query_vec, const uint32_t query_dim, + DISKANN_DLLEXPORT virtual void preprocess_query(const float *query_vec, const size_t query_dim, float *scratch_query_vector) override; }; diff --git a/src/distance.cpp b/src/distance.cpp index 8e68a7b09..120801568 100644 --- a/src/distance.cpp +++ b/src/distance.cpp @@ -41,18 +41,18 @@ template diskann::Metric Distance::get_metric() const return _distance_metric; } -template bool Distance::normalization_required() const +template bool Distance::preprocessing_required() const { return false; } template -void Distance::normalize_data_for_build(T *original_data, const uint32_t orig_dim, const uint32_t num_points) +void Distance::preprocess_base_points(T *original_data, const size_t orig_dim, const size_t num_points) { } template -void Distance::normalize_vector_for_search(const T *query_vec, const uint32_t query_dim, T *scratch_query) +void Distance::preprocess_query(const T *query_vec, const size_t query_dim, T *scratch_query) { std::memcpy(scratch_query, query_vec, query_dim * sizeof(T)); } @@ -570,12 +570,12 @@ uint32_t AVXNormalizedCosineDistanceFloat::post_normalization_dimension(uint32_t { return orig_dimension; } -bool AVXNormalizedCosineDistanceFloat::normalization_required() const +bool AVXNormalizedCosineDistanceFloat::preprocessing_required() const { return true; } -void AVXNormalizedCosineDistanceFloat::normalize_data_for_build(float *original_data, const uint32_t orig_dim, - const uint32_t num_points) +void AVXNormalizedCosineDistanceFloat::preprocess_base_points(float *original_data, const size_t orig_dim, + const size_t num_points) { for (uint32_t i = 0; i < num_points; i++) { @@ -583,7 +583,7 @@ void AVXNormalizedCosineDistanceFloat::normalize_data_for_build(float *original_ } } -void AVXNormalizedCosineDistanceFloat::normalize_vector_for_search(const float *query_vec, const uint32_t query_dim, +void AVXNormalizedCosineDistanceFloat::preprocess_query(const float *query_vec, const size_t query_dim, float *query_scratch) { normalize_and_copy(query_vec, query_dim, query_scratch); diff --git a/src/in_mem_data_store.cpp b/src/in_mem_data_store.cpp index 0e7831b6c..470acc342 100644 --- a/src/in_mem_data_store.cpp +++ b/src/in_mem_data_store.cpp @@ -114,9 +114,9 @@ template void InMemDataStore::populate_data(const data std::memmove(_data + i * _aligned_dim, vectors + i * this->_dim, this->_dim * sizeof(data_t)); } - if (_distance_fn->normalization_required()) + if (_distance_fn->preprocessing_required()) { - _distance_fn->normalize_data_for_build(_data, this->_dim, num_pts); + _distance_fn->preprocess_base_points(_data, this->_dim, num_pts); } } @@ -142,9 +142,9 @@ template void InMemDataStore::populate_data(const std: throw diskann::ANNException(ss.str(), -1); } - if (_distance_fn->normalization_required()) + if (_distance_fn->preprocessing_required()) { - _distance_fn->normalize_data_for_build(_data, this->_dim, this->capacity()); + _distance_fn->preprocess_base_points(_data, this->_dim, this->capacity()); } } @@ -164,9 +164,9 @@ template void InMemDataStore::set_vector(const locatio size_t offset_in_data = loc * _aligned_dim; memset(_data + offset_in_data, 0, _aligned_dim * sizeof(data_t)); memcpy(_data + offset_in_data, vector, this->_dim * sizeof(data_t)); - if (_distance_fn->normalization_required()) + if (_distance_fn->preprocessing_required()) { - _distance_fn->normalize_data_for_build(_data + offset_in_data, _aligned_dim, 1); + _distance_fn->preprocess_base_points(_data + offset_in_data, _aligned_dim, 1); } } diff --git a/src/index.cpp b/src/index.cpp index 92eeff9bf..c0b68c14f 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -848,14 +848,16 @@ std::pair Index::iterate_to_fixed_point( std::vector &id_scratch = scratch->id_scratch(); std::vector &dist_scratch = scratch->dist_scratch(); assert(id_scratch.size() == 0); - T *aligned_query = scratch->aligned_query(); - memcpy(aligned_query, query, _dim * sizeof(T)); - if (_normalize_vecs) - { - normalize((float *)aligned_query, _dim); - } - + //REFACTOR + //T *aligned_query = scratch->aligned_query(); + //memcpy(aligned_query, query, _dim * sizeof(T)); + //if (_normalize_vecs) + //{ + // normalize((float *)aligned_query, _dim); + //} + + T* aligned_query = scratch->aligned_query(); float *query_float = nullptr; float *query_rotated = nullptr; @@ -1989,6 +1991,9 @@ std::pair Index::search(const T *query, con std::shared_lock lock(_update_lock); + if (_distance->preprocessing_required()){ + _distance->preprocess_query(query, _data_store->get_dims(), scratch->aligned_query()); + } auto retval = iterate_to_fixed_point(query, L, init_ids, scratch, false, unused_filter_label, true); NeighborPriorityQueue &best_L_nodes = scratch->best_l_nodes(); From 7375708363cdfdb9390e13c219715e1fb46eaa8c Mon Sep 17 00:00:00 2001 From: Gopal Srinivasa Date: Tue, 18 Apr 2023 22:33:39 +0530 Subject: [PATCH 53/69] Ensuring we call preprocess_query always --- src/index.cpp | 20 +++++++++++--------- src/scratch.cpp | 1 + 2 files changed, 12 insertions(+), 9 deletions(-) diff --git a/src/index.cpp b/src/index.cpp index c0b68c14f..cc74cd52a 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -1991,10 +1991,9 @@ std::pair Index::search(const T *query, con std::shared_lock lock(_update_lock); - if (_distance->preprocessing_required()){ - _distance->preprocess_query(query, _data_store->get_dims(), scratch->aligned_query()); - } - auto retval = iterate_to_fixed_point(query, L, init_ids, scratch, false, unused_filter_label, true); + _distance->preprocess_query(query, _data_store->get_dims(), scratch->aligned_query()); + auto retval = + iterate_to_fixed_point(scratch->aligned_query(), L, init_ids, scratch, false, unused_filter_label, true); NeighborPriorityQueue &best_L_nodes = scratch->best_l_nodes(); @@ -2068,10 +2067,11 @@ std::pair Index::search_with_filters(const } filter_vec.emplace_back(filter_label); - T *aligned_query = scratch->aligned_query(); - memcpy(aligned_query, query, _dim * sizeof(T)); - - auto retval = iterate_to_fixed_point(aligned_query, L, init_ids, scratch, true, filter_vec, true); + //REFACTOR + //T *aligned_query = scratch->aligned_query(); + //memcpy(aligned_query, query, _dim * sizeof(T)); + _distance->preprocess_query(query, _data_store->get_dims(), scratch->aligned_query()); + auto retval = iterate_to_fixed_point(scratch->aligned_query(), L, init_ids, scratch, true, filter_vec, true); auto best_L_nodes = scratch->best_l_nodes(); @@ -2130,7 +2130,9 @@ size_t Index::search_with_tags(const T *query, const uint64_t K const std::vector init_ids = get_init_ids(); const std::vector unused_filter_label; - iterate_to_fixed_point(query, L, init_ids, scratch, false, unused_filter_label, true); + _distance->preprocess_query(query, _data_store->get_dims(), scratch->aligned_query()); + iterate_to_fixed_point(scratch->aligned_query(), L, init_ids, scratch, false, unused_filter_label, true); + NeighborPriorityQueue &best_L_nodes = scratch->best_l_nodes(); assert(best_L_nodes.size() <= L); diff --git a/src/scratch.cpp b/src/scratch.cpp index cc0e74e19..3f8db3bec 100644 --- a/src/scratch.cpp +++ b/src/scratch.cpp @@ -61,6 +61,7 @@ template void InMemQueryScratch::clear() _expanded_nodes_set.clear(); _expanded_nghrs_vec.clear(); _occlude_list_output.clear(); + } template void InMemQueryScratch::resize_for_new_L(uint32_t new_l) From 112d4fe4e9c8ce062b7c4fcfdff2d53534e603ce Mon Sep 17 00:00:00 2001 From: Gopal Srinivasa Date: Tue, 18 Apr 2023 23:14:21 +0530 Subject: [PATCH 54/69] Fixed distance invocations --- .clang-format | 2 +- include/ann_exception.h | 33 +- src/index.cpp | 5557 +++++++++++++++++----------------- tests/build_memory_index.cpp | 345 ++- 4 files changed, 2970 insertions(+), 2967 deletions(-) diff --git a/.clang-format b/.clang-format index ad3192fd6..7ff715f79 100644 --- a/.clang-format +++ b/.clang-format @@ -1,5 +1,5 @@ --- -BasedOnStyle: Microsoft +BasedOnStyle: Google --- Language: Cpp SortIncludes: false diff --git a/include/ann_exception.h b/include/ann_exception.h index 6b81373c1..5d721e31a 100644 --- a/include/ann_exception.h +++ b/include/ann_exception.h @@ -11,24 +11,25 @@ #define __FUNCSIG__ __PRETTY_FUNCTION__ #endif -namespace diskann -{ +namespace diskann { -class ANNException : public std::runtime_error -{ - public: - DISKANN_DLLEXPORT ANNException(const std::string &message, int errorCode); - DISKANN_DLLEXPORT ANNException(const std::string &message, int errorCode, const std::string &funcSig, - const std::string &fileName, uint32_t lineNum); +class ANNException : public std::runtime_error { + public: + DISKANN_DLLEXPORT ANNException(const std::string &message, int errorCode); + DISKANN_DLLEXPORT ANNException(const std::string &message, int errorCode, + const std::string &funcSig, + const std::string &fileName, uint32_t lineNum); - private: - int _errorCode; + private: + int _errorCode; }; -class FileException : public ANNException -{ - public: - DISKANN_DLLEXPORT FileException(const std::string &filename, std::system_error &e, const std::string &funcSig, - const std::string &fileName, uint32_t lineNum); +class FileException : public ANNException { + public: + DISKANN_DLLEXPORT FileException(const std::string &filename, + std::system_error &e, + const std::string &funcSig, + const std::string &fileName, + uint32_t lineNum); }; -} // namespace diskann +} // namespace diskann diff --git a/src/index.cpp b/src/index.cpp index cc74cd52a..4a4959464 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -11,7 +11,8 @@ #include "memory_mapper.h" #include "timer.h" #include "windows_customizations.h" -#if defined(RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS) && defined(DISKANN_BUILD) +#if defined(RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS) && \ + defined(DISKANN_BUILD) #include "gperftools/malloc_extension.h" #endif @@ -22,3055 +23,2925 @@ #define MAX_POINTS_FOR_USING_BITSET 10000000 -namespace diskann -{ +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 IndexWriteParameters &indexParams, const uint32_t initial_search_list_size, - const uint32_t search_threads, const bool enable_tags, const bool concurrent_consolidate, - const bool pq_dist_build, const size_t num_pq_chunks, const bool use_opq) - : Index(m, dim, max_points, dynamic_index, enable_tags, concurrent_consolidate, pq_dist_build, num_pq_chunks, - use_opq, indexParams.num_frozen_points) -{ - _indexingQueueSize = indexParams.search_list_size; - _indexingRange = indexParams.max_degree; - _indexingMaxC = indexParams.max_occlusion_size; - _indexingAlpha = indexParams.alpha; - _filterIndexingQueueSize = indexParams.filter_list_size; - - uint32_t num_threads_indx = indexParams.num_threads; - uint32_t num_scratch_spaces = search_threads + num_threads_indx; - - initialize_query_scratch(num_scratch_spaces, initial_search_list_size, _indexingQueueSize, _indexingRange, - _indexingMaxC, dim); -} - -template -Index::Index(Metric m, const size_t dim, const size_t max_points, 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 size_t num_frozen_pts) - : _dist_metric(m), _dim(dim), _max_points(max_points), _num_frozen_pts(num_frozen_pts), - _dynamic_index(dynamic_index), _enable_tags(enable_tags), _indexingMaxC(DEFAULT_MAXC), _query_scratch(nullptr), - _pq_dist(pq_dist_build), _use_opq(use_opq), _num_pq_chunks(num_pq_chunks), - _delete_set(new tsl::robin_set), _conc_consolidate(concurrent_consolidate) -{ - if (dynamic_index && !enable_tags) - { - throw ANNException("ERROR: Dynamic Indexing must have tags enabled.", -1, __FUNCSIG__, __FILE__, __LINE__); - } - - if (_pq_dist) - { - if (dynamic_index) - throw ANNException("ERROR: Dynamic Indexing not supported with PQ distance based " - "index construction", - -1, __FUNCSIG__, __FILE__, __LINE__); - if (m == diskann::Metric::INNER_PRODUCT) - throw ANNException("ERROR: Inner product metrics not yet supported " - "with PQ distance " - "base index", - -1, __FUNCSIG__, __FILE__, __LINE__); - } - - if (dynamic_index && _num_frozen_pts == 0) - { - _num_frozen_pts = 1; - } - // Sanity check. While logically it is correct, max_points = 0 causes - // downstream problems. - if (_max_points == 0) - { - _max_points = 1; - } - const size_t total_internal_points = _max_points + _num_frozen_pts; - - if (_pq_dist) - { - if (_num_pq_chunks > _dim) - throw diskann::ANNException("ERROR: num_pq_chunks > dim", -1, __FUNCSIG__, __FILE__, __LINE__); - alloc_aligned(((void **)&_pq_data), total_internal_points * _num_pq_chunks * sizeof(char), 8 * sizeof(char)); - std::memset(_pq_data, 0, total_internal_points * _num_pq_chunks * sizeof(char)); - } - - _start = (uint32_t)_max_points; - - _final_graph.resize(total_internal_points); - - // This should come from a factory. - if (m == diskann::Metric::COSINE && std::is_floating_point::value) - { - // This is safe because T is float inside the if block. - this->_distance.reset((Distance *)new AVXNormalizedCosineDistanceFloat()); - this->_normalize_vecs = true; - diskann::cout << "Normalizing vectors and using L2 for cosine " - "AVXNormalizedCosineDistanceFloat()." - << std::endl; - } - else - { - this->_distance.reset((Distance *)get_distance_function(m)); - } - // REFACTOR: TODO This should move to a factory method. - - _data_store = - std::make_unique>((location_t)total_internal_points, _dim, this->_distance); - - _locks = std::vector(total_internal_points); - - if (enable_tags) - { - _location_to_tag.reserve(total_internal_points); - _tag_to_location.reserve(total_internal_points); - } +Index::Index( + Metric m, const size_t dim, const size_t max_points, + const bool dynamic_index, const IndexWriteParameters &indexParams, + const uint32_t initial_search_list_size, const uint32_t search_threads, + const bool enable_tags, const bool concurrent_consolidate, + const bool pq_dist_build, const size_t num_pq_chunks, const bool use_opq) + : Index(m, dim, max_points, dynamic_index, enable_tags, + concurrent_consolidate, pq_dist_build, num_pq_chunks, use_opq, + indexParams.num_frozen_points) { + _indexingQueueSize = indexParams.search_list_size; + _indexingRange = indexParams.max_degree; + _indexingMaxC = indexParams.max_occlusion_size; + _indexingAlpha = indexParams.alpha; + _filterIndexingQueueSize = indexParams.filter_list_size; + + uint32_t num_threads_indx = indexParams.num_threads; + uint32_t num_scratch_spaces = search_threads + num_threads_indx; + + initialize_query_scratch(num_scratch_spaces, initial_search_list_size, + _indexingQueueSize, _indexingRange, _indexingMaxC, + dim); } -template Index::~Index() -{ - // Ensure that no other activity is happening before dtor() - std::unique_lock ul(_update_lock); - std::unique_lock cl(_consolidate_lock); - std::unique_lock tl(_tag_lock); - std::unique_lock dl(_delete_lock); +template +Index::Index(Metric m, const size_t dim, + const size_t max_points, 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 size_t num_frozen_pts) + : _dist_metric(m), + _dim(dim), + _max_points(max_points), + _num_frozen_pts(num_frozen_pts), + _dynamic_index(dynamic_index), + _enable_tags(enable_tags), + _indexingMaxC(DEFAULT_MAXC), + _query_scratch(nullptr), + _pq_dist(pq_dist_build), + _use_opq(use_opq), + _num_pq_chunks(num_pq_chunks), + _delete_set(new tsl::robin_set), + _conc_consolidate(concurrent_consolidate) { + if (dynamic_index && !enable_tags) { + throw ANNException("ERROR: Dynamic Indexing must have tags enabled.", -1, + __FUNCSIG__, __FILE__, __LINE__); + } + + if (_pq_dist) { + if (dynamic_index) + throw ANNException( + "ERROR: Dynamic Indexing not supported with PQ distance based " + "index construction", + -1, __FUNCSIG__, __FILE__, __LINE__); + if (m == diskann::Metric::INNER_PRODUCT) + throw ANNException( + "ERROR: Inner product metrics not yet supported " + "with PQ distance " + "base index", + -1, __FUNCSIG__, __FILE__, __LINE__); + } + + if (dynamic_index && _num_frozen_pts == 0) { + _num_frozen_pts = 1; + } + // Sanity check. While logically it is correct, max_points = 0 causes + // downstream problems. + if (_max_points == 0) { + _max_points = 1; + } + const size_t total_internal_points = _max_points + _num_frozen_pts; + + if (_pq_dist) { + if (_num_pq_chunks > _dim) + throw diskann::ANNException("ERROR: num_pq_chunks > dim", -1, __FUNCSIG__, + __FILE__, __LINE__); + alloc_aligned(((void **)&_pq_data), + total_internal_points * _num_pq_chunks * sizeof(char), + 8 * sizeof(char)); + std::memset(_pq_data, 0, + total_internal_points * _num_pq_chunks * sizeof(char)); + } + + _start = (uint32_t)_max_points; + + _final_graph.resize(total_internal_points); + + // This should come from a factory. + if (m == diskann::Metric::COSINE && std::is_floating_point::value) { + // This is safe because T is float inside the if block. + this->_distance.reset( + (Distance *)new AVXNormalizedCosineDistanceFloat()); + this->_normalize_vecs = true; + diskann::cout << "Normalizing vectors and using L2 for cosine " + "AVXNormalizedCosineDistanceFloat()." + << std::endl; + } else { + this->_distance.reset((Distance *)get_distance_function(m)); + } + // REFACTOR: TODO This should move to a factory method. - for (auto &lock : _locks) - { - LockGuard lg(lock); - } + _data_store = std::make_unique>( + (location_t)total_internal_points, _dim, this->_distance); - // if (this->_distance != nullptr) - //{ - // delete this->_distance; - // this->_distance = nullptr; - // } - // REFACTOR + _locks = std::vector(total_internal_points); - if (_opt_graph != nullptr) - { - delete[] _opt_graph; - } + if (enable_tags) { + _location_to_tag.reserve(total_internal_points); + _tag_to_location.reserve(total_internal_points); + } +} - if (!_query_scratch.empty()) - { - ScratchStoreManager> manager(_query_scratch); - manager.destroy(); - } +template +Index::~Index() { + // Ensure that no other activity is happening before dtor() + std::unique_lock ul(_update_lock); + std::unique_lock cl(_consolidate_lock); + std::unique_lock tl(_tag_lock); + std::unique_lock dl(_delete_lock); + + for (auto &lock : _locks) { + LockGuard lg(lock); + } + + // if (this->_distance != nullptr) + //{ + // delete this->_distance; + // this->_distance = nullptr; + // } + // REFACTOR + + if (_opt_graph != nullptr) { + delete[] _opt_graph; + } + + if (!_query_scratch.empty()) { + ScratchStoreManager> manager(_query_scratch); + manager.destroy(); + } } template -void Index::initialize_query_scratch(uint32_t num_threads, uint32_t search_l, uint32_t indexing_l, - uint32_t r, uint32_t maxc, size_t dim) -{ - for (uint32_t i = 0; i < num_threads; i++) - { - auto scratch = new InMemQueryScratch(search_l, indexing_l, r, maxc, dim, _data_store->get_aligned_dim(), _data_store->get_alignment_factor(), _pq_dist); - _query_scratch.push(scratch); - } +void Index::initialize_query_scratch(uint32_t num_threads, + uint32_t search_l, + uint32_t indexing_l, + uint32_t r, uint32_t maxc, + size_t dim) { + for (uint32_t i = 0; i < num_threads; i++) { + auto scratch = new InMemQueryScratch( + search_l, indexing_l, r, maxc, dim, _data_store->get_aligned_dim(), + _data_store->get_alignment_factor(), _pq_dist); + _query_scratch.push(scratch); + } } -template size_t Index::save_tags(std::string tags_file) -{ - 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 (uint32_t i = 0; i < _nd; i++) - { - TagT tag; - if (_location_to_tag.try_get(i, tag)) - { - tag_data[i] = tag; - } - else - { - // catering to future when tagT can be any type. - std::memset((char *)&tag_data[i], 0, sizeof(TagT)); - } - } - if (_num_frozen_pts > 0) - { - std::memset((char *)&tag_data[_start], 0, sizeof(TagT) * _num_frozen_pts); - } - try - { - tag_bytes_written = save_bin(tags_file, tag_data, _nd + _num_frozen_pts, 1); - } - catch (std::system_error &e) - { - throw FileException(tags_file, e, __FUNCSIG__, __FILE__, __LINE__); - } - delete[] tag_data; - return tag_bytes_written; +template +size_t Index::save_tags(std::string tags_file) { + 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 (uint32_t i = 0; i < _nd; i++) { + TagT tag; + if (_location_to_tag.try_get(i, tag)) { + tag_data[i] = tag; + } else { + // catering to future when tagT can be any type. + std::memset((char *)&tag_data[i], 0, sizeof(TagT)); + } + } + if (_num_frozen_pts > 0) { + std::memset((char *)&tag_data[_start], 0, sizeof(TagT) * _num_frozen_pts); + } + try { + tag_bytes_written = + save_bin(tags_file, tag_data, _nd + _num_frozen_pts, 1); + } catch (std::system_error &e) { + throw FileException(tags_file, e, __FUNCSIG__, __FILE__, __LINE__); + } + delete[] tag_data; + return tag_bytes_written; } -template size_t Index::save_data(std::string data_file) -{ - // Note: at this point, either _nd == _max_points or any frozen points have - // been temporarily moved to _nd, so _nd + _num_frozen_points is the valid - // location limit. - return _data_store->save(data_file, _nd + _num_frozen_pts); +template +size_t Index::save_data(std::string data_file) { + // Note: at this point, either _nd == _max_points or any frozen points have + // been temporarily moved to _nd, so _nd + _num_frozen_points is the valid + // location limit. + return _data_store->save(data_file, _nd + _num_frozen_pts); } // 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 uint32_t) -template size_t Index::save_graph(std::string graph_file) -{ - std::ofstream out; - open_file_to_write(out, graph_file); - - size_t file_offset = 0; // we will use this if we want - out.seekp(file_offset, out.beg); - size_t index_size = 24; - uint32_t max_degree = 0; - out.write((char *)&index_size, sizeof(uint64_t)); - out.write((char *)&_max_observed_degree, sizeof(uint32_t)); - uint32_t ep_u32 = _start; - out.write((char *)&ep_u32, sizeof(uint32_t)); - out.write((char *)&_num_frozen_pts, sizeof(size_t)); - // Note: at this point, either _nd == _max_points or any frozen points have - // been temporarily moved to _nd, so _nd + _num_frozen_points is the valid - // location limit. - for (uint32_t i = 0; i < _nd + _num_frozen_pts; i++) - { - uint32_t GK = (uint32_t)_final_graph[i].size(); - out.write((char *)&GK, sizeof(uint32_t)); - out.write((char *)_final_graph[i].data(), GK * sizeof(uint32_t)); - max_degree = _final_graph[i].size() > max_degree ? (uint32_t)_final_graph[i].size() : max_degree; - index_size += (size_t)(sizeof(uint32_t) * (GK + 1)); - } - out.seekp(file_offset, out.beg); - out.write((char *)&index_size, sizeof(uint64_t)); - out.write((char *)&max_degree, sizeof(uint32_t)); - out.close(); - return index_size; // number of bytes written +template +size_t Index::save_graph(std::string graph_file) { + std::ofstream out; + open_file_to_write(out, graph_file); + + size_t file_offset = 0; // we will use this if we want + out.seekp(file_offset, out.beg); + size_t index_size = 24; + uint32_t max_degree = 0; + out.write((char *)&index_size, sizeof(uint64_t)); + out.write((char *)&_max_observed_degree, sizeof(uint32_t)); + uint32_t ep_u32 = _start; + out.write((char *)&ep_u32, sizeof(uint32_t)); + out.write((char *)&_num_frozen_pts, sizeof(size_t)); + // Note: at this point, either _nd == _max_points or any frozen points have + // been temporarily moved to _nd, so _nd + _num_frozen_points is the valid + // location limit. + for (uint32_t i = 0; i < _nd + _num_frozen_pts; i++) { + uint32_t GK = (uint32_t)_final_graph[i].size(); + out.write((char *)&GK, sizeof(uint32_t)); + out.write((char *)_final_graph[i].data(), GK * sizeof(uint32_t)); + max_degree = _final_graph[i].size() > max_degree + ? (uint32_t)_final_graph[i].size() + : max_degree; + index_size += (size_t)(sizeof(uint32_t) * (GK + 1)); + } + out.seekp(file_offset, out.beg); + out.write((char *)&index_size, sizeof(uint64_t)); + out.write((char *)&max_degree, sizeof(uint32_t)); + out.close(); + return index_size; // number of bytes written } template -size_t Index::save_delete_list(const std::string &filename) -{ - if (_delete_set->size() == 0) - { - return 0; - } - std::unique_ptr delete_list = std::make_unique(_delete_set->size()); - uint32_t i = 0; - for (auto &del : *_delete_set) - { - delete_list[i++] = del; - } - return save_bin(filename, delete_list.get(), _delete_set->size(), 1); +size_t Index::save_delete_list(const std::string &filename) { + if (_delete_set->size() == 0) { + return 0; + } + std::unique_ptr delete_list = + std::make_unique(_delete_set->size()); + uint32_t i = 0; + for (auto &del : *_delete_set) { + delete_list[i++] = del; + } + return save_bin(filename, delete_list.get(), _delete_set->size(), + 1); } template -void Index::save(const char *filename, bool compact_before_save) -{ - diskann::Timer timer; - - std::unique_lock ul(_update_lock); - std::unique_lock cl(_consolidate_lock); - std::unique_lock tl(_tag_lock); - std::unique_lock dl(_delete_lock); - - if (compact_before_save) - { - compact_data(); - compact_frozen_point(); - } - else - { - if (!_data_compacted) - { - throw ANNException("Index save for non-compacted index is not yet implemented", -1, __FUNCSIG__, __FILE__, - __LINE__); - } - } - - if (!_save_as_one_file) - { - if (_filtered_index) - { - if (_label_to_medoid_id.size() > 0) - { - std::ofstream medoid_writer(std::string(filename) + "_labels_to_medoids.txt"); - if (medoid_writer.fail()) - { - throw diskann::ANNException(std::string("Failed to open file ") + filename, -1); - } - for (auto iter : _label_to_medoid_id) - { - medoid_writer << iter.first << ", " << iter.second << std::endl; - } - medoid_writer.close(); - } - - if (_use_universal_label) - { - std::ofstream universal_label_writer(std::string(filename) + "_universal_label.txt"); - assert(universal_label_writer.is_open()); - universal_label_writer << _universal_label << std::endl; - universal_label_writer.close(); - } - - if (_pts_to_labels.size() > 0) - { - std::ofstream label_writer(std::string(filename) + "_labels.txt"); - assert(label_writer.is_open()); - for (uint32_t i = 0; i < _pts_to_labels.size(); i++) - { - for (uint32_t j = 0; j < (_pts_to_labels[i].size() - 1); j++) - { - label_writer << _pts_to_labels[i][j] << ","; - } - if (_pts_to_labels[i].size() != 0) - label_writer << _pts_to_labels[i][_pts_to_labels[i].size() - 1]; - label_writer << std::endl; - } - label_writer.close(); - } - } - - 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 - { - diskann::cout << "Save index in a single file currently not supported. " - "Not saving the index." - << std::endl; - } +void Index::save(const char *filename, + bool compact_before_save) { + diskann::Timer timer; + + std::unique_lock ul(_update_lock); + std::unique_lock cl(_consolidate_lock); + std::unique_lock tl(_tag_lock); + std::unique_lock dl(_delete_lock); + + if (compact_before_save) { + compact_data(); + compact_frozen_point(); + } else { + if (!_data_compacted) { + throw ANNException( + "Index save for non-compacted index is not yet implemented", -1, + __FUNCSIG__, __FILE__, __LINE__); + } + } + + if (!_save_as_one_file) { + if (_filtered_index) { + if (_label_to_medoid_id.size() > 0) { + std::ofstream medoid_writer(std::string(filename) + + "_labels_to_medoids.txt"); + if (medoid_writer.fail()) { + throw diskann::ANNException( + std::string("Failed to open file ") + filename, -1); + } + for (auto iter : _label_to_medoid_id) { + medoid_writer << iter.first << ", " << iter.second << std::endl; + } + medoid_writer.close(); + } + + if (_use_universal_label) { + std::ofstream universal_label_writer(std::string(filename) + + "_universal_label.txt"); + assert(universal_label_writer.is_open()); + universal_label_writer << _universal_label << std::endl; + universal_label_writer.close(); + } + + if (_pts_to_labels.size() > 0) { + std::ofstream label_writer(std::string(filename) + "_labels.txt"); + assert(label_writer.is_open()); + for (uint32_t i = 0; i < _pts_to_labels.size(); i++) { + for (uint32_t j = 0; j < (_pts_to_labels[i].size() - 1); j++) { + label_writer << _pts_to_labels[i][j] << ","; + } + if (_pts_to_labels[i].size() != 0) + label_writer << _pts_to_labels[i][_pts_to_labels[i].size() - 1]; + label_writer << std::endl; + } + label_writer.close(); + } + } + + 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 { + diskann::cout << "Save index in a single file currently not supported. " + "Not saving the index." + << std::endl; + } - // If frozen points were temporarily compacted to _nd, move back to - // _max_points. - reposition_frozen_point_to_end(); + // If frozen points were temporarily compacted to _nd, move back to + // _max_points. + reposition_frozen_point_to_end(); - diskann::cout << "Time taken for save: " << timer.elapsed() / 1000000.0 << "s." << std::endl; + diskann::cout << "Time taken for save: " << timer.elapsed() / 1000000.0 + << "s." << std::endl; } #ifdef EXEC_ENV_OLS template -size_t Index::load_tags(AlignedFileReader &reader) -{ +size_t Index::load_tags(AlignedFileReader &reader) { #else template -size_t Index::load_tags(const std::string tag_filename) -{ - 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__); - } +size_t Index::load_tags(const std::string tag_filename) { + 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__); + } #endif - if (!_enable_tags) - { - diskann::cout << "Tags not loaded as tags not enabled." << std::endl; - return 0; - } + 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; + size_t file_dim, file_num_points; + TagT *tag_data; #ifdef EXEC_ENV_OLS - load_bin(reader, tag_data, file_num_points, file_dim); + load_bin(reader, tag_data, file_num_points, file_dim); #else - load_bin(std::string(tag_filename), tag_data, file_num_points, file_dim); + load_bin(std::string(tag_filename), tag_data, file_num_points, + file_dim); #endif - if (file_dim != 1) - { - std::stringstream stream; - stream << "ERROR: Found " << file_dim << " dimensions for tags," - << "but tag file must have 1 dimension." << std::endl; - diskann::cerr << stream.str() << std::endl; - delete[] tag_data; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } + if (file_dim != 1) { + std::stringstream stream; + stream << "ERROR: Found " << file_dim << " dimensions for tags," + << "but tag file must have 1 dimension." << std::endl; + diskann::cerr << stream.str() << std::endl; + delete[] tag_data; + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } - const size_t num_data_points = file_num_points - _num_frozen_pts; - _location_to_tag.reserve(num_data_points); - _tag_to_location.reserve(num_data_points); - for (uint32_t i = 0; i < (uint32_t)num_data_points; i++) - { - TagT tag = *(tag_data + i); - if (_delete_set->find(i) == _delete_set->end()) - { - _location_to_tag.set(i, tag); - _tag_to_location[tag] = i; - } + const size_t num_data_points = file_num_points - _num_frozen_pts; + _location_to_tag.reserve(num_data_points); + _tag_to_location.reserve(num_data_points); + for (uint32_t i = 0; i < (uint32_t)num_data_points; i++) { + TagT tag = *(tag_data + i); + if (_delete_set->find(i) == _delete_set->end()) { + _location_to_tag.set(i, tag); + _tag_to_location[tag] = i; } - diskann::cout << "Tags loaded." << std::endl; - delete[] tag_data; - return file_num_points; + } + diskann::cout << "Tags loaded." << std::endl; + delete[] tag_data; + return file_num_points; } template #ifdef EXEC_ENV_OLS -size_t Index::load_data(AlignedFileReader &reader) -{ +size_t Index::load_data(AlignedFileReader &reader) { #else -size_t Index::load_data(std::string filename) -{ +size_t Index::load_data(std::string filename) { #endif - size_t file_dim, file_num_points; + size_t file_dim, file_num_points; #ifdef EXEC_ENV_OLS - diskann::get_bin_metadata(reader, file_num_points, file_dim); + diskann::get_bin_metadata(reader, file_num_points, file_dim); #else - if (!file_exists(filename)) - { - std::stringstream stream; - stream << "ERROR: data file " << filename << " does not exist." << std::endl; - diskann::cerr << stream.str() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - diskann::get_bin_metadata(filename, file_num_points, file_dim); + if (!file_exists(filename)) { + std::stringstream stream; + stream << "ERROR: data file " << filename << " does not exist." + << std::endl; + diskann::cerr << stream.str() << std::endl; + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } + diskann::get_bin_metadata(filename, file_num_points, file_dim); #endif - // since we are loading a new dataset, _empty_slots must be cleared - _empty_slots.clear(); + // since we are loading a new dataset, _empty_slots must be cleared + _empty_slots.clear(); - if (file_dim != _dim) - { - std::stringstream stream; - stream << "ERROR: Driver requests loading " << _dim << " dimension," - << "but file has " << file_dim << " dimension." << std::endl; - diskann::cerr << stream.str() << std::endl; - 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; + diskann::cerr << stream.str() << std::endl; + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } - if (file_num_points > _max_points + _num_frozen_pts) - { - // update and tag lock acquired in load() before calling load_data - resize(file_num_points - _num_frozen_pts); - } + if (file_num_points > _max_points + _num_frozen_pts) { + // update and tag lock acquired in load() before calling load_data + resize(file_num_points - _num_frozen_pts); + } #ifdef EXEC_ENV_OLS - // REFACTOR TODO: Must figure out how to support aligned reader in a clean manner. - copy_aligned_data_from_file(reader, _data, file_num_points, file_dim, _aligned_dim); + // REFACTOR TODO: Must figure out how to support aligned reader in a clean + // manner. + copy_aligned_data_from_file(reader, _data, file_num_points, file_dim, + _aligned_dim); #else - _data_store->populate_data(filename, 0U); // offset == 0. + _data_store->populate_data(filename, 0U); // offset == 0. #endif - return file_num_points; + return file_num_points; } #ifdef EXEC_ENV_OLS template -size_t Index::load_delete_set(AlignedFileReader &reader) -{ +size_t Index::load_delete_set(AlignedFileReader &reader) { #else template -size_t Index::load_delete_set(const std::string &filename) -{ +size_t Index::load_delete_set(const std::string &filename) { #endif - std::unique_ptr delete_list; - size_t npts, ndim; + std::unique_ptr delete_list; + size_t npts, ndim; #ifdef EXEC_ENV_OLS - diskann::load_bin(reader, delete_list, npts, ndim); + diskann::load_bin(reader, delete_list, npts, ndim); #else - diskann::load_bin(filename, delete_list, npts, ndim); + diskann::load_bin(filename, delete_list, npts, ndim); #endif - assert(ndim == 1); - for (uint32_t i = 0; i < npts; i++) - { - _delete_set->insert(delete_list[i]); - } - return npts; + assert(ndim == 1); + for (uint32_t i = 0; i < npts; i++) { + _delete_set->insert(delete_list[i]); + } + return npts; } // load the index from file and update the max_degree, cur (navigating // node loc), and _final_graph (adjacency list) template #ifdef EXEC_ENV_OLS -void Index::load(AlignedFileReader &reader, uint32_t num_threads, uint32_t search_l) -{ +void Index::load(AlignedFileReader &reader, + uint32_t num_threads, uint32_t search_l) { #else -void Index::load(const char *filename, uint32_t num_threads, uint32_t search_l) -{ +void Index::load(const char *filename, uint32_t num_threads, + uint32_t search_l) { #endif - std::unique_lock ul(_update_lock); - std::unique_lock cl(_consolidate_lock); - std::unique_lock tl(_tag_lock); - std::unique_lock dl(_delete_lock); + std::unique_lock ul(_update_lock); + std::unique_lock cl(_consolidate_lock); + std::unique_lock tl(_tag_lock); + std::unique_lock dl(_delete_lock); - _has_built = true; + _has_built = true; - size_t tags_file_num_pts = 0, graph_num_pts = 0, data_file_num_pts = 0, label_num_pts = 0; + size_t tags_file_num_pts = 0, graph_num_pts = 0, data_file_num_pts = 0, + label_num_pts = 0; - std::string mem_index_file(filename); - std::string labels_file = mem_index_file + "_labels.txt"; - std::string labels_to_medoids = mem_index_file + "_labels_to_medoids.txt"; - std::string labels_map_file = mem_index_file + "_labels_map.txt"; + std::string mem_index_file(filename); + std::string labels_file = mem_index_file + "_labels.txt"; + std::string labels_to_medoids = mem_index_file + "_labels_to_medoids.txt"; + std::string labels_map_file = mem_index_file + "_labels_map.txt"; - if (!_save_as_one_file) - { - // For DLVS Store, we will not support saving the index in multiple - // files. + if (!_save_as_one_file) { + // For DLVS Store, we will not support saving the index in multiple + // files. #ifndef EXEC_ENV_OLS - 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); + 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); #endif - } - else - { - diskann::cout << "Single index file saving/loading support not yet " - "enabled. Not loading the index." - << std::endl; - return; - } - - 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; - diskann::cerr << stream.str() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - - if (file_exists(labels_file)) - { - _label_map = load_label_map(labels_map_file); - parse_label_file(labels_file, label_num_pts); - assert(label_num_pts == data_file_num_pts); - if (file_exists(labels_to_medoids)) - { - std::ifstream medoid_stream(labels_to_medoids); - std::string line, token; - uint32_t line_cnt = 0; - - _label_to_medoid_id.clear(); - - while (std::getline(medoid_stream, line)) - { - std::istringstream iss(line); - uint32_t cnt = 0; - uint32_t medoid = 0; - LabelT label; - while (std::getline(iss, token, ',')) - { - token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); - token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); - LabelT token_as_num = std::stoul(token); - if (cnt == 0) - label = token_as_num; - else - medoid = token_as_num; - cnt++; - } - _label_to_medoid_id[label] = medoid; - line_cnt++; - } - } - - std::string universal_label_file(filename); - universal_label_file += "_universal_label.txt"; - if (file_exists(universal_label_file)) - { - std::ifstream universal_label_reader(universal_label_file); - universal_label_reader >> _universal_label; - _use_universal_label = true; - universal_label_reader.close(); - } - } - - _nd = data_file_num_pts - _num_frozen_pts; - _empty_slots.clear(); - _empty_slots.reserve(_max_points); - for (auto i = _nd; i < _max_points; i++) - { - _empty_slots.insert((uint32_t)i); - } - - reposition_frozen_point_to_end(); - diskann::cout << "Num frozen points:" << _num_frozen_pts << " _nd: " << _nd << " _start: " << _start - << " size(_location_to_tag): " << _location_to_tag.size() - << " size(_tag_to_location):" << _tag_to_location.size() << " Max points: " << _max_points + } else { + diskann::cout << "Single index file saving/loading support not yet " + "enabled. Not loading the index." << std::endl; + return; + } - // For incremental index, _query_scratch is initialized in the constructor. - // For the bulk index, the params required to initialize _query_scratch - // are known only at load time, hence this check and the call to - // initialize_q_s(). - if (_query_scratch.size() == 0) - { - initialize_query_scratch(num_threads, search_l, search_l, (uint32_t)_max_range_of_loaded_graph, _indexingMaxC, - _dim); - } + 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; + diskann::cerr << stream.str() << std::endl; + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } + + if (file_exists(labels_file)) { + _label_map = load_label_map(labels_map_file); + parse_label_file(labels_file, label_num_pts); + assert(label_num_pts == data_file_num_pts); + if (file_exists(labels_to_medoids)) { + std::ifstream medoid_stream(labels_to_medoids); + std::string line, token; + uint32_t line_cnt = 0; + + _label_to_medoid_id.clear(); + + while (std::getline(medoid_stream, line)) { + std::istringstream iss(line); + uint32_t cnt = 0; + uint32_t medoid = 0; + LabelT label; + while (std::getline(iss, token, ',')) { + token.erase(std::remove(token.begin(), token.end(), '\n'), + token.end()); + token.erase(std::remove(token.begin(), token.end(), '\r'), + token.end()); + LabelT token_as_num = std::stoul(token); + if (cnt == 0) + label = token_as_num; + else + medoid = token_as_num; + cnt++; + } + _label_to_medoid_id[label] = medoid; + line_cnt++; + } + } + + std::string universal_label_file(filename); + universal_label_file += "_universal_label.txt"; + if (file_exists(universal_label_file)) { + std::ifstream universal_label_reader(universal_label_file); + universal_label_reader >> _universal_label; + _use_universal_label = true; + universal_label_reader.close(); + } + } + + _nd = data_file_num_pts - _num_frozen_pts; + _empty_slots.clear(); + _empty_slots.reserve(_max_points); + for (auto i = _nd; i < _max_points; i++) { + _empty_slots.insert((uint32_t)i); + } + + reposition_frozen_point_to_end(); + diskann::cout << "Num frozen points:" << _num_frozen_pts << " _nd: " << _nd + << " _start: " << _start + << " size(_location_to_tag): " << _location_to_tag.size() + << " size(_tag_to_location):" << _tag_to_location.size() + << " Max points: " << _max_points << std::endl; + + // For incremental index, _query_scratch is initialized in the constructor. + // For the bulk index, the params required to initialize _query_scratch + // are known only at load time, hence this check and the call to + // initialize_q_s(). + if (_query_scratch.size() == 0) { + initialize_query_scratch(num_threads, search_l, search_l, + (uint32_t)_max_range_of_loaded_graph, + _indexingMaxC, _dim); + } } #ifndef EXEC_ENV_OLS template -size_t Index::get_graph_num_frozen_points(const std::string &graph_file) -{ - size_t expected_file_size; - uint32_t max_observed_degree, start; - size_t file_frozen_pts; +size_t Index::get_graph_num_frozen_points( + const std::string &graph_file) { + size_t expected_file_size; + uint32_t max_observed_degree, start; + size_t file_frozen_pts; - std::ifstream in; - in.exceptions(std::ios::badbit | std::ios::failbit); + std::ifstream in; + in.exceptions(std::ios::badbit | std::ios::failbit); - in.open(graph_file, std::ios::binary); - in.read((char *)&expected_file_size, sizeof(size_t)); - in.read((char *)&max_observed_degree, sizeof(uint32_t)); - in.read((char *)&start, sizeof(uint32_t)); - in.read((char *)&file_frozen_pts, sizeof(size_t)); + in.open(graph_file, std::ios::binary); + in.read((char *)&expected_file_size, sizeof(size_t)); + in.read((char *)&max_observed_degree, sizeof(uint32_t)); + in.read((char *)&start, sizeof(uint32_t)); + in.read((char *)&file_frozen_pts, sizeof(size_t)); - return file_frozen_pts; + return file_frozen_pts; } #endif #ifdef EXEC_ENV_OLS template -size_t Index::load_graph(AlignedFileReader &reader, size_t expected_num_points) -{ +size_t Index::load_graph(AlignedFileReader &reader, + size_t expected_num_points) { #else template -size_t Index::load_graph(std::string filename, size_t expected_num_points) -{ +size_t Index::load_graph(std::string filename, + size_t expected_num_points) { #endif - size_t expected_file_size; - size_t file_frozen_pts; + size_t expected_file_size; + size_t file_frozen_pts; #ifdef EXEC_ENV_OLS - int header_size = 2 * sizeof(size_t) + 2 * sizeof(uint32_t); - std::unique_ptr header = std::make_unique(header_size); - read_array(reader, header.get(), header_size); - - expected_file_size = *((size_t *)header.get()); - _max_observed_degree = *((uint32_t *)(header.get() + sizeof(size_t))); - _start = *((uint32_t *)(header.get() + sizeof(size_t) + sizeof(uint32_t))); - file_frozen_pts = *((size_t *)(header.get() + sizeof(size_t) + sizeof(uint32_t) + sizeof(uint32_t))); + int header_size = 2 * sizeof(size_t) + 2 * sizeof(uint32_t); + std::unique_ptr header = std::make_unique(header_size); + read_array(reader, header.get(), header_size); + + expected_file_size = *((size_t *)header.get()); + _max_observed_degree = *((uint32_t *)(header.get() + sizeof(size_t))); + _start = *((uint32_t *)(header.get() + sizeof(size_t) + sizeof(uint32_t))); + file_frozen_pts = *((size_t *)(header.get() + sizeof(size_t) + + sizeof(uint32_t) + sizeof(uint32_t))); #else - size_t file_offset = 0; // will need this for single file format support - std::ifstream in; - in.exceptions(std::ios::badbit | std::ios::failbit); - in.open(filename, std::ios::binary); - in.seekg(file_offset, in.beg); - in.read((char *)&expected_file_size, sizeof(size_t)); - in.read((char *)&_max_observed_degree, sizeof(uint32_t)); - in.read((char *)&_start, sizeof(uint32_t)); - in.read((char *)&file_frozen_pts, sizeof(size_t)); - size_t vamana_metadata_size = sizeof(size_t) + sizeof(uint32_t) + sizeof(uint32_t) + sizeof(size_t); + size_t file_offset = 0; // will need this for single file format support + std::ifstream in; + in.exceptions(std::ios::badbit | std::ios::failbit); + in.open(filename, std::ios::binary); + in.seekg(file_offset, in.beg); + in.read((char *)&expected_file_size, sizeof(size_t)); + in.read((char *)&_max_observed_degree, sizeof(uint32_t)); + in.read((char *)&_start, sizeof(uint32_t)); + in.read((char *)&file_frozen_pts, sizeof(size_t)); + size_t vamana_metadata_size = + sizeof(size_t) + sizeof(uint32_t) + sizeof(uint32_t) + sizeof(size_t); #endif - diskann::cout << "From graph header, expected_file_size: " << expected_file_size - << ", _max_observed_degree: " << _max_observed_degree << ", _start: " << _start - << ", file_frozen_pts: " << file_frozen_pts << std::endl; + diskann::cout << "From graph header, expected_file_size: " + << expected_file_size + << ", _max_observed_degree: " << _max_observed_degree + << ", _start: " << _start + << ", file_frozen_pts: " << file_frozen_pts << std::endl; - 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; - } - diskann::cerr << stream.str() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + 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; } + diskann::cerr << stream.str() << std::endl; + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } #ifdef EXEC_ENV_OLS - diskann::cout << "Loading vamana graph from reader..." << std::flush; + diskann::cout << "Loading vamana graph from reader..." << std::flush; #else - diskann::cout << "Loading vamana graph " << filename << "..." << std::flush; + diskann::cout << "Loading vamana graph " << filename << "..." << std::flush; #endif - const size_t expected_max_points = expected_num_points - file_frozen_pts; + const size_t expected_max_points = expected_num_points - file_frozen_pts; - // If user provides more points than max_points - // resize the _final_graph to the larger size. - if (_max_points < expected_max_points) - { - diskann::cout << "Number of points in data: " << expected_max_points - << " is greater than max_points: " << _max_points - << " Setting max points to: " << expected_max_points << std::endl; - _final_graph.resize(expected_max_points + _num_frozen_pts); - _max_points = expected_max_points; - } + // If user provides more points than max_points + // resize the _final_graph to the larger size. + if (_max_points < expected_max_points) { + diskann::cout << "Number of points in data: " << expected_max_points + << " is greater than max_points: " << _max_points + << " Setting max points to: " << expected_max_points + << std::endl; + _final_graph.resize(expected_max_points + _num_frozen_pts); + _max_points = expected_max_points; + } #ifdef EXEC_ENV_OLS - uint32_t nodes_read = 0; - size_t cc = 0; - size_t graph_offset = header_size; - while (nodes_read < expected_num_points) - { - uint32_t k; - read_value(reader, k, graph_offset); - graph_offset += sizeof(uint32_t); - std::vector tmp(k); - tmp.reserve(k); - read_array(reader, tmp.data(), k, graph_offset); - graph_offset += k * sizeof(uint32_t); - cc += k; - _final_graph[nodes_read].swap(tmp); - nodes_read++; - if (nodes_read % 1000000 == 0) - { - diskann::cout << "." << std::flush; - } - if (k > _max_range_of_loaded_graph) - { - _max_range_of_loaded_graph = k; - } - } + uint32_t nodes_read = 0; + size_t cc = 0; + size_t graph_offset = header_size; + while (nodes_read < expected_num_points) { + uint32_t k; + read_value(reader, k, graph_offset); + graph_offset += sizeof(uint32_t); + std::vector tmp(k); + tmp.reserve(k); + read_array(reader, tmp.data(), k, graph_offset); + graph_offset += k * sizeof(uint32_t); + cc += k; + _final_graph[nodes_read].swap(tmp); + nodes_read++; + if (nodes_read % 1000000 == 0) { + diskann::cout << "." << std::flush; + } + if (k > _max_range_of_loaded_graph) { + _max_range_of_loaded_graph = k; + } + } #else - size_t bytes_read = vamana_metadata_size; - size_t cc = 0; - uint32_t nodes_read = 0; - while (bytes_read != expected_file_size) - { - uint32_t k; - in.read((char *)&k, sizeof(uint32_t)); - - if (k == 0) - { - diskann::cerr << "ERROR: Point found with no out-neighbors, point#" << nodes_read << std::endl; - } - - cc += k; - ++nodes_read; - std::vector tmp(k); - tmp.reserve(k); - in.read((char *)tmp.data(), k * sizeof(uint32_t)); - _final_graph[nodes_read - 1].swap(tmp); - bytes_read += sizeof(uint32_t) * ((size_t)k + 1); - if (nodes_read % 10000000 == 0) - diskann::cout << "." << std::flush; - if (k > _max_range_of_loaded_graph) - { - _max_range_of_loaded_graph = k; - } - } + size_t bytes_read = vamana_metadata_size; + size_t cc = 0; + uint32_t nodes_read = 0; + while (bytes_read != expected_file_size) { + uint32_t k; + in.read((char *)&k, sizeof(uint32_t)); + + if (k == 0) { + diskann::cerr << "ERROR: Point found with no out-neighbors, point#" + << nodes_read << std::endl; + } + + cc += k; + ++nodes_read; + std::vector tmp(k); + tmp.reserve(k); + in.read((char *)tmp.data(), k * sizeof(uint32_t)); + _final_graph[nodes_read - 1].swap(tmp); + bytes_read += sizeof(uint32_t) * ((size_t)k + 1); + if (nodes_read % 10000000 == 0) diskann::cout << "." << std::flush; + if (k > _max_range_of_loaded_graph) { + _max_range_of_loaded_graph = k; + } + } #endif - diskann::cout << "done. Index has " << nodes_read << " nodes and " << cc << " out-edges, _start is set to " - << _start << std::endl; - return nodes_read; + diskann::cout << "done. Index has " << nodes_read << " nodes and " << cc + << " out-edges, _start is set to " << _start << std::endl; + return nodes_read; } -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; - } +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; + } - location_t location = _tag_to_location[tag]; - _data_store->get_vector(location, vec); + location_t location = _tag_to_location[tag]; + _data_store->get_vector(location, vec); - return 0; + return 0; } -template uint32_t Index::calculate_entry_point() -{ - // TODO: need to compute medoid with PQ data too, for now sample at random - if (_pq_dist) - { - size_t r = (size_t)rand() * (size_t)RAND_MAX + (size_t)rand(); - return (uint32_t)(r % (size_t)_nd); - } +template +uint32_t Index::calculate_entry_point() { + // TODO: need to compute medoid with PQ data too, for now sample at random + if (_pq_dist) { + size_t r = (size_t)rand() * (size_t)RAND_MAX + (size_t)rand(); + return (uint32_t)(r % (size_t)_nd); + } - // TODO: This function does not support multi-threaded calculation of medoid. - // Must revisit if perf is a concern. - return _data_store->calculate_medoid(); + // TODO: This function does not support multi-threaded calculation of medoid. + // Must revisit if perf is a concern. + return _data_store->calculate_medoid(); } -template std::vector Index::get_init_ids() -{ - std::vector init_ids; - init_ids.reserve(1 + _num_frozen_pts); +template +std::vector Index::get_init_ids() { + std::vector init_ids; + init_ids.reserve(1 + _num_frozen_pts); - init_ids.emplace_back(_start); + init_ids.emplace_back(_start); - for (uint32_t frozen = _max_points; frozen < _max_points + _num_frozen_pts; frozen++) - { - if (frozen != _start) - { - init_ids.emplace_back(frozen); - } + for (uint32_t frozen = _max_points; frozen < _max_points + _num_frozen_pts; + frozen++) { + if (frozen != _start) { + init_ids.emplace_back(frozen); } + } - return init_ids; + return init_ids; } template std::pair Index::iterate_to_fixed_point( - const T *query, const uint32_t Lsize, const std::vector &init_ids, InMemQueryScratch *scratch, - bool use_filter, const std::vector &filter_label, bool search_invocation) -{ - std::vector &expanded_nodes = scratch->pool(); - NeighborPriorityQueue &best_L_nodes = scratch->best_l_nodes(); - best_L_nodes.reserve(Lsize); - tsl::robin_set &inserted_into_pool_rs = scratch->inserted_into_pool_rs(); - boost::dynamic_bitset<> &inserted_into_pool_bs = scratch->inserted_into_pool_bs(); - std::vector &id_scratch = scratch->id_scratch(); - std::vector &dist_scratch = scratch->dist_scratch(); - assert(id_scratch.size() == 0); - - //REFACTOR - //T *aligned_query = scratch->aligned_query(); - //memcpy(aligned_query, query, _dim * sizeof(T)); - //if (_normalize_vecs) - //{ - // normalize((float *)aligned_query, _dim); - //} + const T *query, const uint32_t Lsize, const std::vector &init_ids, + InMemQueryScratch *scratch, bool use_filter, + const std::vector &filter_label, bool search_invocation) { + std::vector &expanded_nodes = scratch->pool(); + NeighborPriorityQueue &best_L_nodes = scratch->best_l_nodes(); + best_L_nodes.reserve(Lsize); + tsl::robin_set &inserted_into_pool_rs = + scratch->inserted_into_pool_rs(); + boost::dynamic_bitset<> &inserted_into_pool_bs = + scratch->inserted_into_pool_bs(); + std::vector &id_scratch = scratch->id_scratch(); + std::vector &dist_scratch = scratch->dist_scratch(); + assert(id_scratch.size() == 0); + + // REFACTOR + // T *aligned_query = scratch->aligned_query(); + // memcpy(aligned_query, query, _dim * sizeof(T)); + // if (_normalize_vecs) + //{ + // normalize((float *)aligned_query, _dim); + // } + + T *aligned_query = scratch->aligned_query(); + + float *query_float = nullptr; + float *query_rotated = nullptr; + float *pq_dists = nullptr; + uint8_t *pq_coord_scratch = nullptr; + // Intialize PQ related scratch to use PQ based distances + if (_pq_dist) { + // Get scratch spaces + PQScratch *pq_query_scratch = scratch->pq_scratch(); + query_float = pq_query_scratch->aligned_query_float; + query_rotated = pq_query_scratch->rotated_query; + pq_dists = pq_query_scratch->aligned_pqtable_dist_scratch; + + // Copy query vector to float and then to "rotated" query + for (size_t d = 0; d < _dim; d++) { + query_float[d] = (float)aligned_query[d]; + } + pq_query_scratch->set(_dim, aligned_query); + + // center the query and rotate if we have a rotation matrix + _pq_table.preprocess_query(query_rotated); + _pq_table.populate_chunk_distances(query_rotated, pq_dists); + + pq_coord_scratch = pq_query_scratch->aligned_pq_coord_scratch; + } + + if (expanded_nodes.size() > 0 || id_scratch.size() > 0) { + throw ANNException("ERROR: Clear scratch space before passing.", -1, + __FUNCSIG__, __FILE__, __LINE__); + } + + // Decide whether to use bitset or robin set to mark visited nodes + auto total_num_points = _max_points + _num_frozen_pts; + bool fast_iterate = total_num_points <= MAX_POINTS_FOR_USING_BITSET; + + if (fast_iterate) { + if (inserted_into_pool_bs.size() < total_num_points) { + // hopefully using 2X will reduce the number of allocations. + auto resize_size = 2 * total_num_points > MAX_POINTS_FOR_USING_BITSET + ? MAX_POINTS_FOR_USING_BITSET + : 2 * total_num_points; + inserted_into_pool_bs.resize(resize_size); + } + } + + // Lambda to determine if a node has been visited + auto is_not_visited = [this, fast_iterate, &inserted_into_pool_bs, + &inserted_into_pool_rs](const uint32_t id) { + return fast_iterate + ? inserted_into_pool_bs[id] == 0 + : inserted_into_pool_rs.find(id) == inserted_into_pool_rs.end(); + }; + + // Lambda to batch compute query<-> node distances in PQ space + auto compute_dists = [this, pq_coord_scratch, pq_dists]( + const std::vector &ids, + std::vector &dists_out) { + diskann::aggregate_coords(ids, this->_pq_data, this->_num_pq_chunks, + pq_coord_scratch); + diskann::pq_dist_lookup(pq_coord_scratch, ids.size(), this->_num_pq_chunks, + pq_dists, dists_out); + }; + + // Initialize the candidate pool with starting points + for (auto id : init_ids) { + if (id >= _max_points + _num_frozen_pts) { + diskann::cerr << "Out of range loc found as an edge : " << id + << std::endl; + throw diskann::ANNException(std::string("Wrong loc") + std::to_string(id), + -1, __FUNCSIG__, __FILE__, __LINE__); + } + + if (use_filter) { + std::vector common_filters; + auto &x = _pts_to_labels[id]; + std::set_intersection(filter_label.begin(), filter_label.end(), x.begin(), + x.end(), std::back_inserter(common_filters)); + if (_use_universal_label) { + if (std::find(filter_label.begin(), filter_label.end(), + _universal_label) != filter_label.end() || + std::find(x.begin(), x.end(), _universal_label) != x.end()) + common_filters.emplace_back(_universal_label); + } + + if (common_filters.size() == 0) continue; + } + + if (is_not_visited(id)) { + if (fast_iterate) { + inserted_into_pool_bs[id] = 1; + } else { + inserted_into_pool_rs.insert(id); + } + + float distance; + if (_pq_dist) { + pq_dist_lookup(pq_coord_scratch, 1, this->_num_pq_chunks, pq_dists, + &distance); + } else { + distance = _data_store->get_distance(aligned_query, id); + } + Neighbor nn = Neighbor(id, distance); + best_L_nodes.insert(nn); + } + } + + uint32_t hops = 0; + uint32_t cmps = 0; + + while (best_L_nodes.has_unexpanded_node()) { + auto nbr = best_L_nodes.closest_unexpanded(); + auto n = nbr.id; + + // Add node to expanded nodes to create pool for prune later + if (!search_invocation) { + if (!use_filter) { + expanded_nodes.emplace_back(nbr); + } else { // in filter based indexing, the same point might invoke + // multiple iterate_to_fixed_points, so need to be careful + // not to add the same item to pool multiple times. + if (std::find(expanded_nodes.begin(), expanded_nodes.end(), nbr) == + expanded_nodes.end()) { + expanded_nodes.emplace_back(nbr); + } + } + } + + // Find which of the nodes in des have not been visited before + id_scratch.clear(); + dist_scratch.clear(); + { + if (_dynamic_index) _locks[n].lock(); + for (auto id : _final_graph[n]) { + assert(id < _max_points + _num_frozen_pts); + + if (use_filter) { + // NOTE: NEED TO CHECK IF THIS CORRECT WITH NEW LOCKS. + std::vector common_filters; + auto &x = _pts_to_labels[id]; + std::set_intersection(filter_label.begin(), filter_label.end(), + x.begin(), x.end(), + std::back_inserter(common_filters)); + if (_use_universal_label) { + if (std::find(filter_label.begin(), filter_label.end(), + _universal_label) != filter_label.end() || + std::find(x.begin(), x.end(), _universal_label) != x.end()) + common_filters.emplace_back(_universal_label); + } + + if (common_filters.size() == 0) continue; + } + + if (is_not_visited(id)) { + id_scratch.push_back(id); + } + } + + if (_dynamic_index) _locks[n].unlock(); + } + + // Mark nodes visited + for (auto id : id_scratch) { + if (fast_iterate) { + inserted_into_pool_bs[id] = 1; + } else { + inserted_into_pool_rs.insert(id); + } + } - T* aligned_query = scratch->aligned_query(); + // Compute distances to unvisited nodes in the expansion + if (_pq_dist) { + assert(dist_scratch.capacity() >= id_scratch.size()); + compute_dists(id_scratch, dist_scratch); + } else { + assert(dist_scratch.size() == 0); + for (size_t m = 0; m < id_scratch.size(); ++m) { + uint32_t id = id_scratch[m]; - float *query_float = nullptr; - float *query_rotated = nullptr; - float *pq_dists = nullptr; - uint8_t *pq_coord_scratch = nullptr; - // Intialize PQ related scratch to use PQ based distances - if (_pq_dist) - { - // Get scratch spaces - PQScratch *pq_query_scratch = scratch->pq_scratch(); - query_float = pq_query_scratch->aligned_query_float; - query_rotated = pq_query_scratch->rotated_query; - pq_dists = pq_query_scratch->aligned_pqtable_dist_scratch; - - // Copy query vector to float and then to "rotated" query - for (size_t d = 0; d < _dim; d++) - { - query_float[d] = (float)aligned_query[d]; + if (m + 1 < id_scratch.size()) { + auto nextn = id_scratch[m + 1]; + _data_store->prefetch_vector(nextn); } - pq_query_scratch->set(_dim, aligned_query); - // center the query and rotate if we have a rotation matrix - _pq_table.preprocess_query(query_rotated); - _pq_table.populate_chunk_distances(query_rotated, pq_dists); - - pq_coord_scratch = pq_query_scratch->aligned_pq_coord_scratch; + dist_scratch.push_back(_data_store->get_distance(aligned_query, id)); + } } + cmps += id_scratch.size(); - if (expanded_nodes.size() > 0 || id_scratch.size() > 0) - { - throw ANNException("ERROR: Clear scratch space before passing.", -1, __FUNCSIG__, __FILE__, __LINE__); + // Insert pairs into the pool of candidates + for (size_t m = 0; m < id_scratch.size(); ++m) { + best_L_nodes.insert(Neighbor(id_scratch[m], dist_scratch[m])); } + } + return std::make_pair(hops, cmps); +} - // Decide whether to use bitset or robin set to mark visited nodes - auto total_num_points = _max_points + _num_frozen_pts; - bool fast_iterate = total_num_points <= MAX_POINTS_FOR_USING_BITSET; +template +void Index::search_for_point_and_prune( + int location, uint32_t Lindex, std::vector &pruned_list, + InMemQueryScratch *scratch, bool use_filter, uint32_t filteredLindex) { + const std::vector init_ids = get_init_ids(); + const std::vector unused_filter_label; - if (fast_iterate) - { - if (inserted_into_pool_bs.size() < total_num_points) - { - // hopefully using 2X will reduce the number of allocations. - auto resize_size = - 2 * total_num_points > MAX_POINTS_FOR_USING_BITSET ? MAX_POINTS_FOR_USING_BITSET : 2 * total_num_points; - inserted_into_pool_bs.resize(resize_size); - } - } + if (!use_filter) { + _data_store->get_vector(location, scratch->aligned_query()); + iterate_to_fixed_point(scratch->aligned_query(), Lindex, init_ids, scratch, + false, unused_filter_label, false); + } else { + std::vector filter_specific_start_nodes; + for (auto &x : _pts_to_labels[location]) + filter_specific_start_nodes.emplace_back(_label_to_medoid_id[x]); - // Lambda to determine if a node has been visited - auto is_not_visited = [this, fast_iterate, &inserted_into_pool_bs, &inserted_into_pool_rs](const uint32_t id) { - return fast_iterate ? inserted_into_pool_bs[id] == 0 - : inserted_into_pool_rs.find(id) == inserted_into_pool_rs.end(); - }; + _data_store->get_vector(location, scratch->aligned_query()); + iterate_to_fixed_point(scratch->aligned_query(), filteredLindex, + filter_specific_start_nodes, scratch, true, + _pts_to_labels[location], false); + } - // Lambda to batch compute query<-> node distances in PQ space - auto compute_dists = [this, pq_coord_scratch, pq_dists](const std::vector &ids, - std::vector &dists_out) { - diskann::aggregate_coords(ids, this->_pq_data, this->_num_pq_chunks, pq_coord_scratch); - diskann::pq_dist_lookup(pq_coord_scratch, ids.size(), this->_num_pq_chunks, pq_dists, dists_out); - }; + auto &pool = scratch->pool(); - // Initialize the candidate pool with starting points - for (auto id : init_ids) - { - if (id >= _max_points + _num_frozen_pts) - { - diskann::cerr << "Out of range loc found as an edge : " << id << std::endl; - throw diskann::ANNException(std::string("Wrong loc") + std::to_string(id), -1, __FUNCSIG__, __FILE__, - __LINE__); - } + for (uint32_t i = 0; i < pool.size(); i++) { + if (pool[i].id == (uint32_t)location) { + pool.erase(pool.begin() + i); + i--; + } + } - if (use_filter) - { - std::vector common_filters; - auto &x = _pts_to_labels[id]; - std::set_intersection(filter_label.begin(), filter_label.end(), x.begin(), x.end(), - std::back_inserter(common_filters)); - if (_use_universal_label) - { - if (std::find(filter_label.begin(), filter_label.end(), _universal_label) != filter_label.end() || - std::find(x.begin(), x.end(), _universal_label) != x.end()) - common_filters.emplace_back(_universal_label); - } + if (pruned_list.size() > 0) { + throw diskann::ANNException("ERROR: non-empty pruned_list passed", -1, + __FUNCSIG__, __FILE__, __LINE__); + } - if (common_filters.size() == 0) - continue; - } + prune_neighbors(location, pool, pruned_list, scratch); - if (is_not_visited(id)) - { - if (fast_iterate) - { - inserted_into_pool_bs[id] = 1; - } - else - { - inserted_into_pool_rs.insert(id); - } + assert(!pruned_list.empty()); + assert(_final_graph.size() == _max_points + _num_frozen_pts); +} - float distance; - if (_pq_dist) - { - pq_dist_lookup(pq_coord_scratch, 1, this->_num_pq_chunks, pq_dists, &distance); - } - else - { - distance = _data_store->get_distance(aligned_query, id); +template +void Index::occlude_list( + const uint32_t location, std::vector &pool, const float alpha, + const uint32_t degree, const uint32_t maxc, std::vector &result, + InMemQueryScratch *scratch, + const tsl::robin_set *const delete_set_ptr) { + if (pool.size() == 0) return; + + // Truncate pool at maxc and initialize scratch spaces + assert(std::is_sorted(pool.begin(), pool.end())); + assert(result.size() == 0); + if (pool.size() > maxc) pool.resize(maxc); + std::vector &occlude_factor = scratch->occlude_factor(); + // occlude_list can be called with the same scratch more than once by + // search_for_point_and_add_link through inter_insert. + occlude_factor.clear(); + // Initialize occlude_factor to pool.size() many 0.0f values for correctness + occlude_factor.insert(occlude_factor.end(), pool.size(), 0.0f); + + float cur_alpha = 1; + while (cur_alpha <= alpha && result.size() < degree) { + // used for MIPS, where we store a value of eps in cur_alpha to + // denote pruned out entries which we can skip in later rounds. + float eps = cur_alpha + 0.01f; + + for (auto iter = pool.begin(); result.size() < degree && iter != pool.end(); + ++iter) { + if (occlude_factor[iter - pool.begin()] > cur_alpha) { + continue; + } + // Set the entry to float::max so that is not considered again + occlude_factor[iter - pool.begin()] = std::numeric_limits::max(); + // Add the entry to the result if its not been deleted, and doesn't + // add a self loop + if (delete_set_ptr == nullptr || + delete_set_ptr->find(iter->id) == delete_set_ptr->end()) { + if (iter->id != location) { + result.push_back(iter->id); + } + } + + // Update occlude factor for points from iter+1 to pool.end() + for (auto iter2 = iter + 1; iter2 != pool.end(); iter2++) { + auto t = iter2 - pool.begin(); + if (occlude_factor[t] > alpha) continue; + + bool prune_allowed = true; + if (_filtered_index) { + uint32_t a = iter->id; + uint32_t b = iter2->id; + for (auto &x : _pts_to_labels[b]) { + if (std::find(_pts_to_labels[a].begin(), _pts_to_labels[a].end(), + x) == _pts_to_labels[a].end()) { + prune_allowed = false; } - Neighbor nn = Neighbor(id, distance); - best_L_nodes.insert(nn); + if (!prune_allowed) break; + } } - } + if (!prune_allowed) continue; - uint32_t hops = 0; - uint32_t cmps = 0; - - while (best_L_nodes.has_unexpanded_node()) - { - auto nbr = best_L_nodes.closest_unexpanded(); - auto n = nbr.id; - - // Add node to expanded nodes to create pool for prune later - if (!search_invocation) - { - if (!use_filter) - { - expanded_nodes.emplace_back(nbr); - } - else - { // in filter based indexing, the same point might invoke - // multiple iterate_to_fixed_points, so need to be careful - // not to add the same item to pool multiple times. - if (std::find(expanded_nodes.begin(), expanded_nodes.end(), nbr) == expanded_nodes.end()) - { - expanded_nodes.emplace_back(nbr); - } - } + float djk = _data_store->get_distance(iter2->id, iter->id); + if (_dist_metric == diskann::Metric::L2 || + _dist_metric == diskann::Metric::COSINE) { + occlude_factor[t] = + (djk == 0) ? std::numeric_limits::max() + : std::max(occlude_factor[t], iter2->distance / djk); + } else if (_dist_metric == diskann::Metric::INNER_PRODUCT) { + // Improvization for flipping max and min dist for MIPS + float x = -iter2->distance; + float y = -djk; + if (y > cur_alpha * x) { + occlude_factor[t] = std::max(occlude_factor[t], eps); + } } + } + } + cur_alpha *= 1.2f; + } +} - // Find which of the nodes in des have not been visited before - id_scratch.clear(); - dist_scratch.clear(); - { - if (_dynamic_index) - _locks[n].lock(); - for (auto id : _final_graph[n]) - { - assert(id < _max_points + _num_frozen_pts); - - if (use_filter) - { - // NOTE: NEED TO CHECK IF THIS CORRECT WITH NEW LOCKS. - std::vector common_filters; - auto &x = _pts_to_labels[id]; - std::set_intersection(filter_label.begin(), filter_label.end(), x.begin(), x.end(), - std::back_inserter(common_filters)); - if (_use_universal_label) - { - if (std::find(filter_label.begin(), filter_label.end(), _universal_label) != - filter_label.end() || - std::find(x.begin(), x.end(), _universal_label) != x.end()) - common_filters.emplace_back(_universal_label); - } - - if (common_filters.size() == 0) - continue; - } - - if (is_not_visited(id)) - { - id_scratch.push_back(id); - } - } - - if (_dynamic_index) - _locks[n].unlock(); - } +template +void Index::prune_neighbors(const uint32_t location, + std::vector &pool, + std::vector &pruned_list, + InMemQueryScratch *scratch) { + prune_neighbors(location, pool, _indexingRange, _indexingMaxC, _indexingAlpha, + pruned_list, scratch); +} - // Mark nodes visited - for (auto id : id_scratch) - { - if (fast_iterate) - { - inserted_into_pool_bs[id] = 1; - } - else - { - inserted_into_pool_rs.insert(id); - } - } +template +void Index::prune_neighbors( + const uint32_t location, std::vector &pool, const uint32_t range, + const uint32_t max_candidate_size, const float alpha, + std::vector &pruned_list, InMemQueryScratch *scratch) { + if (pool.size() == 0) { + // if the pool is empty, behave like a noop + pruned_list.clear(); + return; + } + + _max_observed_degree = (std::max)(_max_observed_degree, range); + + // If using _pq_build, over-write the PQ distances with actual distances + if (_pq_dist) { + for (auto &ngh : pool) + ngh.distance = _data_store->get_distance(ngh.id, location); + } + + // sort the pool based on distance to query and prune it with occlude_list + std::sort(pool.begin(), pool.end()); + pruned_list.clear(); + pruned_list.reserve(range); + + if (pool.begin()->distance == 0) { + diskann::cerr + << "Warning: a candidate with distance 0 found in prune_neighbors" + << std::endl; + } + occlude_list(location, pool, alpha, range, max_candidate_size, pruned_list, + scratch); + assert(pruned_list.size() <= range); + + if (_saturate_graph && alpha > 1) { + for (const auto &node : pool) { + if (pruned_list.size() >= range) break; + if ((std::find(pruned_list.begin(), pruned_list.end(), node.id) == + pruned_list.end()) && + node.id != location) + pruned_list.push_back(node.id); + } + } +} - // Compute distances to unvisited nodes in the expansion - if (_pq_dist) - { - assert(dist_scratch.capacity() >= id_scratch.size()); - compute_dists(id_scratch, dist_scratch); - } - else - { - assert(dist_scratch.size() == 0); - for (size_t m = 0; m < id_scratch.size(); ++m) - { - uint32_t id = id_scratch[m]; - - if (m + 1 < id_scratch.size()) - { - auto nextn = id_scratch[m + 1]; - _data_store->prefetch_vector(nextn); - } - - dist_scratch.push_back(_data_store->get_distance(aligned_query, id)); - } - } - cmps += id_scratch.size(); +template +void Index::inter_insert(uint32_t n, + std::vector &pruned_list, + const uint32_t range, + InMemQueryScratch *scratch) { + const auto &src_pool = pruned_list; + + assert(!src_pool.empty()); + + for (auto des : src_pool) { + // des.loc is the loc of the neighbors of n + assert(des < _max_points + _num_frozen_pts); + // des_pool contains the neighbors of the neighbors of n + std::vector copy_of_neighbors; + bool prune_needed = false; + { + LockGuard guard(_locks[des]); + auto &des_pool = _final_graph[des]; + if (std::find(des_pool.begin(), des_pool.end(), n) == des_pool.end()) { + if (des_pool.size() < (uint64_t)(GRAPH_SLACK_FACTOR * range)) { + des_pool.emplace_back(n); + prune_needed = false; + } else { + copy_of_neighbors.reserve(des_pool.size() + 1); + copy_of_neighbors = des_pool; + copy_of_neighbors.push_back(n); + prune_needed = true; + } + } + } // des lock is released by this point + + if (prune_needed) { + tsl::robin_set dummy_visited(0); + std::vector dummy_pool(0); + + size_t reserveSize = + (size_t)(std::ceil(1.05 * GRAPH_SLACK_FACTOR * range)); + dummy_visited.reserve(reserveSize); + dummy_pool.reserve(reserveSize); + + for (auto cur_nbr : copy_of_neighbors) { + if (dummy_visited.find(cur_nbr) == dummy_visited.end() && + cur_nbr != des) { + float dist = _data_store->get_distance(des, cur_nbr); + dummy_pool.emplace_back(Neighbor(cur_nbr, dist)); + dummy_visited.insert(cur_nbr); + } + } + std::vector new_out_neighbors; + prune_neighbors(des, dummy_pool, new_out_neighbors, scratch); + { + LockGuard guard(_locks[des]); + + _final_graph[des] = new_out_neighbors; + } + } + } +} - // Insert pairs into the pool of candidates - for (size_t m = 0; m < id_scratch.size(); ++m) - { - best_L_nodes.insert(Neighbor(id_scratch[m], dist_scratch[m])); - } - } - return std::make_pair(hops, cmps); +template +void Index::inter_insert(uint32_t n, + std::vector &pruned_list, + InMemQueryScratch *scratch) { + inter_insert(n, pruned_list, _indexingRange, scratch); } template -void Index::search_for_point_and_prune(int location, uint32_t Lindex, - std::vector &pruned_list, - InMemQueryScratch *scratch, bool use_filter, - uint32_t filteredLindex) -{ - const std::vector init_ids = get_init_ids(); - const std::vector unused_filter_label; +void Index::link(IndexWriteParameters ¶meters) { + uint32_t num_threads = parameters.num_threads; + if (num_threads != 0) omp_set_num_threads(num_threads); + + _saturate_graph = parameters.saturate_graph; + + _indexingQueueSize = parameters.search_list_size; + _filterIndexingQueueSize = parameters.filter_list_size; + _indexingRange = parameters.max_degree; + _indexingMaxC = parameters.max_occlusion_size; + _indexingAlpha = parameters.alpha; + + /* visit_order is a vector that is initialized to the entire graph */ + std::vector visit_order; + std::vector pool, tmp; + tsl::robin_set visited; + visit_order.reserve(_nd + _num_frozen_pts); + for (uint32_t i = 0; i < (uint32_t)_nd; i++) { + visit_order.emplace_back(i); + } + + // If there are any frozen points, add them all. + for (uint32_t frozen = _max_points; frozen < _max_points + _num_frozen_pts; + frozen++) { + visit_order.emplace_back(frozen); + } + + // if there are frozen points, the first such one is set to be the _start + if (_num_frozen_pts > 0) + _start = (uint32_t)_max_points; + else + _start = calculate_entry_point(); - if (!use_filter) - { - _data_store->get_vector(location, scratch->aligned_query()); - iterate_to_fixed_point(scratch->aligned_query(), Lindex, init_ids, scratch, false, unused_filter_label, false); - } - else - { - std::vector filter_specific_start_nodes; - for (auto &x : _pts_to_labels[location]) - filter_specific_start_nodes.emplace_back(_label_to_medoid_id[x]); + for (size_t p = 0; p < _nd; p++) { + _final_graph[p].reserve( + (size_t)(std::ceil(_indexingRange * GRAPH_SLACK_FACTOR * 1.05))); + } - _data_store->get_vector(location, scratch->aligned_query()); - iterate_to_fixed_point(scratch->aligned_query(), filteredLindex, filter_specific_start_nodes, scratch, true, - _pts_to_labels[location], false); - } + diskann::Timer link_timer; - auto &pool = scratch->pool(); +#pragma omp parallel for schedule(dynamic, 2048) + for (int64_t node_ctr = 0; node_ctr < (int64_t)(visit_order.size()); + node_ctr++) { + auto node = visit_order[node_ctr]; - for (uint32_t i = 0; i < pool.size(); i++) - { - if (pool[i].id == (uint32_t)location) - { - pool.erase(pool.begin() + i); - i--; - } - } + ScratchStoreManager> manager(_query_scratch); + auto scratch = manager.scratch_space(); - if (pruned_list.size() > 0) + std::vector pruned_list; + if (_filtered_index) { + search_for_point_and_prune(node, _indexingQueueSize, pruned_list, scratch, + _filtered_index, _filterIndexingQueueSize); + } else { + search_for_point_and_prune(node, _indexingQueueSize, pruned_list, + scratch); + } { - throw diskann::ANNException("ERROR: non-empty pruned_list passed", -1, __FUNCSIG__, __FILE__, __LINE__); + LockGuard guard(_locks[node]); + _final_graph[node].reserve( + (size_t)(_indexingRange * GRAPH_SLACK_FACTOR * 1.05)); + _final_graph[node] = pruned_list; + assert(_final_graph[node].size() <= _indexingRange); } - prune_neighbors(location, pool, pruned_list, scratch); + inter_insert(node, pruned_list, scratch); - assert(!pruned_list.empty()); - assert(_final_graph.size() == _max_points + _num_frozen_pts); + if (node_ctr % 100000 == 0) { + diskann::cout << "\r" << (100.0 * node_ctr) / (visit_order.size()) + << "% of index build completed." << std::flush; + } + } + + if (_nd > 0) { + diskann::cout << "Starting final cleanup.." << std::flush; + } +#pragma omp parallel for schedule(dynamic, 2048) + for (int64_t node_ctr = 0; node_ctr < (int64_t)(visit_order.size()); + node_ctr++) { + auto node = visit_order[node_ctr]; + if (_final_graph[node].size() > _indexingRange) { + ScratchStoreManager> manager(_query_scratch); + auto scratch = manager.scratch_space(); + + 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 = _data_store->get_distance(node, cur_nbr); + dummy_pool.emplace_back(Neighbor(cur_nbr, dist)); + dummy_visited.insert(cur_nbr); + } + } + prune_neighbors(node, dummy_pool, new_out_neighbors, scratch); + + _final_graph[node].clear(); + for (auto id : new_out_neighbors) _final_graph[node].emplace_back(id); + } + } + if (_nd > 0) { + diskann::cout << "done. Link time: " + << ((double)link_timer.elapsed() / (double)1000000) << "s" + << std::endl; + } } template -void Index::occlude_list(const uint32_t location, std::vector &pool, const float alpha, - const uint32_t degree, const uint32_t maxc, std::vector &result, - InMemQueryScratch *scratch, - const tsl::robin_set *const delete_set_ptr) -{ - if (pool.size() == 0) - return; +void Index::prune_all_neighbors( + const uint32_t max_degree, const uint32_t max_occlusion_size, + const float alpha) { + const uint32_t range = max_degree; + const uint32_t maxc = max_occlusion_size; - // Truncate pool at maxc and initialize scratch spaces - assert(std::is_sorted(pool.begin(), pool.end())); - assert(result.size() == 0); - if (pool.size() > maxc) - pool.resize(maxc); - std::vector &occlude_factor = scratch->occlude_factor(); - // occlude_list can be called with the same scratch more than once by - // search_for_point_and_add_link through inter_insert. - occlude_factor.clear(); - // Initialize occlude_factor to pool.size() many 0.0f values for correctness - occlude_factor.insert(occlude_factor.end(), pool.size(), 0.0f); + _filtered_index = true; - float cur_alpha = 1; - while (cur_alpha <= alpha && result.size() < degree) - { - // used for MIPS, where we store a value of eps in cur_alpha to - // denote pruned out entries which we can skip in later rounds. - float eps = cur_alpha + 0.01f; - - for (auto iter = pool.begin(); result.size() < degree && iter != pool.end(); ++iter) - { - if (occlude_factor[iter - pool.begin()] > cur_alpha) - { - continue; - } - // Set the entry to float::max so that is not considered again - occlude_factor[iter - pool.begin()] = std::numeric_limits::max(); - // Add the entry to the result if its not been deleted, and doesn't - // add a self loop - if (delete_set_ptr == nullptr || delete_set_ptr->find(iter->id) == delete_set_ptr->end()) - { - if (iter->id != location) - { - result.push_back(iter->id); - } - } + diskann::Timer timer; +#pragma omp parallel for + for (int64_t node = 0; node < (int64_t)(_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; - // Update occlude factor for points from iter+1 to pool.end() - for (auto iter2 = iter + 1; iter2 != pool.end(); iter2++) - { - auto t = iter2 - pool.begin(); - if (occlude_factor[t] > alpha) - continue; - - bool prune_allowed = true; - if (_filtered_index) - { - uint32_t a = iter->id; - uint32_t b = iter2->id; - for (auto &x : _pts_to_labels[b]) - { - if (std::find(_pts_to_labels[a].begin(), _pts_to_labels[a].end(), x) == _pts_to_labels[a].end()) - { - prune_allowed = false; - } - if (!prune_allowed) - break; - } - } - if (!prune_allowed) - continue; - - float djk = _data_store->get_distance(iter2->id, iter->id); - if (_dist_metric == diskann::Metric::L2 || _dist_metric == diskann::Metric::COSINE) - { - occlude_factor[t] = (djk == 0) ? std::numeric_limits::max() - : std::max(occlude_factor[t], iter2->distance / djk); - } - else if (_dist_metric == diskann::Metric::INNER_PRODUCT) - { - // Improvization for flipping max and min dist for MIPS - float x = -iter2->distance; - float y = -djk; - if (y > cur_alpha * x) - { - occlude_factor[t] = std::max(occlude_factor[t], eps); - } - } - } - } - cur_alpha *= 1.2f; - } + ScratchStoreManager> manager(_query_scratch); + auto scratch = manager.scratch_space(); + + for (auto cur_nbr : _final_graph[node]) { + if (dummy_visited.find(cur_nbr) == dummy_visited.end() && + cur_nbr != node) { + float dist = _data_store->get_distance(node, cur_nbr); + dummy_pool.emplace_back(Neighbor(cur_nbr, dist)); + dummy_visited.insert(cur_nbr); + } + } + + prune_neighbors((uint32_t)node, dummy_pool, range, maxc, alpha, + new_out_neighbors, scratch); + _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 < _max_points + _num_frozen_pts; i++) { + if (i < _nd || i >= _max_points) { + const std::vector &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; + } } +// REFACTOR template -void Index::prune_neighbors(const uint32_t location, std::vector &pool, - std::vector &pruned_list, InMemQueryScratch *scratch) -{ - prune_neighbors(location, pool, _indexingRange, _indexingMaxC, _indexingAlpha, pruned_list, scratch); +void Index::set_start_points(const T *data, + size_t data_count) { + std::unique_lock ul(_update_lock); + std::unique_lock tl(_tag_lock); + if (_nd > 0) + throw ANNException("Can not set starting point for a non-empty index", -1, + __FUNCSIG__, __FILE__, __LINE__); + + if (data_count != _num_frozen_pts * _dim) + throw ANNException("Invalid number of points", -1, __FUNCSIG__, __FILE__, + __LINE__); + + // memcpy(_data + _aligned_dim * _max_points, data, _aligned_dim * + // sizeof(T) * _num_frozen_pts); + for (location_t i = _max_points; i < _max_points + _num_frozen_pts; i++) { + _data_store->set_vector(i, data + i * _dim); + } + _has_built = true; + diskann::cout << "Index start points set: #" << _num_frozen_pts << std::endl; } template -void Index::prune_neighbors(const uint32_t location, std::vector &pool, const uint32_t range, - const uint32_t max_candidate_size, const float alpha, - std::vector &pruned_list, InMemQueryScratch *scratch) -{ - if (pool.size() == 0) - { - // if the pool is empty, behave like a noop - pruned_list.clear(); - return; - } +void Index::set_start_points_at_random(T radius, + uint32_t random_seed) { + std::mt19937 gen{random_seed}; + std::normal_distribution<> d{0.0, 1.0}; - _max_observed_degree = (std::max)(_max_observed_degree, range); + std::vector points_data; + points_data.reserve(_dim * _num_frozen_pts); + std::vector real_vec(_dim); - // If using _pq_build, over-write the PQ distances with actual distances - if (_pq_dist) - { - for (auto &ngh : pool) - ngh.distance = _data_store->get_distance(ngh.id, location); + for (size_t frozen_point = 0; frozen_point < _num_frozen_pts; + frozen_point++) { + double norm_sq = 0.0; + for (size_t i = 0; i < _dim; ++i) { + auto r = d(gen); + real_vec[i] = r; + norm_sq += r * r; } - // sort the pool based on distance to query and prune it with occlude_list - std::sort(pool.begin(), pool.end()); - pruned_list.clear(); - pruned_list.reserve(range); - - if (pool.begin()->distance == 0) - { - diskann::cerr << "Warning: a candidate with distance 0 found in prune_neighbors" << std::endl; - } - occlude_list(location, pool, alpha, range, max_candidate_size, pruned_list, scratch); - assert(pruned_list.size() <= range); + const double norm = std::sqrt(norm_sq); + for (auto iter : real_vec) + points_data.push_back(static_cast(iter * radius / norm)); + } - if (_saturate_graph && alpha > 1) - { - for (const auto &node : pool) - { - if (pruned_list.size() >= range) - break; - if ((std::find(pruned_list.begin(), pruned_list.end(), node.id) == pruned_list.end()) && - node.id != location) - pruned_list.push_back(node.id); - } - } + set_start_points(points_data.data(), points_data.size()); } template -void Index::inter_insert(uint32_t n, std::vector &pruned_list, const uint32_t range, - InMemQueryScratch *scratch) -{ - const auto &src_pool = pruned_list; +void Index::build_with_data_populated( + IndexWriteParameters ¶meters, const std::vector &tags) { + diskann::cout << "Starting index build with " << _nd << " points... " + << std::endl; - assert(!src_pool.empty()); + if (_nd < 1) + throw ANNException("Error: Trying to build an index with 0 points", -1, + __FUNCSIG__, __FILE__, __LINE__); - for (auto des : src_pool) - { - // des.loc is the loc of the neighbors of n - assert(des < _max_points + _num_frozen_pts); - // des_pool contains the neighbors of the neighbors of n - std::vector copy_of_neighbors; - bool prune_needed = false; - { - LockGuard guard(_locks[des]); - auto &des_pool = _final_graph[des]; - if (std::find(des_pool.begin(), des_pool.end(), n) == des_pool.end()) - { - if (des_pool.size() < (uint64_t)(GRAPH_SLACK_FACTOR * range)) - { - des_pool.emplace_back(n); - prune_needed = false; - } - else - { - copy_of_neighbors.reserve(des_pool.size() + 1); - copy_of_neighbors = des_pool; - copy_of_neighbors.push_back(n); - prune_needed = true; - } - } - } // des lock is released by this point - - if (prune_needed) - { - tsl::robin_set dummy_visited(0); - std::vector dummy_pool(0); - - size_t reserveSize = (size_t)(std::ceil(1.05 * GRAPH_SLACK_FACTOR * range)); - dummy_visited.reserve(reserveSize); - dummy_pool.reserve(reserveSize); - - for (auto cur_nbr : copy_of_neighbors) - { - if (dummy_visited.find(cur_nbr) == dummy_visited.end() && cur_nbr != des) - { - float dist = _data_store->get_distance(des, cur_nbr); - dummy_pool.emplace_back(Neighbor(cur_nbr, dist)); - dummy_visited.insert(cur_nbr); - } - } - std::vector new_out_neighbors; - prune_neighbors(des, dummy_pool, new_out_neighbors, scratch); - { - LockGuard guard(_locks[des]); - - _final_graph[des] = new_out_neighbors; - } - } - } -} - -template -void Index::inter_insert(uint32_t n, std::vector &pruned_list, InMemQueryScratch *scratch) -{ - inter_insert(n, pruned_list, _indexingRange, scratch); -} - -template -void Index::link(IndexWriteParameters ¶meters) -{ - uint32_t num_threads = parameters.num_threads; - if (num_threads != 0) - omp_set_num_threads(num_threads); - - _saturate_graph = parameters.saturate_graph; - - _indexingQueueSize = parameters.search_list_size; - _filterIndexingQueueSize = parameters.filter_list_size; - _indexingRange = parameters.max_degree; - _indexingMaxC = parameters.max_occlusion_size; - _indexingAlpha = parameters.alpha; - - /* visit_order is a vector that is initialized to the entire graph */ - std::vector visit_order; - std::vector pool, tmp; - tsl::robin_set visited; - visit_order.reserve(_nd + _num_frozen_pts); - for (uint32_t i = 0; i < (uint32_t)_nd; i++) - { - visit_order.emplace_back(i); - } - - // If there are any frozen points, add them all. - for (uint32_t frozen = _max_points; frozen < _max_points + _num_frozen_pts; frozen++) - { - visit_order.emplace_back(frozen); - } - - // if there are frozen points, the first such one is set to be the _start - if (_num_frozen_pts > 0) - _start = (uint32_t)_max_points; - else - _start = calculate_entry_point(); - - for (size_t p = 0; p < _nd; p++) - { - _final_graph[p].reserve((size_t)(std::ceil(_indexingRange * GRAPH_SLACK_FACTOR * 1.05))); - } - - diskann::Timer link_timer; - -#pragma omp parallel for schedule(dynamic, 2048) - for (int64_t node_ctr = 0; node_ctr < (int64_t)(visit_order.size()); node_ctr++) - { - auto node = visit_order[node_ctr]; - - ScratchStoreManager> manager(_query_scratch); - auto scratch = manager.scratch_space(); - - std::vector pruned_list; - if (_filtered_index) - { - search_for_point_and_prune(node, _indexingQueueSize, pruned_list, scratch, _filtered_index, - _filterIndexingQueueSize); - } - else - { - search_for_point_and_prune(node, _indexingQueueSize, pruned_list, scratch); - } - { - LockGuard guard(_locks[node]); - _final_graph[node].reserve((size_t)(_indexingRange * GRAPH_SLACK_FACTOR * 1.05)); - _final_graph[node] = pruned_list; - assert(_final_graph[node].size() <= _indexingRange); - } - - inter_insert(node, pruned_list, scratch); - - if (node_ctr % 100000 == 0) - { - diskann::cout << "\r" << (100.0 * node_ctr) / (visit_order.size()) << "% of index build completed." - << std::flush; - } - } - - if (_nd > 0) - { - diskann::cout << "Starting final cleanup.." << std::flush; - } -#pragma omp parallel for schedule(dynamic, 2048) - for (int64_t node_ctr = 0; node_ctr < (int64_t)(visit_order.size()); node_ctr++) - { - auto node = visit_order[node_ctr]; - if (_final_graph[node].size() > _indexingRange) - { - ScratchStoreManager> manager(_query_scratch); - auto scratch = manager.scratch_space(); - - 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 = _data_store->get_distance(node, cur_nbr); - dummy_pool.emplace_back(Neighbor(cur_nbr, dist)); - dummy_visited.insert(cur_nbr); - } - } - prune_neighbors(node, dummy_pool, new_out_neighbors, scratch); - - _final_graph[node].clear(); - for (auto id : new_out_neighbors) - _final_graph[node].emplace_back(id); - } - } - if (_nd > 0) - { - diskann::cout << "done. Link time: " << ((double)link_timer.elapsed() / (double)1000000) << "s" << std::endl; - } -} - -template -void Index::prune_all_neighbors(const uint32_t max_degree, const uint32_t max_occlusion_size, - const float alpha) -{ - const uint32_t range = max_degree; - const uint32_t maxc = max_occlusion_size; - - _filtered_index = true; - - diskann::Timer timer; -#pragma omp parallel for - for (int64_t node = 0; node < (int64_t)(_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; - - ScratchStoreManager> manager(_query_scratch); - auto scratch = manager.scratch_space(); - - for (auto cur_nbr : _final_graph[node]) - { - if (dummy_visited.find(cur_nbr) == dummy_visited.end() && cur_nbr != node) - { - float dist = _data_store->get_distance(node, cur_nbr); - dummy_pool.emplace_back(Neighbor(cur_nbr, dist)); - dummy_visited.insert(cur_nbr); - } - } - - prune_neighbors((uint32_t)node, dummy_pool, range, maxc, alpha, new_out_neighbors, scratch); - _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 < _max_points + _num_frozen_pts; i++) - { - if (i < _nd || i >= _max_points) - { - const std::vector &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; - } + if (_enable_tags && tags.size() != _nd) { + std::stringstream stream; + stream << "ERROR: Driver requests loading " << _nd << " points from file," + << "but tags vector is of size " << tags.size() << "." << std::endl; + diskann::cerr << stream.str() << std::endl; + 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]] = (uint32_t)i; + _location_to_tag.set(static_cast(i), tags[i]); + } + } + + uint32_t index_R = parameters.max_degree; + uint32_t num_threads_index = parameters.num_threads; + uint32_t index_L = parameters.search_list_size; + uint32_t maxc = parameters.max_occlusion_size; + + if (_query_scratch.size() == 0) { + initialize_query_scratch(5 + num_threads_index, index_L, index_L, index_R, + maxc, _data_store->get_aligned_dim()); + } + + generate_frozen_point(); + link(parameters); + + size_t max = 0, min = SIZE_MAX, 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++; + } + diskann::cout << "Index built with degree: max:" << max + << " avg:" << (float)total / (float)(_nd + _num_frozen_pts) + << " min:" << min << " count(deg<2):" << cnt << std::endl; + + _max_observed_degree = std::max((uint32_t)max, _max_observed_degree); + _has_built = true; } -// REFACTOR template -void Index::set_start_points(const T *data, size_t data_count) -{ - std::unique_lock ul(_update_lock); +void Index::build(const T *data, + const size_t num_points_to_load, + IndexWriteParameters ¶meters, + const std::vector &tags) { + if (num_points_to_load == 0) { + 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); + + { std::unique_lock tl(_tag_lock); - if (_nd > 0) - throw ANNException("Can not set starting point for a non-empty index", -1, __FUNCSIG__, __FILE__, __LINE__); - - if (data_count != _num_frozen_pts * _dim) - throw ANNException("Invalid number of points", -1, __FUNCSIG__, __FILE__, __LINE__); - - // memcpy(_data + _aligned_dim * _max_points, data, _aligned_dim * sizeof(T) * _num_frozen_pts); - for (location_t i = _max_points; i < _max_points + _num_frozen_pts; i++) - { - _data_store->set_vector(i, data + i * _dim); - } - _has_built = true; - diskann::cout << "Index start points set: #" << _num_frozen_pts << std::endl; -} - -template -void Index::set_start_points_at_random(T radius, uint32_t random_seed) -{ - std::mt19937 gen{random_seed}; - std::normal_distribution<> d{0.0, 1.0}; - - std::vector points_data; - points_data.reserve(_dim * _num_frozen_pts); - std::vector real_vec(_dim); + _nd = num_points_to_load; - for (size_t frozen_point = 0; frozen_point < _num_frozen_pts; frozen_point++) - { - double norm_sq = 0.0; - for (size_t i = 0; i < _dim; ++i) - { - auto r = d(gen); - real_vec[i] = r; - norm_sq += r * r; - } + _data_store->populate_data(data, num_points_to_load); - const double norm = std::sqrt(norm_sq); - for (auto iter : real_vec) - points_data.push_back(static_cast(iter * radius / norm)); - } - - set_start_points(points_data.data(), points_data.size()); -} - -template -void Index::build_with_data_populated(IndexWriteParameters ¶meters, const std::vector &tags) -{ - diskann::cout << "Starting index build with " << _nd << " points... " << std::endl; - - if (_nd < 1) - throw ANNException("Error: Trying to build an index with 0 points", -1, __FUNCSIG__, __FILE__, __LINE__); - - if (_enable_tags && tags.size() != _nd) - { - std::stringstream stream; - stream << "ERROR: Driver requests loading " << _nd << " points from file," - << "but tags vector is of size " << tags.size() << "." << std::endl; - diskann::cerr << stream.str() << std::endl; - 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]] = (uint32_t)i; - _location_to_tag.set(static_cast(i), tags[i]); - } - } - - uint32_t index_R = parameters.max_degree; - uint32_t num_threads_index = parameters.num_threads; - uint32_t index_L = parameters.search_list_size; - uint32_t maxc = parameters.max_occlusion_size; - - if (_query_scratch.size() == 0) - { - initialize_query_scratch(5 + num_threads_index, index_L, index_L, index_R, maxc, - _data_store->get_aligned_dim()); - } - - generate_frozen_point(); - link(parameters); - - size_t max = 0, min = SIZE_MAX, 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++; - } - diskann::cout << "Index built with degree: max:" << max << " avg:" << (float)total / (float)(_nd + _num_frozen_pts) - << " min:" << min << " count(deg<2):" << cnt << std::endl; - - _max_observed_degree = std::max((uint32_t)max, _max_observed_degree); - _has_built = true; -} - -template -void Index::build(const T *data, const size_t num_points_to_load, IndexWriteParameters ¶meters, - const std::vector &tags) -{ - if (num_points_to_load == 0) - { - 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); - - { - std::unique_lock tl(_tag_lock); - _nd = num_points_to_load; - - _data_store->populate_data(data, num_points_to_load); - - // REFACTOR - // memcpy((char *)_data, (char *)data, _aligned_dim * _nd * sizeof(T)); - // if (_normalize_vecs) - //{ - // for (size_t i = 0; i < num_points_to_load; i++) - // { - // normalize(_data + _aligned_dim * i, _aligned_dim); - // } - // } - } + // REFACTOR + // memcpy((char *)_data, (char *)data, _aligned_dim * _nd * sizeof(T)); + // if (_normalize_vecs) + //{ + // for (size_t i = 0; i < num_points_to_load; i++) + // { + // normalize(_data + _aligned_dim * i, _aligned_dim); + // } + // } + } - build_with_data_populated(parameters, tags); + build_with_data_populated(parameters, tags); } template -void Index::build(const char *filename, const size_t num_points_to_load, - IndexWriteParameters ¶meters, const std::vector &tags) -{ - std::unique_lock ul(_update_lock); - if (num_points_to_load == 0) - throw ANNException("Do not call build with 0 points", -1, __FUNCSIG__, __FILE__, __LINE__); - - if (!file_exists(filename)) - { - std::stringstream stream; - stream << "ERROR: Data file " << filename << " does not exist." << std::endl; - diskann::cerr << stream.str() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - - size_t file_num_points, file_dim; - if (filename == nullptr) - { - throw diskann::ANNException("Can not build with an empty file", -1, __FUNCSIG__, __FILE__, __LINE__); - } - - diskann::get_bin_metadata(filename, file_num_points, file_dim); - if (file_num_points > _max_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; - - if (_pq_dist) - aligned_free(_pq_data); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - - if (num_points_to_load > file_num_points) - { - std::stringstream stream; - stream << "ERROR: Driver requests loading " << num_points_to_load << " points and file has only " - << file_num_points << " points." << std::endl; - - if (_pq_dist) - aligned_free(_pq_data); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } +void Index::build(const char *filename, + const size_t num_points_to_load, + IndexWriteParameters ¶meters, + const std::vector &tags) { + std::unique_lock ul(_update_lock); + if (num_points_to_load == 0) + throw ANNException("Do not call build with 0 points", -1, __FUNCSIG__, + __FILE__, __LINE__); + + if (!file_exists(filename)) { + std::stringstream stream; + stream << "ERROR: Data file " << filename << " does not exist." + << std::endl; + diskann::cerr << stream.str() << std::endl; + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } + + size_t file_num_points, file_dim; + if (filename == nullptr) { + throw diskann::ANNException("Can not build with an empty file", -1, + __FUNCSIG__, __FILE__, __LINE__); + } + + diskann::get_bin_metadata(filename, file_num_points, file_dim); + if (file_num_points > _max_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; - if (file_dim != _dim) - { - std::stringstream stream; - stream << "ERROR: Driver requests loading " << _dim << " dimension," - << "but file has " << file_dim << " dimension." << std::endl; - diskann::cerr << stream.str() << std::endl; + if (_pq_dist) aligned_free(_pq_data); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } - if (_pq_dist) - aligned_free(_pq_data); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } + if (num_points_to_load > file_num_points) { + std::stringstream stream; + stream << "ERROR: Driver requests loading " << num_points_to_load + << " points and file has only " << file_num_points << " points." + << std::endl; - if (_pq_dist) - { - double p_val = std::min(1.0, ((double)MAX_PQ_TRAINING_SET_SIZE / (double)file_num_points)); + if (_pq_dist) aligned_free(_pq_data); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } - std::string suffix = _use_opq ? "_opq" : "_pq"; - suffix += std::to_string(_num_pq_chunks); - auto pq_pivots_file = std::string(filename) + suffix + "_pivots.bin"; - auto pq_compressed_file = std::string(filename) + suffix + "_compressed.bin"; - generate_quantized_data(std::string(filename), pq_pivots_file, pq_compressed_file, _dist_metric, p_val, - _num_pq_chunks, _use_opq); + if (file_dim != _dim) { + std::stringstream stream; + stream << "ERROR: Driver requests loading " << _dim << " dimension," + << "but file has " << file_dim << " dimension." << std::endl; + diskann::cerr << stream.str() << std::endl; - copy_aligned_data_from_file(pq_compressed_file.c_str(), _pq_data, file_num_points, _num_pq_chunks, - _num_pq_chunks); + if (_pq_dist) aligned_free(_pq_data); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } + + if (_pq_dist) { + double p_val = std::min( + 1.0, ((double)MAX_PQ_TRAINING_SET_SIZE / (double)file_num_points)); + + std::string suffix = _use_opq ? "_opq" : "_pq"; + suffix += std::to_string(_num_pq_chunks); + auto pq_pivots_file = std::string(filename) + suffix + "_pivots.bin"; + auto pq_compressed_file = + std::string(filename) + suffix + "_compressed.bin"; + generate_quantized_data(std::string(filename), pq_pivots_file, + pq_compressed_file, _dist_metric, p_val, + _num_pq_chunks, _use_opq); + + copy_aligned_data_from_file(pq_compressed_file.c_str(), _pq_data, + file_num_points, _num_pq_chunks, + _num_pq_chunks); #ifdef EXEC_ENV_OLS - throw ANNException("load_pq_centroid_bin should not be called when " - "EXEC_ENV_OLS is defined.", - -1, __FUNCSIG__, __FILE__, __LINE__); + throw ANNException( + "load_pq_centroid_bin should not be called when " + "EXEC_ENV_OLS is defined.", + -1, __FUNCSIG__, __FILE__, __LINE__); #else - _pq_table.load_pq_centroid_bin(pq_pivots_file.c_str(), _num_pq_chunks); + _pq_table.load_pq_centroid_bin(pq_pivots_file.c_str(), _num_pq_chunks); #endif - } + } - _data_store->populate_data(filename, 0U); - diskann::cout << "Using only first " << num_points_to_load << " from file.. " << std::endl; + _data_store->populate_data(filename, 0U); + diskann::cout << "Using only first " << num_points_to_load << " from file.. " + << std::endl; - { - std::unique_lock tl(_tag_lock); - _nd = num_points_to_load; - } - build_with_data_populated(parameters, tags); + { + std::unique_lock tl(_tag_lock); + _nd = num_points_to_load; + } + build_with_data_populated(parameters, tags); } template -void Index::build(const char *filename, const size_t num_points_to_load, - IndexWriteParameters ¶meters, const char *tag_filename) -{ - std::vector tags; +void Index::build(const char *filename, + const size_t num_points_to_load, + IndexWriteParameters ¶meters, + const char *tag_filename) { + std::vector tags; - if (_enable_tags) - { - std::unique_lock tl(_tag_lock); - if (tag_filename == nullptr) - { - throw ANNException("Tag filename is null, while _enable_tags is set", -1, __FUNCSIG__, __FILE__, __LINE__); - } - 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, insufficient to populate tags for " << num_points_to_load - << " points to load"; - throw diskann::ANNException(sstream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - for (size_t i = 0; i < num_points_to_load; i++) - { - tags.push_back(tag_data[i]); - } - delete[] tag_data; - } - else - { - throw diskann::ANNException(std::string("Tag file") + tag_filename + " does not exist", -1, __FUNCSIG__, - __FILE__, __LINE__); - } + if (_enable_tags) { + std::unique_lock tl(_tag_lock); + if (tag_filename == nullptr) { + throw ANNException("Tag filename is null, while _enable_tags is set", -1, + __FUNCSIG__, __FILE__, __LINE__); + } 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, insufficient to populate tags for " + << num_points_to_load << " points to load"; + throw diskann::ANNException(sstream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } + for (size_t i = 0; i < num_points_to_load; i++) { + tags.push_back(tag_data[i]); } + delete[] tag_data; + } else { + throw diskann::ANNException( + std::string("Tag file") + tag_filename + " does not exist", -1, + __FUNCSIG__, __FILE__, __LINE__); + } } - build(filename, num_points_to_load, parameters, tags); + } + build(filename, num_points_to_load, parameters, tags); } template -std::unordered_map Index::load_label_map(const std::string &labels_map_file) -{ - std::unordered_map string_to_int_mp; - std::ifstream map_reader(labels_map_file); - std::string line, token; - LabelT token_as_num; - std::string label_str; - while (std::getline(map_reader, line)) - { - std::istringstream iss(line); - getline(iss, token, '\t'); - label_str = token; - getline(iss, token, '\t'); - token_as_num = std::stoul(token); - string_to_int_mp[label_str] = token_as_num; - } - return string_to_int_mp; +std::unordered_map Index::load_label_map( + const std::string &labels_map_file) { + std::unordered_map string_to_int_mp; + std::ifstream map_reader(labels_map_file); + std::string line, token; + LabelT token_as_num; + std::string label_str; + while (std::getline(map_reader, line)) { + std::istringstream iss(line); + getline(iss, token, '\t'); + label_str = token; + getline(iss, token, '\t'); + token_as_num = std::stoul(token); + string_to_int_mp[label_str] = token_as_num; + } + return string_to_int_mp; } template -LabelT Index::get_converted_label(const std::string &raw_label) -{ - if (_label_map.find(raw_label) != _label_map.end()) - { - return _label_map[raw_label]; - } - std::stringstream stream; - stream << "Unable to find label in the Label Map"; - diskann::cerr << stream.str() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); +LabelT Index::get_converted_label( + const std::string &raw_label) { + if (_label_map.find(raw_label) != _label_map.end()) { + return _label_map[raw_label]; + } + std::stringstream stream; + stream << "Unable to find label in the Label Map"; + diskann::cerr << stream.str() << std::endl; + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); } template -void Index::parse_label_file(const std::string &label_file, size_t &num_points) -{ - // Format of Label txt file: filters with comma separators - - std::ifstream infile(label_file); - if (infile.fail()) - { - throw diskann::ANNException(std::string("Failed to open file ") + label_file, -1); - } - - std::string line, token; - uint32_t line_cnt = 0; - - while (std::getline(infile, line)) - { - line_cnt++; - } - _pts_to_labels.resize(line_cnt, std::vector()); - - infile.clear(); - infile.seekg(0, std::ios::beg); - line_cnt = 0; - - while (std::getline(infile, line)) - { - std::istringstream iss(line); - std::vector lbls(0); - getline(iss, token, '\t'); - std::istringstream new_iss(token); - while (getline(new_iss, token, ',')) - { - token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); - token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); - LabelT token_as_num = std::stoul(token); - lbls.push_back(token_as_num); - _labels.insert(token_as_num); - } - if (lbls.size() <= 0) - { - diskann::cout << "No label found"; - exit(-1); - } - std::sort(lbls.begin(), lbls.end()); - _pts_to_labels[line_cnt] = lbls; - line_cnt++; - } - num_points = (size_t)line_cnt; - diskann::cout << "Identified " << _labels.size() << " distinct label(s)" << std::endl; +void Index::parse_label_file(const std::string &label_file, + size_t &num_points) { + // Format of Label txt file: filters with comma separators + + std::ifstream infile(label_file); + if (infile.fail()) { + throw diskann::ANNException( + std::string("Failed to open file ") + label_file, -1); + } + + std::string line, token; + uint32_t line_cnt = 0; + + while (std::getline(infile, line)) { + line_cnt++; + } + _pts_to_labels.resize(line_cnt, std::vector()); + + infile.clear(); + infile.seekg(0, std::ios::beg); + line_cnt = 0; + + while (std::getline(infile, line)) { + std::istringstream iss(line); + std::vector lbls(0); + getline(iss, token, '\t'); + std::istringstream new_iss(token); + while (getline(new_iss, token, ',')) { + token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); + token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); + LabelT token_as_num = std::stoul(token); + lbls.push_back(token_as_num); + _labels.insert(token_as_num); + } + if (lbls.size() <= 0) { + diskann::cout << "No label found"; + exit(-1); + } + std::sort(lbls.begin(), lbls.end()); + _pts_to_labels[line_cnt] = lbls; + line_cnt++; + } + num_points = (size_t)line_cnt; + diskann::cout << "Identified " << _labels.size() << " distinct label(s)" + << std::endl; } template -void Index::set_universal_label(const LabelT &label) -{ - _use_universal_label = true; - _universal_label = label; +void Index::set_universal_label(const LabelT &label) { + _use_universal_label = true; + _universal_label = label; } template -void Index::build_filtered_index(const char *filename, const std::string &label_file, - const size_t num_points_to_load, IndexWriteParameters ¶meters, - const std::vector &tags) -{ - _labels_file = label_file; - _filtered_index = true; - _label_to_medoid_id.clear(); - size_t num_points_labels = 0; - parse_label_file(label_file, - num_points_labels); // determines medoid for each label and identifies the points to label mapping - - std::unordered_map> label_to_points; - - for (typename tsl::robin_set::size_type lbl = 0; lbl < _labels.size(); lbl++) - { - auto itr = _labels.begin(); - std::advance(itr, lbl); - auto &x = *itr; - - std::vector labeled_points; - for (uint32_t point_id = 0; point_id < num_points_to_load; point_id++) - { - bool pt_has_lbl = std::find(_pts_to_labels[point_id].begin(), _pts_to_labels[point_id].end(), x) != - _pts_to_labels[point_id].end(); - - bool pt_has_univ_lbl = - (_use_universal_label && (std::find(_pts_to_labels[point_id].begin(), _pts_to_labels[point_id].end(), - _universal_label) != _pts_to_labels[point_id].end())); - - if (pt_has_lbl || pt_has_univ_lbl) - { - labeled_points.emplace_back(point_id); - } - } - label_to_points[x] = labeled_points; - } - - uint32_t num_cands = 25; - for (auto itr = _labels.begin(); itr != _labels.end(); itr++) - { - uint32_t best_medoid_count = std::numeric_limits::max(); - auto &curr_label = *itr; - uint32_t best_medoid; - auto labeled_points = label_to_points[curr_label]; - for (uint32_t cnd = 0; cnd < num_cands; cnd++) - { - uint32_t cur_cnd = labeled_points[rand() % labeled_points.size()]; - uint32_t cur_cnt = std::numeric_limits::max(); - if (_medoid_counts.find(cur_cnd) == _medoid_counts.end()) - { - _medoid_counts[cur_cnd] = 0; - cur_cnt = 0; - } - else - { - cur_cnt = _medoid_counts[cur_cnd]; - } - if (cur_cnt < best_medoid_count) - { - best_medoid_count = cur_cnt; - best_medoid = cur_cnd; - } - } - _label_to_medoid_id[curr_label] = best_medoid; - _medoid_counts[best_medoid]++; - } - - this->build(filename, num_points_to_load, parameters, tags); +void Index::build_filtered_index( + const char *filename, const std::string &label_file, + const size_t num_points_to_load, IndexWriteParameters ¶meters, + const std::vector &tags) { + _labels_file = label_file; + _filtered_index = true; + _label_to_medoid_id.clear(); + size_t num_points_labels = 0; + parse_label_file( + label_file, + num_points_labels); // determines medoid for each label and identifies + // the points to label mapping + + std::unordered_map> label_to_points; + + for (typename tsl::robin_set::size_type lbl = 0; lbl < _labels.size(); + lbl++) { + auto itr = _labels.begin(); + std::advance(itr, lbl); + auto &x = *itr; + + std::vector labeled_points; + for (uint32_t point_id = 0; point_id < num_points_to_load; point_id++) { + bool pt_has_lbl = std::find(_pts_to_labels[point_id].begin(), + _pts_to_labels[point_id].end(), + x) != _pts_to_labels[point_id].end(); + + bool pt_has_univ_lbl = + (_use_universal_label && + (std::find(_pts_to_labels[point_id].begin(), + _pts_to_labels[point_id].end(), + _universal_label) != _pts_to_labels[point_id].end())); + + if (pt_has_lbl || pt_has_univ_lbl) { + labeled_points.emplace_back(point_id); + } + } + label_to_points[x] = labeled_points; + } + + uint32_t num_cands = 25; + for (auto itr = _labels.begin(); itr != _labels.end(); itr++) { + uint32_t best_medoid_count = std::numeric_limits::max(); + auto &curr_label = *itr; + uint32_t best_medoid; + auto labeled_points = label_to_points[curr_label]; + for (uint32_t cnd = 0; cnd < num_cands; cnd++) { + uint32_t cur_cnd = labeled_points[rand() % labeled_points.size()]; + uint32_t cur_cnt = std::numeric_limits::max(); + if (_medoid_counts.find(cur_cnd) == _medoid_counts.end()) { + _medoid_counts[cur_cnd] = 0; + cur_cnt = 0; + } else { + cur_cnt = _medoid_counts[cur_cnd]; + } + if (cur_cnt < best_medoid_count) { + best_medoid_count = cur_cnt; + best_medoid = cur_cnd; + } + } + _label_to_medoid_id[curr_label] = best_medoid; + _medoid_counts[best_medoid]++; + } + + this->build(filename, num_points_to_load, parameters, tags); } template template -std::pair Index::search(const T *query, const size_t K, const uint32_t L, - IdType *indices, float *distances) -{ - if (K > (uint64_t)L) - { - throw ANNException("Set L to a value of at least K", -1, __FUNCSIG__, __FILE__, __LINE__); - } - - ScratchStoreManager> manager(_query_scratch); - auto scratch = manager.scratch_space(); - - if (L > scratch->get_L()) - { - diskann::cout << "Attempting to expand query scratch_space. Was created " - << "with Lsize: " << scratch->get_L() << " but search L is: " << L << std::endl; - scratch->resize_for_new_L(L); - diskann::cout << "Resize completed. New scratch->L is " << scratch->get_L() << std::endl; - } +std::pair Index::search(const T *query, + const size_t K, + const uint32_t L, + IdType *indices, + float *distances) { + if (K > (uint64_t)L) { + throw ANNException("Set L to a value of at least K", -1, __FUNCSIG__, + __FILE__, __LINE__); + } + + ScratchStoreManager> manager(_query_scratch); + auto scratch = manager.scratch_space(); + + if (L > scratch->get_L()) { + diskann::cout << "Attempting to expand query scratch_space. Was created " + << "with Lsize: " << scratch->get_L() + << " but search L is: " << L << std::endl; + scratch->resize_for_new_L(L); + diskann::cout << "Resize completed. New scratch->L is " << scratch->get_L() + << std::endl; + } - const std::vector unused_filter_label; - const std::vector init_ids = get_init_ids(); + const std::vector unused_filter_label; + const std::vector init_ids = get_init_ids(); - std::shared_lock lock(_update_lock); + std::shared_lock lock(_update_lock); - _distance->preprocess_query(query, _data_store->get_dims(), scratch->aligned_query()); - auto retval = - iterate_to_fixed_point(scratch->aligned_query(), L, init_ids, scratch, false, unused_filter_label, true); + _distance->preprocess_query(query, _data_store->get_dims(), + scratch->aligned_query()); + auto retval = + iterate_to_fixed_point(scratch->aligned_query(), L, init_ids, scratch, + false, unused_filter_label, true); - NeighborPriorityQueue &best_L_nodes = scratch->best_l_nodes(); + NeighborPriorityQueue &best_L_nodes = scratch->best_l_nodes(); - size_t pos = 0; - for (size_t i = 0; i < best_L_nodes.size(); ++i) - { - if (best_L_nodes[i].id < _max_points) - { - // safe because Index uses uint32_t ids internally - // and IDType will be uint32_t or uint64_t - indices[pos] = (IdType)best_L_nodes[i].id; - if (distances != nullptr) - { + size_t pos = 0; + for (size_t i = 0; i < best_L_nodes.size(); ++i) { + if (best_L_nodes[i].id < _max_points) { + // safe because Index uses uint32_t ids internally + // and IDType will be uint32_t or uint64_t + indices[pos] = (IdType)best_L_nodes[i].id; + if (distances != nullptr) { #ifdef EXEC_ENV_OLS - // DLVS expects negative distances - distances[pos] = best_L_nodes[i].distance; + // DLVS expects negative distances + distances[pos] = best_L_nodes[i].distance; #else - distances[pos] = _dist_metric == diskann::Metric::INNER_PRODUCT ? -1 * best_L_nodes[i].distance - : best_L_nodes[i].distance; + distances[pos] = _dist_metric == diskann::Metric::INNER_PRODUCT + ? -1 * best_L_nodes[i].distance + : best_L_nodes[i].distance; #endif - } - pos++; - } - if (pos == K) - break; - } - if (pos < K) - { - diskann::cerr << "Found fewer than K elements for query" << std::endl; + } + pos++; } + if (pos == K) break; + } + if (pos < K) { + diskann::cerr << "Found fewer than K elements for query" << std::endl; + } - return retval; + return retval; } template template -std::pair Index::search_with_filters(const T *query, const LabelT &filter_label, - const size_t K, const uint32_t L, - IdType *indices, float *distances) -{ - if (K > (uint64_t)L) - { - throw ANNException("Set L to a value of at least K", -1, __FUNCSIG__, __FILE__, __LINE__); - } - - ScratchStoreManager> manager(_query_scratch); - auto scratch = manager.scratch_space(); - - if (L > scratch->get_L()) - { - diskann::cout << "Attempting to expand query scratch_space. Was created " - << "with Lsize: " << scratch->get_L() << " but search L is: " << L << std::endl; - scratch->resize_for_new_L(L); - diskann::cout << "Resize completed. New scratch->L is " << scratch->get_L() << std::endl; - } - - std::vector filter_vec; - std::vector init_ids = get_init_ids(); - - std::shared_lock lock(_update_lock); - - if (_label_to_medoid_id.find(filter_label) != _label_to_medoid_id.end()) - { - init_ids.emplace_back(_label_to_medoid_id[filter_label]); - } - else - { - diskann::cout << "No filtered medoid found. exitting " - << std::endl; // RKNOTE: If universal label found start there - throw diskann::ANNException("No filtered medoid found. exitting ", -1); - } - filter_vec.emplace_back(filter_label); - - //REFACTOR - //T *aligned_query = scratch->aligned_query(); - //memcpy(aligned_query, query, _dim * sizeof(T)); - _distance->preprocess_query(query, _data_store->get_dims(), scratch->aligned_query()); - auto retval = iterate_to_fixed_point(scratch->aligned_query(), L, init_ids, scratch, true, filter_vec, true); - - auto best_L_nodes = scratch->best_l_nodes(); - - size_t pos = 0; - for (size_t i = 0; i < best_L_nodes.size(); ++i) - { - if (best_L_nodes[i].id < _max_points) - { - // safe because Index uses uint32_t ids internally - // and IDType will be uint32_t or uint64_t - indices[pos] = (IdType)best_L_nodes[i].id; - if (distances != nullptr) - { +std::pair Index::search_with_filters( + const T *query, const LabelT &filter_label, const size_t K, + const uint32_t L, IdType *indices, float *distances) { + if (K > (uint64_t)L) { + throw ANNException("Set L to a value of at least K", -1, __FUNCSIG__, + __FILE__, __LINE__); + } + + ScratchStoreManager> manager(_query_scratch); + auto scratch = manager.scratch_space(); + + if (L > scratch->get_L()) { + diskann::cout << "Attempting to expand query scratch_space. Was created " + << "with Lsize: " << scratch->get_L() + << " but search L is: " << L << std::endl; + scratch->resize_for_new_L(L); + diskann::cout << "Resize completed. New scratch->L is " << scratch->get_L() + << std::endl; + } + + std::vector filter_vec; + std::vector init_ids = get_init_ids(); + + std::shared_lock lock(_update_lock); + + if (_label_to_medoid_id.find(filter_label) != _label_to_medoid_id.end()) { + init_ids.emplace_back(_label_to_medoid_id[filter_label]); + } else { + diskann::cout << "No filtered medoid found. exitting " + << std::endl; // RKNOTE: If universal label found start there + throw diskann::ANNException("No filtered medoid found. exitting ", -1); + } + filter_vec.emplace_back(filter_label); + + // REFACTOR + // T *aligned_query = scratch->aligned_query(); + // memcpy(aligned_query, query, _dim * sizeof(T)); + _distance->preprocess_query(query, _data_store->get_dims(), + scratch->aligned_query()); + auto retval = iterate_to_fixed_point(scratch->aligned_query(), L, init_ids, + scratch, true, filter_vec, true); + + auto best_L_nodes = scratch->best_l_nodes(); + + size_t pos = 0; + for (size_t i = 0; i < best_L_nodes.size(); ++i) { + if (best_L_nodes[i].id < _max_points) { + // safe because Index uses uint32_t ids internally + // and IDType will be uint32_t or uint64_t + indices[pos] = (IdType)best_L_nodes[i].id; + if (distances != nullptr) { #ifdef EXEC_ENV_OLS - // DLVS expects negative distances - distances[pos] = best_L_nodes[i].distance; + // DLVS expects negative distances + distances[pos] = best_L_nodes[i].distance; #else - distances[pos] = _dist_metric == diskann::Metric::INNER_PRODUCT ? -1 * best_L_nodes[i].distance - : best_L_nodes[i].distance; + distances[pos] = _dist_metric == diskann::Metric::INNER_PRODUCT + ? -1 * best_L_nodes[i].distance + : best_L_nodes[i].distance; #endif - } - pos++; - } - if (pos == K) - break; - } - if (pos < K) - { - diskann::cerr << "Found fewer than K elements for query" << std::endl; + } + pos++; } + if (pos == K) break; + } + if (pos < K) { + diskann::cerr << "Found fewer than K elements for query" << std::endl; + } - return retval; + return retval; } template -size_t Index::search_with_tags(const T *query, const uint64_t K, const uint32_t L, TagT *tags, - float *distances, std::vector &res_vectors) -{ - if (K > (uint64_t)L) - { - throw ANNException("Set L to a value of at least K", -1, __FUNCSIG__, __FILE__, __LINE__); - } - ScratchStoreManager> manager(_query_scratch); - auto scratch = manager.scratch_space(); - - if (L > scratch->get_L()) - { - diskann::cout << "Attempting to expand query scratch_space. Was created " - << "with Lsize: " << scratch->get_L() << " but search L is: " << L << std::endl; - scratch->resize_for_new_L(L); - diskann::cout << "Resize completed. New scratch->L is " << scratch->get_L() << std::endl; - } +size_t Index::search_with_tags(const T *query, + const uint64_t K, + const uint32_t L, TagT *tags, + float *distances, + std::vector &res_vectors) { + if (K > (uint64_t)L) { + throw ANNException("Set L to a value of at least K", -1, __FUNCSIG__, + __FILE__, __LINE__); + } + ScratchStoreManager> manager(_query_scratch); + auto scratch = manager.scratch_space(); + + if (L > scratch->get_L()) { + diskann::cout << "Attempting to expand query scratch_space. Was created " + << "with Lsize: " << scratch->get_L() + << " but search L is: " << L << std::endl; + scratch->resize_for_new_L(L); + diskann::cout << "Resize completed. New scratch->L is " << scratch->get_L() + << std::endl; + } - std::shared_lock ul(_update_lock); + std::shared_lock ul(_update_lock); - const std::vector init_ids = get_init_ids(); - const std::vector unused_filter_label; + const std::vector init_ids = get_init_ids(); + const std::vector unused_filter_label; - _distance->preprocess_query(query, _data_store->get_dims(), scratch->aligned_query()); - iterate_to_fixed_point(scratch->aligned_query(), L, init_ids, scratch, false, unused_filter_label, true); + _distance->preprocess_query(query, _data_store->get_dims(), + scratch->aligned_query()); + iterate_to_fixed_point(scratch->aligned_query(), L, init_ids, scratch, false, + unused_filter_label, true); - NeighborPriorityQueue &best_L_nodes = scratch->best_l_nodes(); - assert(best_L_nodes.size() <= L); + NeighborPriorityQueue &best_L_nodes = scratch->best_l_nodes(); + assert(best_L_nodes.size() <= L); - std::shared_lock tl(_tag_lock); + std::shared_lock tl(_tag_lock); - size_t pos = 0; - for (size_t i = 0; i < best_L_nodes.size(); ++i) - { - auto node = best_L_nodes[i]; + size_t pos = 0; + for (size_t i = 0; i < best_L_nodes.size(); ++i) { + auto node = best_L_nodes[i]; - TagT tag; - if (_location_to_tag.try_get(node.id, tag)) - { - tags[pos] = tag; + TagT tag; + if (_location_to_tag.try_get(node.id, tag)) { + tags[pos] = tag; - if (res_vectors.size() > 0) - { - _data_store->get_vector(node.id, res_vectors[pos]); - } + if (res_vectors.size() > 0) { + _data_store->get_vector(node.id, res_vectors[pos]); + } - if (distances != nullptr) - { + if (distances != nullptr) { #ifdef EXEC_ENV_OLS - distances[pos] = node.distance; // DLVS expects negative distances + distances[pos] = node.distance; // DLVS expects negative distances #else - distances[pos] = _dist_metric == INNER_PRODUCT ? -1 * node.distance : node.distance; + distances[pos] = + _dist_metric == INNER_PRODUCT ? -1 * node.distance : node.distance; #endif - } - pos++; - // If res_vectors.size() < k, clip at the value. - if (pos == K || pos == res_vectors.size()) - break; - } + } + pos++; + // If res_vectors.size() < k, clip at the value. + if (pos == K || pos == res_vectors.size()) break; } + } - return pos; + return pos; } -template size_t Index::get_num_points() -{ - std::shared_lock tl(_tag_lock); - return _nd; +template +size_t Index::get_num_points() { + std::shared_lock tl(_tag_lock); + return _nd; } -template size_t Index::get_max_points() -{ - std::shared_lock tl(_tag_lock); - return _max_points; +template +size_t Index::get_max_points() { + std::shared_lock tl(_tag_lock); + return _max_points; } -template void Index::generate_frozen_point() -{ - if (_num_frozen_pts == 0) - return; - - if (_num_frozen_pts > 1) - { - throw ANNException("More than one frozen point not supported in generate_frozen_point", -1, __FUNCSIG__, - __FILE__, __LINE__); - } - - if (_nd == 0) - { - throw ANNException("ERROR: Can not pick a frozen point since nd=0", -1, __FUNCSIG__, __FILE__, __LINE__); - } - size_t res = calculate_entry_point(); - - if (_pq_dist) - { - // copy the PQ data corresponding to the point returned by - // calculate_entry_point - memcpy(_pq_data + _max_points * _num_pq_chunks, _pq_data + res * _num_pq_chunks, - _num_pq_chunks * DIV_ROUND_UP(NUM_PQ_BITS, 8)); - } - else - { - _data_store->copy_points(res, _max_points, 1); - } +template +void Index::generate_frozen_point() { + if (_num_frozen_pts == 0) return; + + if (_num_frozen_pts > 1) { + throw ANNException( + "More than one frozen point not supported in generate_frozen_point", -1, + __FUNCSIG__, __FILE__, __LINE__); + } + + if (_nd == 0) { + throw ANNException("ERROR: Can not pick a frozen point since nd=0", -1, + __FUNCSIG__, __FILE__, __LINE__); + } + size_t res = calculate_entry_point(); + + if (_pq_dist) { + // copy the PQ data corresponding to the point returned by + // calculate_entry_point + memcpy(_pq_data + _max_points * _num_pq_chunks, + _pq_data + res * _num_pq_chunks, + _num_pq_chunks * DIV_ROUND_UP(NUM_PQ_BITS, 8)); + } else { + _data_store->copy_points(res, _max_points, 1); + } } -template int Index::enable_delete() -{ - assert(_enable_tags); +template +int Index::enable_delete() { + assert(_enable_tags); - if (!_enable_tags) - { - diskann::cerr << "Tags must be instantiated for deletions" << std::endl; - return -2; - } + if (!_enable_tags) { + diskann::cerr << "Tags must be instantiated for deletions" << std::endl; + return -2; + } - std::unique_lock ul(_update_lock); - std::unique_lock tl(_tag_lock); - std::unique_lock dl(_delete_lock); + std::unique_lock ul(_update_lock); + std::unique_lock tl(_tag_lock); + std::unique_lock dl(_delete_lock); - if (_data_compacted) - { - for (uint32_t slot = (uint32_t)_nd; slot < _max_points; ++slot) - { - _empty_slots.insert(slot); - } + if (_data_compacted) { + for (uint32_t slot = (uint32_t)_nd; slot < _max_points; ++slot) { + _empty_slots.insert(slot); } + } - return 0; + return 0; } template -inline void Index::process_delete(const tsl::robin_set &old_delete_set, size_t loc, - const uint32_t range, const uint32_t maxc, const float alpha, - InMemQueryScratch *scratch) -{ - tsl::robin_set &expanded_nodes_set = scratch->expanded_nodes_set(); - std::vector &expanded_nghrs_vec = scratch->expanded_nodes_vec(); - - // If this condition were not true, deadlock could result - assert(old_delete_set.find(loc) == old_delete_set.end()); - - std::vector adj_list; - { - // Acquire and release lock[loc] before acquiring locks for neighbors - std::unique_lock adj_list_lock; - if (_conc_consolidate) - adj_list_lock = std::unique_lock(_locks[loc]); - adj_list = _final_graph[loc]; - } - - bool modify = false; - for (auto ngh : adj_list) - { - if (old_delete_set.find(ngh) == old_delete_set.end()) - { - expanded_nodes_set.insert(ngh); - } - else - { - modify = true; - - std::unique_lock ngh_lock; - if (_conc_consolidate) - ngh_lock = std::unique_lock(_locks[ngh]); - for (auto j : _final_graph[ngh]) - if (j != loc && old_delete_set.find(j) == old_delete_set.end()) - expanded_nodes_set.insert(j); - } - } - - if (modify) - { - if (expanded_nodes_set.size() <= range) - { - std::unique_lock adj_list_lock(_locks[loc]); - _final_graph[loc].clear(); - for (auto &ngh : expanded_nodes_set) - _final_graph[loc].push_back(ngh); - } - else - { - // Create a pool of Neighbor candidates from the expanded_nodes_set - expanded_nghrs_vec.reserve(expanded_nodes_set.size()); - for (auto &ngh : expanded_nodes_set) - { - expanded_nghrs_vec.emplace_back(ngh, _data_store->get_distance(loc, ngh)); - } - std::sort(expanded_nghrs_vec.begin(), expanded_nghrs_vec.end()); - std::vector &occlude_list_output = scratch->occlude_list_output(); - occlude_list(loc, expanded_nghrs_vec, alpha, range, maxc, occlude_list_output, scratch, &old_delete_set); - std::unique_lock adj_list_lock(_locks[loc]); - _final_graph[loc] = occlude_list_output; - } - } +inline void Index::process_delete( + const tsl::robin_set &old_delete_set, size_t loc, + const uint32_t range, const uint32_t maxc, const float alpha, + InMemQueryScratch *scratch) { + tsl::robin_set &expanded_nodes_set = scratch->expanded_nodes_set(); + std::vector &expanded_nghrs_vec = scratch->expanded_nodes_vec(); + + // If this condition were not true, deadlock could result + assert(old_delete_set.find(loc) == old_delete_set.end()); + + std::vector adj_list; + { + // Acquire and release lock[loc] before acquiring locks for neighbors + std::unique_lock adj_list_lock; + if (_conc_consolidate) + adj_list_lock = std::unique_lock(_locks[loc]); + adj_list = _final_graph[loc]; + } + + bool modify = false; + for (auto ngh : adj_list) { + if (old_delete_set.find(ngh) == old_delete_set.end()) { + expanded_nodes_set.insert(ngh); + } else { + modify = true; + + std::unique_lock ngh_lock; + if (_conc_consolidate) + ngh_lock = std::unique_lock(_locks[ngh]); + for (auto j : _final_graph[ngh]) + if (j != loc && old_delete_set.find(j) == old_delete_set.end()) + expanded_nodes_set.insert(j); + } + } + + if (modify) { + if (expanded_nodes_set.size() <= range) { + std::unique_lock adj_list_lock(_locks[loc]); + _final_graph[loc].clear(); + for (auto &ngh : expanded_nodes_set) _final_graph[loc].push_back(ngh); + } else { + // Create a pool of Neighbor candidates from the expanded_nodes_set + expanded_nghrs_vec.reserve(expanded_nodes_set.size()); + for (auto &ngh : expanded_nodes_set) { + expanded_nghrs_vec.emplace_back(ngh, + _data_store->get_distance(loc, ngh)); + } + std::sort(expanded_nghrs_vec.begin(), expanded_nghrs_vec.end()); + std::vector &occlude_list_output = + scratch->occlude_list_output(); + occlude_list(loc, expanded_nghrs_vec, alpha, range, maxc, + occlude_list_output, scratch, &old_delete_set); + std::unique_lock adj_list_lock(_locks[loc]); + _final_graph[loc] = occlude_list_output; + } + } } // Returns number of live points left after consolidation template -consolidation_report Index::consolidate_deletes(const IndexWriteParameters ¶ms) -{ - if (!_enable_tags) - throw diskann::ANNException("Point tag array not instantiated", -1, __FUNCSIG__, __FILE__, __LINE__); +consolidation_report Index::consolidate_deletes( + const IndexWriteParameters ¶ms) { + if (!_enable_tags) + throw diskann::ANNException("Point tag array not instantiated", -1, + __FUNCSIG__, __FILE__, __LINE__); - { - std::shared_lock ul(_update_lock); - std::shared_lock tl(_tag_lock); - std::shared_lock dl(_delete_lock); - if (_empty_slots.size() + _nd != _max_points) - { - std::string err = "#empty slots + nd != max points"; - diskann::cerr << err << std::endl; - throw ANNException(err, -1, __FUNCSIG__, __FILE__, __LINE__); - } - - if (_location_to_tag.size() + _delete_set->size() != _nd) - { - diskann::cerr << "Error: _location_to_tag.size (" << _location_to_tag.size() << ") + _delete_set->size (" - << _delete_set->size() << ") != _nd(" << _nd << ") "; - return consolidation_report(diskann::consolidation_report::status_code::INCONSISTENT_COUNT_ERROR, 0, 0, 0, - 0, 0, 0, 0); - } - - if (_location_to_tag.size() != _tag_to_location.size()) - { - throw diskann::ANNException("_location_to_tag and _tag_to_location not of same size", -1, __FUNCSIG__, - __FILE__, __LINE__); - } - } - - std::unique_lock update_lock(_update_lock, std::defer_lock); - if (!_conc_consolidate) - update_lock.lock(); - - std::unique_lock cl(_consolidate_lock, std::defer_lock); - if (!cl.try_lock()) - { - diskann::cerr << "Consildate delete function failed to acquire consolidate lock" << std::endl; - return consolidation_report(diskann::consolidation_report::status_code::LOCK_FAIL, 0, 0, 0, 0, 0, 0, 0); - } - - diskann::cout << "Starting consolidate_deletes... "; - - std::unique_ptr> old_delete_set(new tsl::robin_set); - { - std::unique_lock dl(_delete_lock); - std::swap(_delete_set, old_delete_set); - } - - if (old_delete_set->find(_start) != old_delete_set->end()) - { - throw diskann::ANNException("ERROR: start node has been deleted", -1, __FUNCSIG__, __FILE__, __LINE__); - } - - const uint32_t range = params.max_degree; - const uint32_t maxc = params.max_occlusion_size; - const float alpha = params.alpha; - const uint32_t num_threads = params.num_threads == 0 ? omp_get_num_threads() : params.num_threads; - - uint32_t num_calls_to_process_delete = 0; - diskann::Timer timer; + { + std::shared_lock ul(_update_lock); + std::shared_lock tl(_tag_lock); + std::shared_lock dl(_delete_lock); + if (_empty_slots.size() + _nd != _max_points) { + std::string err = "#empty slots + nd != max points"; + diskann::cerr << err << std::endl; + throw ANNException(err, -1, __FUNCSIG__, __FILE__, __LINE__); + } + + if (_location_to_tag.size() + _delete_set->size() != _nd) { + diskann::cerr << "Error: _location_to_tag.size (" + << _location_to_tag.size() << ") + _delete_set->size (" + << _delete_set->size() << ") != _nd(" << _nd << ") "; + return consolidation_report( + diskann::consolidation_report::status_code::INCONSISTENT_COUNT_ERROR, + 0, 0, 0, 0, 0, 0, 0); + } + + if (_location_to_tag.size() != _tag_to_location.size()) { + throw diskann::ANNException( + "_location_to_tag and _tag_to_location not of same size", -1, + __FUNCSIG__, __FILE__, __LINE__); + } + } + + std::unique_lock update_lock(_update_lock, + std::defer_lock); + if (!_conc_consolidate) update_lock.lock(); + + std::unique_lock cl(_consolidate_lock, + std::defer_lock); + if (!cl.try_lock()) { + diskann::cerr + << "Consildate delete function failed to acquire consolidate lock" + << std::endl; + return consolidation_report( + diskann::consolidation_report::status_code::LOCK_FAIL, 0, 0, 0, 0, 0, 0, + 0); + } + + diskann::cout << "Starting consolidate_deletes... "; + + std::unique_ptr> old_delete_set( + new tsl::robin_set); + { + std::unique_lock dl(_delete_lock); + std::swap(_delete_set, old_delete_set); + } + + if (old_delete_set->find(_start) != old_delete_set->end()) { + throw diskann::ANNException("ERROR: start node has been deleted", -1, + __FUNCSIG__, __FILE__, __LINE__); + } + + const uint32_t range = params.max_degree; + const uint32_t maxc = params.max_occlusion_size; + const float alpha = params.alpha; + const uint32_t num_threads = + params.num_threads == 0 ? omp_get_num_threads() : params.num_threads; + + uint32_t num_calls_to_process_delete = 0; + diskann::Timer timer; #pragma omp parallel for num_threads(num_threads) schedule(dynamic, 8192) reduction(+ : num_calls_to_process_delete) - for (int64_t loc = 0; loc < (int64_t)_max_points; loc++) - { - if (old_delete_set->find((uint32_t)loc) == old_delete_set->end() && !_empty_slots.is_in_set((uint32_t)loc)) - { - ScratchStoreManager> manager(_query_scratch); - auto scratch = manager.scratch_space(); - process_delete(*old_delete_set, loc, range, maxc, alpha, scratch); - num_calls_to_process_delete += 1; - } - } - for (int64_t loc = _max_points; loc < (int64_t)(_max_points + _num_frozen_pts); loc++) - { - ScratchStoreManager> manager(_query_scratch); - auto scratch = manager.scratch_space(); - process_delete(*old_delete_set, loc, range, maxc, alpha, scratch); - num_calls_to_process_delete += 1; - } + for (int64_t loc = 0; loc < (int64_t)_max_points; loc++) { + if (old_delete_set->find((uint32_t)loc) == old_delete_set->end() && + !_empty_slots.is_in_set((uint32_t)loc)) { + ScratchStoreManager> manager(_query_scratch); + auto scratch = manager.scratch_space(); + process_delete(*old_delete_set, loc, range, maxc, alpha, scratch); + num_calls_to_process_delete += 1; + } + } + for (int64_t loc = _max_points; + loc < (int64_t)(_max_points + _num_frozen_pts); loc++) { + ScratchStoreManager> manager(_query_scratch); + auto scratch = manager.scratch_space(); + process_delete(*old_delete_set, loc, range, maxc, alpha, scratch); + num_calls_to_process_delete += 1; + } - std::unique_lock tl(_tag_lock); - size_t ret_nd = release_locations(*old_delete_set); - size_t max_points = _max_points; - size_t empty_slots_size = _empty_slots.size(); + std::unique_lock tl(_tag_lock); + size_t ret_nd = release_locations(*old_delete_set); + size_t max_points = _max_points; + size_t empty_slots_size = _empty_slots.size(); - std::shared_lock dl(_delete_lock); - size_t delete_set_size = _delete_set->size(); - size_t old_delete_set_size = old_delete_set->size(); + std::shared_lock dl(_delete_lock); + size_t delete_set_size = _delete_set->size(); + size_t old_delete_set_size = old_delete_set->size(); - if (!_conc_consolidate) - { - update_lock.unlock(); - } + if (!_conc_consolidate) { + update_lock.unlock(); + } - double duration = timer.elapsed() / 1000000.0; - diskann::cout << " done in " << duration << " seconds." << std::endl; - return consolidation_report(diskann::consolidation_report::status_code::SUCCESS, ret_nd, max_points, - empty_slots_size, old_delete_set_size, delete_set_size, num_calls_to_process_delete, - duration); + double duration = timer.elapsed() / 1000000.0; + diskann::cout << " done in " << duration << " seconds." << std::endl; + return consolidation_report( + diskann::consolidation_report::status_code::SUCCESS, ret_nd, max_points, + empty_slots_size, old_delete_set_size, delete_set_size, + num_calls_to_process_delete, duration); } -template void Index::compact_frozen_point() -{ - if (_nd < _max_points && _num_frozen_pts > 0) - { - reposition_points(_max_points, _nd, _num_frozen_pts); - _start = (uint32_t)_nd; - } +template +void Index::compact_frozen_point() { + if (_nd < _max_points && _num_frozen_pts > 0) { + reposition_points(_max_points, _nd, _num_frozen_pts); + _start = (uint32_t)_nd; + } } // Should be called after acquiring _update_lock -template void Index::compact_data() -{ - if (!_dynamic_index) - throw ANNException("Can not compact a non-dynamic index", -1, __FUNCSIG__, __FILE__, __LINE__); - - if (_data_compacted) - { - diskann::cerr << "Warning! Calling compact_data() when _data_compacted is true!" << std::endl; - return; - } - - if (_delete_set->size() > 0) - { - throw ANNException("Can not compact data when index has non-empty _delete_set of " - "size: " + - std::to_string(_delete_set->size()), - -1, __FUNCSIG__, __FILE__, __LINE__); - } - - diskann::Timer timer; - - std::vector new_location = std::vector(_max_points + _num_frozen_pts, UINT32_MAX); - - uint32_t new_counter = 0; - std::set empty_locations; - for (uint32_t old_location = 0; old_location < _max_points; old_location++) - { - if (_location_to_tag.contains(old_location)) - { - new_location[old_location] = new_counter; - new_counter++; - } - else - { - empty_locations.insert(old_location); - } - } - for (uint32_t old_location = _max_points; old_location < _max_points + _num_frozen_pts; old_location++) - { - new_location[old_location] = old_location; - } - - // If start node is removed, throw an exception - if (_start < _max_points && !_location_to_tag.contains(_start)) - { - throw diskann::ANNException("ERROR: Start node deleted.", -1, __FUNCSIG__, __FILE__, __LINE__); - } - - size_t num_dangling = 0; - for (uint32_t old = 0; old < _max_points + _num_frozen_pts; ++old) - { - std::vector new_adj_list; - - if ((new_location[old] < _max_points) // If point continues to exist - || (old >= _max_points && old < _max_points + _num_frozen_pts)) - { - new_adj_list.reserve(_final_graph[old].size()); - for (auto ngh_iter : _final_graph[old]) - { - if (empty_locations.find(ngh_iter) != empty_locations.end()) - { - ++num_dangling; - diskann::cerr << "Error in compact_data(). _final_graph[" << old << "] has neighbor " << ngh_iter - << " which is a location not associated with any tag." << std::endl; - } - else - { - new_adj_list.push_back(new_location[ngh_iter]); - } - } - _final_graph[old].swap(new_adj_list); - - // Move the data and adj list to the correct position - if (new_location[old] != old) - { - assert(new_location[old] < old); - _final_graph[new_location[old]].swap(_final_graph[old]); - - _data_store->copy_points(old, new_location[old], 1); - } - } - else - { - _final_graph[old].clear(); - } - } - diskann::cerr << "#dangling references after data compaction: " << num_dangling << std::endl; - - _tag_to_location.clear(); - for (auto pos = _location_to_tag.find_first(); pos.is_valid(); pos = _location_to_tag.find_next(pos)) - { - const auto tag = _location_to_tag.get(pos); - _tag_to_location[tag] = new_location[pos._key]; - } - _location_to_tag.clear(); - for (const auto &iter : _tag_to_location) - { - _location_to_tag.set(iter.second, iter.first); - } - - for (size_t old = _nd; old < _max_points; ++old) - { - _final_graph[old].clear(); - } - _empty_slots.clear(); - for (auto i = _nd; i < _max_points; i++) - { - _empty_slots.insert((uint32_t)i); - } - _data_compacted = true; - diskann::cout << "Time taken for compact_data: " << timer.elapsed() / 1000000. << "s." << std::endl; +template +void Index::compact_data() { + if (!_dynamic_index) + throw ANNException("Can not compact a non-dynamic index", -1, __FUNCSIG__, + __FILE__, __LINE__); + + if (_data_compacted) { + diskann::cerr + << "Warning! Calling compact_data() when _data_compacted is true!" + << std::endl; + return; + } + + if (_delete_set->size() > 0) { + throw ANNException( + "Can not compact data when index has non-empty _delete_set of " + "size: " + + std::to_string(_delete_set->size()), + -1, __FUNCSIG__, __FILE__, __LINE__); + } + + diskann::Timer timer; + + std::vector new_location = + std::vector(_max_points + _num_frozen_pts, UINT32_MAX); + + uint32_t new_counter = 0; + std::set empty_locations; + for (uint32_t old_location = 0; old_location < _max_points; old_location++) { + if (_location_to_tag.contains(old_location)) { + new_location[old_location] = new_counter; + new_counter++; + } else { + empty_locations.insert(old_location); + } + } + for (uint32_t old_location = _max_points; + old_location < _max_points + _num_frozen_pts; old_location++) { + new_location[old_location] = old_location; + } + + // If start node is removed, throw an exception + if (_start < _max_points && !_location_to_tag.contains(_start)) { + throw diskann::ANNException("ERROR: Start node deleted.", -1, __FUNCSIG__, + __FILE__, __LINE__); + } + + size_t num_dangling = 0; + for (uint32_t old = 0; old < _max_points + _num_frozen_pts; ++old) { + std::vector new_adj_list; + + if ((new_location[old] < _max_points) // If point continues to exist + || (old >= _max_points && old < _max_points + _num_frozen_pts)) { + new_adj_list.reserve(_final_graph[old].size()); + for (auto ngh_iter : _final_graph[old]) { + if (empty_locations.find(ngh_iter) != empty_locations.end()) { + ++num_dangling; + diskann::cerr << "Error in compact_data(). _final_graph[" << old + << "] has neighbor " << ngh_iter + << " which is a location not associated with any tag." + << std::endl; + } else { + new_adj_list.push_back(new_location[ngh_iter]); + } + } + _final_graph[old].swap(new_adj_list); + + // Move the data and adj list to the correct position + if (new_location[old] != old) { + assert(new_location[old] < old); + _final_graph[new_location[old]].swap(_final_graph[old]); + + _data_store->copy_points(old, new_location[old], 1); + } + } else { + _final_graph[old].clear(); + } + } + diskann::cerr << "#dangling references after data compaction: " + << num_dangling << std::endl; + + _tag_to_location.clear(); + for (auto pos = _location_to_tag.find_first(); pos.is_valid(); + pos = _location_to_tag.find_next(pos)) { + const auto tag = _location_to_tag.get(pos); + _tag_to_location[tag] = new_location[pos._key]; + } + _location_to_tag.clear(); + for (const auto &iter : _tag_to_location) { + _location_to_tag.set(iter.second, iter.first); + } + + for (size_t old = _nd; old < _max_points; ++old) { + _final_graph[old].clear(); + } + _empty_slots.clear(); + for (auto i = _nd; i < _max_points; i++) { + _empty_slots.insert((uint32_t)i); + } + _data_compacted = true; + diskann::cout << "Time taken for compact_data: " << timer.elapsed() / 1000000. + << "s." << std::endl; } // // Caller must hold unique _tag_lock and _delete_lock before calling this // -template int Index::reserve_location() -{ - if (_nd >= _max_points) - { - return -1; - } - uint32_t location; - if (_data_compacted && _empty_slots.is_empty()) - { - // This code path is encountered when enable_delete hasn't been - // called yet, so no points have been deleted and _empty_slots - // hasn't been filled in. In that case, just keep assigning - // consecutive locations. - location = (uint32_t)_nd; - } - else - { - assert(_empty_slots.size() != 0); - assert(_empty_slots.size() + _nd == _max_points); +template +int Index::reserve_location() { + if (_nd >= _max_points) { + return -1; + } + uint32_t location; + if (_data_compacted && _empty_slots.is_empty()) { + // This code path is encountered when enable_delete hasn't been + // called yet, so no points have been deleted and _empty_slots + // hasn't been filled in. In that case, just keep assigning + // consecutive locations. + location = (uint32_t)_nd; + } else { + assert(_empty_slots.size() != 0); + assert(_empty_slots.size() + _nd == _max_points); + + location = _empty_slots.pop_any(); + _delete_set->erase(location); + } + + ++_nd; + return location; +} - location = _empty_slots.pop_any(); - _delete_set->erase(location); - } +template +size_t Index::release_location(int location) { + if (_empty_slots.is_in_set(location)) + throw ANNException( + "Trying to release location, but location already in empty slots", -1, + __FUNCSIG__, __FILE__, __LINE__); + _empty_slots.insert(location); - ++_nd; - return location; + _nd--; + return _nd; } -template size_t Index::release_location(int location) -{ +template +size_t Index::release_locations( + const tsl::robin_set &locations) { + for (auto location : locations) { if (_empty_slots.is_in_set(location)) - throw ANNException("Trying to release location, but location already in empty slots", -1, __FUNCSIG__, __FILE__, - __LINE__); + throw ANNException( + "Trying to release location, but location " + "already in empty slots", + -1, __FUNCSIG__, __FILE__, __LINE__); _empty_slots.insert(location); _nd--; - return _nd; -} - -template -size_t Index::release_locations(const tsl::robin_set &locations) -{ - for (auto location : locations) - { - if (_empty_slots.is_in_set(location)) - throw ANNException("Trying to release location, but location " - "already in empty slots", - -1, __FUNCSIG__, __FILE__, __LINE__); - _empty_slots.insert(location); - - _nd--; - } + } - if (_empty_slots.size() + _nd != _max_points) - throw ANNException("#empty slots + nd != max points", -1, __FUNCSIG__, __FILE__, __LINE__); + if (_empty_slots.size() + _nd != _max_points) + throw ANNException("#empty slots + nd != max points", -1, __FUNCSIG__, + __FILE__, __LINE__); - return _nd; + return _nd; } template -void Index::reposition_points(uint32_t old_location_start, uint32_t new_location_start, - uint32_t num_locations) -{ - if (num_locations == 0 || old_location_start == new_location_start) - { - return; - } - - // Update pointers to the moved nodes. Note: the computation is correct even - // when new_location_start < old_location_start given the C++ uint32_t - // integer arithmetic rules. - const uint32_t location_delta = new_location_start - old_location_start; - - for (uint32_t i = 0; i < _max_points + _num_frozen_pts; i++) - for (auto &loc : _final_graph[i]) - if (loc >= old_location_start && loc < old_location_start + num_locations) - loc += location_delta; - - // 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; - - // Move the adjacency lists. Make sure that overlapping ranges are handled - // correctly. - if (new_location_start < old_location_start) - { - // New location before the old location: copy the entries in order - // to avoid modifying locations that are yet to be copied. - for (uint32_t loc_offset = 0; loc_offset < num_locations; loc_offset++) - { - assert(_final_graph[new_location_start + loc_offset].empty()); - _final_graph[new_location_start + loc_offset].swap(_final_graph[old_location_start + loc_offset]); - } - - // 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 - { - // Old location after the new location: copy from the end of the range - // to avoid modifying locations that are yet to be copied. - for (uint32_t loc_offset = num_locations; loc_offset > 0; loc_offset--) - { - assert(_final_graph[new_location_start + loc_offset - 1u].empty()); - _final_graph[new_location_start + loc_offset - 1u].swap(_final_graph[old_location_start + loc_offset - 1u]); - } - - // 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; - } - } - _data_store->reposition_points(old_location_start, new_location_start, num_locations); +void Index::reposition_points(uint32_t old_location_start, + uint32_t new_location_start, + uint32_t num_locations) { + if (num_locations == 0 || old_location_start == new_location_start) { + return; + } + + // Update pointers to the moved nodes. Note: the computation is correct even + // when new_location_start < old_location_start given the C++ uint32_t + // integer arithmetic rules. + const uint32_t location_delta = new_location_start - old_location_start; + + for (uint32_t i = 0; i < _max_points + _num_frozen_pts; i++) + for (auto &loc : _final_graph[i]) + if (loc >= old_location_start && loc < old_location_start + num_locations) + loc += location_delta; + + // 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; + + // Move the adjacency lists. Make sure that overlapping ranges are handled + // correctly. + if (new_location_start < old_location_start) { + // New location before the old location: copy the entries in order + // to avoid modifying locations that are yet to be copied. + for (uint32_t loc_offset = 0; loc_offset < num_locations; loc_offset++) { + assert(_final_graph[new_location_start + loc_offset].empty()); + _final_graph[new_location_start + loc_offset].swap( + _final_graph[old_location_start + loc_offset]); + } + + // 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 { + // Old location after the new location: copy from the end of the range + // to avoid modifying locations that are yet to be copied. + for (uint32_t loc_offset = num_locations; loc_offset > 0; loc_offset--) { + assert(_final_graph[new_location_start + loc_offset - 1u].empty()); + _final_graph[new_location_start + loc_offset - 1u].swap( + _final_graph[old_location_start + loc_offset - 1u]); + } + + // 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; + } + } + _data_store->reposition_points(old_location_start, new_location_start, + num_locations); } -template void Index::reposition_frozen_point_to_end() -{ - if (_num_frozen_pts == 0) - return; +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; - } + if (_nd == _max_points) { + diskann::cout + << "Not repositioning frozen point as it is already at the end." + << std::endl; + return; + } - reposition_points((uint32_t)_nd, (uint32_t)_max_points, (uint32_t)_num_frozen_pts); - _start = (uint32_t)_max_points; + reposition_points((uint32_t)_nd, (uint32_t)_max_points, + (uint32_t)_num_frozen_pts); + _start = (uint32_t)_max_points; } -template void Index::resize(size_t new_max_points) -{ - const size_t new_internal_points = new_max_points + _num_frozen_pts; - auto start = std::chrono::high_resolution_clock::now(); - assert(_empty_slots.size() == 0); // should not resize if there are empty slots. - - _data_store->resize(new_internal_points); - _final_graph.resize(new_internal_points); - _locks = std::vector(new_internal_points); - - if (_num_frozen_pts != 0) - { - reposition_points((uint32_t)_max_points, (uint32_t)new_max_points, (uint32_t)_num_frozen_pts); - _start = (uint32_t)new_max_points; - } - - _max_points = new_max_points; - _empty_slots.reserve(_max_points); - for (auto i = _nd; i < _max_points; i++) - { - _empty_slots.insert((uint32_t)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::resize(size_t new_max_points) { + const size_t new_internal_points = new_max_points + _num_frozen_pts; + auto start = std::chrono::high_resolution_clock::now(); + assert(_empty_slots.size() == + 0); // should not resize if there are empty slots. + + _data_store->resize(new_internal_points); + _final_graph.resize(new_internal_points); + _locks = std::vector(new_internal_points); + + if (_num_frozen_pts != 0) { + reposition_points((uint32_t)_max_points, (uint32_t)new_max_points, + (uint32_t)_num_frozen_pts); + _start = (uint32_t)new_max_points; + } + + _max_points = new_max_points; + _empty_slots.reserve(_max_points); + for (auto i = _nd; i < _max_points; i++) { + _empty_slots.insert((uint32_t)i); + } + + auto stop = std::chrono::high_resolution_clock::now(); + diskann::cout << "Resizing took: " + << std::chrono::duration(stop - start).count() << "s" + << std::endl; } template -int Index::insert_point(const T *point, const TagT tag) -{ - assert(_has_built); - if (tag == static_cast(0)) - { - throw diskann::ANNException("Do not insert point with tag 0. That is " - "reserved for points hidden " - "from the user.", - -1, __FUNCSIG__, __FILE__, __LINE__); - } - - std::shared_lock shared_ul(_update_lock); - std::unique_lock tl(_tag_lock); - std::unique_lock dl(_delete_lock); +int Index::insert_point(const T *point, const TagT tag) { + assert(_has_built); + if (tag == static_cast(0)) { + throw diskann::ANNException( + "Do not insert point with tag 0. That is " + "reserved for points hidden " + "from the user.", + -1, __FUNCSIG__, __FILE__, __LINE__); + } + + std::shared_lock shared_ul(_update_lock); + std::unique_lock tl(_tag_lock); + std::unique_lock dl(_delete_lock); + + // Find a vacant location in the data array to insert the new point + auto location = reserve_location(); + if (location == -1) { +#if EXPAND_IF_FULL + dl.unlock(); + tl.unlock(); + shared_ul.unlock(); - // Find a vacant location in the data array to insert the new point - auto location = reserve_location(); - if (location == -1) { -#if EXPAND_IF_FULL - dl.unlock(); - tl.unlock(); - shared_ul.unlock(); - - { - std::unique_lock ul(_update_lock); - tl.lock(); - dl.lock(); - - if (_nd >= _max_points) - { - auto new_max_points = (size_t)(_max_points * INDEX_GROWTH_FACTOR); - resize(new_max_points); - } + std::unique_lock ul(_update_lock); + tl.lock(); + dl.lock(); - dl.unlock(); - tl.unlock(); - ul.unlock(); - } + if (_nd >= _max_points) { + auto new_max_points = (size_t)(_max_points * INDEX_GROWTH_FACTOR); + resize(new_max_points); + } - shared_ul.lock(); - tl.lock(); - dl.lock(); + dl.unlock(); + tl.unlock(); + ul.unlock(); + } - location = reserve_location(); - if (location == -1) - { - throw diskann::ANNException("Cannot reserve location even after " - "expanding graph. Terminating.", - -1, __FUNCSIG__, __FILE__, __LINE__); - } + shared_ul.lock(); + tl.lock(); + dl.lock(); + + location = reserve_location(); + if (location == -1) { + throw diskann::ANNException( + "Cannot reserve location even after " + "expanding graph. Terminating.", + -1, __FUNCSIG__, __FILE__, __LINE__); + } #else - return -1; + return -1; #endif + } + dl.unlock(); + + // Insert tag and mapping to location + if (_enable_tags) { + if (_tag_to_location.find(tag) != _tag_to_location.end()) { + release_location(location); + return -1; } - dl.unlock(); - // Insert tag and mapping to location - if (_enable_tags) - { - if (_tag_to_location.find(tag) != _tag_to_location.end()) - { - release_location(location); - return -1; - } + _tag_to_location[tag] = location; + _location_to_tag.set(location, tag); + } + tl.unlock(); - _tag_to_location[tag] = location; - _location_to_tag.set(location, tag); - } - tl.unlock(); + _data_store->set_vector(location, point); - _data_store->set_vector(location, point); + // Find and add appropriate graph edges + ScratchStoreManager> manager(_query_scratch); + auto scratch = manager.scratch_space(); + std::vector pruned_list; + if (_filtered_index) { + search_for_point_and_prune(location, _indexingQueueSize, pruned_list, + scratch, true, _filterIndexingQueueSize); + } else { + search_for_point_and_prune(location, _indexingQueueSize, pruned_list, + scratch); + } + { + std::shared_lock tlock(_tag_lock, std::defer_lock); + if (_conc_consolidate) tlock.lock(); - // Find and add appropriate graph edges - ScratchStoreManager> manager(_query_scratch); - auto scratch = manager.scratch_space(); - std::vector pruned_list; - if (_filtered_index) - { - search_for_point_and_prune(location, _indexingQueueSize, pruned_list, scratch, true, _filterIndexingQueueSize); - } - else - { - search_for_point_and_prune(location, _indexingQueueSize, pruned_list, scratch); - } - { - std::shared_lock tlock(_tag_lock, std::defer_lock); - if (_conc_consolidate) - tlock.lock(); - - LockGuard guard(_locks[location]); - _final_graph[location].clear(); - _final_graph[location].reserve((size_t)(_indexingRange * GRAPH_SLACK_FACTOR * 1.05)); - - for (auto link : pruned_list) - { - if (_conc_consolidate) - if (!_location_to_tag.contains(link)) - continue; - _final_graph[location].emplace_back(link); - } - assert(_final_graph[location].size() <= _indexingRange); + LockGuard guard(_locks[location]); + _final_graph[location].clear(); + _final_graph[location].reserve( + (size_t)(_indexingRange * GRAPH_SLACK_FACTOR * 1.05)); - if (_conc_consolidate) - tlock.unlock(); + for (auto link : pruned_list) { + if (_conc_consolidate) + if (!_location_to_tag.contains(link)) continue; + _final_graph[location].emplace_back(link); } + assert(_final_graph[location].size() <= _indexingRange); - inter_insert(location, pruned_list, scratch); + if (_conc_consolidate) tlock.unlock(); + } - return 0; + inter_insert(location, pruned_list, scratch); + + return 0; } -template int Index::lazy_delete(const TagT &tag) -{ - std::shared_lock ul(_update_lock); - std::unique_lock tl(_tag_lock); - std::unique_lock dl(_delete_lock); - _data_compacted = false; +template +int Index::lazy_delete(const TagT &tag) { + std::shared_lock ul(_update_lock); + std::unique_lock tl(_tag_lock); + std::unique_lock dl(_delete_lock); + _data_compacted = false; - if (_tag_to_location.find(tag) == _tag_to_location.end()) - { - diskann::cerr << "Delete tag not found " << tag << std::endl; - return -1; - } - assert(_tag_to_location[tag] < _max_points); + if (_tag_to_location.find(tag) == _tag_to_location.end()) { + diskann::cerr << "Delete tag not found " << tag << std::endl; + return -1; + } + assert(_tag_to_location[tag] < _max_points); - const auto location = _tag_to_location[tag]; - _delete_set->insert(location); - _location_to_tag.erase(location); - _tag_to_location.erase(tag); + const auto location = _tag_to_location[tag]; + _delete_set->insert(location); + _location_to_tag.erase(location); + _tag_to_location.erase(tag); - return 0; + return 0; } template -void Index::lazy_delete(const std::vector &tags, std::vector &failed_tags) -{ - if (failed_tags.size() > 0) - { - throw ANNException("failed_tags should be passed as an empty list", -1, __FUNCSIG__, __FILE__, __LINE__); - } - std::shared_lock ul(_update_lock); - std::unique_lock tl(_tag_lock); - std::unique_lock dl(_delete_lock); - _data_compacted = false; - - for (auto tag : tags) - { - if (_tag_to_location.find(tag) == _tag_to_location.end()) - { - failed_tags.push_back(tag); - } - else - { - const auto location = _tag_to_location[tag]; - _delete_set->insert(location); - _location_to_tag.erase(location); - _tag_to_location.erase(tag); - } - } +void Index::lazy_delete(const std::vector &tags, + std::vector &failed_tags) { + if (failed_tags.size() > 0) { + throw ANNException("failed_tags should be passed as an empty list", -1, + __FUNCSIG__, __FILE__, __LINE__); + } + std::shared_lock ul(_update_lock); + std::unique_lock tl(_tag_lock); + std::unique_lock dl(_delete_lock); + _data_compacted = false; + + for (auto tag : tags) { + if (_tag_to_location.find(tag) == _tag_to_location.end()) { + failed_tags.push_back(tag); + } else { + const auto location = _tag_to_location[tag]; + _delete_set->insert(location); + _location_to_tag.erase(location); + _tag_to_location.erase(tag); + } + } } -template bool Index::is_index_saved() -{ - return _is_saved; +template +bool Index::is_index_saved() { + return _is_saved; } template -void Index::get_active_tags(tsl::robin_set &active_tags) -{ - active_tags.clear(); - std::shared_lock tl(_tag_lock); - for (auto iter : _tag_to_location) - { - active_tags.insert(iter.first); - } +void Index::get_active_tags( + tsl::robin_set &active_tags) { + active_tags.clear(); + std::shared_lock tl(_tag_lock); + for (auto iter : _tag_to_location) { + active_tags.insert(iter.first); + } } -template void Index::print_status() -{ - std::shared_lock ul(_update_lock); - std::shared_lock cl(_consolidate_lock); - std::shared_lock tl(_tag_lock); - std::shared_lock dl(_delete_lock); - - 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 << std::endl; - diskann::cout << "---------------------------------------------------------" - "------------" - << std::endl; +template +void Index::print_status() { + std::shared_lock ul(_update_lock); + std::shared_lock cl(_consolidate_lock); + std::shared_lock tl(_tag_lock); + std::shared_lock dl(_delete_lock); + + 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 + << std::endl; + diskann::cout << "---------------------------------------------------------" + "------------" + << std::endl; } -template void Index::count_nodes_at_bfs_levels() -{ - std::unique_lock ul(_update_lock); +template +void Index::count_nodes_at_bfs_levels() { + std::unique_lock ul(_update_lock); - boost::dynamic_bitset<> visited(_max_points + _num_frozen_pts); + boost::dynamic_bitset<> visited(_max_points + _num_frozen_pts); - size_t MAX_BFS_LEVELS = 32; - auto bfs_sets = new tsl::robin_set[MAX_BFS_LEVELS]; + size_t MAX_BFS_LEVELS = 32; + auto bfs_sets = new tsl::robin_set[MAX_BFS_LEVELS]; - bfs_sets[0].insert(_start); - visited.set(_start); + bfs_sets[0].insert(_start); + visited.set(_start); - for (uint32_t i = _max_points; i < _max_points + _num_frozen_pts; ++i) - { - if (i != _start) - { - bfs_sets[0].insert(i); - visited.set(i); - } + for (uint32_t i = _max_points; i < _max_points + _num_frozen_pts; ++i) { + if (i != _start) { + bfs_sets[0].insert(i); + visited.set(i); } + } - for (size_t l = 0; l < MAX_BFS_LEVELS - 1; ++l) - { - diskann::cout << "Number of nodes at BFS level " << l << " is " << bfs_sets[l].size() << std::endl; - if (bfs_sets[l].size() == 0) - break; - for (auto node : bfs_sets[l]) - { - for (auto nghbr : _final_graph[node]) - { - if (!visited.test(nghbr)) - { - visited.set(nghbr); - bfs_sets[l + 1].insert(nghbr); - } - } + for (size_t l = 0; l < MAX_BFS_LEVELS - 1; ++l) { + diskann::cout << "Number of nodes at BFS level " << l << " is " + << bfs_sets[l].size() << std::endl; + if (bfs_sets[l].size() == 0) break; + for (auto node : bfs_sets[l]) { + for (auto nghbr : _final_graph[node]) { + if (!visited.test(nghbr)) { + visited.set(nghbr); + bfs_sets[l + 1].insert(nghbr); } + } } + } - delete[] bfs_sets; + delete[] bfs_sets; } -// REFACTOR: This should be an OptimizedDataStore class, dummy impl here for compiling sake -// template void Index::optimize_index_layout() +// REFACTOR: This should be an OptimizedDataStore class, dummy impl here for +// compiling sake template void +// Index::optimize_index_layout() //{ // use after build or load //} // REFACTOR: This should be an OptimizedDataStore class -template void Index::optimize_index_layout() -{ // use after build or load - if (_dynamic_index) - { - throw diskann::ANNException("Optimize_index_layout not implemented for dyanmic indices", -1, __FUNCSIG__, - __FILE__, __LINE__); - } - - float *cur_vec = new float[_data_store->get_aligned_dim()]; - std::memset(cur_vec, 0, _data_store->get_aligned_dim() * sizeof(float)); - _data_len = (_data_store->get_aligned_dim() + 1) * sizeof(float); - _neighbor_len = (_max_observed_degree + 1) * sizeof(uint32_t); - _node_size = _data_len + _neighbor_len; - _opt_graph = new char[_node_size * _nd]; - DistanceFastL2 *dist_fast = (DistanceFastL2 *)_data_store->get_dist_fn(); - for (uint32_t i = 0; i < _nd; i++) - { - char *cur_node_offset = _opt_graph + i * _node_size; - _data_store->get_vector(i, (T *)cur_vec); - float cur_norm = dist_fast->norm((T *)cur_vec, _data_store->get_aligned_dim()); - std::memcpy(cur_node_offset, &cur_norm, sizeof(float)); - std::memcpy(cur_node_offset + sizeof(float), cur_vec, _data_len - sizeof(float)); - - cur_node_offset += _data_len; - uint32_t k = _final_graph[i].size(); - std::memcpy(cur_node_offset, &k, sizeof(uint32_t)); - std::memcpy(cur_node_offset + sizeof(uint32_t), _final_graph[i].data(), k * sizeof(uint32_t)); - std::vector().swap(_final_graph[i]); - } - _final_graph.clear(); - _final_graph.shrink_to_fit(); - delete[] cur_vec; -} - -// REFACTOR: once optimized layout becomes its own Data+Graph store, we should just invoke regular search +template +void Index::optimize_index_layout() { // use after build or load + if (_dynamic_index) { + throw diskann::ANNException( + "Optimize_index_layout not implemented for dyanmic indices", -1, + __FUNCSIG__, __FILE__, __LINE__); + } + + float *cur_vec = new float[_data_store->get_aligned_dim()]; + std::memset(cur_vec, 0, _data_store->get_aligned_dim() * sizeof(float)); + _data_len = (_data_store->get_aligned_dim() + 1) * sizeof(float); + _neighbor_len = (_max_observed_degree + 1) * sizeof(uint32_t); + _node_size = _data_len + _neighbor_len; + _opt_graph = new char[_node_size * _nd]; + DistanceFastL2 *dist_fast = + (DistanceFastL2 *)_data_store->get_dist_fn(); + for (uint32_t i = 0; i < _nd; i++) { + char *cur_node_offset = _opt_graph + i * _node_size; + _data_store->get_vector(i, (T *)cur_vec); + float cur_norm = + dist_fast->norm((T *)cur_vec, _data_store->get_aligned_dim()); + std::memcpy(cur_node_offset, &cur_norm, sizeof(float)); + std::memcpy(cur_node_offset + sizeof(float), cur_vec, + _data_len - sizeof(float)); + + cur_node_offset += _data_len; + uint32_t k = _final_graph[i].size(); + std::memcpy(cur_node_offset, &k, sizeof(uint32_t)); + std::memcpy(cur_node_offset + sizeof(uint32_t), _final_graph[i].data(), + k * sizeof(uint32_t)); + std::vector().swap(_final_graph[i]); + } + _final_graph.clear(); + _final_graph.shrink_to_fit(); + delete[] cur_vec; +} + +// REFACTOR: once optimized layout becomes its own Data+Graph store, we should +// just invoke regular search // template -// void Index::search_with_optimized_layout(const T *query, size_t K, size_t L, uint32_t *indices) +// void Index::search_with_optimized_layout(const T *query, +// size_t K, size_t L, uint32_t *indices) //{ //} template -void Index::search_with_optimized_layout(const T *query, size_t K, size_t L, uint32_t *indices) -{ - DistanceFastL2 *dist_fast = (DistanceFastL2 *)_data_store->get_dist_fn(); - - NeighborPriorityQueue retset(L); - std::vector init_ids(L); - - boost::dynamic_bitset<> flags{_nd, 0}; - uint32_t tmp_l = 0; - uint32_t *neighbors = (uint32_t *)(_opt_graph + _node_size * _start + _data_len); - uint32_t MaxM_ep = *neighbors; +void Index::search_with_optimized_layout(const T *query, + size_t K, size_t L, + uint32_t *indices) { + DistanceFastL2 *dist_fast = + (DistanceFastL2 *)_data_store->get_dist_fn(); + + NeighborPriorityQueue retset(L); + std::vector init_ids(L); + + boost::dynamic_bitset<> flags{_nd, 0}; + uint32_t tmp_l = 0; + uint32_t *neighbors = + (uint32_t *)(_opt_graph + _node_size * _start + _data_len); + uint32_t 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) { + uint32_t id = rand() % _nd; + if (flags[id]) continue; + flags[id] = true; + init_ids[tmp_l] = id; + tmp_l++; + } + + for (uint32_t i = 0; i < init_ids.size(); i++) { + uint32_t id = init_ids[i]; + if (id >= _nd) continue; + _mm_prefetch(_opt_graph + _node_size * id, _MM_HINT_T0); + } + L = 0; + for (uint32_t i = 0; i < init_ids.size(); i++) { + uint32_t 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, + (uint32_t)_data_store->get_aligned_dim()); + retset.insert(Neighbor(id, dist)); + flags[id] = true; + L++; + } + + while (retset.has_unexpanded_node()) { + auto nbr = retset.closest_unexpanded(); + auto n = nbr.id; + _mm_prefetch(_opt_graph + _node_size * n + _data_len, _MM_HINT_T0); + neighbors = (uint32_t *)(_opt_graph + _node_size * n + _data_len); + uint32_t MaxM = *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) - { - uint32_t id = rand() % _nd; - if (flags[id]) - continue; - flags[id] = true; - init_ids[tmp_l] = id; - tmp_l++; - } - - for (uint32_t i = 0; i < init_ids.size(); i++) - { - uint32_t id = init_ids[i]; - if (id >= _nd) - continue; - _mm_prefetch(_opt_graph + _node_size * id, _MM_HINT_T0); - } - L = 0; - for (uint32_t i = 0; i < init_ids.size(); i++) - { - uint32_t 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, (uint32_t)_data_store->get_aligned_dim()); - retset.insert(Neighbor(id, dist)); - flags[id] = true; - L++; - } - - while (retset.has_unexpanded_node()) - { - auto nbr = retset.closest_unexpanded(); - auto n = nbr.id; - _mm_prefetch(_opt_graph + _node_size * n + _data_len, _MM_HINT_T0); - neighbors = (uint32_t *)(_opt_graph + _node_size * n + _data_len); - uint32_t MaxM = *neighbors; - neighbors++; - for (uint32_t m = 0; m < MaxM; ++m) - _mm_prefetch(_opt_graph + _node_size * neighbors[m], _MM_HINT_T0); - for (uint32_t m = 0; m < MaxM; ++m) - { - uint32_t id = neighbors[m]; - if (flags[id]) - continue; - flags[id] = 1; - T *data = (T *)(_opt_graph + _node_size * id); - float norm = *data; - data++; - float dist = dist_fast->compare(query, data, norm, (uint32_t)_data_store->get_aligned_dim()); - Neighbor nn(id, dist); - retset.insert(nn); - } - } - - for (size_t i = 0; i < K; i++) - { - indices[i] = retset[i].id; - } + for (uint32_t m = 0; m < MaxM; ++m) + _mm_prefetch(_opt_graph + _node_size * neighbors[m], _MM_HINT_T0); + for (uint32_t m = 0; m < MaxM; ++m) { + uint32_t id = neighbors[m]; + if (flags[id]) continue; + flags[id] = 1; + T *data = (T *)(_opt_graph + _node_size * id); + float norm = *data; + data++; + float dist = dist_fast->compare(query, data, norm, + (uint32_t)_data_store->get_aligned_dim()); + Neighbor nn(id, dist); + retset.insert(nn); + } + } + + for (size_t i = 0; i < K; i++) { + indices[i] = retset[i].id; + } } /* Internals of the library */ -template const float Index::INDEX_GROWTH_FACTOR = 1.5f; +template +const float Index::INDEX_GROWTH_FACTOR = 1.5f; // EXPORTS template DISKANN_DLLEXPORT class Index; @@ -3099,132 +2970,252 @@ template DISKANN_DLLEXPORT class Index; template DISKANN_DLLEXPORT class Index; template DISKANN_DLLEXPORT class Index; -template DISKANN_DLLEXPORT std::pair Index::search( - const float *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const float *query, const size_t K, const uint32_t L, uint32_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const uint8_t *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const uint8_t *query, const size_t K, const uint32_t L, uint32_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const int8_t *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const int8_t *query, const size_t K, const uint32_t L, uint32_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search(const float *query, + const size_t K, + const uint32_t L, + uint64_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search(const float *query, + const size_t K, + const uint32_t L, + uint32_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search(const uint8_t *query, + const size_t K, + const uint32_t L, + uint64_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search(const uint8_t *query, + const size_t K, + const uint32_t L, + uint32_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search(const int8_t *query, + const size_t K, + const uint32_t L, + uint64_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search(const int8_t *query, + const size_t K, + const uint32_t L, + uint32_t *indices, + float *distances); // TagT==uint32_t -template DISKANN_DLLEXPORT std::pair Index::search( - const float *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const float *query, const size_t K, const uint32_t L, uint32_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const uint8_t *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const uint8_t *query, const size_t K, const uint32_t L, uint32_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const int8_t *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const int8_t *query, const size_t K, const uint32_t L, uint32_t *indices, float *distances); - -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint64_t>(const float *query, const uint32_t &filter_label, const size_t K, const uint32_t L, uint64_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint32_t>(const float *query, const uint32_t &filter_label, const size_t K, const uint32_t L, uint32_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint64_t>(const uint8_t *query, const uint32_t &filter_label, const size_t K, const uint32_t L, uint64_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint32_t>(const uint8_t *query, const uint32_t &filter_label, const size_t K, const uint32_t L, uint32_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint64_t>(const int8_t *query, const uint32_t &filter_label, const size_t K, const uint32_t L, uint64_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint32_t>(const int8_t *query, const uint32_t &filter_label, const size_t K, const uint32_t L, uint32_t *indices, - float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search(const float *query, + const size_t K, + const uint32_t L, + uint64_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search(const float *query, + const size_t K, + const uint32_t L, + uint32_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search(const uint8_t *query, + const size_t K, + const uint32_t L, + uint64_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search(const uint8_t *query, + const size_t K, + const uint32_t L, + uint32_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search(const int8_t *query, + const size_t K, + const uint32_t L, + uint64_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search(const int8_t *query, + const size_t K, + const uint32_t L, + uint32_t *indices, + float *distances); + +template DISKANN_DLLEXPORT std::pair +Index::search_with_filters( + const float *query, const uint32_t &filter_label, const size_t K, + const uint32_t L, uint64_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search_with_filters( + const float *query, const uint32_t &filter_label, const size_t K, + const uint32_t L, uint32_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search_with_filters( + const uint8_t *query, const uint32_t &filter_label, const size_t K, + const uint32_t L, uint64_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search_with_filters( + const uint8_t *query, const uint32_t &filter_label, const size_t K, + const uint32_t L, uint32_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search_with_filters( + const int8_t *query, const uint32_t &filter_label, const size_t K, + const uint32_t L, uint64_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search_with_filters( + const int8_t *query, const uint32_t &filter_label, const size_t K, + const uint32_t L, uint32_t *indices, float *distances); // TagT==uint32_t -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint64_t>(const float *query, const uint32_t &filter_label, const size_t K, const uint32_t L, uint64_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint32_t>(const float *query, const uint32_t &filter_label, const size_t K, const uint32_t L, uint32_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint64_t>(const uint8_t *query, const uint32_t &filter_label, const size_t K, const uint32_t L, uint64_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint32_t>(const uint8_t *query, const uint32_t &filter_label, const size_t K, const uint32_t L, uint32_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint64_t>(const int8_t *query, const uint32_t &filter_label, const size_t K, const uint32_t L, uint64_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint32_t>(const int8_t *query, const uint32_t &filter_label, const size_t K, const uint32_t L, uint32_t *indices, - float *distances); - -template DISKANN_DLLEXPORT std::pair Index::search( - const float *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const float *query, const size_t K, const uint32_t L, uint32_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const uint8_t *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const uint8_t *query, const size_t K, const uint32_t L, uint32_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const int8_t *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const int8_t *query, const size_t K, const uint32_t L, uint32_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search_with_filters( + const float *query, const uint32_t &filter_label, const size_t K, + const uint32_t L, uint64_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search_with_filters( + const float *query, const uint32_t &filter_label, const size_t K, + const uint32_t L, uint32_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search_with_filters( + const uint8_t *query, const uint32_t &filter_label, const size_t K, + const uint32_t L, uint64_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search_with_filters( + const uint8_t *query, const uint32_t &filter_label, const size_t K, + const uint32_t L, uint32_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search_with_filters( + const int8_t *query, const uint32_t &filter_label, const size_t K, + const uint32_t L, uint64_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search_with_filters( + const int8_t *query, const uint32_t &filter_label, const size_t K, + const uint32_t L, uint32_t *indices, float *distances); + +template DISKANN_DLLEXPORT std::pair +Index::search(const float *query, + const size_t K, + const uint32_t L, + uint64_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search(const float *query, + const size_t K, + const uint32_t L, + uint32_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search(const uint8_t *query, + const size_t K, + const uint32_t L, + uint64_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search(const uint8_t *query, + const size_t K, + const uint32_t L, + uint32_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search(const int8_t *query, + const size_t K, + const uint32_t L, + uint64_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search(const int8_t *query, + const size_t K, + const uint32_t L, + uint32_t *indices, + float *distances); // TagT==uint32_t -template DISKANN_DLLEXPORT std::pair Index::search( - const float *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const float *query, const size_t K, const uint32_t L, uint32_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const uint8_t *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const uint8_t *query, const size_t K, const uint32_t L, uint32_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const int8_t *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair Index::search( - const int8_t *query, const size_t K, const uint32_t L, uint32_t *indices, float *distances); - -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint64_t>(const float *query, const uint16_t &filter_label, const size_t K, const uint32_t L, uint64_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint32_t>(const float *query, const uint16_t &filter_label, const size_t K, const uint32_t L, uint32_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint64_t>(const uint8_t *query, const uint16_t &filter_label, const size_t K, const uint32_t L, uint64_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint32_t>(const uint8_t *query, const uint16_t &filter_label, const size_t K, const uint32_t L, uint32_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint64_t>(const int8_t *query, const uint16_t &filter_label, const size_t K, const uint32_t L, uint64_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint32_t>(const int8_t *query, const uint16_t &filter_label, const size_t K, const uint32_t L, uint32_t *indices, - float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search(const float *query, + const size_t K, + const uint32_t L, + uint64_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search(const float *query, + const size_t K, + const uint32_t L, + uint32_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search(const uint8_t *query, + const size_t K, + const uint32_t L, + uint64_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search(const uint8_t *query, + const size_t K, + const uint32_t L, + uint32_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search(const int8_t *query, + const size_t K, + const uint32_t L, + uint64_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search(const int8_t *query, + const size_t K, + const uint32_t L, + uint32_t *indices, + float *distances); + +template DISKANN_DLLEXPORT std::pair +Index::search_with_filters( + const float *query, const uint16_t &filter_label, const size_t K, + const uint32_t L, uint64_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search_with_filters( + const float *query, const uint16_t &filter_label, const size_t K, + const uint32_t L, uint32_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search_with_filters( + const uint8_t *query, const uint16_t &filter_label, const size_t K, + const uint32_t L, uint64_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search_with_filters( + const uint8_t *query, const uint16_t &filter_label, const size_t K, + const uint32_t L, uint32_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search_with_filters( + const int8_t *query, const uint16_t &filter_label, const size_t K, + const uint32_t L, uint64_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search_with_filters( + const int8_t *query, const uint16_t &filter_label, const size_t K, + const uint32_t L, uint32_t *indices, float *distances); // TagT==uint32_t -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint64_t>(const float *query, const uint16_t &filter_label, const size_t K, const uint32_t L, uint64_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint32_t>(const float *query, const uint16_t &filter_label, const size_t K, const uint32_t L, uint32_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint64_t>(const uint8_t *query, const uint16_t &filter_label, const size_t K, const uint32_t L, uint64_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint32_t>(const uint8_t *query, const uint16_t &filter_label, const size_t K, const uint32_t L, uint32_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint64_t>(const int8_t *query, const uint16_t &filter_label, const size_t K, const uint32_t L, uint64_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair Index::search_with_filters< - uint32_t>(const int8_t *query, const uint16_t &filter_label, const size_t K, const uint32_t L, uint32_t *indices, - float *distances); - -} // namespace diskann +template DISKANN_DLLEXPORT std::pair +Index::search_with_filters( + const float *query, const uint16_t &filter_label, const size_t K, + const uint32_t L, uint64_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search_with_filters( + const float *query, const uint16_t &filter_label, const size_t K, + const uint32_t L, uint32_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search_with_filters( + const uint8_t *query, const uint16_t &filter_label, const size_t K, + const uint32_t L, uint64_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search_with_filters( + const uint8_t *query, const uint16_t &filter_label, const size_t K, + const uint32_t L, uint32_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search_with_filters( + const int8_t *query, const uint16_t &filter_label, const size_t K, + const uint32_t L, uint64_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair +Index::search_with_filters( + const int8_t *query, const uint16_t &filter_label, const size_t K, + const uint32_t L, uint32_t *indices, float *distances); + +} // namespace diskann diff --git a/tests/build_memory_index.cpp b/tests/build_memory_index.cpp index 3712350c3..5146afa60 100644 --- a/tests/build_memory_index.cpp +++ b/tests/build_memory_index.cpp @@ -21,185 +21,196 @@ namespace po = boost::program_options; template -int build_in_memory_index(const diskann::Metric &metric, const std::string &data_path, const uint32_t R, - const uint32_t L, const float alpha, const std::string &save_path, const uint32_t num_threads, - const bool use_pq_build, const size_t num_pq_bytes, const bool use_opq, - const std::string &label_file, const std::string &universal_label, const uint32_t Lf) -{ - diskann::IndexWriteParameters paras = diskann::IndexWriteParametersBuilder(L, R) - .with_filter_list_size(Lf) - .with_alpha(alpha) - .with_saturate_graph(false) - .with_num_threads(num_threads) - .build(); - std::string labels_file_to_use = save_path + "_label_formatted.txt"; - std::string mem_labels_int_map_file = save_path + "_labels_map.txt"; +int build_in_memory_index(const diskann::Metric &metric, + const std::string &data_path, const uint32_t R, + const uint32_t L, const float alpha, + const std::string &save_path, + const uint32_t num_threads, const bool use_pq_build, + const size_t num_pq_bytes, const bool use_opq, + const std::string &label_file, + const std::string &universal_label, + const uint32_t Lf) { + diskann::IndexWriteParameters paras = + diskann::IndexWriteParametersBuilder(L, R) + .with_filter_list_size(Lf) + .with_alpha(alpha) + .with_saturate_graph(false) + .with_num_threads(num_threads) + .build(); + std::string labels_file_to_use = save_path + "_label_formatted.txt"; + std::string mem_labels_int_map_file = save_path + "_labels_map.txt"; - size_t data_num, data_dim; - diskann::get_bin_metadata(data_path, data_num, data_dim); + size_t data_num, data_dim; + diskann::get_bin_metadata(data_path, data_num, data_dim); - diskann::Index index(metric, data_dim, data_num, false, false, false, use_pq_build, num_pq_bytes, - use_opq); - auto s = std::chrono::high_resolution_clock::now(); - if (label_file == "") - { - index.build(data_path.c_str(), data_num, paras); + diskann::Index index(metric, data_dim, data_num, false, + false, false, use_pq_build, + num_pq_bytes, use_opq); + auto s = std::chrono::high_resolution_clock::now(); + if (label_file == "") { + index.build(data_path.c_str(), data_num, paras); + } else { + convert_labels_string_to_int(label_file, labels_file_to_use, + mem_labels_int_map_file, universal_label); + if (universal_label != "") { + LabelT unv_label_as_num = 0; + index.set_universal_label(unv_label_as_num); } - else - { - convert_labels_string_to_int(label_file, labels_file_to_use, mem_labels_int_map_file, universal_label); - if (universal_label != "") - { - LabelT unv_label_as_num = 0; - index.set_universal_label(unv_label_as_num); - } - index.build_filtered_index(data_path.c_str(), labels_file_to_use, data_num, paras); - } - std::chrono::duration diff = std::chrono::high_resolution_clock::now() - s; + index.build_filtered_index(data_path.c_str(), labels_file_to_use, data_num, + paras); + } + std::chrono::duration diff = + std::chrono::high_resolution_clock::now() - s; - std::cout << "Indexing time: " << diff.count() << "\n"; - index.save(save_path.c_str()); - if (label_file != "") - std::remove(labels_file_to_use.c_str()); - return 0; + std::cout << "Indexing time: " << diff.count() << "\n"; + index.save(save_path.c_str()); + if (label_file != "") std::remove(labels_file_to_use.c_str()); + return 0; } -int main(int argc, char **argv) -{ - std::string data_type, dist_fn, data_path, index_path_prefix, label_file, universal_label, label_type; - uint32_t num_threads, R, L, Lf, build_PQ_bytes; - float alpha; - bool use_pq_build, use_opq; +int main(int argc, char **argv) { + std::string data_type, dist_fn, data_path, index_path_prefix, label_file, + universal_label, label_type; + uint32_t num_threads, R, L, Lf, build_PQ_bytes; + float alpha; + bool use_pq_build, use_opq; - po::options_description desc{"Arguments"}; - try - { - desc.add_options()("help,h", "Print information on arguments"); - desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); - desc.add_options()("dist_fn", po::value(&dist_fn)->required(), - "distance function "); - desc.add_options()("data_path", po::value(&data_path)->required(), - "Input data file in bin format"); - desc.add_options()("index_path_prefix", po::value(&index_path_prefix)->required(), - "Path prefix for saving index file components"); - desc.add_options()("max_degree,R", po::value(&R)->default_value(64), "Maximum graph degree"); - desc.add_options()("Lbuild,L", po::value(&L)->default_value(100), - "Build complexity, higher value results in better graphs"); - desc.add_options()("alpha", po::value(&alpha)->default_value(1.2f), - "alpha controls density and diameter of graph, set " - "1 for sparse graph, " - "1.2 or 1.4 for denser graphs with lower diameter"); - desc.add_options()("num_threads,T", po::value(&num_threads)->default_value(omp_get_num_procs()), - "Number of threads used for building index (defaults to " - "omp_get_num_procs())"); - desc.add_options()("build_PQ_bytes", po::value(&build_PQ_bytes)->default_value(0), - "Number of PQ bytes to build the index; 0 for full precision " - "build"); - desc.add_options()("use_opq", po::bool_switch()->default_value(false), - "Set true for OPQ compression while using PQ " - "distance comparisons for " - "building the index, and false for PQ compression"); - desc.add_options()("label_file", po::value(&label_file)->default_value(""), - "Input label file in txt format for Filtered Index search. " - "The file should contain comma separated filters for each node " - "with each line corresponding to a graph node"); - desc.add_options()("universal_label", po::value(&universal_label)->default_value(""), - "Universal label, if using it, only in conjunction with " - "labels_file"); - desc.add_options()("FilteredLbuild,Lf", po::value(&Lf)->default_value(0), - "Build complexity for filtered points, higher value " - "results in better graphs"); - desc.add_options()("label_type", po::value(&label_type)->default_value("uint"), - "Storage type of Labels , default value is uint which " - "will consume memory 4 bytes per filter"); + po::options_description desc{"Arguments"}; + try { + desc.add_options()("help,h", "Print information on arguments"); + desc.add_options()("data_type", + po::value(&data_type)->required(), + "data type "); + desc.add_options()("dist_fn", po::value(&dist_fn)->required(), + "distance function "); + desc.add_options()("data_path", + po::value(&data_path)->required(), + "Input data file in bin format"); + desc.add_options()("index_path_prefix", + po::value(&index_path_prefix)->required(), + "Path prefix for saving index file components"); + desc.add_options()("max_degree,R", + po::value(&R)->default_value(64), + "Maximum graph degree"); + desc.add_options()( + "Lbuild,L", po::value(&L)->default_value(100), + "Build complexity, higher value results in better graphs"); + desc.add_options()("alpha", po::value(&alpha)->default_value(1.2f), + "alpha controls density and diameter of graph, set " + "1 for sparse graph, " + "1.2 or 1.4 for denser graphs with lower diameter"); + desc.add_options()( + "num_threads,T", + po::value(&num_threads)->default_value(omp_get_num_procs()), + "Number of threads used for building index (defaults to " + "omp_get_num_procs())"); + desc.add_options()( + "build_PQ_bytes", + po::value(&build_PQ_bytes)->default_value(0), + "Number of PQ bytes to build the index; 0 for full precision " + "build"); + desc.add_options()("use_opq", po::bool_switch()->default_value(false), + "Set true for OPQ compression while using PQ " + "distance comparisons for " + "building the index, and false for PQ compression"); + desc.add_options()( + "label_file", po::value(&label_file)->default_value(""), + "Input label file in txt format for Filtered Index search. " + "The file should contain comma separated filters for each node " + "with each line corresponding to a graph node"); + desc.add_options()( + "universal_label", + po::value(&universal_label)->default_value(""), + "Universal label, if using it, only in conjunction with " + "labels_file"); + desc.add_options()("FilteredLbuild,Lf", + po::value(&Lf)->default_value(0), + "Build complexity for filtered points, higher value " + "results in better graphs"); + desc.add_options()( + "label_type", + po::value(&label_type)->default_value("uint"), + "Storage type of Labels , default value is uint which " + "will consume memory 4 bytes per filter"); - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) - { - std::cout << desc; - return 0; - } - po::notify(vm); - use_pq_build = (build_PQ_bytes > 0); - use_opq = vm["use_opq"].as(); - } - catch (const std::exception &ex) - { - std::cerr << ex.what() << '\n'; - return -1; + po::variables_map vm; + po::store(po::parse_command_line(argc, argv, desc), vm); + if (vm.count("help")) { + std::cout << desc; + return 0; } + po::notify(vm); + use_pq_build = (build_PQ_bytes > 0); + use_opq = vm["use_opq"].as(); + } catch (const std::exception &ex) { + std::cerr << ex.what() << '\n'; + return -1; + } - diskann::Metric metric; - if (dist_fn == std::string("mips")) - { - metric = diskann::Metric::INNER_PRODUCT; - } - else if (dist_fn == std::string("l2")) - { - metric = diskann::Metric::L2; - } - else if (dist_fn == std::string("cosine")) - { - metric = diskann::Metric::COSINE; - } - else - { - std::cout << "Unsupported distance function. Currently only L2/ Inner " - "Product/Cosine are supported." + diskann::Metric metric; + if (dist_fn == std::string("mips")) { + metric = diskann::Metric::INNER_PRODUCT; + } else if (dist_fn == std::string("l2")) { + metric = diskann::Metric::L2; + } else if (dist_fn == std::string("cosine")) { + metric = diskann::Metric::COSINE; + } else { + std::cout << "Unsupported distance function. Currently only L2/ Inner " + "Product/Cosine are supported." + << std::endl; + return -1; + } + + try { + diskann::cout << "Starting index build with R: " << R << " Lbuild: " << L + << " alpha: " << alpha << " #threads: " << num_threads + << std::endl; + if (label_file != "" && label_type == "ushort") { + if (data_type == std::string("int8")) + return build_in_memory_index( + metric, data_path, R, L, alpha, index_path_prefix, num_threads, + use_pq_build, build_PQ_bytes, use_opq, label_file, universal_label, + Lf); + else if (data_type == std::string("uint8")) + return build_in_memory_index( + metric, data_path, R, L, alpha, index_path_prefix, num_threads, + use_pq_build, build_PQ_bytes, use_opq, label_file, universal_label, + Lf); + else if (data_type == std::string("float")) + return build_in_memory_index( + metric, data_path, R, L, alpha, index_path_prefix, num_threads, + use_pq_build, build_PQ_bytes, use_opq, label_file, universal_label, + Lf); + else { + std::cout << "Unsupported type. Use one of int8, uint8 or float." << std::endl; return -1; - } - - try - { - diskann::cout << "Starting index build with R: " << R << " Lbuild: " << L << " alpha: " << alpha - << " #threads: " << num_threads << std::endl; - if (label_file != "" && label_type == "ushort") - { - if (data_type == std::string("int8")) - return build_in_memory_index( - metric, data_path, R, L, alpha, index_path_prefix, num_threads, use_pq_build, build_PQ_bytes, - use_opq, label_file, universal_label, Lf); - else if (data_type == std::string("uint8")) - return build_in_memory_index( - metric, data_path, R, L, alpha, index_path_prefix, num_threads, use_pq_build, build_PQ_bytes, - use_opq, label_file, universal_label, Lf); - else if (data_type == std::string("float")) - return build_in_memory_index( - metric, data_path, R, L, alpha, index_path_prefix, num_threads, use_pq_build, build_PQ_bytes, - use_opq, label_file, universal_label, Lf); - else - { - std::cout << "Unsupported type. Use one of int8, uint8 or float." << std::endl; - return -1; - } - } - else - { - if (data_type == std::string("int8")) - return build_in_memory_index(metric, data_path, R, L, alpha, index_path_prefix, num_threads, - use_pq_build, build_PQ_bytes, use_opq, label_file, universal_label, - Lf); - else if (data_type == std::string("uint8")) - return build_in_memory_index(metric, data_path, R, L, alpha, index_path_prefix, num_threads, - use_pq_build, build_PQ_bytes, use_opq, label_file, - universal_label, Lf); - else if (data_type == std::string("float")) - return build_in_memory_index(metric, data_path, R, L, alpha, index_path_prefix, num_threads, - use_pq_build, build_PQ_bytes, use_opq, label_file, universal_label, - Lf); - else - { - std::cout << "Unsupported type. Use one of int8, uint8 or float." << std::endl; - return -1; - } - } - } - catch (const std::exception &e) - { - std::cout << std::string(e.what()) << std::endl; - diskann::cerr << "Index build failed." << std::endl; + } + } else { + if (data_type == std::string("int8")) + return build_in_memory_index( + metric, data_path, R, L, alpha, index_path_prefix, num_threads, + use_pq_build, build_PQ_bytes, use_opq, label_file, universal_label, + Lf); + else if (data_type == std::string("uint8")) + return build_in_memory_index( + metric, data_path, R, L, alpha, index_path_prefix, num_threads, + use_pq_build, build_PQ_bytes, use_opq, label_file, universal_label, + Lf); + else if (data_type == std::string("float")) + return build_in_memory_index( + metric, data_path, R, L, alpha, index_path_prefix, num_threads, + use_pq_build, build_PQ_bytes, use_opq, label_file, universal_label, + Lf); + else { + std::cout << "Unsupported type. Use one of int8, uint8 or float." + << std::endl; return -1; + } } + } catch (const std::exception &e) { + std::cout << std::string(e.what()) << std::endl; + diskann::cerr << "Index build failed." << std::endl; + return -1; + } } From 3f73901aba66b8a82c683fffd0b0e30db70c0385 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Wed, 19 Apr 2023 01:02:16 +0530 Subject: [PATCH 55/69] fixed cosine bug, clang-formatted --- include/abstract_data_store.h | 174 +- include/abstract_graph_store.h | 28 +- include/aligned_file_reader.h | 139 +- include/boost_dynamic_bitset_fwd.h | 9 +- include/cached_io.h | 348 ++-- include/concurrent_queue.h | 175 +- include/cosine_similarity.h | 392 ++-- include/defaults.h | 10 +- include/disk_utils.h | 98 +- include/distance.h | 370 ++-- include/exceptions.h | 15 +- include/filter_utils.h | 275 +-- include/in_mem_data_store.h | 122 +- include/in_mem_graph_store.h | 20 +- include/index.h | 793 ++++---- include/linux_aligned_file_reader.h | 56 +- include/locking.h | 5 +- include/logger.h | 15 +- include/logger_impl.h | 65 +- include/math_utils.h | 48 +- include/memory_mapper.h | 34 +- include/natural_number_map.h | 130 +- include/natural_number_set.h | 57 +- include/neighbor.h | 191 +- include/parameters.h | 216 ++- include/partition.h | 39 +- include/percentile_stats.h | 78 +- include/pq.h | 174 +- include/pq_flash_index.h | 356 ++-- include/scratch.h | 321 ++- include/simd_utils.h | 137 +- include/timer.h | 46 +- include/types.h | 5 +- include/utils.h | 1716 ++++++++--------- include/windows_aligned_file_reader.h | 64 +- include/windows_slim_lock.h | 76 +- src/abstract_data_store.cpp | 47 +- src/ann_exception.cpp | 43 +- src/disk_utils.cpp | 2571 +++++++++++++------------ src/distance.cpp | 1092 +++++------ src/filter_utils.cpp | 452 ++--- src/in_mem_data_store.cpp | 564 +++--- src/in_mem_graph_store.cpp | 28 +- src/linux_aligned_file_reader.cpp | 333 ++-- src/logger.cpp | 112 +- src/math_utils.cpp | 668 ++++--- src/memory_mapper.cpp | 146 +- src/natural_number_map.cpp | 132 +- src/natural_number_set.cpp | 76 +- src/partition.cpp | 1121 +++++------ src/pq.cpp | 1739 +++++++++-------- src/pq_flash_index.cpp | 2521 ++++++++++++------------ src/scratch.cpp | 189 +- src/utils.cpp | 814 ++++---- src/windows_aligned_file_reader.cpp | 284 ++- 55 files changed, 9728 insertions(+), 10001 deletions(-) diff --git a/include/abstract_data_store.h b/include/abstract_data_store.h index 4c59efd2f..6cf83e693 100644 --- a/include/abstract_data_store.h +++ b/include/abstract_data_store.h @@ -9,87 +9,97 @@ #include "types.h" #include "windows_customizations.h" -namespace diskann -{ - -template class AbstractDataStore -{ - public: - AbstractDataStore(const location_t capacity, const size_t dim); - - // Return number of points returned - virtual location_t load(const std::string &filename) = 0; - - // Why does store take num_pts? Since store only has capacity, but we allow resizing - // we can end up in a situation where the store has spare capacity. To optimize disk - // utilization, we pass the number of points that are "true" points, so that the store - // can discard the empty locations before saving. - virtual size_t save(const std::string &filename, const location_t num_pts) = 0; - - DISKANN_DLLEXPORT virtual location_t capacity() const; - - DISKANN_DLLEXPORT virtual size_t get_dims() const; - - // Implementers can choose to return _dim if they are not - // concerned about memory alignment. - // Returns _dim aligned to a 8-byte value. Used for allocating - // aligned memory for efficiency/simplicity of code. - virtual size_t get_aligned_dim() const = 0; - - // populate the store with vectors (either from a pointer or bin file), - // potentially after normalizing the vectors if the metric deems so - virtual void populate_data(const data_t *vectors, const location_t num_pts) = 0; - virtual void populate_data(const std::string &filename, const size_t offset) = 0; - - // reverse of populate, save the first num_pts many points back to bin file - virtual void save_data_to_bin(const std::string &filename, const location_t num_pts) = 0; - - // Returns the updated capacity of the datastore. Clients should check - // if resize actually changed the capacity to new_num_points before - // proceeding with operations. See the code below: - // auto new_capcity = data_store->resize(new_num_points); - // if ( new_capacity >= new_num_points) { - // //PROCEED - // else - // //ERROR. - virtual location_t resize(const location_t new_num_points); - - // operations on vectors - virtual void get_vector(const location_t i, data_t *dest) const = 0; - virtual void set_vector(const location_t i, const data_t *const vector) = 0; - virtual void prefetch_vector(const location_t loc) = 0; - - // internal shuffle operations to move around vectors - virtual void reposition_points(const location_t start_loc, const location_t end_loc, - const location_t num_points) = 0; - virtual void copy_points(const location_t from_loc, const location_t to_loc, const location_t num_points) = 0; - - // metric specific operations - - virtual float get_distance(const data_t *query, const location_t loc) const = 0; - virtual void get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, - float *distances) const = 0; - virtual float get_distance(const location_t loc1, const location_t loc2) const = 0; - - // stats of the data stored in store - // Returns the point in the dataset that is closest to the mean of all points in the dataset - virtual location_t calculate_medoid() const = 0; - - //search helpers - virtual size_t get_alignment_factor() const = 0; - - protected: - // Expand the datastore to new_num_points. Returns the new capacity created, which should be == new_num_points - // in the normal case. Implementers can also return _capacity to indicate that there are not implementing this - // method. - virtual location_t expand(const location_t new_num_points) = 0; - - // Shrink the datastore to new_num_points. It is NOT an error if shrink doesn't reduce the capacity - // so callers need to check this correctly. See also for "default" implementation - virtual location_t shrink(const location_t new_num_points) = 0; - - location_t _capacity; - size_t _dim; +namespace diskann { + +template +class AbstractDataStore { + public: + AbstractDataStore(const location_t capacity, const size_t dim); + + // Return number of points returned + virtual location_t load(const std::string &filename) = 0; + + // Why does store take num_pts? Since store only has capacity, but we allow + // resizing we can end up in a situation where the store has spare capacity. + // To optimize disk utilization, we pass the number of points that are "true" + // points, so that the store can discard the empty locations before saving. + virtual size_t save(const std::string &filename, + const location_t num_pts) = 0; + + DISKANN_DLLEXPORT virtual location_t capacity() const; + + DISKANN_DLLEXPORT virtual size_t get_dims() const; + + // Implementers can choose to return _dim if they are not + // concerned about memory alignment. + // Returns _dim aligned to a 8-byte value. Used for allocating + // aligned memory for efficiency/simplicity of code. + virtual size_t get_aligned_dim() const = 0; + + // populate the store with vectors (either from a pointer or bin file), + // potentially after normalizing the vectors if the metric deems so + virtual void populate_data(const data_t *vectors, + const location_t num_pts) = 0; + virtual void populate_data(const std::string &filename, + const size_t offset) = 0; + + // reverse of populate, save the first num_pts many points back to bin file + virtual void save_data_to_bin(const std::string &filename, + const location_t num_pts) = 0; + + // Returns the updated capacity of the datastore. Clients should check + // if resize actually changed the capacity to new_num_points before + // proceeding with operations. See the code below: + // auto new_capcity = data_store->resize(new_num_points); + // if ( new_capacity >= new_num_points) { + // //PROCEED + // else + // //ERROR. + virtual location_t resize(const location_t new_num_points); + + // operations on vectors + virtual void get_vector(const location_t i, data_t *dest) const = 0; + virtual void set_vector(const location_t i, const data_t *const vector) = 0; + virtual void prefetch_vector(const location_t loc) = 0; + + // internal shuffle operations to move around vectors + virtual void reposition_points(const location_t start_loc, + const location_t end_loc, + const location_t num_points) = 0; + virtual void copy_points(const location_t from_loc, const location_t to_loc, + const location_t num_points) = 0; + + // metric specific operations + + virtual float get_distance(const data_t *query, + const location_t loc) const = 0; + virtual void get_distance(const data_t *query, const location_t *locations, + const uint32_t location_count, + float *distances) const = 0; + virtual float get_distance(const location_t loc1, + const location_t loc2) const = 0; + + // stats of the data stored in store + // Returns the point in the dataset that is closest to the mean of all points + // in the dataset + virtual location_t calculate_medoid() const = 0; + + // search helpers + virtual size_t get_alignment_factor() const = 0; + + protected: + // Expand the datastore to new_num_points. Returns the new capacity created, + // which should be == new_num_points in the normal case. Implementers can also + // return _capacity to indicate that there are not implementing this method. + virtual location_t expand(const location_t new_num_points) = 0; + + // Shrink the datastore to new_num_points. It is NOT an error if shrink + // doesn't reduce the capacity so callers need to check this correctly. See + // also for "default" implementation + virtual location_t shrink(const location_t new_num_points) = 0; + + location_t _capacity; + size_t _dim; }; -} // namespace diskann +} // namespace diskann diff --git a/include/abstract_graph_store.h b/include/abstract_graph_store.h index f7735b79a..6f0eecdd7 100644 --- a/include/abstract_graph_store.h +++ b/include/abstract_graph_store.h @@ -8,24 +8,22 @@ #include "types.h" -namespace diskann -{ +namespace diskann { -class AbstractGraphStore -{ - public: - AbstractGraphStore(const size_t max_pts) : _capacity(max_pts) - { - } +class AbstractGraphStore { + public: + AbstractGraphStore(const size_t max_pts) : _capacity(max_pts) {} - virtual int load(const std::string &index_path_prefix) = 0; - virtual int store(const std::string &index_path_prefix) = 0; + virtual int load(const std::string &index_path_prefix) = 0; + virtual int store(const std::string &index_path_prefix) = 0; - virtual void get_adj_list(const location_t i, std::vector &neighbors) = 0; - virtual void set_adj_list(const location_t i, std::vector &neighbors) = 0; + virtual void get_adj_list(const location_t i, + std::vector &neighbors) = 0; + virtual void set_adj_list(const location_t i, + std::vector &neighbors) = 0; - private: - size_t _capacity; + private: + size_t _capacity; }; -} // namespace diskann +} // namespace diskann diff --git a/include/aligned_file_reader.h b/include/aligned_file_reader.h index f5e2af5c3..cf5d5fd78 100644 --- a/include/aligned_file_reader.h +++ b/include/aligned_file_reader.h @@ -18,11 +18,10 @@ typedef io_context_t IOContext; #include #ifndef USE_BING_INFRA -struct IOContext -{ - HANDLE fhandle = NULL; - HANDLE iocp = NULL; - std::vector reqs; +struct IOContext { + HANDLE fhandle = NULL; + HANDLE iocp = NULL; + std::vector reqs; }; #else #include "IDiskPriorityIO.h" @@ -32,32 +31,25 @@ struct IOContext // errors. // Because of such callous copying, we have to use ptr->atomic instead // of atomic, as atomic is not copyable. -struct IOContext -{ - enum Status - { - READ_WAIT = 0, - READ_SUCCESS, - READ_FAILED, - PROCESS_COMPLETE - }; - - std::shared_ptr m_pDiskIO = nullptr; - std::shared_ptr> m_pRequests; - std::shared_ptr> m_pRequestsStatus; - - // waitonaddress on this memory to wait for IO completion signal - // reader should signal this memory after IO completion - // TODO: WindowsAlignedFileReader can be modified to take advantage of this - // and can largely share code with the file reader for Bing. - mutable volatile long m_completeCount = 0; - - IOContext() - : m_pRequestsStatus(new std::vector()), m_pRequests(new std::vector()) - { - (*m_pRequestsStatus).reserve(MAX_IO_DEPTH); - (*m_pRequests).reserve(MAX_IO_DEPTH); - } +struct IOContext { + enum Status { READ_WAIT = 0, READ_SUCCESS, READ_FAILED, PROCESS_COMPLETE }; + + std::shared_ptr m_pDiskIO = nullptr; + std::shared_ptr > m_pRequests; + std::shared_ptr > m_pRequestsStatus; + + // waitonaddress on this memory to wait for IO completion signal + // reader should signal this memory after IO completion + // TODO: WindowsAlignedFileReader can be modified to take advantage of this + // and can largely share code with the file reader for Bing. + mutable volatile long m_completeCount = 0; + + IOContext() + : m_pRequestsStatus(new std::vector()), + m_pRequests(new std::vector()) { + (*m_pRequestsStatus).reserve(MAX_IO_DEPTH); + (*m_pRequests).reserve(MAX_IO_DEPTH); + } }; #endif @@ -71,50 +63,47 @@ struct IOContext #include "utils.h" // NOTE :: all 3 fields must be 512-aligned -struct AlignedRead -{ - uint64_t offset; // where to read from - uint64_t len; // how much to read - void *buf; // where to read into - - AlignedRead() : offset(0), len(0), buf(nullptr) - { - } - - AlignedRead(uint64_t offset, uint64_t len, void *buf) : offset(offset), len(len), buf(buf) - { - assert(IS_512_ALIGNED(offset)); - assert(IS_512_ALIGNED(len)); - assert(IS_512_ALIGNED(buf)); - // assert(malloc_usable_size(buf) >= len); - } +struct AlignedRead { + uint64_t offset; // where to read from + uint64_t len; // how much to read + void *buf; // where to read into + + AlignedRead() : offset(0), len(0), buf(nullptr) {} + + AlignedRead(uint64_t offset, uint64_t len, void *buf) + : offset(offset), len(len), buf(buf) { + assert(IS_512_ALIGNED(offset)); + assert(IS_512_ALIGNED(len)); + assert(IS_512_ALIGNED(buf)); + // assert(malloc_usable_size(buf) >= len); + } }; -class AlignedFileReader -{ - protected: - tsl::robin_map ctx_map; - std::mutex ctx_mut; - - public: - // returns the thread-specific context - // returns (io_context_t)(-1) if thread is not registered - virtual IOContext &get_ctx() = 0; - - virtual ~AlignedFileReader(){}; - - // register thread-id for a context - virtual void register_thread() = 0; - // 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 close() = 0; - - // process batch of aligned requests in parallel - // NOTE :: blocking call - virtual void read(std::vector &read_reqs, IOContext &ctx, bool async = false) = 0; +class AlignedFileReader { + protected: + tsl::robin_map ctx_map; + std::mutex ctx_mut; + + public: + // returns the thread-specific context + // returns (io_context_t)(-1) if thread is not registered + virtual IOContext &get_ctx() = 0; + + virtual ~AlignedFileReader(){}; + + // register thread-id for a context + virtual void register_thread() = 0; + // 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 close() = 0; + + // process batch of aligned requests in parallel + // NOTE :: blocking call + virtual void read(std::vector &read_reqs, IOContext &ctx, + bool async = false) = 0; }; diff --git a/include/boost_dynamic_bitset_fwd.h b/include/boost_dynamic_bitset_fwd.h index 5aebb2bc2..d0dd72c13 100644 --- a/include/boost_dynamic_bitset_fwd.h +++ b/include/boost_dynamic_bitset_fwd.h @@ -3,9 +3,10 @@ #pragma once -namespace boost -{ +namespace boost { #ifndef BOOST_DYNAMIC_BITSET_FWD_HPP -template > class dynamic_bitset; +template > +class dynamic_bitset; #endif -} // namespace boost +} // namespace boost diff --git a/include/cached_io.h b/include/cached_io.h index daef2f2f7..ac31283a8 100644 --- a/include/cached_io.h +++ b/include/cached_io.h @@ -11,207 +11,175 @@ #include "ann_exception.h" // sequential cached reads -class cached_ifstream -{ - public: - cached_ifstream() - { +class cached_ifstream { + public: + cached_ifstream() {} + cached_ifstream(const std::string &filename, uint64_t cacheSize) + : cache_size(cacheSize), cur_off(0) { + reader.exceptions(std::ifstream::failbit | std::ifstream::badbit); + this->open(filename, cache_size); + } + ~cached_ifstream() { + delete[] cache_buf; + reader.close(); + } + + void open(const std::string &filename, uint64_t cacheSize) { + this->cur_off = 0; + + try { + reader.open(filename, std::ios::binary | std::ios::ate); + fsize = reader.tellg(); + reader.seekg(0, std::ios::beg); + assert(reader.is_open()); + assert(cacheSize > 0); + cacheSize = (std::min)(cacheSize, fsize); + this->cache_size = cacheSize; + cache_buf = new char[cacheSize]; + reader.read(cache_buf, cacheSize); + diskann::cout << "Opened: " << filename.c_str() << ", size: " << fsize + << ", cache_size: " << cacheSize << std::endl; + } catch (std::system_error &e) { + throw diskann::FileException(filename, e, __FUNCSIG__, __FILE__, + __LINE__); } - cached_ifstream(const std::string &filename, uint64_t cacheSize) : cache_size(cacheSize), cur_off(0) - { - reader.exceptions(std::ifstream::failbit | std::ifstream::badbit); - this->open(filename, cache_size); - } - ~cached_ifstream() - { - delete[] cache_buf; - reader.close(); - } - - void open(const std::string &filename, uint64_t cacheSize) - { - this->cur_off = 0; - - try - { - reader.open(filename, std::ios::binary | std::ios::ate); - fsize = reader.tellg(); - reader.seekg(0, std::ios::beg); - assert(reader.is_open()); - assert(cacheSize > 0); - cacheSize = (std::min)(cacheSize, fsize); - this->cache_size = cacheSize; - cache_buf = new char[cacheSize]; - reader.read(cache_buf, cacheSize); - diskann::cout << "Opened: " << filename.c_str() << ", size: " << fsize << ", cache_size: " << cacheSize - << std::endl; - } - catch (std::system_error &e) - { - throw diskann::FileException(filename, e, __FUNCSIG__, __FILE__, __LINE__); - } - } - - size_t get_file_size() - { - return fsize; - } - - void read(char *read_buf, uint64_t n_bytes) - { - assert(cache_buf != nullptr); - assert(read_buf != nullptr); - - if (n_bytes <= (cache_size - cur_off)) - { - // case 1: cache contains all data - memcpy(read_buf, cache_buf + cur_off, n_bytes); - cur_off += n_bytes; - } - else - { - // case 2: cache contains some data - uint64_t cached_bytes = cache_size - cur_off; - if (n_bytes - cached_bytes > fsize - reader.tellg()) - { - std::stringstream stream; - stream << "Reading beyond end of file" << std::endl; - stream << "n_bytes: " << n_bytes << " cached_bytes: " << cached_bytes << " fsize: " << fsize - << " current pos:" << reader.tellg() << std::endl; - diskann::cout << stream.str() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - memcpy(read_buf, cache_buf + cur_off, cached_bytes); - - // go to disk and fetch more data - reader.read(read_buf + cached_bytes, n_bytes - cached_bytes); - // reset cur off - cur_off = cache_size; - - uint64_t size_left = fsize - reader.tellg(); - - if (size_left >= cache_size) - { - reader.read(cache_buf, cache_size); - cur_off = 0; - } - // note that if size_left < cache_size, then cur_off = cache_size, - // so subsequent reads will all be directly from file - } + } + + size_t get_file_size() { return fsize; } + + void read(char *read_buf, uint64_t n_bytes) { + assert(cache_buf != nullptr); + assert(read_buf != nullptr); + + if (n_bytes <= (cache_size - cur_off)) { + // case 1: cache contains all data + memcpy(read_buf, cache_buf + cur_off, n_bytes); + cur_off += n_bytes; + } else { + // case 2: cache contains some data + uint64_t cached_bytes = cache_size - cur_off; + if (n_bytes - cached_bytes > fsize - reader.tellg()) { + std::stringstream stream; + stream << "Reading beyond end of file" << std::endl; + stream << "n_bytes: " << n_bytes << " cached_bytes: " << cached_bytes + << " fsize: " << fsize << " current pos:" << reader.tellg() + << std::endl; + diskann::cout << stream.str() << std::endl; + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } + memcpy(read_buf, cache_buf + cur_off, cached_bytes); + + // go to disk and fetch more data + reader.read(read_buf + cached_bytes, n_bytes - cached_bytes); + // reset cur off + cur_off = cache_size; + + uint64_t size_left = fsize - reader.tellg(); + + if (size_left >= cache_size) { + reader.read(cache_buf, cache_size); + cur_off = 0; + } + // note that if size_left < cache_size, then cur_off = cache_size, + // so subsequent reads will all be directly from file } - - private: - // underlying ifstream - std::ifstream reader; - // # bytes to cache in one shot read - uint64_t cache_size = 0; - // underlying buf for cache - char *cache_buf = nullptr; - // offset into cache_buf for cur_pos - uint64_t cur_off = 0; - // file size - uint64_t fsize = 0; + } + + private: + // underlying ifstream + std::ifstream reader; + // # bytes to cache in one shot read + uint64_t cache_size = 0; + // underlying buf for cache + char *cache_buf = nullptr; + // offset into cache_buf for cur_pos + uint64_t cur_off = 0; + // file size + uint64_t fsize = 0; }; // sequential cached writes -class cached_ofstream -{ - public: - cached_ofstream(const std::string &filename, uint64_t cache_size) : cache_size(cache_size), cur_off(0) - { - writer.exceptions(std::ifstream::failbit | std::ifstream::badbit); - try - { - writer.open(filename, std::ios::binary); - assert(writer.is_open()); - assert(cache_size > 0); - cache_buf = new char[cache_size]; - diskann::cout << "Opened: " << filename.c_str() << ", cache_size: " << cache_size << std::endl; - } - catch (std::system_error &e) - { - throw diskann::FileException(filename, e, __FUNCSIG__, __FILE__, __LINE__); - } +class cached_ofstream { + public: + cached_ofstream(const std::string &filename, uint64_t cache_size) + : cache_size(cache_size), cur_off(0) { + writer.exceptions(std::ifstream::failbit | std::ifstream::badbit); + try { + writer.open(filename, std::ios::binary); + assert(writer.is_open()); + assert(cache_size > 0); + cache_buf = new char[cache_size]; + diskann::cout << "Opened: " << filename.c_str() + << ", cache_size: " << cache_size << std::endl; + } catch (std::system_error &e) { + throw diskann::FileException(filename, e, __FUNCSIG__, __FILE__, + __LINE__); } + } - ~cached_ofstream() - { - this->close(); - } + ~cached_ofstream() { this->close(); } - void close() - { - // dump any remaining data in memory - if (cur_off > 0) - { - this->flush_cache(); - } - - if (cache_buf != nullptr) - { - delete[] cache_buf; - cache_buf = nullptr; - } - - if (writer.is_open()) - writer.close(); - diskann::cout << "Finished writing " << fsize << "B" << std::endl; + void close() { + // dump any remaining data in memory + if (cur_off > 0) { + this->flush_cache(); } - size_t get_file_size() - { - return fsize; - } - // writes n_bytes from write_buf to the underlying ofstream/cache - void write(char *write_buf, uint64_t n_bytes) - { - assert(cache_buf != nullptr); - if (n_bytes <= (cache_size - cur_off)) - { - // case 1: cache can take all data - memcpy(cache_buf + cur_off, write_buf, n_bytes); - cur_off += n_bytes; - } - else - { - // case 2: cache cant take all data - // go to disk and write existing cache data - writer.write(cache_buf, cur_off); - fsize += cur_off; - // write the new data to disk - writer.write(write_buf, n_bytes); - fsize += n_bytes; - // memset all cache data and reset cur_off - memset(cache_buf, 0, cache_size); - cur_off = 0; - } + if (cache_buf != nullptr) { + delete[] cache_buf; + cache_buf = nullptr; } - void flush_cache() - { - assert(cache_buf != nullptr); - writer.write(cache_buf, cur_off); - fsize += cur_off; - memset(cache_buf, 0, cache_size); - cur_off = 0; + if (writer.is_open()) writer.close(); + diskann::cout << "Finished writing " << fsize << "B" << std::endl; + } + + size_t get_file_size() { return fsize; } + // writes n_bytes from write_buf to the underlying ofstream/cache + void write(char *write_buf, uint64_t n_bytes) { + assert(cache_buf != nullptr); + if (n_bytes <= (cache_size - cur_off)) { + // case 1: cache can take all data + memcpy(cache_buf + cur_off, write_buf, n_bytes); + cur_off += n_bytes; + } else { + // case 2: cache cant take all data + // go to disk and write existing cache data + writer.write(cache_buf, cur_off); + fsize += cur_off; + // write the new data to disk + writer.write(write_buf, n_bytes); + fsize += n_bytes; + // memset all cache data and reset cur_off + memset(cache_buf, 0, cache_size); + cur_off = 0; } - - void reset() - { - flush_cache(); - writer.seekp(0); - } - - private: - // underlying ofstream - std::ofstream writer; - // # bytes to cache for one shot write - uint64_t cache_size = 0; - // underlying buf for cache - char *cache_buf = nullptr; - // offset into cache_buf for cur_pos - uint64_t cur_off = 0; - - // file size - uint64_t fsize = 0; + } + + void flush_cache() { + assert(cache_buf != nullptr); + writer.write(cache_buf, cur_off); + fsize += cur_off; + memset(cache_buf, 0, cache_size); + cur_off = 0; + } + + void reset() { + flush_cache(); + writer.seekp(0); + } + + private: + // underlying ofstream + std::ofstream writer; + // # bytes to cache for one shot write + uint64_t cache_size = 0; + // underlying buf for cache + char *cache_buf = nullptr; + // offset into cache_buf for cur_pos + uint64_t cur_off = 0; + + // file size + uint64_t fsize = 0; }; diff --git a/include/concurrent_queue.h b/include/concurrent_queue.h index 1e57bbf0f..a375a7466 100644 --- a/include/concurrent_queue.h +++ b/include/concurrent_queue.h @@ -11,122 +11,91 @@ #include #include -namespace diskann -{ +namespace diskann { -template class ConcurrentQueue -{ - typedef std::chrono::microseconds chrono_us_t; - typedef std::unique_lock mutex_locker; +template +class ConcurrentQueue { + typedef std::chrono::microseconds chrono_us_t; + typedef std::unique_lock mutex_locker; - std::queue q; - std::mutex mut; - std::mutex push_mut; - std::mutex pop_mut; - std::condition_variable push_cv; - std::condition_variable pop_cv; - T null_T; + std::queue q; + std::mutex mut; + std::mutex push_mut; + std::mutex pop_mut; + std::condition_variable push_cv; + std::condition_variable pop_cv; + T null_T; - public: - ConcurrentQueue() - { - } + public: + ConcurrentQueue() {} - ConcurrentQueue(T nullT) - { - this->null_T = nullT; - } + ConcurrentQueue(T nullT) { this->null_T = nullT; } - ~ConcurrentQueue() - { - this->push_cv.notify_all(); - this->pop_cv.notify_all(); - } + ~ConcurrentQueue() { + this->push_cv.notify_all(); + this->pop_cv.notify_all(); + } - // queue stats - uint64_t size() - { - mutex_locker lk(this->mut); - uint64_t ret = q.size(); - lk.unlock(); - return ret; - } + // queue stats + uint64_t size() { + mutex_locker lk(this->mut); + uint64_t ret = q.size(); + lk.unlock(); + return ret; + } - bool empty() - { - return (this->size() == 0); - } + bool empty() { return (this->size() == 0); } - // PUSH BACK - void push(T &new_val) - { - mutex_locker lk(this->mut); - this->q.push(new_val); - lk.unlock(); - } + // PUSH BACK + void push(T &new_val) { + mutex_locker lk(this->mut); + this->q.push(new_val); + lk.unlock(); + } - template void insert(Iterator iter_begin, Iterator iter_end) - { - mutex_locker lk(this->mut); - for (Iterator it = iter_begin; it != iter_end; it++) - { - this->q.push(*it); - } - lk.unlock(); + template + void insert(Iterator iter_begin, Iterator iter_end) { + mutex_locker lk(this->mut); + for (Iterator it = iter_begin; it != iter_end; it++) { + this->q.push(*it); } + lk.unlock(); + } - // POP FRONT - T pop() - { - mutex_locker lk(this->mut); - if (this->q.empty()) - { - lk.unlock(); - return this->null_T; - } - else - { - T ret = this->q.front(); - this->q.pop(); - // diskann::cout << "thread_id: " << std::this_thread::get_id() << - // ", ctx: " - // << ret.ctx << "\n"; - lk.unlock(); - return ret; - } + // POP FRONT + T pop() { + mutex_locker lk(this->mut); + if (this->q.empty()) { + lk.unlock(); + return this->null_T; + } else { + T ret = this->q.front(); + this->q.pop(); + // diskann::cout << "thread_id: " << std::this_thread::get_id() << + // ", ctx: " + // << ret.ctx << "\n"; + lk.unlock(); + return ret; } + } - // register for notifications - void wait_for_push_notify(chrono_us_t wait_time = chrono_us_t{10}) - { - mutex_locker lk(this->push_mut); - this->push_cv.wait_for(lk, wait_time); - lk.unlock(); - } + // register for notifications + void wait_for_push_notify(chrono_us_t wait_time = chrono_us_t{10}) { + mutex_locker lk(this->push_mut); + this->push_cv.wait_for(lk, wait_time); + lk.unlock(); + } - void wait_for_pop_notify(chrono_us_t wait_time = chrono_us_t{10}) - { - mutex_locker lk(this->pop_mut); - this->pop_cv.wait_for(lk, wait_time); - lk.unlock(); - } + void wait_for_pop_notify(chrono_us_t wait_time = chrono_us_t{10}) { + mutex_locker lk(this->pop_mut); + this->pop_cv.wait_for(lk, wait_time); + lk.unlock(); + } - // just notify functions - void push_notify_one() - { - this->push_cv.notify_one(); - } - void push_notify_all() - { - this->push_cv.notify_all(); - } - void pop_notify_one() - { - this->pop_cv.notify_one(); - } - void pop_notify_all() - { - this->pop_cv.notify_all(); - } + // just notify functions + void push_notify_one() { this->push_cv.notify_one(); } + void push_notify_all() { this->push_cv.notify_all(); } + void pop_notify_one() { this->pop_cv.notify_one(); } + void pop_notify_all() { this->pop_cv.notify_all(); } }; -} // namespace diskann +} // namespace diskann diff --git a/include/cosine_similarity.h b/include/cosine_similarity.h index dc51f6c0a..5c2ec3d50 100644 --- a/include/cosine_similarity.h +++ b/include/cosine_similarity.h @@ -38,202 +38,194 @@ extern bool Avx2SupportedCPU; * */ -namespace diskann -{ +namespace diskann { using namespace std; #define PORTABLE_ALIGN16 __declspec(align(16)) -static float NormScalarProductSIMD2(const int8_t *pVect1, const int8_t *pVect2, uint32_t qty) -{ - if (Avx2SupportedCPU) - { - __m256 cos, p1Len, p2Len; - cos = p1Len = p2Len = _mm256_setzero_ps(); - while (qty >= 32) - { - __m256i rx = _mm256_load_si256((__m256i *)pVect1), ry = _mm256_load_si256((__m256i *)pVect2); - cos = _mm256_add_ps(cos, _mm256_mul_epi8(rx, ry)); - p1Len = _mm256_add_ps(p1Len, _mm256_mul_epi8(rx, rx)); - p2Len = _mm256_add_ps(p2Len, _mm256_mul_epi8(ry, ry)); - pVect1 += 32; - pVect2 += 32; - qty -= 32; - } - while (qty > 0) - { - __m128i rx = _mm_load_si128((__m128i *)pVect1), ry = _mm_load_si128((__m128i *)pVect2); - cos = _mm256_add_ps(cos, _mm256_mul32_pi8(rx, ry)); - p1Len = _mm256_add_ps(p1Len, _mm256_mul32_pi8(rx, rx)); - p2Len = _mm256_add_ps(p2Len, _mm256_mul32_pi8(ry, ry)); - pVect1 += 4; - pVect2 += 4; - qty -= 4; - } - cos = _mm256_hadd_ps(_mm256_hadd_ps(cos, cos), cos); - p1Len = _mm256_hadd_ps(_mm256_hadd_ps(p1Len, p1Len), p1Len); - p2Len = _mm256_hadd_ps(_mm256_hadd_ps(p2Len, p2Len), p2Len); - float denominator = max(numeric_limits::min() * 2, sqrt(p1Len.m256_f32[0] + p1Len.m256_f32[4]) * - sqrt(p2Len.m256_f32[0] + p2Len.m256_f32[4])); - float cosine = (cos.m256_f32[0] + cos.m256_f32[4]) / denominator; - - return max(float(-1), min(float(1), cosine)); +static float NormScalarProductSIMD2(const int8_t *pVect1, const int8_t *pVect2, + uint32_t qty) { + if (Avx2SupportedCPU) { + __m256 cos, p1Len, p2Len; + cos = p1Len = p2Len = _mm256_setzero_ps(); + while (qty >= 32) { + __m256i rx = _mm256_load_si256((__m256i *)pVect1), + ry = _mm256_load_si256((__m256i *)pVect2); + cos = _mm256_add_ps(cos, _mm256_mul_epi8(rx, ry)); + p1Len = _mm256_add_ps(p1Len, _mm256_mul_epi8(rx, rx)); + p2Len = _mm256_add_ps(p2Len, _mm256_mul_epi8(ry, ry)); + pVect1 += 32; + pVect2 += 32; + qty -= 32; } - - __m128 cos, p1Len, p2Len; - cos = p1Len = p2Len = _mm_setzero_ps(); - __m128i rx, ry; - while (qty >= 16) - { - rx = _mm_load_si128((__m128i *)pVect1); - ry = _mm_load_si128((__m128i *)pVect2); - cos = _mm_add_ps(cos, _mm_mul_epi8(rx, ry)); - p1Len = _mm_add_ps(p1Len, _mm_mul_epi8(rx, rx)); - p2Len = _mm_add_ps(p2Len, _mm_mul_epi8(ry, ry)); - pVect1 += 16; - pVect2 += 16; - qty -= 16; - } - while (qty > 0) - { - rx = _mm_load_si128((__m128i *)pVect1); - ry = _mm_load_si128((__m128i *)pVect2); - cos = _mm_add_ps(cos, _mm_mul32_pi8(rx, ry)); - p1Len = _mm_add_ps(p1Len, _mm_mul32_pi8(rx, rx)); - p2Len = _mm_add_ps(p2Len, _mm_mul32_pi8(ry, ry)); - pVect1 += 4; - pVect2 += 4; - qty -= 4; + while (qty > 0) { + __m128i rx = _mm_load_si128((__m128i *)pVect1), + ry = _mm_load_si128((__m128i *)pVect2); + cos = _mm256_add_ps(cos, _mm256_mul32_pi8(rx, ry)); + p1Len = _mm256_add_ps(p1Len, _mm256_mul32_pi8(rx, rx)); + p2Len = _mm256_add_ps(p2Len, _mm256_mul32_pi8(ry, ry)); + pVect1 += 4; + pVect2 += 4; + qty -= 4; } - cos = _mm_hadd_ps(_mm_hadd_ps(cos, cos), cos); - p1Len = _mm_hadd_ps(_mm_hadd_ps(p1Len, p1Len), p1Len); - p2Len = _mm_hadd_ps(_mm_hadd_ps(p2Len, p2Len), p2Len); - float norm1 = p1Len.m128_f32[0]; - float norm2 = p2Len.m128_f32[0]; - - static const float eps = numeric_limits::min() * 2; - - if (norm1 < eps) - { /* - * This shouldn't normally happen for this space, but - * if it does, we don't want to get NANs - */ - if (norm2 < eps) - { - return 1; - } - return 0; + cos = _mm256_hadd_ps(_mm256_hadd_ps(cos, cos), cos); + p1Len = _mm256_hadd_ps(_mm256_hadd_ps(p1Len, p1Len), p1Len); + p2Len = _mm256_hadd_ps(_mm256_hadd_ps(p2Len, p2Len), p2Len); + float denominator = max(numeric_limits::min() * 2, + sqrt(p1Len.m256_f32[0] + p1Len.m256_f32[4]) * + sqrt(p2Len.m256_f32[0] + p2Len.m256_f32[4])); + float cosine = (cos.m256_f32[0] + cos.m256_f32[4]) / denominator; + + return max(float(-1), min(float(1), cosine)); + } + + __m128 cos, p1Len, p2Len; + cos = p1Len = p2Len = _mm_setzero_ps(); + __m128i rx, ry; + while (qty >= 16) { + rx = _mm_load_si128((__m128i *)pVect1); + ry = _mm_load_si128((__m128i *)pVect2); + cos = _mm_add_ps(cos, _mm_mul_epi8(rx, ry)); + p1Len = _mm_add_ps(p1Len, _mm_mul_epi8(rx, rx)); + p2Len = _mm_add_ps(p2Len, _mm_mul_epi8(ry, ry)); + pVect1 += 16; + pVect2 += 16; + qty -= 16; + } + while (qty > 0) { + rx = _mm_load_si128((__m128i *)pVect1); + ry = _mm_load_si128((__m128i *)pVect2); + cos = _mm_add_ps(cos, _mm_mul32_pi8(rx, ry)); + p1Len = _mm_add_ps(p1Len, _mm_mul32_pi8(rx, rx)); + p2Len = _mm_add_ps(p2Len, _mm_mul32_pi8(ry, ry)); + pVect1 += 4; + pVect2 += 4; + qty -= 4; + } + cos = _mm_hadd_ps(_mm_hadd_ps(cos, cos), cos); + p1Len = _mm_hadd_ps(_mm_hadd_ps(p1Len, p1Len), p1Len); + p2Len = _mm_hadd_ps(_mm_hadd_ps(p2Len, p2Len), p2Len); + float norm1 = p1Len.m128_f32[0]; + float norm2 = p2Len.m128_f32[0]; + + static const float eps = numeric_limits::min() * 2; + + if (norm1 < eps) { /* + * This shouldn't normally happen for this space, but + * if it does, we don't want to get NANs + */ + if (norm2 < eps) { + return 1; } - /* - * Sometimes due to rounding errors, we get values > 1 or < -1. - * This throws off other functions that use scalar product, e.g., acos - */ - return max(float(-1), min(float(1), cos.m128_f32[0] / sqrt(norm1) / sqrt(norm2))); + return 0; + } + /* + * Sometimes due to rounding errors, we get values > 1 or < -1. + * This throws off other functions that use scalar product, e.g., acos + */ + return max(float(-1), + min(float(1), cos.m128_f32[0] / sqrt(norm1) / sqrt(norm2))); } -static float NormScalarProductSIMD(const float *pVect1, const float *pVect2, uint32_t qty) -{ - // Didn't get significant performance gain compared with 128bit version. - static const float eps = numeric_limits::min() * 2; - - if (Avx2SupportedCPU) - { - uint32_t qty8 = qty / 8; - - const float *pEnd1 = pVect1 + 8 * qty8; - const float *pEnd2 = pVect1 + qty; - - __m256 v1, v2; - __m256 sum_prod = _mm256_set_ps(0, 0, 0, 0, 0, 0, 0, 0); - __m256 sum_square1 = sum_prod; - __m256 sum_square2 = sum_prod; - - while (pVect1 < pEnd1) - { - v1 = _mm256_loadu_ps(pVect1); - pVect1 += 8; - v2 = _mm256_loadu_ps(pVect2); - pVect2 += 8; - sum_prod = _mm256_add_ps(sum_prod, _mm256_mul_ps(v1, v2)); - sum_square1 = _mm256_add_ps(sum_square1, _mm256_mul_ps(v1, v1)); - sum_square2 = _mm256_add_ps(sum_square2, _mm256_mul_ps(v2, v2)); - } - - float PORTABLE_ALIGN16 TmpResProd[8]; - float PORTABLE_ALIGN16 TmpResSquare1[8]; - float PORTABLE_ALIGN16 TmpResSquare2[8]; - - _mm256_store_ps(TmpResProd, sum_prod); - _mm256_store_ps(TmpResSquare1, sum_square1); - _mm256_store_ps(TmpResSquare2, sum_square2); - - float sum = 0.0f; - float norm1 = 0.0f; - float norm2 = 0.0f; - for (uint32_t i = 0; i < 8; ++i) - { - sum += TmpResProd[i]; - norm1 += TmpResSquare1[i]; - norm2 += TmpResSquare2[i]; - } - - while (pVect1 < pEnd2) - { - sum += (*pVect1) * (*pVect2); - norm1 += (*pVect1) * (*pVect1); - norm2 += (*pVect2) * (*pVect2); - - ++pVect1; - ++pVect2; - } - - if (norm1 < eps) - { - return norm2 < eps ? 1.0f : 0.0f; - } - - return max(float(-1), min(float(1), sum / sqrt(norm1) / sqrt(norm2))); +static float NormScalarProductSIMD(const float *pVect1, const float *pVect2, + uint32_t qty) { + // Didn't get significant performance gain compared with 128bit version. + static const float eps = numeric_limits::min() * 2; + + if (Avx2SupportedCPU) { + uint32_t qty8 = qty / 8; + + const float *pEnd1 = pVect1 + 8 * qty8; + const float *pEnd2 = pVect1 + qty; + + __m256 v1, v2; + __m256 sum_prod = _mm256_set_ps(0, 0, 0, 0, 0, 0, 0, 0); + __m256 sum_square1 = sum_prod; + __m256 sum_square2 = sum_prod; + + while (pVect1 < pEnd1) { + v1 = _mm256_loadu_ps(pVect1); + pVect1 += 8; + v2 = _mm256_loadu_ps(pVect2); + pVect2 += 8; + sum_prod = _mm256_add_ps(sum_prod, _mm256_mul_ps(v1, v2)); + sum_square1 = _mm256_add_ps(sum_square1, _mm256_mul_ps(v1, v1)); + sum_square2 = _mm256_add_ps(sum_square2, _mm256_mul_ps(v2, v2)); } - __m128 v1, v2; - __m128 sum_prod = _mm_set1_ps(0); - __m128 sum_square1 = sum_prod; - __m128 sum_square2 = sum_prod; - - while (qty >= 4) - { - v1 = _mm_loadu_ps(pVect1); - pVect1 += 4; - v2 = _mm_loadu_ps(pVect2); - pVect2 += 4; - sum_prod = _mm_add_ps(sum_prod, _mm_mul_ps(v1, v2)); - sum_square1 = _mm_add_ps(sum_square1, _mm_mul_ps(v1, v1)); - sum_square2 = _mm_add_ps(sum_square2, _mm_mul_ps(v2, v2)); - - qty -= 4; + float PORTABLE_ALIGN16 TmpResProd[8]; + float PORTABLE_ALIGN16 TmpResSquare1[8]; + float PORTABLE_ALIGN16 TmpResSquare2[8]; + + _mm256_store_ps(TmpResProd, sum_prod); + _mm256_store_ps(TmpResSquare1, sum_square1); + _mm256_store_ps(TmpResSquare2, sum_square2); + + float sum = 0.0f; + float norm1 = 0.0f; + float norm2 = 0.0f; + for (uint32_t i = 0; i < 8; ++i) { + sum += TmpResProd[i]; + norm1 += TmpResSquare1[i]; + norm2 += TmpResSquare2[i]; } - float sum = sum_prod.m128_f32[0] + sum_prod.m128_f32[1] + sum_prod.m128_f32[2] + sum_prod.m128_f32[3]; - float norm1 = sum_square1.m128_f32[0] + sum_square1.m128_f32[1] + sum_square1.m128_f32[2] + sum_square1.m128_f32[3]; - float norm2 = sum_square2.m128_f32[0] + sum_square2.m128_f32[1] + sum_square2.m128_f32[2] + sum_square2.m128_f32[3]; + while (pVect1 < pEnd2) { + sum += (*pVect1) * (*pVect2); + norm1 += (*pVect1) * (*pVect1); + norm2 += (*pVect2) * (*pVect2); + + ++pVect1; + ++pVect2; + } - if (norm1 < eps) - { - return norm2 < eps ? 1.0f : 0.0f; + if (norm1 < eps) { + return norm2 < eps ? 1.0f : 0.0f; } return max(float(-1), min(float(1), sum / sqrt(norm1) / sqrt(norm2))); + } + + __m128 v1, v2; + __m128 sum_prod = _mm_set1_ps(0); + __m128 sum_square1 = sum_prod; + __m128 sum_square2 = sum_prod; + + while (qty >= 4) { + v1 = _mm_loadu_ps(pVect1); + pVect1 += 4; + v2 = _mm_loadu_ps(pVect2); + pVect2 += 4; + sum_prod = _mm_add_ps(sum_prod, _mm_mul_ps(v1, v2)); + sum_square1 = _mm_add_ps(sum_square1, _mm_mul_ps(v1, v1)); + sum_square2 = _mm_add_ps(sum_square2, _mm_mul_ps(v2, v2)); + + qty -= 4; + } + + float sum = sum_prod.m128_f32[0] + sum_prod.m128_f32[1] + + sum_prod.m128_f32[2] + sum_prod.m128_f32[3]; + float norm1 = sum_square1.m128_f32[0] + sum_square1.m128_f32[1] + + sum_square1.m128_f32[2] + sum_square1.m128_f32[3]; + float norm2 = sum_square2.m128_f32[0] + sum_square2.m128_f32[1] + + sum_square2.m128_f32[2] + sum_square2.m128_f32[3]; + + if (norm1 < eps) { + return norm2 < eps ? 1.0f : 0.0f; + } + + return max(float(-1), min(float(1), sum / sqrt(norm1) / sqrt(norm2))); } -static float NormScalarProductSIMD2(const float *pVect1, const float *pVect2, uint32_t qty) -{ - return NormScalarProductSIMD(pVect1, pVect2, qty); +static float NormScalarProductSIMD2(const float *pVect1, const float *pVect2, + uint32_t qty) { + return NormScalarProductSIMD(pVect1, pVect2, qty); } -template static float CosineSimilarity2(const T *p1, const T *p2, uint32_t qty) -{ - return std::max(0.0f, 1.0f - NormScalarProductSIMD2(p1, p2, qty)); +template +static float CosineSimilarity2(const T *p1, const T *p2, uint32_t qty) { + return std::max(0.0f, 1.0f - NormScalarProductSIMD2(p1, p2, qty)); } // static template float CosineSimilarity2<__int8>(const __int8* pVect1, @@ -242,22 +234,19 @@ template static float CosineSimilarity2(const T *p1, const T *p2, uint // static template float CosineSimilarity2(const float* pVect1, // const float* pVect2, size_t qty); -template static void CosineSimilarityNormalize(T *pVector, uint32_t qty) -{ - T sum = 0; - for (uint32_t i = 0; i < qty; ++i) - { - sum += pVector[i] * pVector[i]; - } - sum = 1 / sqrt(sum); - if (sum == 0) - { - sum = numeric_limits::min(); - } - for (uint32_t i = 0; i < qty; ++i) - { - pVector[i] *= sum; - } +template +static void CosineSimilarityNormalize(T *pVector, uint32_t qty) { + T sum = 0; + for (uint32_t i = 0; i < qty; ++i) { + sum += pVector[i] * pVector[i]; + } + sum = 1 / sqrt(sum); + if (sum == 0) { + sum = numeric_limits::min(); + } + for (uint32_t i = 0; i < qty; ++i) { + pVector[i] *= sum; + } } // template static void CosineSimilarityNormalize(float* pVector, @@ -265,19 +254,22 @@ template static void CosineSimilarityNormalize(T *pVector, uint32_t qt // template static void CosineSimilarityNormalize(double* pVector, // size_t qty); -template <> void CosineSimilarityNormalize(__int8 * /*pVector*/, uint32_t /*qty*/) -{ - throw std::runtime_error("For int8 type vector, you can not use cosine distance!"); +template <> +void CosineSimilarityNormalize(__int8 * /*pVector*/, uint32_t /*qty*/) { + throw std::runtime_error( + "For int8 type vector, you can not use cosine distance!"); } -template <> void CosineSimilarityNormalize(__int16 * /*pVector*/, uint32_t /*qty*/) -{ - throw std::runtime_error("For int16 type vector, you can not use cosine distance!"); +template <> +void CosineSimilarityNormalize(__int16 * /*pVector*/, uint32_t /*qty*/) { + throw std::runtime_error( + "For int16 type vector, you can not use cosine distance!"); } -template <> void CosineSimilarityNormalize(int * /*pVector*/, uint32_t /*qty*/) -{ - throw std::runtime_error("For int type vector, you can not use cosine distance!"); +template <> +void CosineSimilarityNormalize(int * /*pVector*/, uint32_t /*qty*/) { + throw std::runtime_error( + "For int type vector, you can not use cosine distance!"); } -} // namespace diskann +} // namespace diskann #endif \ No newline at end of file diff --git a/include/defaults.h b/include/defaults.h index 6b14bc2b3..87e9323e8 100644 --- a/include/defaults.h +++ b/include/defaults.h @@ -4,10 +4,8 @@ #pragma once #include -namespace diskann -{ -namespace defaults -{ +namespace diskann { +namespace defaults { const float ALPHA = 1.2f; const uint32_t NUM_THREADS = 0; const uint32_t NUM_ROUNDS = 2; @@ -20,5 +18,5 @@ const uint32_t MAX_DEGREE = 64; const uint32_t BUILD_LIST_SIZE = 100; const uint32_t SATURATE_GRAPH = false; const uint32_t SEARCH_LIST_SIZE = 100; -} // namespace defaults -} // namespace diskann +} // namespace defaults +} // namespace diskann diff --git a/include/disk_utils.h b/include/disk_utils.h index 08f046dcd..65d0a51aa 100644 --- a/include/disk_utils.h +++ b/include/disk_utils.h @@ -31,8 +31,7 @@ typedef int FileHandle; #include "utils.h" #include "windows_customizations.h" -namespace diskann -{ +namespace diskann { const size_t MAX_SAMPLE_POINTS_FOR_WARMUP = 100000; const double PQ_TRAINING_SET_FRACTION = 0.1; const double SPACE_FOR_CACHED_NODES_IN_GB = 0.25; @@ -41,68 +40,87 @@ const uint32_t NUM_NODES_TO_CACHE = 250000; const uint32_t WARMUP_L = 20; const uint32_t NUM_KMEANS_REPS = 12; -template class PQFlashIndex; +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 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 size_t calculate_num_pq_chunks(double final_index_ram_limit, + size_t points_num, + uint32_t dim); -DISKANN_DLLEXPORT void read_idmap(const std::string &fname, std::vector &ivecs); +DISKANN_DLLEXPORT void read_idmap(const std::string &fname, + std::vector &ivecs); #ifdef EXEC_ENV_OLS template -DISKANN_DLLEXPORT T *load_warmup(MemoryMappedFiles &files, const std::string &cache_warmup_file, uint64_t &warmup_num, - uint64_t warmup_dim, uint64_t warmup_aligned_dim); +DISKANN_DLLEXPORT T *load_warmup(MemoryMappedFiles &files, + const std::string &cache_warmup_file, + uint64_t &warmup_num, uint64_t warmup_dim, + uint64_t warmup_aligned_dim); #else template -DISKANN_DLLEXPORT T *load_warmup(const std::string &cache_warmup_file, uint64_t &warmup_num, uint64_t warmup_dim, +DISKANN_DLLEXPORT T *load_warmup(const std::string &cache_warmup_file, + uint64_t &warmup_num, uint64_t warmup_dim, uint64_t warmup_aligned_dim); #endif -DISKANN_DLLEXPORT int merge_shards(const std::string &vamana_prefix, const std::string &vamana_suffix, - const std::string &idmaps_prefix, const std::string &idmaps_suffix, - const uint64_t nshards, uint32_t max_degree, const std::string &output_vamana, - const std::string &medoids_file, bool use_filters = false, - const std::string &labels_to_medoids_file = std::string("")); +DISKANN_DLLEXPORT int merge_shards( + const std::string &vamana_prefix, const std::string &vamana_suffix, + const std::string &idmaps_prefix, const std::string &idmaps_suffix, + const uint64_t nshards, uint32_t max_degree, + const std::string &output_vamana, const std::string &medoids_file, + bool use_filters = false, + const std::string &labels_to_medoids_file = std::string("")); -DISKANN_DLLEXPORT void extract_shard_labels(const std::string &in_label_file, const std::string &shard_ids_bin, - const std::string &shard_label_file); +DISKANN_DLLEXPORT void extract_shard_labels( + const std::string &in_label_file, const std::string &shard_ids_bin, + const std::string &shard_label_file); template -DISKANN_DLLEXPORT std::string preprocess_base_file(const std::string &infile, const std::string &indexPrefix, - diskann::Metric &distMetric); +DISKANN_DLLEXPORT std::string preprocess_base_file( + const std::string &infile, const std::string &indexPrefix, + diskann::Metric &distMetric); template -DISKANN_DLLEXPORT int build_merged_vamana_index(std::string base_file, diskann::Metric _compareMetric, uint32_t L, - uint32_t R, double sampling_rate, double ram_budget, - std::string mem_index_path, std::string medoids_file, - std::string centroids_file, size_t build_pq_bytes, bool use_opq, - uint32_t num_threads, bool use_filters = false, - const std::string &label_file = std::string(""), - const std::string &labels_to_medoids_file = std::string(""), - const std::string &universal_label = "", const uint32_t Lf = 0); +DISKANN_DLLEXPORT int build_merged_vamana_index( + std::string base_file, diskann::Metric _compareMetric, uint32_t L, + uint32_t R, double sampling_rate, double ram_budget, + std::string mem_index_path, std::string medoids_file, + std::string centroids_file, size_t build_pq_bytes, bool use_opq, + uint32_t num_threads, bool use_filters = false, + const std::string &label_file = std::string(""), + const std::string &labels_to_medoids_file = std::string(""), + const std::string &universal_label = "", const uint32_t Lf = 0); template -DISKANN_DLLEXPORT uint32_t optimize_beamwidth(std::unique_ptr> &_pFlashIndex, - T *tuning_sample, uint64_t tuning_sample_num, - uint64_t tuning_sample_aligned_dim, uint32_t L, uint32_t nthreads, - uint32_t start_bw = 2); +DISKANN_DLLEXPORT uint32_t optimize_beamwidth( + std::unique_ptr> &_pFlashIndex, + T *tuning_sample, uint64_t tuning_sample_num, + uint64_t tuning_sample_aligned_dim, uint32_t L, uint32_t nthreads, + uint32_t start_bw = 2); template DISKANN_DLLEXPORT int build_disk_index( - const char *dataFilePath, const char *indexFilePath, const char *indexBuildParameters, - diskann::Metric _compareMetric, bool use_opq = false, - const std::string &codebook_prefix = "", // default is empty for no codebook pass in + const char *dataFilePath, const char *indexFilePath, + const char *indexBuildParameters, diskann::Metric _compareMetric, + bool use_opq = false, + const std::string &codebook_prefix = + "", // default is empty for no codebook pass in bool use_filters = false, - const std::string &label_file = std::string(""), // default is empty string for no label_file - const std::string &universal_label = "", const uint32_t filter_threshold = 0, - const uint32_t Lf = 0); // default is empty string for no universal label + const std::string &label_file = + std::string(""), // default is empty string for no label_file + const std::string &universal_label = "", + const uint32_t filter_threshold = 0, + const uint32_t Lf = 0); // default is empty string for no universal label template -DISKANN_DLLEXPORT void create_disk_layout(const std::string base_file, const std::string mem_index_file, - const std::string output_file, - const std::string reorder_data_file = std::string("")); +DISKANN_DLLEXPORT void create_disk_layout( + const std::string base_file, const std::string mem_index_file, + const std::string output_file, + const std::string reorder_data_file = std::string("")); -} // namespace diskann +} // namespace diskann diff --git a/include/distance.h b/include/distance.h index da820a4f1..b76a3ec30 100644 --- a/include/distance.h +++ b/include/distance.h @@ -2,234 +2,210 @@ #include "windows_customizations.h" #include -namespace diskann -{ -enum Metric -{ - L2 = 0, - INNER_PRODUCT = 1, - COSINE = 2, - FAST_L2 = 3 +namespace diskann { +enum Metric { L2 = 0, INNER_PRODUCT = 1, COSINE = 2, FAST_L2 = 3 }; + +template +class Distance { + public: + DISKANN_DLLEXPORT Distance(diskann::Metric dist_metric) + : _distance_metric(dist_metric) {} + + // distance comparison function + DISKANN_DLLEXPORT virtual float compare(const T *a, const T *b, + uint32_t length) const = 0; + + // Needed only for COSINE-BYTE and INNER_PRODUCT-BYTE + DISKANN_DLLEXPORT virtual float compare(const T *a, const T *b, + const float normA, const float normB, + uint32_t length) const; + + // For MIPS, normalization adds an extra dimension to the vectors. + // This function lets callers know if the normalization process + // changes the dimension. + DISKANN_DLLEXPORT virtual uint32_t post_normalization_dimension( + uint32_t orig_dimension) const; + + DISKANN_DLLEXPORT virtual diskann::Metric get_metric() const; + + // This is for efficiency. If no normalization is required, the callers + // can simply ignore the normalize_data_for_build() function. + DISKANN_DLLEXPORT virtual bool preprocessing_required() const; + + // Check the preprocessing_required() function before calling this. + // Clients can call the function like this: + // + // if (metric->preprocessing_required()){ + // T* normalized_data_batch; + // Split data into batches of batch_size and for each, call: + // metric->preprocess_base_points(data_batch, batch_size); + // + // TODO: This does not take into account the case for SSD inner product + // where the dimensions change after normalization. + DISKANN_DLLEXPORT virtual void preprocess_base_points( + T *original_data, const size_t orig_dim, const size_t num_points); + + // Invokes normalization for a single vector during search. The scratch space + // has to be created by the caller keeping track of the fact that + // normalization might change the dimension of the query vector. + DISKANN_DLLEXPORT virtual void preprocess_query(const T *query_vec, + const size_t query_dim, + T *scratch_query); + + // If an algorithm has a requirement that some data be aligned to a certain + // boundary it can use this function to indicate that requirement. Currently, + // we are setting it to 8 because that works well for AVX2. If we have AVX512 + // implementations of distance algos, they might have to set this to 16 + // (depending on how they are implemented) + DISKANN_DLLEXPORT virtual size_t get_required_alignment() const; + + // Providing a default implementation for the virtual destructor because we + // don't expect most metric implementations to need it. + DISKANN_DLLEXPORT virtual ~Distance(); + + protected: + diskann::Metric _distance_metric; + size_t _alignment_factor = 8; }; -template class Distance -{ - public: - DISKANN_DLLEXPORT Distance(diskann::Metric dist_metric) : _distance_metric(dist_metric) - { - } - - // distance comparison function - DISKANN_DLLEXPORT virtual float compare(const T *a, const T *b, uint32_t length) const = 0; - - // Needed only for COSINE-BYTE and INNER_PRODUCT-BYTE - DISKANN_DLLEXPORT virtual float compare(const T *a, const T *b, const float normA, const float normB, - uint32_t length) const; - - // For MIPS, normalization adds an extra dimension to the vectors. - // This function lets callers know if the normalization process - // changes the dimension. - DISKANN_DLLEXPORT virtual uint32_t post_normalization_dimension(uint32_t orig_dimension) const; - - DISKANN_DLLEXPORT virtual diskann::Metric get_metric() const; - - // This is for efficiency. If no normalization is required, the callers - // can simply ignore the normalize_data_for_build() function. - DISKANN_DLLEXPORT virtual bool preprocessing_required() const; - - // Check the preprocessing_required() function before calling this. - // Clients can call the function like this: - // - // if (metric->preprocessing_required()){ - // T* normalized_data_batch; - // Split data into batches of batch_size and for each, call: - // metric->preprocess_base_points(data_batch, batch_size); - // - // TODO: This does not take into account the case for SSD inner product - // where the dimensions change after normalization. - DISKANN_DLLEXPORT virtual void preprocess_base_points(T *original_data, const size_t orig_dim, - const size_t num_points); - - // Invokes normalization for a single vector during search. The scratch space - // has to be created by the caller keeping track of the fact that normalization - // might change the dimension of the query vector. - DISKANN_DLLEXPORT virtual void preprocess_query(const T *query_vec, const size_t query_dim, - T *scratch_query); - - //If an algorithm has a requirement that some data be aligned to a certain boundary - //it can use this function to indicate that requirement. Currently, we are setting it to 8 - //because that works well for AVX2. If we have AVX512 implementations of distance algos, - //they might have to set this to 16 (depending on how they are implemented) - DISKANN_DLLEXPORT virtual size_t get_required_alignment() const; - - // Providing a default implementation for the virtual destructor because we don't - // expect most metric implementations to need it. - DISKANN_DLLEXPORT virtual ~Distance(); - - protected: - diskann::Metric _distance_metric; - size_t _alignment_factor = 8; +class DistanceCosineInt8 : public Distance { + public: + DistanceCosineInt8() : Distance(diskann::Metric::COSINE) {} + DISKANN_DLLEXPORT virtual float compare(const int8_t *a, const int8_t *b, + uint32_t length) const; }; -class DistanceCosineInt8 : public Distance -{ - public: - DistanceCosineInt8() : Distance(diskann::Metric::COSINE) - { - } - DISKANN_DLLEXPORT virtual float compare(const int8_t *a, const int8_t *b, uint32_t length) const; -}; - -class DistanceL2Int8 : public Distance -{ - public: - DistanceL2Int8() : Distance(diskann::Metric::L2) - { - } - DISKANN_DLLEXPORT virtual float compare(const int8_t *a, const int8_t *b, uint32_t size) const; +class DistanceL2Int8 : public Distance { + public: + DistanceL2Int8() : Distance(diskann::Metric::L2) {} + DISKANN_DLLEXPORT virtual float compare(const int8_t *a, const int8_t *b, + uint32_t size) const; }; // AVX implementations. Borrowed from HNSW code. -class AVXDistanceL2Int8 : public Distance -{ - public: - AVXDistanceL2Int8() : Distance(diskann::Metric::L2) - { - } - DISKANN_DLLEXPORT virtual float compare(const int8_t *a, const int8_t *b, uint32_t length) const; +class AVXDistanceL2Int8 : public Distance { + public: + AVXDistanceL2Int8() : Distance(diskann::Metric::L2) {} + DISKANN_DLLEXPORT virtual float compare(const int8_t *a, const int8_t *b, + uint32_t length) const; }; -class DistanceCosineFloat : public Distance -{ - public: - DistanceCosineFloat() : Distance(diskann::Metric::COSINE) - { - } - DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t length) const; +class DistanceCosineFloat : public Distance { + public: + DistanceCosineFloat() : Distance(diskann::Metric::COSINE) {} + DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, + uint32_t length) const; }; -class DistanceL2Float : public Distance -{ - public: - DistanceL2Float() : Distance(diskann::Metric::L2) - { - } +class DistanceL2Float : public Distance { + public: + DistanceL2Float() : Distance(diskann::Metric::L2) {} #ifdef _WINDOWS - DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t size) const; + DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, + uint32_t size) const; #else - DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t size) const __attribute__((hot)); + DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, + uint32_t size) const + __attribute__((hot)); #endif }; -class AVXDistanceL2Float : public Distance -{ - public: - AVXDistanceL2Float() : Distance(diskann::Metric::L2) - { - } - DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t length) const; +class AVXDistanceL2Float : public Distance { + public: + AVXDistanceL2Float() : Distance(diskann::Metric::L2) {} + DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, + uint32_t length) const; }; -template class SlowDistanceL2 : public Distance -{ - public: - SlowDistanceL2() : Distance(diskann::Metric::L2) - { - } - DISKANN_DLLEXPORT virtual float compare(const T *a, const T *b, uint32_t length) const; +template +class SlowDistanceL2 : public Distance { + public: + SlowDistanceL2() : Distance(diskann::Metric::L2) {} + DISKANN_DLLEXPORT virtual float compare(const T *a, const T *b, + uint32_t length) const; }; -class SlowDistanceCosineUInt8 : public Distance -{ - public: - SlowDistanceCosineUInt8() : Distance(diskann::Metric::COSINE) - { - } - DISKANN_DLLEXPORT virtual float compare(const uint8_t *a, const uint8_t *b, uint32_t length) const; +class SlowDistanceCosineUInt8 : public Distance { + public: + SlowDistanceCosineUInt8() : Distance(diskann::Metric::COSINE) {} + DISKANN_DLLEXPORT virtual float compare(const uint8_t *a, const uint8_t *b, + uint32_t length) const; }; -class DistanceL2UInt8 : public Distance -{ - public: - DistanceL2UInt8() : Distance(diskann::Metric::L2) - { - } - DISKANN_DLLEXPORT virtual float compare(const uint8_t *a, const uint8_t *b, uint32_t size) const; +class DistanceL2UInt8 : public Distance { + public: + DistanceL2UInt8() : Distance(diskann::Metric::L2) {} + DISKANN_DLLEXPORT virtual float compare(const uint8_t *a, const uint8_t *b, + uint32_t size) const; }; -template class DistanceInnerProduct : public Distance -{ - public: - DistanceInnerProduct() : Distance(diskann::Metric::INNER_PRODUCT) - { - } - - DistanceInnerProduct(diskann::Metric metric) : Distance(metric) - { - } - inline float inner_product(const T *a, const T *b, unsigned size) const; - - inline float compare(const T *a, const T *b, unsigned size) const - { - float result = inner_product(a, b, size); - // if (result < 0) - // return std::numeric_limits::max(); - // else - return -result; - } +template +class DistanceInnerProduct : public Distance { + public: + DistanceInnerProduct() : Distance(diskann::Metric::INNER_PRODUCT) {} + + DistanceInnerProduct(diskann::Metric metric) : Distance(metric) {} + inline float inner_product(const T *a, const T *b, unsigned size) const; + + inline float compare(const T *a, const T *b, unsigned size) const { + float result = inner_product(a, b, size); + // if (result < 0) + // return std::numeric_limits::max(); + // else + return -result; + } }; -template class DistanceFastL2 : public DistanceInnerProduct -{ - // currently defined only for float. - // templated for future use. - public: - DistanceFastL2() : DistanceInnerProduct(diskann::Metric::FAST_L2) - { - } - float norm(const T *a, unsigned size) const; - float compare(const T *a, const T *b, float norm, unsigned size) const; +template +class DistanceFastL2 : public DistanceInnerProduct { + // currently defined only for float. + // templated for future use. + public: + DistanceFastL2() : DistanceInnerProduct(diskann::Metric::FAST_L2) {} + float norm(const T *a, unsigned size) const; + float compare(const T *a, const T *b, float norm, unsigned size) const; }; -class AVXDistanceInnerProductFloat : public Distance -{ - public: - AVXDistanceInnerProductFloat() : Distance(diskann::Metric::INNER_PRODUCT) - { - } - DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t length) const; +class AVXDistanceInnerProductFloat : public Distance { + public: + AVXDistanceInnerProductFloat() + : Distance(diskann::Metric::INNER_PRODUCT) {} + DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, + uint32_t length) const; }; -class AVXNormalizedCosineDistanceFloat : public Distance -{ - private: - AVXDistanceInnerProductFloat _innerProduct; - - protected: - void normalize_and_copy(const float *a, uint32_t length, float *a_norm) const; - - public: - AVXNormalizedCosineDistanceFloat() : Distance(diskann::Metric::COSINE) - { - } - DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t length) const - { - // Inner product returns negative values to indicate distance. - // This will ensure that cosine is between -1 and 1. - return 1.0f + _innerProduct.compare(a, b, length); - } - DISKANN_DLLEXPORT virtual uint32_t post_normalization_dimension(uint32_t orig_dimension) const override; - - DISKANN_DLLEXPORT virtual bool preprocessing_required() const; - - DISKANN_DLLEXPORT virtual void preprocess_base_points(float *original_data, const size_t orig_dim, - const size_t num_points) override; - - DISKANN_DLLEXPORT virtual void preprocess_query(const float *query_vec, const size_t query_dim, - float *scratch_query_vector) override; +class AVXNormalizedCosineDistanceFloat : public Distance { + private: + AVXDistanceInnerProductFloat _innerProduct; + + protected: + void normalize_and_copy(const float *a, uint32_t length, float *a_norm) const; + + public: + AVXNormalizedCosineDistanceFloat() + : Distance(diskann::Metric::COSINE) {} + DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, + uint32_t length) const { + // Inner product returns negative values to indicate distance. + // This will ensure that cosine is between -1 and 1. + return 1.0f + _innerProduct.compare(a, b, length); + } + DISKANN_DLLEXPORT virtual uint32_t post_normalization_dimension( + uint32_t orig_dimension) const override; + + DISKANN_DLLEXPORT virtual bool preprocessing_required() const; + + DISKANN_DLLEXPORT virtual void preprocess_base_points( + float *original_data, const size_t orig_dim, + const size_t num_points) override; + + DISKANN_DLLEXPORT virtual void preprocess_query( + const float *query_vec, const size_t query_dim, + float *scratch_query_vector) override; }; -template Distance *get_distance_function(Metric m); +template +Distance *get_distance_function(Metric m); -} // namespace diskann +} // namespace diskann diff --git a/include/exceptions.h b/include/exceptions.h index 99e4e7361..b23dae72c 100644 --- a/include/exceptions.h +++ b/include/exceptions.h @@ -4,14 +4,11 @@ #pragma once #include -namespace diskann -{ +namespace diskann { -class NotImplementedException : public std::logic_error -{ - public: - NotImplementedException() : std::logic_error("Function not yet implemented.") - { - } +class NotImplementedException : public std::logic_error { + public: + NotImplementedException() + : std::logic_error("Function not yet implemented.") {} }; -} // namespace diskann +} // namespace diskann diff --git a/include/filter_utils.h b/include/filter_utils.h index b7114d264..073543f1f 100644 --- a/include/filter_utils.h +++ b/include/filter_utils.h @@ -44,24 +44,32 @@ typedef tsl::robin_set label_set; typedef std::string path; // structs for returning multiple items from a function -typedef std::tuple, tsl::robin_map, tsl::robin_set> +typedef std::tuple, + tsl::robin_map, + tsl::robin_set> parse_label_file_return_values; -typedef std::tuple>, uint64_t> load_label_index_return_values; +typedef std::tuple>, uint64_t> + load_label_index_return_values; -namespace diskann -{ +namespace diskann { template -DISKANN_DLLEXPORT void generate_label_indices(path input_data_path, path final_index_path_prefix, label_set all_labels, - unsigned R, unsigned L, float alpha, unsigned num_threads); +DISKANN_DLLEXPORT void generate_label_indices(path input_data_path, + path final_index_path_prefix, + label_set all_labels, unsigned R, + unsigned L, float alpha, + unsigned num_threads); -DISKANN_DLLEXPORT load_label_index_return_values load_label_index(path label_index_path, - uint32_t label_number_of_points); +DISKANN_DLLEXPORT load_label_index_return_values +load_label_index(path label_index_path, uint32_t label_number_of_points); -DISKANN_DLLEXPORT parse_label_file_return_values parse_label_file(path label_data_path, std::string universal_label); +DISKANN_DLLEXPORT parse_label_file_return_values +parse_label_file(path label_data_path, std::string universal_label); template -DISKANN_DLLEXPORT tsl::robin_map> generate_label_specific_vector_files_compat( - path input_data_path, tsl::robin_map labels_to_number_of_points, +DISKANN_DLLEXPORT tsl::robin_map> +generate_label_specific_vector_files_compat( + path input_data_path, + tsl::robin_map labels_to_number_of_points, std::vector point_ids_to_labels, label_set all_labels); /* @@ -75,139 +83,138 @@ DISKANN_DLLEXPORT tsl::robin_map> generate_la * input_data_path + "_" + label */ template -inline tsl::robin_map> generate_label_specific_vector_files( - path input_data_path, tsl::robin_map labels_to_number_of_points, - std::vector point_ids_to_labels, label_set all_labels) -{ - auto file_writing_timer = std::chrono::high_resolution_clock::now(); - diskann::MemoryMapper input_data(input_data_path); - char *input_start = input_data.getBuf(); - - uint32_t number_of_points, dimension; - std::memcpy(&number_of_points, input_start, sizeof(uint32_t)); - std::memcpy(&dimension, input_start + sizeof(uint32_t), sizeof(uint32_t)); - const uint32_t VECTOR_SIZE = dimension * sizeof(T); - const size_t METADATA = 2 * sizeof(uint32_t); - if (number_of_points != point_ids_to_labels.size()) - { - std::cerr << "Error: number of points in labels file and data file differ." << std::endl; - throw; +inline tsl::robin_map> +generate_label_specific_vector_files( + path input_data_path, + tsl::robin_map labels_to_number_of_points, + std::vector point_ids_to_labels, label_set all_labels) { + auto file_writing_timer = std::chrono::high_resolution_clock::now(); + diskann::MemoryMapper input_data(input_data_path); + char *input_start = input_data.getBuf(); + + uint32_t number_of_points, dimension; + std::memcpy(&number_of_points, input_start, sizeof(uint32_t)); + std::memcpy(&dimension, input_start + sizeof(uint32_t), sizeof(uint32_t)); + const uint32_t VECTOR_SIZE = dimension * sizeof(T); + const size_t METADATA = 2 * sizeof(uint32_t); + if (number_of_points != point_ids_to_labels.size()) { + std::cerr << "Error: number of points in labels file and data file differ." + << std::endl; + throw; + } + + tsl::robin_map label_to_iovec_map; + tsl::robin_map label_to_curr_iovec; + tsl::robin_map> label_id_to_orig_id; + + // setup iovec list for each label + for (const auto &lbl : all_labels) { + iovec *label_iovecs = + (iovec *)malloc(labels_to_number_of_points[lbl] * sizeof(iovec)); + if (label_iovecs == nullptr) { + throw; } - - tsl::robin_map label_to_iovec_map; - tsl::robin_map label_to_curr_iovec; - tsl::robin_map> label_id_to_orig_id; - - // setup iovec list for each label - for (const auto &lbl : all_labels) - { - iovec *label_iovecs = (iovec *)malloc(labels_to_number_of_points[lbl] * sizeof(iovec)); - if (label_iovecs == nullptr) - { - throw; - } - label_to_iovec_map[lbl] = label_iovecs; - label_to_curr_iovec[lbl] = 0; - label_id_to_orig_id[lbl].reserve(labels_to_number_of_points[lbl]); + label_to_iovec_map[lbl] = label_iovecs; + label_to_curr_iovec[lbl] = 0; + label_id_to_orig_id[lbl].reserve(labels_to_number_of_points[lbl]); + } + + // each point added to corresponding per-label iovec list + for (uint32_t point_id = 0; point_id < number_of_points; point_id++) { + char *curr_point = input_start + METADATA + (VECTOR_SIZE * point_id); + iovec curr_iovec; + + curr_iovec.iov_base = curr_point; + curr_iovec.iov_len = VECTOR_SIZE; + for (const auto &lbl : point_ids_to_labels[point_id]) { + *(label_to_iovec_map[lbl] + label_to_curr_iovec[lbl]) = curr_iovec; + label_to_curr_iovec[lbl]++; + label_id_to_orig_id[lbl].push_back(point_id); } - - // each point added to corresponding per-label iovec list - for (uint32_t point_id = 0; point_id < number_of_points; point_id++) - { - char *curr_point = input_start + METADATA + (VECTOR_SIZE * point_id); - iovec curr_iovec; - - curr_iovec.iov_base = curr_point; - curr_iovec.iov_len = VECTOR_SIZE; - for (const auto &lbl : point_ids_to_labels[point_id]) - { - *(label_to_iovec_map[lbl] + label_to_curr_iovec[lbl]) = curr_iovec; - label_to_curr_iovec[lbl]++; - label_id_to_orig_id[lbl].push_back(point_id); - } + } + + // write each label iovec to resp. file + for (const auto &lbl : all_labels) { + int label_input_data_fd; + path curr_label_input_data_path(input_data_path + "_" + lbl); + uint32_t curr_num_pts = labels_to_number_of_points[lbl]; + + label_input_data_fd = + open(curr_label_input_data_path.c_str(), + O_CREAT | O_WRONLY | O_TRUNC | O_APPEND, (mode_t)0644); + if (label_input_data_fd == -1) throw; + + // write metadata + uint32_t metadata[2] = {curr_num_pts, dimension}; + int return_value = + write(label_input_data_fd, metadata, sizeof(uint32_t) * 2); + if (return_value == -1) { + throw; } - // write each label iovec to resp. file - for (const auto &lbl : all_labels) - { - int label_input_data_fd; - path curr_label_input_data_path(input_data_path + "_" + lbl); - uint32_t curr_num_pts = labels_to_number_of_points[lbl]; - - label_input_data_fd = - open(curr_label_input_data_path.c_str(), O_CREAT | O_WRONLY | O_TRUNC | O_APPEND, (mode_t)0644); - if (label_input_data_fd == -1) - throw; - - // write metadata - uint32_t metadata[2] = {curr_num_pts, dimension}; - int return_value = write(label_input_data_fd, metadata, sizeof(uint32_t) * 2); - if (return_value == -1) - { - throw; - } - - // limits on number of iovec structs per writev means we need to perform - // multiple writevs - size_t i = 0; - while (curr_num_pts > IOV_MAX) - { - return_value = writev(label_input_data_fd, (label_to_iovec_map[lbl] + (IOV_MAX * i)), IOV_MAX); - if (return_value == -1) - { - close(label_input_data_fd); - throw; - } - curr_num_pts -= IOV_MAX; - i += 1; - } - return_value = writev(label_input_data_fd, (label_to_iovec_map[lbl] + (IOV_MAX * i)), curr_num_pts); - if (return_value == -1) - { - close(label_input_data_fd); - throw; - } - - free(label_to_iovec_map[lbl]); + // limits on number of iovec structs per writev means we need to perform + // multiple writevs + size_t i = 0; + while (curr_num_pts > IOV_MAX) { + return_value = writev(label_input_data_fd, + (label_to_iovec_map[lbl] + (IOV_MAX * i)), IOV_MAX); + if (return_value == -1) { close(label_input_data_fd); + throw; + } + curr_num_pts -= IOV_MAX; + i += 1; + } + return_value = + writev(label_input_data_fd, (label_to_iovec_map[lbl] + (IOV_MAX * i)), + curr_num_pts); + if (return_value == -1) { + close(label_input_data_fd); + throw; } - std::chrono::duration file_writing_time = std::chrono::high_resolution_clock::now() - file_writing_timer; - std::cout << "generated " << all_labels.size() << " label-specific vector files for index building in time " - << file_writing_time.count() << "\n" - << std::endl; + free(label_to_iovec_map[lbl]); + close(label_input_data_fd); + } + + std::chrono::duration file_writing_time = + std::chrono::high_resolution_clock::now() - file_writing_timer; + std::cout << "generated " << all_labels.size() + << " label-specific vector files for index building in time " + << file_writing_time.count() << "\n" + << std::endl; - return label_id_to_orig_id; + return label_id_to_orig_id; } -inline std::vector loadTags(const std::string &tags_file, const std::string &base_file) -{ - const bool tags_enabled = tags_file.empty() ? false : true; - std::vector location_to_tag; - if (tags_enabled) - { - size_t tag_file_ndims, tag_file_npts; - std::uint32_t *tag_data; - diskann::load_bin(tags_file, tag_data, tag_file_npts, tag_file_ndims); - if (tag_file_ndims != 1) - { - diskann::cerr << "tags file error" << std::endl; - throw diskann::ANNException("tag file error", -1, __FUNCSIG__, __FILE__, __LINE__); - } - - // check if the point count match - size_t base_file_npts, base_file_ndims; - diskann::get_bin_metadata(base_file, base_file_npts, base_file_ndims); - if (base_file_npts != tag_file_npts) - { - diskann::cerr << "point num in tags file mismatch" << std::endl; - throw diskann::ANNException("point num in tags file mismatch", -1, __FUNCSIG__, __FILE__, __LINE__); - } - - location_to_tag.assign(tag_data, tag_data + tag_file_npts); - delete[] tag_data; +inline std::vector loadTags(const std::string &tags_file, + const std::string &base_file) { + const bool tags_enabled = tags_file.empty() ? false : true; + std::vector location_to_tag; + if (tags_enabled) { + size_t tag_file_ndims, tag_file_npts; + std::uint32_t *tag_data; + diskann::load_bin(tags_file, tag_data, tag_file_npts, + tag_file_ndims); + if (tag_file_ndims != 1) { + diskann::cerr << "tags file error" << std::endl; + throw diskann::ANNException("tag file error", -1, __FUNCSIG__, __FILE__, + __LINE__); } - return location_to_tag; + + // check if the point count match + size_t base_file_npts, base_file_ndims; + diskann::get_bin_metadata(base_file, base_file_npts, base_file_ndims); + if (base_file_npts != tag_file_npts) { + diskann::cerr << "point num in tags file mismatch" << std::endl; + throw diskann::ANNException("point num in tags file mismatch", -1, + __FUNCSIG__, __FILE__, __LINE__); + } + + location_to_tag.assign(tag_data, tag_data + tag_file_npts); + delete[] tag_data; + } + return location_to_tag; } -} // namespace diskann +} // namespace diskann diff --git a/include/in_mem_data_store.h b/include/in_mem_data_store.h index 60dd78313..6314b90bc 100644 --- a/include/in_mem_data_store.h +++ b/include/in_mem_data_store.h @@ -15,67 +15,77 @@ #include "natural_number_map.h" #include "natural_number_set.h" -namespace diskann -{ -template class InMemDataStore : public AbstractDataStore -{ - public: - InMemDataStore(const location_t capacity, const size_t dim, std::shared_ptr> distance_metric); - virtual ~InMemDataStore(); - - virtual location_t load(const std::string &filename) override; - virtual size_t save(const std::string &filename, const location_t num_points) override; - - virtual size_t get_aligned_dim() const override; - - - // Populate internal data from unaligned data while doing alignment and any normalization that is required. - virtual void populate_data(const data_t *vectors, const location_t num_pts) override; - virtual void populate_data(const std::string &filename, const size_t offset) override; - - virtual void save_data_to_bin(const std::string &filename, const location_t num_pts) override; - - virtual void get_vector(const location_t i, data_t *target) const override; - virtual void set_vector(const location_t i, const data_t *const vector) override; - virtual void prefetch_vector(const location_t loc) override; - - virtual void reposition_points(const location_t old_location_start, const location_t new_location_start, - const location_t num_points) override; - virtual void copy_points(const location_t from_loc, const location_t to_loc, const location_t num_points) override; - - virtual float get_distance(const data_t *query, const location_t loc) const override; - virtual float get_distance(const location_t loc1, const location_t loc2) const override; - virtual void get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, - float *distances) const override; - - virtual location_t calculate_medoid() const override; - - virtual Distance *get_dist_fn(); - - virtual size_t get_alignment_factor() const override; - - protected: - virtual location_t expand(const location_t new_size) override; - virtual location_t shrink(const location_t new_size) override; - - virtual location_t load_impl(const std::string &filename); +namespace diskann { +template +class InMemDataStore : public AbstractDataStore { + public: + InMemDataStore(const location_t capacity, const size_t dim, + std::shared_ptr> distance_metric); + virtual ~InMemDataStore(); + + virtual location_t load(const std::string &filename) override; + virtual size_t save(const std::string &filename, + const location_t num_points) override; + + virtual size_t get_aligned_dim() const override; + + // Populate internal data from unaligned data while doing alignment and any + // normalization that is required. + virtual void populate_data(const data_t *vectors, + const location_t num_pts) override; + virtual void populate_data(const std::string &filename, + const size_t offset) override; + + virtual void save_data_to_bin(const std::string &filename, + const location_t num_pts) override; + + virtual void get_vector(const location_t i, data_t *target) const override; + virtual void set_vector(const location_t i, + const data_t *const vector) override; + virtual void prefetch_vector(const location_t loc) override; + + virtual void reposition_points(const location_t old_location_start, + const location_t new_location_start, + const location_t num_points) override; + virtual void copy_points(const location_t from_loc, const location_t to_loc, + const location_t num_points) override; + + virtual float get_distance(const data_t *query, + const location_t loc) const override; + virtual float get_distance(const location_t loc1, + const location_t loc2) const override; + virtual void get_distance(const data_t *query, const location_t *locations, + const uint32_t location_count, + float *distances) const override; + + virtual location_t calculate_medoid() const override; + + virtual Distance *get_dist_fn(); + + virtual size_t get_alignment_factor() const override; + + protected: + virtual location_t expand(const location_t new_size) override; + virtual location_t shrink(const location_t new_size) override; + + virtual location_t load_impl(const std::string &filename); #ifdef EXEC_ENV_OLS - virtual location_t load_impl(AlignedFileReader &reader); + virtual location_t load_impl(AlignedFileReader &reader); #endif - private: - data_t *_data = nullptr; + private: + data_t *_data = nullptr; - size_t _aligned_dim; + size_t _aligned_dim; - // It may seem weird to put distance metric along with the data store class, but - // this gives us perf benefits as the datastore can do distance computations during - // search and compute norms of vectors internally without have to copy - // data back and forth. - std::shared_ptr> _distance_fn; + // It may seem weird to put distance metric along with the data store class, + // but this gives us perf benefits as the datastore can do distance + // computations during search and compute norms of vectors internally without + // have to copy data back and forth. + std::shared_ptr> _distance_fn; - // in case we need to save vector norms for optimization - std::shared_ptr _pre_computed_norms; + // in case we need to save vector norms for optimization + std::shared_ptr _pre_computed_norms; }; -} // namespace diskann \ No newline at end of file +} // namespace diskann \ No newline at end of file diff --git a/include/in_mem_graph_store.h b/include/in_mem_graph_store.h index 98a9e4dc5..a5a907b38 100644 --- a/include/in_mem_graph_store.h +++ b/include/in_mem_graph_store.h @@ -5,19 +5,17 @@ #include "abstract_graph_store.h" -namespace diskann -{ +namespace diskann { -class InMemGraphStore : public AbstractGraphStore -{ - public: - InMemGraphStore(const size_t max_pts); +class InMemGraphStore : public AbstractGraphStore { + public: + InMemGraphStore(const size_t max_pts); - int load(const std::string &index_path_prefix); - int store(const std::string &index_path_prefix); + int load(const std::string &index_path_prefix); + int store(const std::string &index_path_prefix); - void get_adj_list(const location_t i, std::vector &neighbors); - void set_adj_list(const location_t i, std::vector &neighbors); + void get_adj_list(const location_t i, std::vector &neighbors); + void set_adj_list(const location_t i, std::vector &neighbors); }; -} // namespace diskann +} // namespace diskann diff --git a/include/index.h b/include/index.h index d32314383..aa6415e3a 100644 --- a/include/index.h +++ b/include/index.h @@ -24,394 +24,437 @@ #define EXPAND_IF_FULL 0 #define DEFAULT_MAXC 750 -namespace diskann -{ - -inline double estimate_ram_usage(size_t size, uint32_t dim, uint32_t datasize, uint32_t degree) -{ - double size_of_data = ((double)size) * ROUND_UP(dim, 8) * datasize; - double size_of_graph = ((double)size) * degree * sizeof(uint32_t) * GRAPH_SLACK_FACTOR; - double size_of_locks = ((double)size) * sizeof(non_recursive_mutex); - double size_of_outer_vector = ((double)size) * sizeof(ptrdiff_t); - - return OVERHEAD_FACTOR * (size_of_data + size_of_graph + size_of_locks + size_of_outer_vector); +namespace diskann { + +inline double estimate_ram_usage(size_t size, uint32_t dim, uint32_t datasize, + uint32_t degree) { + double size_of_data = ((double)size) * ROUND_UP(dim, 8) * datasize; + double size_of_graph = + ((double)size) * degree * sizeof(uint32_t) * GRAPH_SLACK_FACTOR; + double size_of_locks = ((double)size) * sizeof(non_recursive_mutex); + double size_of_outer_vector = ((double)size) * sizeof(ptrdiff_t); + + return OVERHEAD_FACTOR * + (size_of_data + size_of_graph + size_of_locks + size_of_outer_vector); } -struct consolidation_report -{ - enum status_code - { - SUCCESS = 0, - FAIL = 1, - LOCK_FAIL = 2, - INCONSISTENT_COUNT_ERROR = 3 - }; - status_code _status; - size_t _active_points, _max_points, _empty_slots, _slots_released, _delete_set_size, _num_calls_to_process_delete; - double _time; - - consolidation_report(status_code status, size_t active_points, size_t max_points, size_t empty_slots, - size_t slots_released, size_t delete_set_size, size_t num_calls_to_process_delete, - double time_secs) - : _status(status), _active_points(active_points), _max_points(max_points), _empty_slots(empty_slots), - _slots_released(slots_released), _delete_set_size(delete_set_size), - _num_calls_to_process_delete(num_calls_to_process_delete), _time(time_secs) - { - } +struct consolidation_report { + enum status_code { + SUCCESS = 0, + FAIL = 1, + LOCK_FAIL = 2, + INCONSISTENT_COUNT_ERROR = 3 + }; + status_code _status; + size_t _active_points, _max_points, _empty_slots, _slots_released, + _delete_set_size, _num_calls_to_process_delete; + double _time; + + consolidation_report(status_code status, size_t active_points, + size_t max_points, size_t empty_slots, + size_t slots_released, size_t delete_set_size, + size_t num_calls_to_process_delete, double time_secs) + : _status(status), + _active_points(active_points), + _max_points(max_points), + _empty_slots(empty_slots), + _slots_released(slots_released), + _delete_set_size(delete_set_size), + _num_calls_to_process_delete(num_calls_to_process_delete), + _time(time_secs) {} }; -template class Index -{ - /************************************************************************** - * - * Public functions acquire one or more of _update_lock, _consolidate_lock, - * _tag_lock, _delete_lock before calling protected functions which DO NOT - * acquire these locks. They might acquire locks on _locks[i] - * - **************************************************************************/ - - public: - // Constructor for Bulk operations and for creating the index object solely - // for loading a prexisting index. - DISKANN_DLLEXPORT Index(Metric m, const size_t dim, const size_t max_points = 1, 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 size_t num_frozen_pts = 0); - - // Constructor for incremental index - DISKANN_DLLEXPORT Index(Metric m, const size_t dim, const size_t max_points, const bool dynamic_index, - const IndexWriteParameters &indexParameters, const uint32_t initial_search_list_size, - const uint32_t search_threads, 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); - - DISKANN_DLLEXPORT ~Index(); - - // Saves graph, data, metadata and associated tags. - DISKANN_DLLEXPORT void save(const char *filename, bool compact_before_save = false); - - // Load functions +template +class Index { + /************************************************************************** + * + * Public functions acquire one or more of _update_lock, _consolidate_lock, + * _tag_lock, _delete_lock before calling protected functions which DO NOT + * acquire these locks. They might acquire locks on _locks[i] + * + **************************************************************************/ + + public: + // Constructor for Bulk operations and for creating the index object solely + // for loading a prexisting index. + DISKANN_DLLEXPORT Index( + Metric m, const size_t dim, const size_t max_points = 1, + 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 size_t num_frozen_pts = 0); + + // Constructor for incremental index + DISKANN_DLLEXPORT Index( + Metric m, const size_t dim, const size_t max_points, + const bool dynamic_index, const IndexWriteParameters &indexParameters, + const uint32_t initial_search_list_size, const uint32_t search_threads, + 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); + + DISKANN_DLLEXPORT ~Index(); + + // Saves graph, data, metadata and associated tags. + DISKANN_DLLEXPORT void save(const char *filename, + bool compact_before_save = false); + + // Load functions #ifdef EXEC_ENV_OLS - DISKANN_DLLEXPORT void load(AlignedFileReader &reader, uint32_t num_threads, uint32_t search_l); + DISKANN_DLLEXPORT void load(AlignedFileReader &reader, uint32_t num_threads, + uint32_t search_l); #else - // Reads the number of frozen points from graph's metadata file section. - DISKANN_DLLEXPORT static size_t get_graph_num_frozen_points(const std::string &graph_file); + // Reads the number of frozen points from graph's metadata file section. + DISKANN_DLLEXPORT static size_t get_graph_num_frozen_points( + const std::string &graph_file); - DISKANN_DLLEXPORT void load(const char *index_file, uint32_t num_threads, uint32_t search_l); + DISKANN_DLLEXPORT void load(const char *index_file, uint32_t num_threads, + uint32_t search_l); #endif - // get some private variables - DISKANN_DLLEXPORT size_t get_num_points(); - DISKANN_DLLEXPORT size_t get_max_points(); - - // Batch build from a file. Optionally pass tags vector. - DISKANN_DLLEXPORT void build(const char *filename, const size_t num_points_to_load, - IndexWriteParameters ¶meters, const std::vector &tags = std::vector()); - - // Batch build from a file. Optionally pass tags file. - DISKANN_DLLEXPORT void build(const char *filename, const size_t num_points_to_load, - IndexWriteParameters ¶meters, const char *tag_filename); - - // Batch build from a data array, which must pad vectors to aligned_dim - DISKANN_DLLEXPORT void build(const T *data, const size_t num_points_to_load, IndexWriteParameters ¶meters, + // get some private variables + DISKANN_DLLEXPORT size_t get_num_points(); + DISKANN_DLLEXPORT size_t get_max_points(); + + // Batch build from a file. Optionally pass tags vector. + DISKANN_DLLEXPORT void build( + const char *filename, const size_t num_points_to_load, + IndexWriteParameters ¶meters, + const std::vector &tags = std::vector()); + + // Batch build from a file. Optionally pass tags file. + DISKANN_DLLEXPORT void build(const char *filename, + const size_t num_points_to_load, + IndexWriteParameters ¶meters, + const char *tag_filename); + + // Batch build from a data array, which must pad vectors to aligned_dim + DISKANN_DLLEXPORT void build(const T *data, const size_t num_points_to_load, + IndexWriteParameters ¶meters, + const std::vector &tags); + + // Filtered Support + DISKANN_DLLEXPORT void build_filtered_index( + const char *filename, const std::string &label_file, + const size_t num_points_to_load, IndexWriteParameters ¶meters, + const std::vector &tags = std::vector()); + + DISKANN_DLLEXPORT void set_universal_label(const LabelT &label); + + // Get converted integer label from string to int map (_label_map) + DISKANN_DLLEXPORT LabelT get_converted_label(const std::string &raw_label); + + // Set starting point of an index before inserting any points incrementally. + // The data count should be equal to _num_frozen_pts * _aligned_dim. + DISKANN_DLLEXPORT void set_start_points(const T *data, size_t data_count); + // Set starting points to random points on a sphere of certain radius. + // A fixed random seed can be specified for scenarios where it's important + // to have higher consistency between index builds. + DISKANN_DLLEXPORT void set_start_points_at_random(T radius, + uint32_t random_seed = 0); + + // For FastL2 search on a static index, we interleave the data with graph + DISKANN_DLLEXPORT void optimize_index_layout(); + + // For FastL2 search on optimized layout + DISKANN_DLLEXPORT void search_with_optimized_layout(const T *query, size_t K, + size_t L, + uint32_t *indices); + + // Added search overload that takes L as parameter, so that we + // can customize L on a per-query basis without tampering with "Parameters" + template + DISKANN_DLLEXPORT std::pair search( + const T *query, const size_t K, const uint32_t L, IDType *indices, + float *distances = nullptr); + + // Initialize space for res_vectors before calling. + DISKANN_DLLEXPORT size_t search_with_tags(const T *query, const uint64_t K, + const uint32_t L, TagT *tags, + float *distances, + std::vector &res_vectors); + + // Filter support search + template + DISKANN_DLLEXPORT std::pair search_with_filters( + const T *query, const LabelT &filter_label, const size_t K, + const uint32_t L, IndexType *indices, float *distances); + + // Will fail if tag already in the index or if tag=0. + DISKANN_DLLEXPORT int insert_point(const T *point, const TagT tag); + + // call this before issuing deletions to sets relevant flags + DISKANN_DLLEXPORT int enable_delete(); + + // Record deleted point now and restructure graph later. Return -1 if tag + // not found, 0 if OK. + DISKANN_DLLEXPORT int lazy_delete(const TagT &tag); + + // Record deleted points now and restructure graph later. Add to failed_tags + // if tag not found. + DISKANN_DLLEXPORT void lazy_delete(const std::vector &tags, + std::vector &failed_tags); + + // Call after a series of lazy deletions + // Returns number of live points left after consolidation + // If _conc_consolidates is set in the ctor, then this call can be invoked + // alongside inserts and lazy deletes, else it acquires _update_lock + DISKANN_DLLEXPORT consolidation_report + consolidate_deletes(const IndexWriteParameters ¶meters); + + DISKANN_DLLEXPORT void prune_all_neighbors(const uint32_t max_degree, + const uint32_t max_occlusion, + const float alpha); + + DISKANN_DLLEXPORT bool is_index_saved(); + + // 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_points(uint32_t old_location_start, + uint32_t new_location_start, + uint32_t num_locations); + + // DISKANN_DLLEXPORT void save_index_as_one_file(bool flag); + + DISKANN_DLLEXPORT void get_active_tags(tsl::robin_set &active_tags); + + // memory should be allocated for vec before calling this function + DISKANN_DLLEXPORT int get_vector_by_tag(TagT &tag, T *vec); + + DISKANN_DLLEXPORT void print_status(); + + DISKANN_DLLEXPORT void count_nodes_at_bfs_levels(); + + // 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: + // No copy/assign. + Index(const Index &) = delete; + Index &operator=(const Index &) = delete; + + // Use after _data and _nd have been populated + // Acquire exclusive _update_lock before calling + void build_with_data_populated(IndexWriteParameters ¶meters, const std::vector &tags); - // Filtered Support - DISKANN_DLLEXPORT void build_filtered_index(const char *filename, const std::string &label_file, - const size_t num_points_to_load, IndexWriteParameters ¶meters, - const std::vector &tags = std::vector()); - - DISKANN_DLLEXPORT void set_universal_label(const LabelT &label); - - // Get converted integer label from string to int map (_label_map) - DISKANN_DLLEXPORT LabelT get_converted_label(const std::string &raw_label); - - // Set starting point of an index before inserting any points incrementally. - // The data count should be equal to _num_frozen_pts * _aligned_dim. - DISKANN_DLLEXPORT void set_start_points(const T *data, size_t data_count); - // Set starting points to random points on a sphere of certain radius. - // A fixed random seed can be specified for scenarios where it's important - // to have higher consistency between index builds. - DISKANN_DLLEXPORT void set_start_points_at_random(T radius, uint32_t random_seed = 0); - - // For FastL2 search on a static index, we interleave the data with graph - DISKANN_DLLEXPORT void optimize_index_layout(); - - // For FastL2 search on optimized layout - DISKANN_DLLEXPORT void search_with_optimized_layout(const T *query, size_t K, size_t L, uint32_t *indices); - - // Added search overload that takes L as parameter, so that we - // can customize L on a per-query basis without tampering with "Parameters" - template - DISKANN_DLLEXPORT std::pair search(const T *query, const size_t K, const uint32_t L, - IDType *indices, float *distances = nullptr); - - // Initialize space for res_vectors before calling. - DISKANN_DLLEXPORT size_t search_with_tags(const T *query, const uint64_t K, const uint32_t L, TagT *tags, - float *distances, std::vector &res_vectors); - - // Filter support search - template - DISKANN_DLLEXPORT std::pair search_with_filters(const T *query, const LabelT &filter_label, - const size_t K, const uint32_t L, - IndexType *indices, float *distances); - - // Will fail if tag already in the index or if tag=0. - DISKANN_DLLEXPORT int insert_point(const T *point, const TagT tag); - - // call this before issuing deletions to sets relevant flags - DISKANN_DLLEXPORT int enable_delete(); - - // Record deleted point now and restructure graph later. Return -1 if tag - // not found, 0 if OK. - DISKANN_DLLEXPORT int lazy_delete(const TagT &tag); - - // Record deleted points now and restructure graph later. Add to failed_tags - // if tag not found. - DISKANN_DLLEXPORT void lazy_delete(const std::vector &tags, std::vector &failed_tags); - - // Call after a series of lazy deletions - // Returns number of live points left after consolidation - // If _conc_consolidates is set in the ctor, then this call can be invoked - // alongside inserts and lazy deletes, else it acquires _update_lock - DISKANN_DLLEXPORT consolidation_report consolidate_deletes(const IndexWriteParameters ¶meters); - - DISKANN_DLLEXPORT void prune_all_neighbors(const uint32_t max_degree, const uint32_t max_occlusion, - const float alpha); - - DISKANN_DLLEXPORT bool is_index_saved(); - - // 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_points(uint32_t old_location_start, uint32_t new_location_start, - uint32_t num_locations); - - // DISKANN_DLLEXPORT void save_index_as_one_file(bool flag); - - DISKANN_DLLEXPORT void get_active_tags(tsl::robin_set &active_tags); - - // memory should be allocated for vec before calling this function - DISKANN_DLLEXPORT int get_vector_by_tag(TagT &tag, T *vec); - - DISKANN_DLLEXPORT void print_status(); - - DISKANN_DLLEXPORT void count_nodes_at_bfs_levels(); - - // 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: - // No copy/assign. - Index(const Index &) = delete; - Index &operator=(const Index &) = delete; - - // Use after _data and _nd have been populated - // Acquire exclusive _update_lock before calling - void build_with_data_populated(IndexWriteParameters ¶meters, const std::vector &tags); - - // generates 1 frozen point that will never be deleted from the graph - // This is not visible to the user - void generate_frozen_point(); - - // determines navigating node of the graph by calculating medoid of datafopt - uint32_t calculate_entry_point(); - - void parse_label_file(const std::string &label_file, size_t &num_pts_labels); - - std::unordered_map load_label_map(const std::string &map_file); - - // Returns the locations of start point and frozen points suitable for use - // with iterate_to_fixed_point. - std::vector get_init_ids(); - - std::pair iterate_to_fixed_point(const T *node_coords, const uint32_t Lindex, - const std::vector &init_ids, - InMemQueryScratch *scratch, bool use_filter, - const std::vector &filters, bool search_invocation); - - void search_for_point_and_prune(int location, uint32_t Lindex, std::vector &pruned_list, - InMemQueryScratch *scratch, bool use_filter = false, - uint32_t filteredLindex = 0); - - void prune_neighbors(const uint32_t location, std::vector &pool, std::vector &pruned_list, - InMemQueryScratch *scratch); - - void prune_neighbors(const uint32_t location, std::vector &pool, const uint32_t range, - const uint32_t max_candidate_size, const float alpha, std::vector &pruned_list, - InMemQueryScratch *scratch); - - // Prunes candidates in @pool to a shorter list @result - // @pool must be sorted before calling - void occlude_list(const uint32_t location, std::vector &pool, const float alpha, const uint32_t degree, - const uint32_t maxc, std::vector &result, InMemQueryScratch *scratch, - const tsl::robin_set *const delete_set_ptr = nullptr); - - // add reverse links from all the visited nodes to node n. - void inter_insert(uint32_t n, std::vector &pruned_list, const uint32_t range, - InMemQueryScratch *scratch); - - void inter_insert(uint32_t n, std::vector &pruned_list, InMemQueryScratch *scratch); - - // Acquire exclusive _update_lock before calling - void link(IndexWriteParameters ¶meters); - - // Acquire exclusive _tag_lock and _delete_lock before calling - int reserve_location(); - - // Acquire exclusive _tag_lock before calling - size_t release_location(int location); - size_t release_locations(const tsl::robin_set &locations); - - // Resize the index when no slots are left for insertion. - // Acquire exclusive _update_lock and _tag_lock before calling. - void resize(size_t new_max_points); - - // Acquire unique lock on _update_lock, _consolidate_lock, _tag_lock - // and _delete_lock before calling these functions. - // Renumber nodes, update tag and location maps and compact the - // graph, mode = _consolidated_order in case of lazy deletion and - // _compacted_order in case of eager deletion - DISKANN_DLLEXPORT void compact_data(); - DISKANN_DLLEXPORT void compact_frozen_point(); - - // Remove deleted nodes from adjacency list of node loc - // Replace removed neighbors with second order neighbors. - // Also acquires _locks[i] for i = loc and out-neighbors of loc. - void process_delete(const tsl::robin_set &old_delete_set, size_t loc, const uint32_t range, - const uint32_t maxc, const float alpha, InMemQueryScratch *scratch); - - void initialize_query_scratch(uint32_t num_threads, uint32_t search_l, uint32_t indexing_l, uint32_t r, - uint32_t maxc, size_t dim); - - // Do not call without acquiring appropriate locks - // call public member functions save and load to invoke these. - DISKANN_DLLEXPORT size_t save_graph(std::string filename); - DISKANN_DLLEXPORT size_t save_data(std::string filename); - DISKANN_DLLEXPORT size_t save_tags(std::string filename); - DISKANN_DLLEXPORT size_t save_delete_list(const std::string &filename); + // generates 1 frozen point that will never be deleted from the graph + // This is not visible to the user + void generate_frozen_point(); + + // determines navigating node of the graph by calculating medoid of datafopt + uint32_t calculate_entry_point(); + + void parse_label_file(const std::string &label_file, size_t &num_pts_labels); + + std::unordered_map load_label_map( + const std::string &map_file); + + // Returns the locations of start point and frozen points suitable for use + // with iterate_to_fixed_point. + std::vector get_init_ids(); + + std::pair iterate_to_fixed_point( + const T *node_coords, const uint32_t Lindex, + const std::vector &init_ids, InMemQueryScratch *scratch, + bool use_filter, const std::vector &filters, + bool search_invocation); + + void search_for_point_and_prune(int location, uint32_t Lindex, + std::vector &pruned_list, + InMemQueryScratch *scratch, + bool use_filter = false, + uint32_t filteredLindex = 0); + + void prune_neighbors(const uint32_t location, std::vector &pool, + std::vector &pruned_list, + InMemQueryScratch *scratch); + + void prune_neighbors(const uint32_t location, std::vector &pool, + const uint32_t range, const uint32_t max_candidate_size, + const float alpha, std::vector &pruned_list, + InMemQueryScratch *scratch); + + // Prunes candidates in @pool to a shorter list @result + // @pool must be sorted before calling + void occlude_list( + const uint32_t location, std::vector &pool, const float alpha, + const uint32_t degree, const uint32_t maxc, std::vector &result, + InMemQueryScratch *scratch, + const tsl::robin_set *const delete_set_ptr = nullptr); + + // add reverse links from all the visited nodes to node n. + void inter_insert(uint32_t n, std::vector &pruned_list, + const uint32_t range, InMemQueryScratch *scratch); + + void inter_insert(uint32_t n, std::vector &pruned_list, + InMemQueryScratch *scratch); + + // Acquire exclusive _update_lock before calling + void link(IndexWriteParameters ¶meters); + + // Acquire exclusive _tag_lock and _delete_lock before calling + int reserve_location(); + + // Acquire exclusive _tag_lock before calling + size_t release_location(int location); + size_t release_locations(const tsl::robin_set &locations); + + // Resize the index when no slots are left for insertion. + // Acquire exclusive _update_lock and _tag_lock before calling. + void resize(size_t new_max_points); + + // Acquire unique lock on _update_lock, _consolidate_lock, _tag_lock + // and _delete_lock before calling these functions. + // Renumber nodes, update tag and location maps and compact the + // graph, mode = _consolidated_order in case of lazy deletion and + // _compacted_order in case of eager deletion + DISKANN_DLLEXPORT void compact_data(); + DISKANN_DLLEXPORT void compact_frozen_point(); + + // Remove deleted nodes from adjacency list of node loc + // Replace removed neighbors with second order neighbors. + // Also acquires _locks[i] for i = loc and out-neighbors of loc. + void process_delete(const tsl::robin_set &old_delete_set, + size_t loc, const uint32_t range, const uint32_t maxc, + const float alpha, InMemQueryScratch *scratch); + + void initialize_query_scratch(uint32_t num_threads, uint32_t search_l, + uint32_t indexing_l, uint32_t r, uint32_t maxc, + size_t dim); + + // Do not call without acquiring appropriate locks + // call public member functions save and load to invoke these. + DISKANN_DLLEXPORT size_t save_graph(std::string filename); + DISKANN_DLLEXPORT size_t save_data(std::string filename); + DISKANN_DLLEXPORT size_t save_tags(std::string filename); + DISKANN_DLLEXPORT size_t save_delete_list(const std::string &filename); #ifdef EXEC_ENV_OLS - DISKANN_DLLEXPORT size_t load_graph(AlignedFileReader &reader, size_t expected_num_points); - DISKANN_DLLEXPORT size_t load_data(AlignedFileReader &reader); - DISKANN_DLLEXPORT size_t load_tags(AlignedFileReader &reader); - DISKANN_DLLEXPORT size_t load_delete_set(AlignedFileReader &reader); + DISKANN_DLLEXPORT size_t load_graph(AlignedFileReader &reader, + size_t expected_num_points); + DISKANN_DLLEXPORT size_t load_data(AlignedFileReader &reader); + DISKANN_DLLEXPORT size_t load_tags(AlignedFileReader &reader); + DISKANN_DLLEXPORT size_t load_delete_set(AlignedFileReader &reader); #else - DISKANN_DLLEXPORT size_t load_graph(const std::string filename, size_t expected_num_points); - DISKANN_DLLEXPORT size_t load_data(std::string filename0); - DISKANN_DLLEXPORT size_t load_tags(const std::string tag_file_name); - DISKANN_DLLEXPORT size_t load_delete_set(const std::string &filename); + DISKANN_DLLEXPORT size_t load_graph(const std::string filename, + size_t expected_num_points); + DISKANN_DLLEXPORT size_t load_data(std::string filename0); + DISKANN_DLLEXPORT size_t load_tags(const std::string tag_file_name); + DISKANN_DLLEXPORT size_t load_delete_set(const std::string &filename); #endif - private: - // Distance functions - Metric _dist_metric = diskann::L2; - std::shared_ptr> _distance; - - // Data - std::unique_ptr> _data_store; - char *_opt_graph = nullptr; - - // Graph related data structures - std::vector> _final_graph; - - // Dimensions - size_t _dim = 0; - 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 - - // _num_frozen_pts is the number of points which are used as initial candidates - // when iterating to closest point(s). These are not visible externally and won't - // be returned by search. At least 1 frozen point is needed for a dynamic index. - // The frozen points have consecutive locations. See also _start below. - size_t _num_frozen_pts = 0; - size_t _max_range_of_loaded_graph = 0; - size_t _node_size; - size_t _data_len; - size_t _neighbor_len; - - uint32_t _max_observed_degree = 0; - // Start point of the search. When _num_frozen_pts is greater than zero, - // this is the location of the first frozen point. Otherwise, this is a - // location of one of the points in index. - uint32_t _start = 0; - - bool _has_built = false; - bool _saturate_graph = false; - bool _save_as_one_file = false; // plan to support in next version - bool _dynamic_index = false; - bool _enable_tags = false; - bool _normalize_vecs = false; // Using normalied L2 for cosine. - - // Filter Support - - bool _filtered_index = false; - std::vector> _pts_to_labels; - tsl::robin_set _labels; - std::string _labels_file; - std::unordered_map _label_to_medoid_id; - std::unordered_map _medoid_counts; - bool _use_universal_label = false; - LabelT _universal_label = 0; - uint32_t _filterIndexingQueueSize; - std::unordered_map _label_map; - - // Indexing parameters - uint32_t _indexingQueueSize; - uint32_t _indexingRange; - uint32_t _indexingMaxC; - float _indexingAlpha; - - // Query scratch data structures - ConcurrentQueue *> _query_scratch; - - // Flags for PQ based distance calculation - bool _pq_dist = false; - bool _use_opq = false; - size_t _num_pq_chunks = 0; - uint8_t *_pq_data = nullptr; - bool _pq_generated = false; - FixedChunkPQTable _pq_table; - - // - // Data structures, locks and flags for dynamic indexing and tags - // - - // lazy_delete removes entry from _location_to_tag and _tag_to_location. If - // _location_to_tag does not resolve a location, infer that it was deleted. - tsl::sparse_map _tag_to_location; - natural_number_map _location_to_tag; - - // _empty_slots has unallocated slots and those freed by consolidate_delete. - // _delete_set has locations marked deleted by lazy_delete. Will not be - // immediately available for insert. consolidate_delete will release these - // slots to _empty_slots. - natural_number_set _empty_slots; - std::unique_ptr> _delete_set; - - bool _data_compacted = true; // true if data has been compacted - bool _is_saved = false; // Checking if the index is already saved. - bool _conc_consolidate = false; // use _lock while searching - - // Acquire locks in the order below when acquiring multiple locks - std::shared_timed_mutex // RW mutex between save/load (exclusive lock) and - _update_lock; // search/inserts/deletes/consolidate (shared lock) - std::shared_timed_mutex // Ensure only one consolidate or compact_data is - _consolidate_lock; // ever active - std::shared_timed_mutex // RW lock for _tag_to_location, - _tag_lock; // _location_to_tag, _empty_slots, _nd, _max_points - std::shared_timed_mutex // RW Lock on _delete_set and _data_compacted - _delete_lock; // variable - - // Per node lock, cardinality=_max_points - std::vector _locks; - - static const float INDEX_GROWTH_FACTOR; + private: + // Distance functions + Metric _dist_metric = diskann::L2; + std::shared_ptr> _distance; + + // Data + std::unique_ptr> _data_store; + char *_opt_graph = nullptr; + + // Graph related data structures + std::vector> _final_graph; + + // Dimensions + size_t _dim = 0; + 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 + + // _num_frozen_pts is the number of points which are used as initial + // candidates when iterating to closest point(s). These are not visible + // externally and won't be returned by search. At least 1 frozen point is + // needed for a dynamic index. The frozen points have consecutive locations. + // See also _start below. + size_t _num_frozen_pts = 0; + size_t _max_range_of_loaded_graph = 0; + size_t _node_size; + size_t _data_len; + size_t _neighbor_len; + + uint32_t _max_observed_degree = 0; + // Start point of the search. When _num_frozen_pts is greater than zero, + // this is the location of the first frozen point. Otherwise, this is a + // location of one of the points in index. + uint32_t _start = 0; + + bool _has_built = false; + bool _saturate_graph = false; + bool _save_as_one_file = false; // plan to support in next version + bool _dynamic_index = false; + bool _enable_tags = false; + bool _normalize_vecs = false; // Using normalied L2 for cosine. + + // Filter Support + + bool _filtered_index = false; + std::vector> _pts_to_labels; + tsl::robin_set _labels; + std::string _labels_file; + std::unordered_map _label_to_medoid_id; + std::unordered_map _medoid_counts; + bool _use_universal_label = false; + LabelT _universal_label = 0; + uint32_t _filterIndexingQueueSize; + std::unordered_map _label_map; + + // Indexing parameters + uint32_t _indexingQueueSize; + uint32_t _indexingRange; + uint32_t _indexingMaxC; + float _indexingAlpha; + + // Query scratch data structures + ConcurrentQueue *> _query_scratch; + + // Flags for PQ based distance calculation + bool _pq_dist = false; + bool _use_opq = false; + size_t _num_pq_chunks = 0; + uint8_t *_pq_data = nullptr; + bool _pq_generated = false; + FixedChunkPQTable _pq_table; + + // + // Data structures, locks and flags for dynamic indexing and tags + // + + // lazy_delete removes entry from _location_to_tag and _tag_to_location. If + // _location_to_tag does not resolve a location, infer that it was deleted. + tsl::sparse_map _tag_to_location; + natural_number_map _location_to_tag; + + // _empty_slots has unallocated slots and those freed by consolidate_delete. + // _delete_set has locations marked deleted by lazy_delete. Will not be + // immediately available for insert. consolidate_delete will release these + // slots to _empty_slots. + natural_number_set _empty_slots; + std::unique_ptr> _delete_set; + + bool _data_compacted = true; // true if data has been compacted + bool _is_saved = false; // Checking if the index is already saved. + bool _conc_consolidate = false; // use _lock while searching + + // Acquire locks in the order below when acquiring multiple locks + std::shared_timed_mutex // RW mutex between save/load (exclusive lock) and + _update_lock; // search/inserts/deletes/consolidate (shared lock) + std::shared_timed_mutex // Ensure only one consolidate or compact_data is + _consolidate_lock; // ever active + std::shared_timed_mutex // RW lock for _tag_to_location, + _tag_lock; // _location_to_tag, _empty_slots, _nd, _max_points + std::shared_timed_mutex // RW Lock on _delete_set and _data_compacted + _delete_lock; // variable + + // Per node lock, cardinality=_max_points + std::vector _locks; + + static const float INDEX_GROWTH_FACTOR; }; -} // namespace diskann +} // namespace diskann diff --git a/include/linux_aligned_file_reader.h b/include/linux_aligned_file_reader.h index 7620e3194..fd78cac79 100644 --- a/include/linux_aligned_file_reader.h +++ b/include/linux_aligned_file_reader.h @@ -6,34 +6,34 @@ #include "aligned_file_reader.h" -class LinuxAlignedFileReader : public AlignedFileReader -{ - private: - uint64_t file_sz; - FileHandle file_desc; - io_context_t bad_ctx = (io_context_t)-1; - - public: - LinuxAlignedFileReader(); - ~LinuxAlignedFileReader(); - - IOContext &get_ctx(); - - // register thread-id for a context - void register_thread(); - - // 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 close(); - - // process batch of aligned requests in parallel - // NOTE :: blocking call - void read(std::vector &read_reqs, IOContext &ctx, bool async = false); +class LinuxAlignedFileReader : public AlignedFileReader { + private: + uint64_t file_sz; + FileHandle file_desc; + io_context_t bad_ctx = (io_context_t)-1; + + public: + LinuxAlignedFileReader(); + ~LinuxAlignedFileReader(); + + IOContext &get_ctx(); + + // register thread-id for a context + void register_thread(); + + // 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 close(); + + // process batch of aligned requests in parallel + // NOTE :: blocking call + void read(std::vector &read_reqs, IOContext &ctx, + bool async = false); }; #endif diff --git a/include/locking.h b/include/locking.h index 2a70f4ffa..6f0c51477 100644 --- a/include/locking.h +++ b/include/locking.h @@ -7,8 +7,7 @@ #include "windows_slim_lock.h" #endif -namespace diskann -{ +namespace diskann { #ifdef _WINDOWS using non_recursive_mutex = windows_exclusive_slim_lock; using LockGuard = windows_exclusive_slim_lock_guard; @@ -16,4 +15,4 @@ using LockGuard = windows_exclusive_slim_lock_guard; using non_recursive_mutex = std::mutex; using LockGuard = std::lock_guard; #endif -} // namespace diskann +} // namespace diskann diff --git a/include/logger.h b/include/logger.h index 28a9f619c..5fe04ac0c 100644 --- a/include/logger.h +++ b/include/logger.h @@ -6,19 +6,14 @@ #include #include "windows_customizations.h" -namespace diskann -{ +namespace diskann { DISKANN_DLLEXPORT extern std::basic_ostream cout; DISKANN_DLLEXPORT extern std::basic_ostream cerr; -enum class DISKANN_DLLEXPORT LogLevel -{ - LL_Info = 0, - LL_Error, - LL_Count -}; +enum class DISKANN_DLLEXPORT LogLevel { LL_Info = 0, LL_Error, LL_Count }; #ifdef EXEC_ENV_OLS -DISKANN_DLLEXPORT void SetCustomLogger(std::function logger); +DISKANN_DLLEXPORT void SetCustomLogger( + std::function logger); #endif -} // namespace diskann +} // namespace diskann diff --git a/include/logger_impl.h b/include/logger_impl.h index af6108c06..4196f32bf 100644 --- a/include/logger_impl.h +++ b/include/logger_impl.h @@ -9,32 +9,29 @@ #include "ann_exception.h" #include "logger.h" -namespace diskann -{ -class ANNStreamBuf : public std::basic_streambuf -{ - public: - DISKANN_DLLEXPORT explicit ANNStreamBuf(FILE *fp); - DISKANN_DLLEXPORT ~ANNStreamBuf(); - - DISKANN_DLLEXPORT bool is_open() const - { - return true; // because stdout and stderr are always open. - } - DISKANN_DLLEXPORT void close(); - DISKANN_DLLEXPORT virtual int underflow(); - DISKANN_DLLEXPORT virtual int overflow(int c); - DISKANN_DLLEXPORT virtual int sync(); - - private: - FILE *_fp; - char *_buf; - int _bufIndex; - std::mutex _mutex; - LogLevel _logLevel; - - int flush(); - void logImpl(char *str, int numchars); +namespace diskann { +class ANNStreamBuf : public std::basic_streambuf { + public: + DISKANN_DLLEXPORT explicit ANNStreamBuf(FILE *fp); + DISKANN_DLLEXPORT ~ANNStreamBuf(); + + DISKANN_DLLEXPORT bool is_open() const { + return true; // because stdout and stderr are always open. + } + DISKANN_DLLEXPORT void close(); + DISKANN_DLLEXPORT virtual int underflow(); + DISKANN_DLLEXPORT virtual int overflow(int c); + DISKANN_DLLEXPORT virtual int sync(); + + private: + FILE *_fp; + char *_buf; + int _bufIndex; + std::mutex _mutex; + LogLevel _logLevel; + + int flush(); + void logImpl(char *str, int numchars); // Why the two buffer-sizes? If we are running normally, we are basically // interacting with a character output system, so we short-circuit the @@ -51,15 +48,15 @@ class ANNStreamBuf : public std::basic_streambuf // 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 = 1024; #else - // Allocating an arbitrarily small buffer here because the overflow() and - // other function implementations push the BUFFER_SIZE chars into the - // buffer before flushing to fwrite. - static const int BUFFER_SIZE = 4; + // Allocating an arbitrarily small buffer here because the overflow() and + // other function implementations push the BUFFER_SIZE chars into the + // buffer before flushing to fwrite. + static const int BUFFER_SIZE = 4; #endif - ANNStreamBuf(const ANNStreamBuf &); - ANNStreamBuf &operator=(const ANNStreamBuf &); + ANNStreamBuf(const ANNStreamBuf &); + ANNStreamBuf &operator=(const ANNStreamBuf &); }; -} // namespace diskann +} // namespace diskann diff --git a/include/math_utils.h b/include/math_utils.h index 83d189f70..333c00c2c 100644 --- a/include/math_utils.h +++ b/include/math_utils.h @@ -6,17 +6,18 @@ #include "common_includes.h" #include "utils.h" -namespace math_utils -{ +namespace math_utils { float calc_distance(float *vec_1, float *vec_2, size_t dim); // compute l2-squared norms of data stored in row major num_points * dim, // needs // to be pre-allocated -void compute_vecs_l2sq(float *vecs_l2sq, float *data, const size_t num_points, const size_t dim); +void compute_vecs_l2sq(float *vecs_l2sq, float *data, const size_t num_points, + const size_t dim); -void rotate_data_randomly(float *data, size_t num_points, size_t dim, float *rot_mat, float *&new_mat, +void rotate_data_randomly(float *data, size_t num_points, size_t dim, + float *rot_mat, float *&new_mat, bool transpose_rot = false); // calculate closest center to data of num_points * dim (row major) @@ -28,10 +29,11 @@ void rotate_data_randomly(float *data, size_t num_points, size_t dim, float *rot // squared distances // Ideally used only by compute_closest_centers -void compute_closest_centers_in_block(const float *const data, const size_t num_points, const size_t dim, - const float *const centers, const size_t num_centers, - const float *const docs_l2sq, const float *const centers_l2sq, - uint32_t *center_index, float *const dist_matrix, size_t k = 1); +void compute_closest_centers_in_block( + const float *const data, const size_t num_points, const size_t dim, + const float *const centers, const size_t num_centers, + const float *const docs_l2sq, const float *const centers_l2sq, + uint32_t *center_index, float *const dist_matrix, size_t k = 1); // Given data in num_points * new_dim row major // Pivots stored in full_pivot_data as k * new_dim row major @@ -43,21 +45,23 @@ void compute_closest_centers_in_block(const float *const data, const size_t num_ // those // values -void compute_closest_centers(float *data, size_t num_points, size_t dim, float *pivot_data, size_t num_centers, - size_t k, uint32_t *closest_centers_ivf, std::vector *inverted_index = NULL, +void compute_closest_centers(float *data, size_t num_points, size_t dim, + float *pivot_data, size_t num_centers, size_t k, + uint32_t *closest_centers_ivf, + std::vector *inverted_index = NULL, float *pts_norms_squared = NULL); // if to_subtract is 1, will subtract nearest center from each row. Else will // add. Output will be in data_load iself. // Nearest centers need to be provided in closst_centers. -void process_residuals(float *data_load, size_t num_points, size_t dim, float *cur_pivot_data, size_t num_centers, +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); -} // namespace math_utils +} // namespace math_utils -namespace kmeans -{ +namespace kmeans { // run Lloyds one iteration // Given data in row major num_points * dim, and centers in row major @@ -67,7 +71,8 @@ namespace kmeans // If closest_centers == NULL, will allocate memory and return. // Similarly, if closest_docs == NULL, will allocate memory and return. -float lloyds_iter(float *data, size_t num_points, size_t dim, float *centers, size_t num_centers, float *docs_l2sq, +float lloyds_iter(float *data, size_t num_points, size_t dim, float *centers, + size_t num_centers, float *docs_l2sq, std::vector *closest_docs, uint32_t *&closest_center); // Run Lloyds until max_reps or stopping criterion @@ -76,12 +81,15 @@ float lloyds_iter(float *data, size_t num_points, size_t dim, float *centers, si // vector [num_centers], and closest_center = new size_t[num_points] // Final centers are output in centers as row major num_centers * dim // -float run_lloyds(float *data, size_t num_points, size_t dim, float *centers, const size_t num_centers, - const size_t max_reps, std::vector *closest_docs, uint32_t *closest_center); +float run_lloyds(float *data, size_t num_points, size_t dim, float *centers, + const size_t num_centers, const size_t max_reps, + std::vector *closest_docs, uint32_t *closest_center); // assumes already memory allocated for pivot_data as new // float[num_centers*dim] and select randomly num_centers points as pivots -void selecting_pivots(float *data, size_t num_points, size_t dim, float *pivot_data, size_t num_centers); +void selecting_pivots(float *data, size_t num_points, size_t dim, + float *pivot_data, size_t num_centers); -void kmeanspp_selecting_pivots(float *data, size_t num_points, size_t dim, float *pivot_data, size_t num_centers); -} // namespace kmeans +void kmeanspp_selecting_pivots(float *data, size_t num_points, size_t dim, + float *pivot_data, size_t num_centers); +} // namespace kmeans diff --git a/include/memory_mapper.h b/include/memory_mapper.h index 75faca1bb..bf78d1687 100644 --- a/include/memory_mapper.h +++ b/include/memory_mapper.h @@ -15,29 +15,27 @@ #endif #include -namespace diskann -{ -class MemoryMapper -{ - private: +namespace diskann { +class MemoryMapper { + private: #ifndef _WINDOWS - int _fd; + int _fd; #else - HANDLE _bareFile; - HANDLE _fd; + HANDLE _bareFile; + HANDLE _fd; #endif - char *_buf; - size_t _fileSize; - const char *_fileName; + char *_buf; + size_t _fileSize; + const char *_fileName; - public: - MemoryMapper(const char *filename); - MemoryMapper(const std::string &filename); + public: + MemoryMapper(const char *filename); + MemoryMapper(const std::string &filename); - char *getBuf(); - size_t getFileSize(); + char *getBuf(); + size_t getFileSize(); - ~MemoryMapper(); + ~MemoryMapper(); }; -} // namespace diskann +} // namespace diskann diff --git a/include/natural_number_map.h b/include/natural_number_map.h index 820ac3fdf..6868e5551 100644 --- a/include/natural_number_map.h +++ b/include/natural_number_map.h @@ -9,8 +9,7 @@ #include -namespace diskann -{ +namespace diskann { // A map whose key is a natural number (from 0 onwards) and maps to a value. // Made as both memory and performance efficient map for scenario such as // DiskANN location-to-tag map. There, the pool of numbers is consecutive from @@ -22,68 +21,67 @@ namespace diskann // Thread-safety: this class is not thread-safe in general. // Exception: multiple read-only operations are safe on the object only if // there are no writers to it in parallel. -template class natural_number_map -{ - public: - static_assert(std::is_trivial::value, "Key must be a trivial type"); - // Some of the class member prototypes are done with this assumption to - // minimize verbosity since it's the only use case. - static_assert(std::is_trivial::value, "Value must be a trivial type"); - - // Represents a reference to a element in the map. Used while iterating - // over map entries. - struct position - { - size_t _key; - // The number of keys that were enumerated when iterating through the - // map so far. Used to early-terminate enumeration when ithere are no - // more entries in the map. - size_t _keys_already_enumerated; - - // Returns whether it's valid to access the element at this position in - // the map. - bool is_valid() const; - }; - - natural_number_map(); - - void reserve(size_t count); - size_t size() const; - - void set(Key key, Value value); - void erase(Key key); - - bool contains(Key key) const; - bool try_get(Key key, Value &value) const; - - // Returns the value at the specified position. Prerequisite: position is - // valid. - Value get(const position &pos) const; - - // Finds the first element in the map, if any. Invalidated by changes in the - // map. - position find_first() const; - - // Finds the next element in the map after the specified position. - // Invalidated by changes in the map. - position find_next(const position &after_position) const; - - void clear(); - - private: - // Number of entries in the map. Not the same as size() of the - // _values_vector below. - size_t _size; - - // Array of values. The key is the index of the value. - std::vector _values_vector; - - // Values that are in the set have the corresponding bit index set - // to 1. - // - // Use a pointer here to allow for forward declaration of dynamic_bitset - // in public headers to avoid making boost a dependency for clients - // of DiskANN. - std::unique_ptr> _values_bitset; +template +class natural_number_map { + public: + static_assert(std::is_trivial::value, "Key must be a trivial type"); + // Some of the class member prototypes are done with this assumption to + // minimize verbosity since it's the only use case. + static_assert(std::is_trivial::value, "Value must be a trivial type"); + + // Represents a reference to a element in the map. Used while iterating + // over map entries. + struct position { + size_t _key; + // The number of keys that were enumerated when iterating through the + // map so far. Used to early-terminate enumeration when ithere are no + // more entries in the map. + size_t _keys_already_enumerated; + + // Returns whether it's valid to access the element at this position in + // the map. + bool is_valid() const; + }; + + natural_number_map(); + + void reserve(size_t count); + size_t size() const; + + void set(Key key, Value value); + void erase(Key key); + + bool contains(Key key) const; + bool try_get(Key key, Value &value) const; + + // Returns the value at the specified position. Prerequisite: position is + // valid. + Value get(const position &pos) const; + + // Finds the first element in the map, if any. Invalidated by changes in the + // map. + position find_first() const; + + // Finds the next element in the map after the specified position. + // Invalidated by changes in the map. + position find_next(const position &after_position) const; + + void clear(); + + private: + // Number of entries in the map. Not the same as size() of the + // _values_vector below. + size_t _size; + + // Array of values. The key is the index of the value. + std::vector _values_vector; + + // Values that are in the set have the corresponding bit index set + // to 1. + // + // Use a pointer here to allow for forward declaration of dynamic_bitset + // in public headers to avoid making boost a dependency for clients + // of DiskANN. + std::unique_ptr> _values_bitset; }; -} // namespace diskann +} // namespace diskann diff --git a/include/natural_number_set.h b/include/natural_number_set.h index ec5b827e6..eda3de5c3 100644 --- a/include/natural_number_set.h +++ b/include/natural_number_set.h @@ -8,8 +8,7 @@ #include "boost_dynamic_bitset_fwd.h" -namespace diskann -{ +namespace diskann { // A set of natural numbers (from 0 onwards). Made for scenario where the // pool of numbers is consecutive from zero to some max value and very // efficient methods for "add to set", "get any value from set", "is in set" @@ -20,31 +19,31 @@ namespace diskann // Thread-safety: this class is not thread-safe in general. // Exception: multiple read-only operations (e.g. is_in_set, empty, size) are // safe on the object only if there are no writers to it in parallel. -template class natural_number_set -{ - public: - static_assert(std::is_trivial::value, "Identifier must be a trivial type"); - - natural_number_set(); - - bool is_empty() const; - void reserve(size_t count); - void insert(T id); - T pop_any(); - void clear(); - size_t size() const; - bool is_in_set(T id) const; - - private: - // Values that are currently in set. - std::vector _values_vector; - - // Values that are in the set have the corresponding bit index set - // to 1. - // - // Use a pointer here to allow for forward declaration of dynamic_bitset - // in public headers to avoid making boost a dependency for clients - // of DiskANN. - std::unique_ptr> _values_bitset; +template +class natural_number_set { + public: + static_assert(std::is_trivial::value, "Identifier must be a trivial type"); + + natural_number_set(); + + bool is_empty() const; + void reserve(size_t count); + void insert(T id); + T pop_any(); + void clear(); + size_t size() const; + bool is_in_set(T id) const; + + private: + // Values that are currently in set. + std::vector _values_vector; + + // Values that are in the set have the corresponding bit index set + // to 1. + // + // Use a pointer here to allow for forward declaration of dynamic_bitset + // in public headers to avoid making boost a dependency for clients + // of DiskANN. + std::unique_ptr> _values_bitset; }; -} // namespace diskann +} // namespace diskann diff --git a/include/neighbor.h b/include/neighbor.h index d7c0c25ed..b6b66eb97 100644 --- a/include/neighbor.h +++ b/include/neighbor.h @@ -8,145 +8,106 @@ #include #include "utils.h" -namespace diskann -{ +namespace diskann { -struct Neighbor -{ - unsigned id; - float distance; - bool expanded; +struct Neighbor { + unsigned id; + float distance; + bool expanded; - Neighbor() = default; + Neighbor() = default; - Neighbor(unsigned id, float distance) : id{id}, distance{distance}, expanded(false) - { - } + Neighbor(unsigned id, float distance) + : id{id}, distance{distance}, expanded(false) {} - inline bool operator<(const Neighbor &other) const - { - return distance < other.distance || (distance == other.distance && id < other.id); - } + inline bool operator<(const Neighbor &other) const { + return distance < other.distance || + (distance == other.distance && id < other.id); + } - inline bool operator==(const Neighbor &other) const - { - return (id == other.id); - } + inline bool operator==(const Neighbor &other) const { + return (id == other.id); + } }; // Invariant: after every `insert` and `closest_unexpanded()`, `_cur` points to // the first Neighbor which is unexpanded. -class NeighborPriorityQueue -{ - public: - NeighborPriorityQueue() : _size(0), _capacity(0), _cur(0) - { +class NeighborPriorityQueue { + public: + NeighborPriorityQueue() : _size(0), _capacity(0), _cur(0) {} + + explicit NeighborPriorityQueue(size_t capacity) + : _size(0), _capacity(capacity), _cur(0), _data(capacity + 1) {} + + // Inserts the item ordered into the set up to the sets capacity. + // The item will be dropped if it is the same id as an exiting + // set item or it has a greated distance than the final + // item in the set. The set cursor that is used to pop() the + // next item will be set to the lowest index of an uncheck item + void insert(const Neighbor &nbr) { + if (_size == _capacity && _data[_size - 1] < nbr) { + return; } - explicit NeighborPriorityQueue(size_t capacity) : _size(0), _capacity(capacity), _cur(0), _data(capacity + 1) - { + size_t lo = 0, hi = _size; + while (lo < hi) { + size_t mid = (lo + hi) >> 1; + if (nbr < _data[mid]) { + hi = mid; + // Make sure the same id isn't inserted into the set + } else if (_data[mid].id == nbr.id) { + return; + } else { + lo = mid + 1; + } } - // Inserts the item ordered into the set up to the sets capacity. - // The item will be dropped if it is the same id as an exiting - // set item or it has a greated distance than the final - // item in the set. The set cursor that is used to pop() the - // next item will be set to the lowest index of an uncheck item - void insert(const Neighbor &nbr) - { - if (_size == _capacity && _data[_size - 1] < nbr) - { - return; - } - - size_t lo = 0, hi = _size; - while (lo < hi) - { - size_t mid = (lo + hi) >> 1; - if (nbr < _data[mid]) - { - hi = mid; - // Make sure the same id isn't inserted into the set - } - else if (_data[mid].id == nbr.id) - { - return; - } - else - { - lo = mid + 1; - } - } - - if (lo < _capacity) - { - std::memmove(&_data[lo + 1], &_data[lo], (_size - lo) * sizeof(Neighbor)); - } - _data[lo] = {nbr.id, nbr.distance}; - if (_size < _capacity) - { - _size++; - } - if (lo < _cur) - { - _cur = lo; - } + if (lo < _capacity) { + std::memmove(&_data[lo + 1], &_data[lo], (_size - lo) * sizeof(Neighbor)); } - - Neighbor closest_unexpanded() - { - _data[_cur].expanded = true; - size_t pre = _cur; - while (_cur < _size && _data[_cur].expanded) - { - _cur++; - } - return _data[pre]; + _data[lo] = {nbr.id, nbr.distance}; + if (_size < _capacity) { + _size++; } - - bool has_unexpanded_node() const - { - return _cur < _size; + if (lo < _cur) { + _cur = lo; } + } - size_t size() const - { - return _size; + Neighbor closest_unexpanded() { + _data[_cur].expanded = true; + size_t pre = _cur; + while (_cur < _size && _data[_cur].expanded) { + _cur++; } + return _data[pre]; + } - size_t capacity() const - { - return _capacity; - } + bool has_unexpanded_node() const { return _cur < _size; } - void reserve(size_t capacity) - { - if (capacity + 1 > _data.size()) - { - _data.resize(capacity + 1); - } - _capacity = capacity; - } + size_t size() const { return _size; } - Neighbor &operator[](size_t i) - { - return _data[i]; - } + size_t capacity() const { return _capacity; } - Neighbor operator[](size_t i) const - { - return _data[i]; + void reserve(size_t capacity) { + if (capacity + 1 > _data.size()) { + _data.resize(capacity + 1); } + _capacity = capacity; + } - void clear() - { - _size = 0; - _cur = 0; - } + Neighbor &operator[](size_t i) { return _data[i]; } + + Neighbor operator[](size_t i) const { return _data[i]; } + + void clear() { + _size = 0; + _cur = 0; + } - private: - size_t _size, _capacity, _cur; - std::vector _data; + private: + size_t _size, _capacity, _cur; + std::vector _data; }; -} // namespace diskann +} // namespace diskann diff --git a/include/parameters.h b/include/parameters.h index 36f8a94e0..1a42c29dd 100644 --- a/include/parameters.h +++ b/include/parameters.h @@ -8,117 +8,121 @@ #include "defaults.h" -namespace diskann -{ +namespace diskann { class IndexWriteParameters { - public: - const uint32_t search_list_size; // L - const uint32_t max_degree; // R - const bool saturate_graph; - const uint32_t max_occlusion_size; // C - const float alpha; - const uint32_t num_rounds; - const uint32_t num_threads; - const uint32_t filter_list_size; // Lf - const uint32_t num_frozen_points; - - private: - IndexWriteParameters(const uint32_t search_list_size, const uint32_t max_degree, const bool saturate_graph, - const uint32_t max_occlusion_size, const float alpha, const uint32_t num_rounds, - const uint32_t num_threads, const uint32_t filter_list_size, const uint32_t num_frozen_points) - : search_list_size(search_list_size), max_degree(max_degree), saturate_graph(saturate_graph), - max_occlusion_size(max_occlusion_size), alpha(alpha), num_rounds(num_rounds), num_threads(num_threads), - filter_list_size(filter_list_size), num_frozen_points(num_frozen_points) - { - } - - friend class IndexWriteParametersBuilder; + public: + const uint32_t search_list_size; // L + const uint32_t max_degree; // R + const bool saturate_graph; + const uint32_t max_occlusion_size; // C + const float alpha; + const uint32_t num_rounds; + const uint32_t num_threads; + const uint32_t filter_list_size; // Lf + const uint32_t num_frozen_points; + + private: + IndexWriteParameters(const uint32_t search_list_size, + const uint32_t max_degree, const bool saturate_graph, + const uint32_t max_occlusion_size, const float alpha, + const uint32_t num_rounds, const uint32_t num_threads, + const uint32_t filter_list_size, + const uint32_t num_frozen_points) + : search_list_size(search_list_size), + max_degree(max_degree), + saturate_graph(saturate_graph), + max_occlusion_size(max_occlusion_size), + alpha(alpha), + num_rounds(num_rounds), + num_threads(num_threads), + filter_list_size(filter_list_size), + num_frozen_points(num_frozen_points) {} + + friend class IndexWriteParametersBuilder; }; -class IndexWriteParametersBuilder -{ - /** - * Fluent builder pattern to keep track of the 7 non-default properties - * and their order. The basic ctor was getting unwieldy. - */ - public: - IndexWriteParametersBuilder(const uint32_t search_list_size, // L - const uint32_t max_degree // R - ) - : _search_list_size(search_list_size), _max_degree(max_degree) - { - } - - IndexWriteParametersBuilder &with_max_occlusion_size(const uint32_t max_occlusion_size) - { - _max_occlusion_size = max_occlusion_size; - return *this; - } - - IndexWriteParametersBuilder &with_saturate_graph(const bool saturate_graph) - { - _saturate_graph = saturate_graph; - return *this; - } - - IndexWriteParametersBuilder &with_alpha(const float alpha) - { - _alpha = alpha; - return *this; - } - - IndexWriteParametersBuilder &with_num_rounds(const uint32_t num_rounds) - { - _num_rounds = num_rounds; - return *this; - } - - IndexWriteParametersBuilder &with_num_threads(const uint32_t num_threads) - { - _num_threads = num_threads; - return *this; - } - - IndexWriteParametersBuilder &with_filter_list_size(const uint32_t filter_list_size) - { - _filter_list_size = filter_list_size; - return *this; - } - - IndexWriteParametersBuilder &with_num_frozen_points(const uint32_t num_frozen_points) - { - _num_frozen_points = num_frozen_points; - return *this; - } - - IndexWriteParameters build() const - { - return IndexWriteParameters(_search_list_size, _max_degree, _saturate_graph, _max_occlusion_size, _alpha, - _num_rounds, _num_threads, _filter_list_size, _num_frozen_points); - } - - IndexWriteParametersBuilder(const IndexWriteParameters &wp) - : _search_list_size(wp.search_list_size), _max_degree(wp.max_degree), - _max_occlusion_size(wp.max_occlusion_size), _saturate_graph(wp.saturate_graph), _alpha(wp.alpha), - _num_rounds(wp.num_rounds), _filter_list_size(wp.filter_list_size), _num_frozen_points(wp.num_frozen_points) - { - } - IndexWriteParametersBuilder(const IndexWriteParametersBuilder &) = delete; - IndexWriteParametersBuilder &operator=(const IndexWriteParametersBuilder &) = delete; - - private: - uint32_t _search_list_size{}; - uint32_t _max_degree{}; - uint32_t _max_occlusion_size{defaults::MAX_OCCLUSION_SIZE}; - bool _saturate_graph{defaults::SATURATE_GRAPH}; - float _alpha{defaults::ALPHA}; - uint32_t _num_rounds{defaults::NUM_ROUNDS}; - uint32_t _num_threads{defaults::NUM_THREADS}; - uint32_t _filter_list_size{defaults::FILTER_LIST_SIZE}; - uint32_t _num_frozen_points{defaults::NUM_FROZEN_POINTS}; +class IndexWriteParametersBuilder { + /** + * Fluent builder pattern to keep track of the 7 non-default properties + * and their order. The basic ctor was getting unwieldy. + */ + public: + IndexWriteParametersBuilder(const uint32_t search_list_size, // L + const uint32_t max_degree // R + ) + : _search_list_size(search_list_size), _max_degree(max_degree) {} + + IndexWriteParametersBuilder &with_max_occlusion_size( + const uint32_t max_occlusion_size) { + _max_occlusion_size = max_occlusion_size; + return *this; + } + + IndexWriteParametersBuilder &with_saturate_graph(const bool saturate_graph) { + _saturate_graph = saturate_graph; + return *this; + } + + IndexWriteParametersBuilder &with_alpha(const float alpha) { + _alpha = alpha; + return *this; + } + + IndexWriteParametersBuilder &with_num_rounds(const uint32_t num_rounds) { + _num_rounds = num_rounds; + return *this; + } + + IndexWriteParametersBuilder &with_num_threads(const uint32_t num_threads) { + _num_threads = num_threads; + return *this; + } + + IndexWriteParametersBuilder &with_filter_list_size( + const uint32_t filter_list_size) { + _filter_list_size = filter_list_size; + return *this; + } + + IndexWriteParametersBuilder &with_num_frozen_points( + const uint32_t num_frozen_points) { + _num_frozen_points = num_frozen_points; + return *this; + } + + IndexWriteParameters build() const { + return IndexWriteParameters(_search_list_size, _max_degree, _saturate_graph, + _max_occlusion_size, _alpha, _num_rounds, + _num_threads, _filter_list_size, + _num_frozen_points); + } + + IndexWriteParametersBuilder(const IndexWriteParameters &wp) + : _search_list_size(wp.search_list_size), + _max_degree(wp.max_degree), + _max_occlusion_size(wp.max_occlusion_size), + _saturate_graph(wp.saturate_graph), + _alpha(wp.alpha), + _num_rounds(wp.num_rounds), + _filter_list_size(wp.filter_list_size), + _num_frozen_points(wp.num_frozen_points) {} + IndexWriteParametersBuilder(const IndexWriteParametersBuilder &) = delete; + IndexWriteParametersBuilder &operator=(const IndexWriteParametersBuilder &) = + delete; + + private: + uint32_t _search_list_size{}; + uint32_t _max_degree{}; + uint32_t _max_occlusion_size{defaults::MAX_OCCLUSION_SIZE}; + bool _saturate_graph{defaults::SATURATE_GRAPH}; + float _alpha{defaults::ALPHA}; + uint32_t _num_rounds{defaults::NUM_ROUNDS}; + uint32_t _num_threads{defaults::NUM_THREADS}; + uint32_t _filter_list_size{defaults::FILTER_LIST_SIZE}; + uint32_t _num_frozen_points{defaults::NUM_FROZEN_POINTS}; }; -} // namespace diskann +} // namespace diskann diff --git a/include/partition.h b/include/partition.h index c2c4c76ad..a5af890be 100644 --- a/include/partition.h +++ b/include/partition.h @@ -16,34 +16,45 @@ #include "windows_customizations.h" template -void gen_random_slice(const std::string base_file, const std::string output_prefix, double sampling_rate); +void gen_random_slice(const std::string base_file, + const std::string output_prefix, double sampling_rate); template -void gen_random_slice(const std::string data_file, double p_val, float *&sampled_data, size_t &slice_size, - size_t &ndims); +void gen_random_slice(const std::string data_file, double p_val, + float *&sampled_data, size_t &slice_size, size_t &ndims); template -void gen_random_slice(const T *inputdata, size_t npts, size_t ndims, double p_val, float *&sampled_data, - size_t &slice_size); +void gen_random_slice(const T *inputdata, size_t npts, size_t ndims, + double p_val, float *&sampled_data, size_t &slice_size); -int estimate_cluster_sizes(float *test_data_float, size_t num_test, float *pivots, const size_t num_centers, - const size_t dim, const size_t k_base, std::vector &cluster_sizes); +int estimate_cluster_sizes(float *test_data_float, size_t num_test, + float *pivots, const size_t num_centers, + const size_t dim, const size_t k_base, + std::vector &cluster_sizes); template -int shard_data_into_clusters(const std::string data_file, float *pivots, const size_t num_centers, const size_t dim, +int shard_data_into_clusters(const std::string data_file, float *pivots, + const size_t num_centers, const size_t dim, const size_t k_base, std::string prefix_path); template -int shard_data_into_clusters_only_ids(const std::string data_file, float *pivots, const size_t num_centers, - const size_t dim, const size_t k_base, std::string prefix_path); +int shard_data_into_clusters_only_ids(const std::string data_file, + float *pivots, const size_t num_centers, + const size_t dim, const size_t k_base, + std::string prefix_path); template -int retrieve_shard_data_from_ids(const std::string data_file, std::string idmap_filename, std::string data_filename); +int retrieve_shard_data_from_ids(const std::string data_file, + std::string idmap_filename, + std::string data_filename); template -int partition(const std::string data_file, const float sampling_rate, size_t num_centers, size_t max_k_means_reps, +int partition(const std::string data_file, const float sampling_rate, + size_t num_centers, size_t max_k_means_reps, const std::string prefix_path, size_t k_base); template -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); +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); diff --git a/include/percentile_stats.h b/include/percentile_stats.h index 793257577..c7a70fe16 100644 --- a/include/percentile_stats.h +++ b/include/percentile_stats.h @@ -16,50 +16,48 @@ #include "distance.h" #include "parameters.h" -namespace diskann -{ -struct QueryStats -{ - float total_us = 0; // total time to process query in micros - float io_us = 0; // total time spent in IO - float cpu_us = 0; // total time spent in CPU - - unsigned n_4k = 0; // # of 4kB reads - unsigned n_8k = 0; // # of 8kB reads - unsigned n_12k = 0; // # of 12kB reads - unsigned n_ios = 0; // total # of IOs issued - unsigned read_size = 0; // total # of bytes read - unsigned n_cmps_saved = 0; // # cmps saved - unsigned n_cmps = 0; // # cmps - unsigned n_cache_hits = 0; // # cache_hits - unsigned n_hops = 0; // # search hops +namespace diskann { +struct QueryStats { + float total_us = 0; // total time to process query in micros + float io_us = 0; // total time spent in IO + float cpu_us = 0; // total time spent in CPU + + unsigned n_4k = 0; // # of 4kB reads + unsigned n_8k = 0; // # of 8kB reads + unsigned n_12k = 0; // # of 12kB reads + unsigned n_ios = 0; // total # of IOs issued + unsigned read_size = 0; // total # of bytes read + unsigned n_cmps_saved = 0; // # cmps saved + unsigned n_cmps = 0; // # cmps + unsigned n_cache_hits = 0; // # cache_hits + unsigned n_hops = 0; // # search hops }; template -inline T get_percentile_stats(QueryStats *stats, uint64_t len, float percentile, - const std::function &member_fn) -{ - std::vector vals(len); - for (uint64_t i = 0; i < len; i++) - { - vals[i] = member_fn(stats[i]); - } - - std::sort(vals.begin(), vals.end(), [](const T &left, const T &right) { return left < right; }); - - auto retval = vals[(uint64_t)(percentile * len)]; - vals.clear(); - return retval; +inline T get_percentile_stats( + QueryStats *stats, uint64_t len, float percentile, + const std::function &member_fn) { + std::vector vals(len); + for (uint64_t i = 0; i < len; i++) { + vals[i] = member_fn(stats[i]); + } + + std::sort(vals.begin(), vals.end(), + [](const T &left, const T &right) { return left < right; }); + + auto retval = vals[(uint64_t)(percentile * len)]; + vals.clear(); + return retval; } template -inline double get_mean_stats(QueryStats *stats, uint64_t len, const std::function &member_fn) -{ - double avg = 0; - for (uint64_t i = 0; i < len; i++) - { - avg += (double)member_fn(stats[i]); - } - return avg / len; +inline double get_mean_stats( + QueryStats *stats, uint64_t len, + const std::function &member_fn) { + double avg = 0; + for (uint64_t i = 0; i < len; i++) { + avg += (double)member_fn(stats[i]); + } + return avg / len; } -} // namespace diskann +} // namespace diskann diff --git a/include/pq.h b/include/pq.h index acfa1b30a..f26f566aa 100644 --- a/include/pq.h +++ b/include/pq.h @@ -12,114 +12,134 @@ #define MAX_PQ_TRAINING_SET_SIZE 256000 #define MAX_PQ_CHUNKS 512 -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(); +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); + 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); + void load_pq_centroid_bin(const char *pq_table_file, size_t num_chunks); #endif - uint32_t get_num_chunks(); + uint32_t get_num_chunks(); - void preprocess_query(float *query_vec); + void preprocess_query(float *query_vec); - // assumes pre-processed query - void populate_chunk_distances(const float *query_vec, float *dist_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 l2_distance(const float *query_vec, uint8_t *base_vec); - float inner_product(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); + // 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 populate_chunk_inner_products(const float *query_vec, float *dist_vec); }; -template struct PQScratch -{ - float *aligned_pqtable_dist_scratch = nullptr; // MUST BE AT LEAST [256 * NCHUNKS] - float *aligned_dist_scratch = nullptr; // MUST BE AT LEAST diskann MAX_DEGREE - uint8_t *aligned_pq_coord_scratch = nullptr; // MUST BE AT LEAST [N_CHUNKS * MAX_DEGREE] - float *rotated_query = nullptr; - float *aligned_query_float = nullptr; - - PQScratch(size_t graph_degree, size_t aligned_dim) - { - diskann::alloc_aligned((void **)&aligned_pq_coord_scratch, - (size_t)graph_degree * (size_t)MAX_PQ_CHUNKS * sizeof(uint8_t), 256); - diskann::alloc_aligned((void **)&aligned_pqtable_dist_scratch, 256 * (size_t)MAX_PQ_CHUNKS * sizeof(float), - 256); - diskann::alloc_aligned((void **)&aligned_dist_scratch, (size_t)graph_degree * sizeof(float), 256); - diskann::alloc_aligned((void **)&aligned_query_float, aligned_dim * sizeof(float), 8 * sizeof(float)); - diskann::alloc_aligned((void **)&rotated_query, aligned_dim * sizeof(float), 8 * sizeof(float)); - - memset(aligned_query_float, 0, aligned_dim * sizeof(float)); - memset(rotated_query, 0, aligned_dim * sizeof(float)); - } - - void set(size_t dim, T *query, const float norm = 1.0f) - { - for (size_t d = 0; d < dim; ++d) - { - if (norm != 1.0f) - rotated_query[d] = aligned_query_float[d] = static_cast(query[d]) / norm; - else - rotated_query[d] = aligned_query_float[d] = static_cast(query[d]); - } +template +struct PQScratch { + float *aligned_pqtable_dist_scratch = + nullptr; // MUST BE AT LEAST [256 * NCHUNKS] + float *aligned_dist_scratch = nullptr; // MUST BE AT LEAST diskann MAX_DEGREE + uint8_t *aligned_pq_coord_scratch = + nullptr; // MUST BE AT LEAST [N_CHUNKS * MAX_DEGREE] + float *rotated_query = nullptr; + float *aligned_query_float = nullptr; + + PQScratch(size_t graph_degree, size_t aligned_dim) { + diskann::alloc_aligned( + (void **)&aligned_pq_coord_scratch, + (size_t)graph_degree * (size_t)MAX_PQ_CHUNKS * sizeof(uint8_t), 256); + diskann::alloc_aligned((void **)&aligned_pqtable_dist_scratch, + 256 * (size_t)MAX_PQ_CHUNKS * sizeof(float), 256); + diskann::alloc_aligned((void **)&aligned_dist_scratch, + (size_t)graph_degree * sizeof(float), 256); + diskann::alloc_aligned((void **)&aligned_query_float, + aligned_dim * sizeof(float), 8 * sizeof(float)); + diskann::alloc_aligned((void **)&rotated_query, aligned_dim * sizeof(float), + 8 * sizeof(float)); + + memset(aligned_query_float, 0, aligned_dim * sizeof(float)); + memset(rotated_query, 0, aligned_dim * sizeof(float)); + } + + void set(size_t dim, T *query, const float norm = 1.0f) { + for (size_t d = 0; d < dim; ++d) { + if (norm != 1.0f) + rotated_query[d] = aligned_query_float[d] = + static_cast(query[d]) / norm; + else + rotated_query[d] = aligned_query_float[d] = + static_cast(query[d]); } + } }; -void aggregate_coords(const std::vector &ids, const uint8_t *all_coords, const uint64_t ndims, uint8_t *out); +void aggregate_coords(const std::vector &ids, + const uint8_t *all_coords, const uint64_t ndims, + uint8_t *out); -void pq_dist_lookup(const uint8_t *pq_ids, const size_t n_pts, const size_t pq_nchunks, const float *pq_dists, +void pq_dist_lookup(const uint8_t *pq_ids, const size_t n_pts, + const size_t pq_nchunks, const float *pq_dists, std::vector &dists_out); // Need to replace calls to these with calls to vector& based functions above -void aggregate_coords(const unsigned *ids, const uint64_t n_ids, const uint8_t *all_coords, const uint64_t ndims, +void aggregate_coords(const unsigned *ids, const uint64_t n_ids, + const uint8_t *all_coords, const uint64_t ndims, uint8_t *out); -void pq_dist_lookup(const uint8_t *pq_ids, const size_t n_pts, const size_t pq_nchunks, const float *pq_dists, +void pq_dist_lookup(const uint8_t *pq_ids, const size_t n_pts, + const size_t pq_nchunks, const float *pq_dists, float *dists_out); -DISKANN_DLLEXPORT int generate_pq_pivots(const float *const 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, bool make_zero_mean = false); +DISKANN_DLLEXPORT int generate_pq_pivots( + const float *const 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, bool make_zero_mean = false); -DISKANN_DLLEXPORT int generate_opq_pivots(const float *train_data, size_t num_train, unsigned dim, unsigned num_centers, - unsigned num_pq_chunks, std::string opq_pivots_path, +DISKANN_DLLEXPORT int generate_opq_pivots(const float *train_data, + size_t num_train, unsigned dim, + unsigned num_centers, + unsigned num_pq_chunks, + std::string opq_pivots_path, bool make_zero_mean = false); template -int generate_pq_data_from_pivots(const std::string &data_file, unsigned num_centers, unsigned num_pq_chunks, - const std::string &pq_pivots_path, const std::string &pq_compressed_vectors_path, +int generate_pq_data_from_pivots(const std::string &data_file, + unsigned num_centers, unsigned num_pq_chunks, + const std::string &pq_pivots_path, + const std::string &pq_compressed_vectors_path, bool use_opq = false); template -void generate_disk_quantized_data(const std::string &data_file_to_use, const std::string &disk_pq_pivots_path, - const std::string &disk_pq_compressed_vectors_path, - const diskann::Metric compareMetric, const double p_val, size_t &disk_pq_dims); +void generate_disk_quantized_data( + const std::string &data_file_to_use, const std::string &disk_pq_pivots_path, + const std::string &disk_pq_compressed_vectors_path, + const diskann::Metric compareMetric, const double p_val, + size_t &disk_pq_dims); template -void generate_quantized_data(const std::string &data_file_to_use, const std::string &pq_pivots_path, - const std::string &pq_compressed_vectors_path, const diskann::Metric compareMetric, - const double p_val, const uint64_t num_pq_chunks, const bool use_opq, +void generate_quantized_data(const std::string &data_file_to_use, + const std::string &pq_pivots_path, + const std::string &pq_compressed_vectors_path, + const diskann::Metric compareMetric, + const double p_val, const uint64_t num_pq_chunks, + const bool use_opq, const std::string &codebook_prefix = ""); -} // namespace diskann +} // namespace diskann diff --git a/include/pq_flash_index.h b/include/pq_flash_index.h index e11c0dc31..df5dc5a99 100644 --- a/include/pq_flash_index.h +++ b/include/pq_flash_index.h @@ -18,191 +18,209 @@ #define FULL_PRECISION_REORDER_MULTIPLIER 3 -namespace diskann -{ +namespace diskann { -template class PQFlashIndex -{ - public: - DISKANN_DLLEXPORT PQFlashIndex(std::shared_ptr &fileReader, - diskann::Metric metric = diskann::Metric::L2); - DISKANN_DLLEXPORT ~PQFlashIndex(); +template +class PQFlashIndex { + public: + DISKANN_DLLEXPORT PQFlashIndex(std::shared_ptr &fileReader, + diskann::Metric metric = diskann::Metric::L2); + DISKANN_DLLEXPORT ~PQFlashIndex(); #ifdef EXEC_ENV_OLS - DISKANN_DLLEXPORT int load(diskann::MemoryMappedFiles &files, uint32_t num_threads, const char *index_prefix); + DISKANN_DLLEXPORT int load(diskann::MemoryMappedFiles &files, + uint32_t num_threads, const char *index_prefix); #else - // load compressed data, and obtains the handle to the disk-resident index - DISKANN_DLLEXPORT int load(uint32_t num_threads, const char *index_prefix); + // load compressed data, and obtains the handle to the disk-resident index + DISKANN_DLLEXPORT int load(uint32_t num_threads, const char *index_prefix); #endif #ifdef EXEC_ENV_OLS - DISKANN_DLLEXPORT int load_from_separate_paths(diskann::MemoryMappedFiles &files, uint32_t num_threads, - const char *index_filepath, const char *pivots_filepath, - const char *compressed_filepath); + DISKANN_DLLEXPORT int load_from_separate_paths( + diskann::MemoryMappedFiles &files, uint32_t num_threads, + const char *index_filepath, const char *pivots_filepath, + const char *compressed_filepath); #else - DISKANN_DLLEXPORT int load_from_separate_paths(uint32_t num_threads, const char *index_filepath, - const char *pivots_filepath, const char *compressed_filepath); + DISKANN_DLLEXPORT int load_from_separate_paths( + uint32_t num_threads, const char *index_filepath, + const char *pivots_filepath, const char *compressed_filepath); #endif - DISKANN_DLLEXPORT void load_cache_list(std::vector &node_list); + DISKANN_DLLEXPORT void load_cache_list(std::vector &node_list); #ifdef EXEC_ENV_OLS - DISKANN_DLLEXPORT void generate_cache_list_from_sample_queries(MemoryMappedFiles &files, std::string sample_bin, - uint64_t l_search, uint64_t beamwidth, - uint64_t num_nodes_to_cache, uint32_t nthreads, - std::vector &node_list); + DISKANN_DLLEXPORT void generate_cache_list_from_sample_queries( + MemoryMappedFiles &files, std::string sample_bin, uint64_t l_search, + uint64_t beamwidth, uint64_t num_nodes_to_cache, uint32_t nthreads, + std::vector &node_list); #else - DISKANN_DLLEXPORT void generate_cache_list_from_sample_queries(std::string sample_bin, uint64_t l_search, - uint64_t beamwidth, uint64_t num_nodes_to_cache, - uint32_t num_threads, - std::vector &node_list); + DISKANN_DLLEXPORT void generate_cache_list_from_sample_queries( + std::string sample_bin, uint64_t l_search, uint64_t beamwidth, + uint64_t num_nodes_to_cache, uint32_t num_threads, + std::vector &node_list); #endif - DISKANN_DLLEXPORT void cache_bfs_levels(uint64_t num_nodes_to_cache, std::vector &node_list, - const bool shuffle = false); - - DISKANN_DLLEXPORT void cached_beam_search(const T *query, const uint64_t k_search, const uint64_t l_search, - uint64_t *res_ids, float *res_dists, const uint64_t beam_width, - const bool use_reorder_data = false, QueryStats *stats = nullptr); - - DISKANN_DLLEXPORT void cached_beam_search(const T *query, const uint64_t k_search, const uint64_t l_search, - uint64_t *res_ids, float *res_dists, const uint64_t beam_width, - const bool use_filter, const LabelT &filter_label, - const bool use_reorder_data = false, QueryStats *stats = nullptr); - - DISKANN_DLLEXPORT void cached_beam_search(const T *query, const uint64_t k_search, const uint64_t l_search, - uint64_t *res_ids, float *res_dists, const uint64_t beam_width, - const uint32_t io_limit, const bool use_reorder_data = false, - QueryStats *stats = nullptr); - - DISKANN_DLLEXPORT void cached_beam_search(const T *query, const uint64_t k_search, const uint64_t l_search, - uint64_t *res_ids, float *res_dists, const uint64_t beam_width, - const bool use_filter, const LabelT &filter_label, - const uint32_t io_limit, const bool use_reorder_data = false, - QueryStats *stats = nullptr); - - DISKANN_DLLEXPORT LabelT get_converted_label(const std::string &filter_label); - - DISKANN_DLLEXPORT uint32_t range_search(const T *query1, const double range, const uint64_t min_l_search, - const uint64_t max_l_search, std::vector &indices, - std::vector &distances, const uint64_t min_beam_width, - QueryStats *stats = nullptr); - - DISKANN_DLLEXPORT uint64_t get_data_dim(); - - std::shared_ptr &reader; - - DISKANN_DLLEXPORT diskann::Metric get_metric(); - - protected: - DISKANN_DLLEXPORT void use_medoids_data_as_centroids(); - DISKANN_DLLEXPORT void setup_thread_data(uint64_t nthreads, uint64_t visited_reserve = 4096); - - DISKANN_DLLEXPORT void set_universal_label(const LabelT &label); - - private: - DISKANN_DLLEXPORT inline bool point_has_label(uint32_t point_id, uint32_t label_id); - std::unordered_map load_label_map(const std::string &map_file); - DISKANN_DLLEXPORT void parse_label_file(const std::string &map_file, size_t &num_pts_labels); - DISKANN_DLLEXPORT void get_label_file_metadata(std::string map_file, uint32_t &num_pts, uint32_t &num_total_labels); - DISKANN_DLLEXPORT inline int32_t get_filter_number(const LabelT &filter_label); - - // index info - // nhood of node `i` is in sector: [i / nnodes_per_sector] - // offset in sector: [(i % nnodes_per_sector) * max_node_len] - // nnbrs of node `i`: *(unsigned*) (buf) - // nbrs of node `i`: ((unsigned*)buf) + 1 - - uint64_t max_node_len = 0, nnodes_per_sector = 0, max_degree = 0; - - // Data used for searching with re-order vectors - uint64_t ndims_reorder_vecs = 0, reorder_data_start_sector = 0, nvecs_per_sector = 0; - - diskann::Metric metric = diskann::Metric::L2; - - // used only for inner product search to re-scale the result value - // (due to the pre-processing of base during index build) - float max_base_norm = 0.0f; - - // data info - uint64_t num_points = 0; - uint64_t num_frozen_points = 0; - uint64_t frozen_location = 0; - uint64_t data_dim = 0; - uint64_t disk_data_dim = 0; // will be different from data_dim only if we use - // PQ for disk data (very large dimensionality) - uint64_t aligned_dim = 0; - uint64_t disk_bytes_per_point = 0; - - std::string disk_index_file; - std::vector> node_visit_counter; - - // PQ data - // n_chunks = # of chunks ndims is split into - // data: char * n_chunks - // chunk_size = chunk size of each dimension chunk - // pq_tables = float* [[2^8 * [chunk_size]] * n_chunks] - uint8_t *data = nullptr; - uint64_t n_chunks; - FixedChunkPQTable pq_table; - - // distance comparator - std::shared_ptr> dist_cmp; - std::shared_ptr> dist_cmp_float; - - // for very large datasets: we use PQ even for the disk resident index - bool use_disk_index_pq = false; - uint64_t disk_pq_n_chunks = 0; - FixedChunkPQTable disk_pq_table; - - // medoid/start info - - // graph has one entry point by default, - // we can optionally have multiple starting points - uint32_t *medoids = nullptr; - // defaults to 1 - size_t num_medoids; - // by default, it is empty. If there are multiple - // centroids, we pick the medoid corresponding to the - // closest centroid as the starting point of search - float *centroid_data = nullptr; - - // nhood_cache - unsigned *nhood_cache_buf = nullptr; - tsl::robin_map> nhood_cache; - - // coord_cache - T *coord_cache_buf = nullptr; - tsl::robin_map coord_cache; - - // thread-specific scratch - ConcurrentQueue *> thread_data; - uint64_t max_nthreads; - bool load_flag = false; - bool count_visited_nodes = false; - bool reorder_data_exists = false; - uint64_t reoreder_data_offset = 0; - - // filter support - uint32_t *_pts_to_label_offsets = nullptr; - uint32_t *_pts_to_labels = nullptr; - tsl::robin_set _labels; - std::unordered_map _filter_to_medoid_id; - bool _use_universal_label; - uint32_t _universal_filter_num; - std::vector _filter_list; - tsl::robin_set _dummy_pts; - tsl::robin_set _has_dummy_pts; - tsl::robin_map _dummy_to_real_map; - tsl::robin_map> _real_to_dummy_map; - std::unordered_map _label_map; + DISKANN_DLLEXPORT void cache_bfs_levels(uint64_t num_nodes_to_cache, + std::vector &node_list, + const bool shuffle = false); + + DISKANN_DLLEXPORT void cached_beam_search( + const T *query, const uint64_t k_search, const uint64_t l_search, + uint64_t *res_ids, float *res_dists, const uint64_t beam_width, + const bool use_reorder_data = false, QueryStats *stats = nullptr); + + DISKANN_DLLEXPORT void cached_beam_search( + const T *query, const uint64_t k_search, const uint64_t l_search, + uint64_t *res_ids, float *res_dists, const uint64_t beam_width, + const bool use_filter, const LabelT &filter_label, + const bool use_reorder_data = false, QueryStats *stats = nullptr); + + DISKANN_DLLEXPORT void cached_beam_search( + const T *query, const uint64_t k_search, const uint64_t l_search, + uint64_t *res_ids, float *res_dists, const uint64_t beam_width, + const uint32_t io_limit, const bool use_reorder_data = false, + QueryStats *stats = nullptr); + + DISKANN_DLLEXPORT void cached_beam_search( + const T *query, const uint64_t k_search, const uint64_t l_search, + uint64_t *res_ids, float *res_dists, const uint64_t beam_width, + const bool use_filter, const LabelT &filter_label, + const uint32_t io_limit, const bool use_reorder_data = false, + QueryStats *stats = nullptr); + + DISKANN_DLLEXPORT LabelT get_converted_label(const std::string &filter_label); + + DISKANN_DLLEXPORT uint32_t range_search(const T *query1, const double range, + const uint64_t min_l_search, + const uint64_t max_l_search, + std::vector &indices, + std::vector &distances, + const uint64_t min_beam_width, + QueryStats *stats = nullptr); + + DISKANN_DLLEXPORT uint64_t get_data_dim(); + + std::shared_ptr &reader; + + DISKANN_DLLEXPORT diskann::Metric get_metric(); + + protected: + DISKANN_DLLEXPORT void use_medoids_data_as_centroids(); + DISKANN_DLLEXPORT void setup_thread_data(uint64_t nthreads, + uint64_t visited_reserve = 4096); + + DISKANN_DLLEXPORT void set_universal_label(const LabelT &label); + + private: + DISKANN_DLLEXPORT inline bool point_has_label(uint32_t point_id, + uint32_t label_id); + std::unordered_map load_label_map( + const std::string &map_file); + DISKANN_DLLEXPORT void parse_label_file(const std::string &map_file, + size_t &num_pts_labels); + DISKANN_DLLEXPORT void get_label_file_metadata(std::string map_file, + uint32_t &num_pts, + uint32_t &num_total_labels); + DISKANN_DLLEXPORT inline int32_t get_filter_number( + const LabelT &filter_label); + + // index info + // nhood of node `i` is in sector: [i / nnodes_per_sector] + // offset in sector: [(i % nnodes_per_sector) * max_node_len] + // nnbrs of node `i`: *(unsigned*) (buf) + // nbrs of node `i`: ((unsigned*)buf) + 1 + + uint64_t max_node_len = 0, nnodes_per_sector = 0, max_degree = 0; + + // Data used for searching with re-order vectors + uint64_t ndims_reorder_vecs = 0, reorder_data_start_sector = 0, + nvecs_per_sector = 0; + + diskann::Metric metric = diskann::Metric::L2; + + // used only for inner product search to re-scale the result value + // (due to the pre-processing of base during index build) + float max_base_norm = 0.0f; + + // data info + uint64_t num_points = 0; + uint64_t num_frozen_points = 0; + uint64_t frozen_location = 0; + uint64_t data_dim = 0; + uint64_t disk_data_dim = 0; // will be different from data_dim only if we use + // PQ for disk data (very large dimensionality) + uint64_t aligned_dim = 0; + uint64_t disk_bytes_per_point = 0; + + std::string disk_index_file; + std::vector> node_visit_counter; + + // PQ data + // n_chunks = # of chunks ndims is split into + // data: char * n_chunks + // chunk_size = chunk size of each dimension chunk + // pq_tables = float* [[2^8 * [chunk_size]] * n_chunks] + uint8_t *data = nullptr; + uint64_t n_chunks; + FixedChunkPQTable pq_table; + + // distance comparator + std::shared_ptr> dist_cmp; + std::shared_ptr> dist_cmp_float; + + // for very large datasets: we use PQ even for the disk resident index + bool use_disk_index_pq = false; + uint64_t disk_pq_n_chunks = 0; + FixedChunkPQTable disk_pq_table; + + // medoid/start info + + // graph has one entry point by default, + // we can optionally have multiple starting points + uint32_t *medoids = nullptr; + // defaults to 1 + size_t num_medoids; + // by default, it is empty. If there are multiple + // centroids, we pick the medoid corresponding to the + // closest centroid as the starting point of search + float *centroid_data = nullptr; + + // nhood_cache + unsigned *nhood_cache_buf = nullptr; + tsl::robin_map> nhood_cache; + + // coord_cache + T *coord_cache_buf = nullptr; + tsl::robin_map coord_cache; + + // thread-specific scratch + ConcurrentQueue *> thread_data; + uint64_t max_nthreads; + bool load_flag = false; + bool count_visited_nodes = false; + bool reorder_data_exists = false; + uint64_t reoreder_data_offset = 0; + + // filter support + uint32_t *_pts_to_label_offsets = nullptr; + uint32_t *_pts_to_labels = nullptr; + tsl::robin_set _labels; + std::unordered_map _filter_to_medoid_id; + bool _use_universal_label; + uint32_t _universal_filter_num; + std::vector _filter_list; + tsl::robin_set _dummy_pts; + tsl::robin_set _has_dummy_pts; + tsl::robin_map _dummy_to_real_map; + tsl::robin_map> _real_to_dummy_map; + std::unordered_map _label_map; #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 - // on how big the header can be. - static const int HEADER_SIZE = SECTOR_LEN; - char *getHeaderBytes(); + // Set to a larger value than the actual header to accommodate + // any additions we make to the header. This is an outer limit + // on how big the header can be. + static const int HEADER_SIZE = SECTOR_LEN; + char *getHeaderBytes(); #endif }; -} // namespace diskann +} // namespace diskann diff --git a/include/scratch.h b/include/scratch.h index 9425e1885..0875275d3 100644 --- a/include/scratch.h +++ b/include/scratch.h @@ -25,210 +25,167 @@ #define SECTOR_LEN (size_t)4096 #define MAX_N_SECTOR_READS 128 -namespace diskann -{ - +namespace diskann { // // Scratch space for in-memory index based search // -template class InMemQueryScratch -{ - public: - ~InMemQueryScratch(); - //REFACTOR TODO: move all parameters to a new class. - InMemQueryScratch(uint32_t search_l, uint32_t indexing_l, uint32_t r, uint32_t maxc, size_t dim, size_t aligned_dim, size_t alignment_factor, - bool init_pq_scratch = false); - void resize_for_new_L(uint32_t new_search_l); - void clear(); - - inline uint32_t get_L() - { - return _L; - } - inline uint32_t get_R() - { - return _R; - } - inline uint32_t get_maxc() - { - return _maxc; - } - inline T *aligned_query() - { - return _aligned_query; - } - inline PQScratch *pq_scratch() - { - return _pq_scratch; - } - inline std::vector &pool() - { - return _pool; - } - inline NeighborPriorityQueue &best_l_nodes() - { - return _best_l_nodes; - } - inline std::vector &occlude_factor() - { - return _occlude_factor; - } - inline tsl::robin_set &inserted_into_pool_rs() - { - return _inserted_into_pool_rs; - } - inline boost::dynamic_bitset<> &inserted_into_pool_bs() - { - return *_inserted_into_pool_bs; - } - inline std::vector &id_scratch() - { - return _id_scratch; - } - inline std::vector &dist_scratch() - { - return _dist_scratch; - } - inline tsl::robin_set &expanded_nodes_set() - { - return _expanded_nodes_set; - } - inline std::vector &expanded_nodes_vec() - { - return _expanded_nghrs_vec; - } - inline std::vector &occlude_list_output() - { - return _occlude_list_output; - } - - private: - uint32_t _L; - uint32_t _R; - uint32_t _maxc; - - T *_aligned_query = nullptr; - - PQScratch *_pq_scratch = nullptr; - - // _pool stores all neighbors explored from best_L_nodes. - // Usually around L+R, but could be higher. - // Initialized to 3L+R for some slack, expands as needed. - std::vector _pool; - - // _best_l_nodes is reserved for storing best L entries - // Underlying storage is L+1 to support inserts - NeighborPriorityQueue _best_l_nodes; - - // _occlude_factor.size() >= pool.size() in occlude_list function - // _pool is clipped to maxc in occlude_list before affecting _occlude_factor - // _occlude_factor is initialized to maxc size - std::vector _occlude_factor; - - // Capacity initialized to 20L - tsl::robin_set _inserted_into_pool_rs; - - // Use a pointer here to allow for forward declaration of dynamic_bitset - // in public headers to avoid making boost a dependency for clients - // of DiskANN. - boost::dynamic_bitset<> *_inserted_into_pool_bs; - - // _id_scratch.size() must be > R*GRAPH_SLACK_FACTOR for iterate_to_fp - std::vector _id_scratch; - - // _dist_scratch must be > R*GRAPH_SLACK_FACTOR for iterate_to_fp - // _dist_scratch should be at least the size of id_scratch - std::vector _dist_scratch; - - // Buffers used in process delete, capacity increases as needed - tsl::robin_set _expanded_nodes_set; - std::vector _expanded_nghrs_vec; - std::vector _occlude_list_output; +template +class InMemQueryScratch { + public: + ~InMemQueryScratch(); + // REFACTOR TODO: move all parameters to a new class. + InMemQueryScratch(uint32_t search_l, uint32_t indexing_l, uint32_t r, + uint32_t maxc, size_t dim, size_t aligned_dim, + size_t alignment_factor, bool init_pq_scratch = false); + void resize_for_new_L(uint32_t new_search_l); + void clear(); + + inline uint32_t get_L() { return _L; } + inline uint32_t get_R() { return _R; } + inline uint32_t get_maxc() { return _maxc; } + inline T *aligned_query() { return _aligned_query; } + inline PQScratch *pq_scratch() { return _pq_scratch; } + inline std::vector &pool() { return _pool; } + inline NeighborPriorityQueue &best_l_nodes() { return _best_l_nodes; } + inline std::vector &occlude_factor() { return _occlude_factor; } + inline tsl::robin_set &inserted_into_pool_rs() { + return _inserted_into_pool_rs; + } + inline boost::dynamic_bitset<> &inserted_into_pool_bs() { + return *_inserted_into_pool_bs; + } + inline std::vector &id_scratch() { return _id_scratch; } + inline std::vector &dist_scratch() { return _dist_scratch; } + inline tsl::robin_set &expanded_nodes_set() { + return _expanded_nodes_set; + } + inline std::vector &expanded_nodes_vec() { + return _expanded_nghrs_vec; + } + inline std::vector &occlude_list_output() { + return _occlude_list_output; + } + + private: + uint32_t _L; + uint32_t _R; + uint32_t _maxc; + + T *_aligned_query = nullptr; + + PQScratch *_pq_scratch = nullptr; + + // _pool stores all neighbors explored from best_L_nodes. + // Usually around L+R, but could be higher. + // Initialized to 3L+R for some slack, expands as needed. + std::vector _pool; + + // _best_l_nodes is reserved for storing best L entries + // Underlying storage is L+1 to support inserts + NeighborPriorityQueue _best_l_nodes; + + // _occlude_factor.size() >= pool.size() in occlude_list function + // _pool is clipped to maxc in occlude_list before affecting _occlude_factor + // _occlude_factor is initialized to maxc size + std::vector _occlude_factor; + + // Capacity initialized to 20L + tsl::robin_set _inserted_into_pool_rs; + + // Use a pointer here to allow for forward declaration of dynamic_bitset + // in public headers to avoid making boost a dependency for clients + // of DiskANN. + boost::dynamic_bitset<> *_inserted_into_pool_bs; + + // _id_scratch.size() must be > R*GRAPH_SLACK_FACTOR for iterate_to_fp + std::vector _id_scratch; + + // _dist_scratch must be > R*GRAPH_SLACK_FACTOR for iterate_to_fp + // _dist_scratch should be at least the size of id_scratch + std::vector _dist_scratch; + + // Buffers used in process delete, capacity increases as needed + tsl::robin_set _expanded_nodes_set; + std::vector _expanded_nghrs_vec; + std::vector _occlude_list_output; }; // // Scratch space for SSD index based search // -template class SSDQueryScratch -{ - public: - T *coord_scratch = nullptr; // MUST BE AT LEAST [MAX_N_CMPS * data_dim] - size_t coord_idx = 0; // index of next [data_dim] scratch to use +template +class SSDQueryScratch { + public: + T *coord_scratch = nullptr; // MUST BE AT LEAST [MAX_N_CMPS * data_dim] + size_t coord_idx = 0; // index of next [data_dim] scratch to use - char *sector_scratch = nullptr; // MUST BE AT LEAST [MAX_N_SECTOR_READS * SECTOR_LEN] - size_t sector_idx = 0; // index of next [SECTOR_LEN] scratch to use + char *sector_scratch = + nullptr; // MUST BE AT LEAST [MAX_N_SECTOR_READS * SECTOR_LEN] + size_t sector_idx = 0; // index of next [SECTOR_LEN] scratch to use - T *aligned_query_T = nullptr; + T *aligned_query_T = nullptr; - PQScratch *_pq_scratch; + PQScratch *_pq_scratch; - tsl::robin_set visited; - NeighborPriorityQueue retset; - std::vector full_retset; + tsl::robin_set visited; + NeighborPriorityQueue retset; + std::vector full_retset; - SSDQueryScratch(size_t aligned_dim, size_t visited_reserve); - ~SSDQueryScratch(); + SSDQueryScratch(size_t aligned_dim, size_t visited_reserve); + ~SSDQueryScratch(); - void reset(); + void reset(); }; -template class SSDThreadData -{ - public: - SSDQueryScratch scratch; - IOContext ctx; +template +class SSDThreadData { + public: + SSDQueryScratch scratch; + IOContext ctx; - SSDThreadData(size_t aligned_dim, size_t visited_reserve); - void clear(); + SSDThreadData(size_t aligned_dim, size_t visited_reserve); + void clear(); }; // // Class to avoid the hassle of pushing and popping the query scratch. // -template class ScratchStoreManager -{ - public: - ScratchStoreManager(ConcurrentQueue &query_scratch) : _scratch_pool(query_scratch) - { - _scratch = query_scratch.pop(); - while (_scratch == nullptr) - { - query_scratch.wait_for_push_notify(); - _scratch = query_scratch.pop(); - } - } - T *scratch_space() - { - return _scratch; - } - - ~ScratchStoreManager() - { - _scratch->clear(); - _scratch_pool.push(_scratch); - _scratch_pool.push_notify_all(); - } - - void destroy() - { - while (!_scratch_pool.empty()) - { - auto scratch = _scratch_pool.pop(); - while (scratch == nullptr) - { - _scratch_pool.wait_for_push_notify(); - scratch = _scratch_pool.pop(); - } - delete scratch; - } - } - - private: - T *_scratch; - ConcurrentQueue &_scratch_pool; - ScratchStoreManager(const ScratchStoreManager &); - ScratchStoreManager &operator=(const ScratchStoreManager &); +template +class ScratchStoreManager { + public: + ScratchStoreManager(ConcurrentQueue &query_scratch) + : _scratch_pool(query_scratch) { + _scratch = query_scratch.pop(); + while (_scratch == nullptr) { + query_scratch.wait_for_push_notify(); + _scratch = query_scratch.pop(); + } + } + T *scratch_space() { return _scratch; } + + ~ScratchStoreManager() { + _scratch->clear(); + _scratch_pool.push(_scratch); + _scratch_pool.push_notify_all(); + } + + void destroy() { + while (!_scratch_pool.empty()) { + auto scratch = _scratch_pool.pop(); + while (scratch == nullptr) { + _scratch_pool.wait_for_push_notify(); + scratch = _scratch_pool.pop(); + } + delete scratch; + } + } + + private: + T *_scratch; + ConcurrentQueue &_scratch_pool; + ScratchStoreManager(const ScratchStoreManager &); + ScratchStoreManager &operator=(const ScratchStoreManager &); }; -} // namespace diskann +} // namespace diskann diff --git a/include/simd_utils.h b/include/simd_utils.h index 4b0736998..fb94c35c2 100644 --- a/include/simd_utils.h +++ b/include/simd_utils.h @@ -9,98 +9,97 @@ #include #endif -namespace diskann -{ -static inline __m256 _mm256_mul_epi8(__m256i X) -{ - __m256i zero = _mm256_setzero_si256(); +namespace diskann { +static inline __m256 _mm256_mul_epi8(__m256i X) { + __m256i zero = _mm256_setzero_si256(); - __m256i sign_x = _mm256_cmpgt_epi8(zero, X); + __m256i sign_x = _mm256_cmpgt_epi8(zero, X); - __m256i xlo = _mm256_unpacklo_epi8(X, sign_x); - __m256i xhi = _mm256_unpackhi_epi8(X, sign_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))); + 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); +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))); + 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); +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))); + 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(); +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 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); + __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))); + 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_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 __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(); +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 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); + __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))); + 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 __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); +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 +} // namespace diskann diff --git a/include/timer.h b/include/timer.h index 963927bd7..d187a4e3f 100644 --- a/include/timer.h +++ b/include/timer.h @@ -3,37 +3,27 @@ #include -namespace diskann -{ -class Timer -{ - typedef std::chrono::high_resolution_clock _clock; - std::chrono::time_point<_clock> check_point; +namespace diskann { +class Timer { + typedef std::chrono::high_resolution_clock _clock; + std::chrono::time_point<_clock> check_point; - public: - Timer() : check_point(_clock::now()) - { - } + public: + Timer() : check_point(_clock::now()) {} - void reset() - { - check_point = _clock::now(); - } + void reset() { check_point = _clock::now(); } - long long elapsed() const - { - return std::chrono::duration_cast(_clock::now() - check_point).count(); - } + long long elapsed() const { + return std::chrono::duration_cast(_clock::now() - + check_point) + .count(); + } - float elapsed_seconds() const - { - return (float)elapsed() / 1000000.0; - } + float elapsed_seconds() const { return (float)elapsed() / 1000000.0; } - std::string elapsed_seconds_for_step(const std::string &step) const - { - return std::string("Time for ") + step + std::string(": ") + std::to_string(elapsed_seconds()) + - std::string(" seconds"); - } + std::string elapsed_seconds_for_step(const std::string &step) const { + return std::string("Time for ") + step + std::string(": ") + + std::to_string(elapsed_seconds()) + std::string(" seconds"); + } }; -} // namespace diskann +} // namespace diskann diff --git a/include/types.h b/include/types.h index ea04cd34d..9b2794474 100644 --- a/include/types.h +++ b/include/types.h @@ -6,7 +6,6 @@ #include #include -namespace diskann -{ +namespace diskann { typedef uint32_t location_t; -} // namespace diskann \ No newline at end of file +} // namespace diskann \ No newline at end of file diff --git a/include/utils.h b/include/utils.h index 1b9cce924..0d4aa8722 100644 --- a/include/utils.h +++ b/include/utils.h @@ -36,7 +36,8 @@ typedef int FileHandle; // taken from // https://github.com/Microsoft/BLAS-on-flash/blob/master/include/utils.h // round up X to the nearest multiple of Y -#define ROUND_UP(X, Y) ((((uint64_t)(X) / (Y)) + ((uint64_t)(X) % (Y) != 0)) * (Y)) +#define ROUND_UP(X, Y) \ + ((((uint64_t)(X) / (Y)) + ((uint64_t)(X) % (Y) != 0)) * (Y)) #define DIV_ROUND_UP(X, Y) (((uint64_t)(X) / (Y)) + ((uint64_t)(X) % (Y) != 0)) @@ -47,750 +48,730 @@ 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 +#define METADATA_SIZE \ + 4096 // all metadata of individual sub-component files is written in first + // 4KB for unified files #define BUFFER_SIZE_FOR_CACHED_IO (size_t)1024 * (size_t)1048576 #define PBSTR "||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||" #define PBWIDTH 60 -inline bool file_exists(const std::string &name, bool dirCheck = false) -{ - int val; +inline bool file_exists(const std::string &name, bool dirCheck = false) { + int val; #ifndef _WINDOWS - struct stat buffer; - val = stat(name.c_str(), &buffer); + struct stat buffer; + val = stat(name.c_str(), &buffer); #else - // It is the 21st century but Windows API still thinks in 32-bit terms. - // Turns out calling stat() on a file > 4GB results in errno = 132 - // (OVERFLOW). How silly is this!? So calling _stat64() - struct _stat64 buffer; - val = _stat64(name.c_str(), &buffer); + // It is the 21st century but Windows API still thinks in 32-bit terms. + // Turns out calling stat() on a file > 4GB results in errno = 132 + // (OVERFLOW). How silly is this!? So calling _stat64() + struct _stat64 buffer; + val = _stat64(name.c_str(), &buffer); #endif - if (val != 0) - { - switch (errno) - { - case EINVAL: - diskann::cout << "Invalid argument passed to stat()" << std::endl; - break; - case ENOENT: - // file is not existing, not an issue, so we won't cout anything. - 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; + if (val != 0) { + switch (errno) { + case EINVAL: + diskann::cout << "Invalid argument passed to stat()" << std::endl; + break; + case ENOENT: + // file is not existing, not an issue, so we won't cout anything. + 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 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()) - { - char buff[1024]; +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()) { + char buff[1024]; #ifdef _WINDOWS - auto ret = std::to_string(strerror_s(buff, 1024, errno)); + auto ret = std::to_string(strerror_s(buff, 1024, errno)); #else - auto ret = std::string(strerror_r(errno, buff, 1024)); + auto ret = std::string(strerror_r(errno, buff, 1024)); #endif - auto message = std::string("Failed to open file") + filename + " for write because " + buff + ", ret=" + ret; - diskann::cerr << message << std::endl; - throw diskann::ANNException(message, -1); - } + auto message = std::string("Failed to open file") + filename + + " for write because " + buff + ", ret=" + ret; + diskann::cerr << message << std::endl; + throw diskann::ANNException(message, -1); + } } -inline size_t get_file_size(const std::string &fname) -{ - std::ifstream reader(fname, std::ios::binary | std::ios::ate); - if (!reader.fail() && reader.is_open()) - { - size_t end_pos = reader.tellg(); - reader.close(); - return end_pos; - } - else - { - diskann::cerr << "Could not open file: " << fname << std::endl; - return 0; - } +inline size_t get_file_size(const std::string &fname) { + std::ifstream reader(fname, std::ios::binary | std::ios::ate); + if (!reader.fail() && reader.is_open()) { + size_t 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; +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; + } } -inline void convert_labels_string_to_int(const std::string &inFileName, const std::string &outFileName, - const std::string &mapFileName, const std::string &unv_label) -{ - std::unordered_map string_int_map; - std::ofstream label_writer(outFileName); - std::ifstream label_reader(inFileName); - if (unv_label != "") - string_int_map[unv_label] = 0; - std::string line, token; - while (std::getline(label_reader, line)) - { - std::istringstream new_iss(line); - std::vector lbls; - while (getline(new_iss, token, ',')) - { - token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); - token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); - if (string_int_map.find(token) == string_int_map.end()) - { - uint32_t nextId = (uint32_t)string_int_map.size() + 1; - string_int_map[token] = nextId; - } - lbls.push_back(string_int_map[token]); - } - if (lbls.size() <= 0) - { - std::cout << "No label found"; - exit(-1); - } - for (size_t j = 0; j < lbls.size(); j++) - { - if (j != lbls.size() - 1) - label_writer << lbls[j] << ","; - else - label_writer << lbls[j] << std::endl; - } +inline void convert_labels_string_to_int(const std::string &inFileName, + const std::string &outFileName, + const std::string &mapFileName, + const std::string &unv_label) { + std::unordered_map string_int_map; + std::ofstream label_writer(outFileName); + std::ifstream label_reader(inFileName); + if (unv_label != "") string_int_map[unv_label] = 0; + std::string line, token; + while (std::getline(label_reader, line)) { + std::istringstream new_iss(line); + std::vector lbls; + while (getline(new_iss, token, ',')) { + token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); + token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); + if (string_int_map.find(token) == string_int_map.end()) { + uint32_t nextId = (uint32_t)string_int_map.size() + 1; + string_int_map[token] = nextId; + } + lbls.push_back(string_int_map[token]); } - label_writer.close(); - - std::ofstream map_writer(mapFileName); - for (auto mp : string_int_map) - { - map_writer << mp.first << "\t" << mp.second << std::endl; + if (lbls.size() <= 0) { + std::cout << "No label found"; + exit(-1); + } + for (size_t j = 0; j < lbls.size(); j++) { + if (j != lbls.size() - 1) + label_writer << lbls[j] << ","; + else + label_writer << lbls[j] << std::endl; } - map_writer.close(); + } + label_writer.close(); + + std::ofstream map_writer(mapFileName); + for (auto mp : string_int_map) { + map_writer << mp.first << "\t" << mp.second << std::endl; + } + map_writer.close(); } #ifdef EXEC_ENV_OLS class AlignedFileReader; #endif -namespace diskann -{ +namespace diskann { static const size_t MAX_SIZE_OF_STREAMBUF = 2LL * 1024 * 1024 * 1024; -inline void print_error_and_terminate(std::stringstream &error_stream) -{ - diskann::cerr << error_stream.str() << std::endl; - throw diskann::ANNException(error_stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); +inline void print_error_and_terminate(std::stringstream &error_stream) { + diskann::cerr << error_stream.str() << std::endl; + throw diskann::ANNException(error_stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); } -inline void report_memory_allocation_failure() -{ - std::stringstream stream; - stream << "Memory Allocation Failed."; - print_error_and_terminate(stream); +inline void report_memory_allocation_failure() { + std::stringstream stream; + stream << "Memory Allocation Failed."; + print_error_and_terminate(stream); } -inline void report_misalignment_of_requested_size(size_t align) -{ - std::stringstream stream; - stream << "Requested memory size is not a multiple of " << align << ". Can not be allocated."; - print_error_and_terminate(stream); +inline void report_misalignment_of_requested_size(size_t align) { + std::stringstream stream; + stream << "Requested memory size is not a multiple of " << align + << ". Can not be allocated."; + print_error_and_terminate(stream); } -inline void alloc_aligned(void **ptr, size_t size, size_t align) -{ - *ptr = nullptr; - if (IS_ALIGNED(size, align) == 0) - report_misalignment_of_requested_size(align); +inline void alloc_aligned(void **ptr, size_t size, size_t align) { + *ptr = nullptr; + if (IS_ALIGNED(size, align) == 0) + report_misalignment_of_requested_size(align); #ifndef _WINDOWS - *ptr = ::aligned_alloc(align, size); + *ptr = ::aligned_alloc(align, size); #else - *ptr = ::_aligned_malloc(size, align); // note the swapped arguments! + *ptr = ::_aligned_malloc(size, align); // note the swapped arguments! #endif - if (*ptr == nullptr) - report_memory_allocation_failure(); + if (*ptr == nullptr) report_memory_allocation_failure(); } -inline void realloc_aligned(void **ptr, size_t size, size_t align) -{ - if (IS_ALIGNED(size, align) == 0) - report_misalignment_of_requested_size(align); +inline void realloc_aligned(void **ptr, size_t size, size_t align) { + if (IS_ALIGNED(size, align) == 0) + report_misalignment_of_requested_size(align); #ifdef _WINDOWS - *ptr = ::_aligned_realloc(*ptr, size, align); + *ptr = ::_aligned_realloc(*ptr, size, align); #else - diskann::cerr << "No aligned realloc on GCC. Must malloc and mem_align, " - "left it out for now." - << std::endl; + diskann::cerr << "No aligned realloc on GCC. Must malloc and mem_align, " + "left it out for now." + << std::endl; #endif - if (*ptr == nullptr) - report_memory_allocation_failure(); + if (*ptr == nullptr) report_memory_allocation_failure(); } -inline void check_stop(std::string arnd) -{ - int brnd; - diskann::cout << arnd << std::endl; - std::cin >> brnd; +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; - } +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; + } #ifndef _WINDOWS - free(ptr); + free(ptr); #else - ::_aligned_free(ptr); + ::_aligned_free(ptr); #endif } -inline void GenRandom(std::mt19937 &rng, unsigned *addr, unsigned size, unsigned N) -{ - for (unsigned i = 0; i < size; ++i) - { - addr[i] = rng() % (N - size); - } +inline void GenRandom(std::mt19937 &rng, unsigned *addr, unsigned size, + unsigned N) { + for (unsigned i = 0; i < size; ++i) { + addr[i] = rng() % (N - size); + } - std::sort(addr, addr + size); - for (unsigned i = 1; i < size; ++i) - { - if (addr[i] <= addr[i - 1]) - { - addr[i] = addr[i - 1] + 1; - } - } - unsigned off = rng() % N; - for (unsigned i = 0; i < size; ++i) - { - addr[i] = (addr[i] + off) % N; + std::sort(addr, addr + size); + for (unsigned i = 1; i < size; ++i) { + if (addr[i] <= addr[i - 1]) { + addr[i] = addr[i - 1] + 1; } + } + unsigned off = rng() % N; + for (unsigned i = 0; i < size; ++i) { + addr[i] = (addr[i] + off) % N; + } } // get_bin_metadata functions START -inline void get_bin_metadata_impl(std::basic_istream &reader, 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; - ncols = ncols_32; +inline void get_bin_metadata_impl(std::basic_istream &reader, + 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; + ncols = ncols_32; } #ifdef EXEC_ENV_OLS -inline void get_bin_metadata(MemoryMappedFiles &files, const std::string &bin_file, size_t &nrows, 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, offset); - - 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; +inline void get_bin_metadata(MemoryMappedFiles &files, + const std::string &bin_file, size_t &nrows, + 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, offset); + + 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 offset = 0) -{ - std::ifstream reader(bin_file.c_str(), std::ios::binary); - get_bin_metadata_impl(reader, nrows, ncols, offset); +inline void get_bin_metadata(const std::string &bin_file, size_t &nrows, + size_t &ncols, size_t offset = 0) { + std::ifstream reader(bin_file.c_str(), std::ios::binary); + get_bin_metadata_impl(reader, nrows, ncols, offset); } // get_bin_metadata functions END -template inline std::string getValues(T *data, size_t num) -{ - std::stringstream stream; - stream << "["; - for (size_t i = 0; i < num; i++) - { - stream << std::to_string(data[i]) << ","; - } - stream << "]" << std::endl; - - return stream.str(); +template +inline std::string getValues(T *data, size_t num) { + std::stringstream stream; + stream << "["; + for (size_t i = 0; i < num; i++) { + stream << std::to_string(data[i]) << ","; + } + stream << "]" << std::endl; + + return stream.str(); } // load_bin functions START template -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; +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; - dim = (unsigned)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; + dim = (unsigned)dim_i32; - std::cout << "Metadata: #pts = " << npts << ", #dims = " << dim << "..." << std::endl; + std::cout << "Metadata: #pts = " << npts << ", #dims = " << dim << "..." + << std::endl; - data = new T[npts * dim]; - reader.read((char *)data, npts * dim * sizeof(T)); + data = new T[npts * dim]; + reader.read((char *)data, npts * dim * sizeof(T)); } #ifdef EXEC_ENV_OLS template -inline void load_bin(MemoryMappedFiles &files, const std::string &bin_file, 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 *)((char *)fc._content + offset); - t_npts = *(contentAsIntPtr); - t_dim = *(contentAsIntPtr + 1); - - npts = t_npts; - dim = t_dim; - - data = (T *)((char *)fc._content + offset + 2 * sizeof(uint32_t)); // No need to copy! +inline void load_bin(MemoryMappedFiles &files, const std::string &bin_file, + 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 *)((char *)fc._content + offset); + t_npts = *(contentAsIntPtr); + t_dim = *(contentAsIntPtr + 1); + + npts = t_npts; + dim = t_dim; + + data = (T *)((char *)fc._content + offset + + 2 * sizeof(uint32_t)); // No need to copy! } -DISKANN_DLLEXPORT void get_bin_metadata(AlignedFileReader &reader, size_t &npts, size_t &ndim, size_t offset = 0); +DISKANN_DLLEXPORT void get_bin_metadata(AlignedFileReader &reader, size_t &npts, + size_t &ndim, size_t offset = 0); template -DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, T *&data, size_t &npts, size_t &ndim, size_t offset = 0); +DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, T *&data, + size_t &npts, size_t &ndim, size_t offset = 0); template -DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, std::unique_ptr &data, size_t &npts, size_t &ndim, - size_t offset = 0); +DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, + std::unique_ptr &data, size_t &npts, + size_t &ndim, size_t offset = 0); template -DISKANN_DLLEXPORT void copy_aligned_data_from_file(AlignedFileReader &reader, T *&data, size_t &npts, size_t &dim, - const size_t &rounded_dim, size_t offset = 0); +DISKANN_DLLEXPORT void copy_aligned_data_from_file(AlignedFileReader &reader, + T *&data, size_t &npts, + size_t &dim, + const size_t &rounded_dim, + size_t offset = 0); // Unlike load_bin, assumes that data is already allocated 'size' entries template -DISKANN_DLLEXPORT void read_array(AlignedFileReader &reader, T *data, size_t size, size_t offset = 0); +DISKANN_DLLEXPORT void read_array(AlignedFileReader &reader, T *data, + size_t size, size_t offset = 0); -template DISKANN_DLLEXPORT void read_value(AlignedFileReader &reader, T &value, size_t offset = 0); +template +DISKANN_DLLEXPORT void read_value(AlignedFileReader &reader, T &value, + size_t offset = 0); #endif template -inline void load_bin(const std::string &bin_file, T *&data, size_t &npts, size_t &dim, size_t offset = 0) -{ - diskann::cout << "Reading bin file " << bin_file.c_str() << " ..." << std::endl; - std::ifstream reader; - reader.exceptions(std::ifstream::failbit | std::ifstream::badbit); - - try - { - diskann::cout << "Opening bin file " << bin_file.c_str() << "... " << std::endl; - reader.open(bin_file, std::ios::binary | std::ios::ate); - reader.seekg(0); - load_bin_impl(reader, data, npts, dim, offset); - } - catch (std::system_error &e) - { - throw FileException(bin_file, e, __FUNCSIG__, __FILE__, __LINE__); - } - diskann::cout << "done." << std::endl; +inline void load_bin(const std::string &bin_file, T *&data, size_t &npts, + size_t &dim, size_t offset = 0) { + diskann::cout << "Reading bin file " << bin_file.c_str() << " ..." + << std::endl; + std::ifstream reader; + reader.exceptions(std::ifstream::failbit | std::ifstream::badbit); + + try { + diskann::cout << "Opening bin file " << bin_file.c_str() << "... " + << std::endl; + reader.open(bin_file, std::ios::binary | std::ios::ate); + reader.seekg(0); + load_bin_impl(reader, data, npts, dim, offset); + } catch (std::system_error &e) { + throw FileException(bin_file, e, __FUNCSIG__, __FILE__, __LINE__); + } + diskann::cout << "done." << std::endl; } -inline void wait_for_keystroke() -{ - int a; - std::cout << "Press any number to continue.." << std::endl; - std::cin >> a; +inline void wait_for_keystroke() { + int a; + std::cout << "Press any number to continue.." << std::endl; + std::cin >> a; } // load_bin functions END -inline void load_truthset(const std::string &bin_file, uint32_t *&ids, float *&dists, size_t &npts, size_t &dim) -{ - size_t read_blk_size = 64 * 1024 * 1024; - cached_ifstream reader(bin_file, read_blk_size); - diskann::cout << "Reading truthset file " << bin_file.c_str() << " ..." << std::endl; - size_t actual_file_size = reader.get_file_size(); - - 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; - - diskann::cout << "Metadata: #pts = " << npts << ", #dims = " << dim << "... " << std::endl; - - int truthset_type = -1; // 1 means truthset has ids and distances, 2 means - // only ids, -1 is error - size_t expected_file_size_with_dists = 2 * npts * dim * sizeof(uint32_t) + 2 * sizeof(uint32_t); - - if (actual_file_size == expected_file_size_with_dists) - truthset_type = 1; - - size_t expected_file_size_just_ids = npts * dim * sizeof(uint32_t) + 2 * sizeof(uint32_t); - - if (actual_file_size == expected_file_size_just_ids) - truthset_type = 2; - - if (truthset_type == -1) - { - std::stringstream stream; - stream << "Error. File size mismatch. File should have bin format, with " - "npts followed by ngt followed by npts*ngt ids and optionally " - "followed by npts*ngt distance values; actual size: " - << actual_file_size << ", expected: " << expected_file_size_with_dists << " or " - << expected_file_size_just_ids; - diskann::cout << stream.str(); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } +inline void load_truthset(const std::string &bin_file, uint32_t *&ids, + float *&dists, size_t &npts, size_t &dim) { + size_t read_blk_size = 64 * 1024 * 1024; + cached_ifstream reader(bin_file, read_blk_size); + diskann::cout << "Reading truthset file " << bin_file.c_str() << " ..." + << std::endl; + size_t actual_file_size = reader.get_file_size(); - ids = new uint32_t[npts * dim]; - reader.read((char *)ids, npts * dim * sizeof(uint32_t)); + 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; - if (truthset_type == 1) - { - dists = new float[npts * dim]; - reader.read((char *)dists, npts * dim * sizeof(float)); - } -} + diskann::cout << "Metadata: #pts = " << npts << ", #dims = " << dim << "... " + << std::endl; -inline void prune_truthset_for_range(const std::string &bin_file, float range, - std::vector> &groundtruth, size_t &npts) -{ - size_t read_blk_size = 64 * 1024 * 1024; - cached_ifstream reader(bin_file, read_blk_size); - diskann::cout << "Reading truthset file " << bin_file.c_str() << "... " << std::endl; - size_t actual_file_size = reader.get_file_size(); - - int npts_i32, dim_i32; - reader.read((char *)&npts_i32, sizeof(int)); - reader.read((char *)&dim_i32, sizeof(int)); - npts = (unsigned)npts_i32; - uint64_t dim = (unsigned)dim_i32; - uint32_t *ids; - float *dists; - - diskann::cout << "Metadata: #pts = " << npts << ", #dims = " << dim << "... " << std::endl; - - int truthset_type = -1; // 1 means truthset has ids and distances, 2 means - // only ids, -1 is error - size_t expected_file_size_with_dists = 2 * npts * dim * sizeof(uint32_t) + 2 * sizeof(uint32_t); - - if (actual_file_size == expected_file_size_with_dists) - truthset_type = 1; - - if (truthset_type == -1) - { - std::stringstream stream; - stream << "Error. File size mismatch. File should have bin format, with " - "npts followed by ngt followed by npts*ngt ids and optionally " - "followed by npts*ngt distance values; actual size: " - << actual_file_size << ", expected: " << expected_file_size_with_dists; - diskann::cout << stream.str(); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } + int truthset_type = -1; // 1 means truthset has ids and distances, 2 means + // only ids, -1 is error + size_t expected_file_size_with_dists = + 2 * npts * dim * sizeof(uint32_t) + 2 * sizeof(uint32_t); - ids = new uint32_t[npts * dim]; - reader.read((char *)ids, npts * dim * sizeof(uint32_t)); + if (actual_file_size == expected_file_size_with_dists) truthset_type = 1; - if (truthset_type == 1) - { - dists = new float[npts * dim]; - reader.read((char *)dists, npts * dim * sizeof(float)); - } - float min_dist = std::numeric_limits::max(); - float max_dist = 0; - groundtruth.resize(npts); - for (uint32_t i = 0; i < npts; i++) - { - groundtruth[i].clear(); - for (uint32_t j = 0; j < dim; j++) - { - if (dists[i * dim + j] <= range) - { - groundtruth[i].emplace_back(ids[i * dim + j]); - } - min_dist = min_dist > dists[i * dim + j] ? dists[i * dim + j] : min_dist; - max_dist = max_dist < dists[i * dim + j] ? dists[i * dim + j] : max_dist; - } - // std::cout<> &groundtruth, - uint64_t >_num) -{ - size_t read_blk_size = 64 * 1024 * 1024; - cached_ifstream reader(bin_file, read_blk_size); - diskann::cout << "Reading truthset file " << bin_file.c_str() << "... " << std::flush; - size_t actual_file_size = reader.get_file_size(); + if (actual_file_size == expected_file_size_just_ids) truthset_type = 2; - int nptsuint32_t, totaluint32_t; - reader.read((char *)&nptsuint32_t, sizeof(int)); - reader.read((char *)&totaluint32_t, sizeof(int)); + if (truthset_type == -1) { + std::stringstream stream; + stream << "Error. File size mismatch. File should have bin format, with " + "npts followed by ngt followed by npts*ngt ids and optionally " + "followed by npts*ngt distance values; actual size: " + << actual_file_size + << ", expected: " << expected_file_size_with_dists << " or " + << expected_file_size_just_ids; + diskann::cout << stream.str(); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } + + ids = new uint32_t[npts * dim]; + reader.read((char *)ids, npts * dim * sizeof(uint32_t)); + + if (truthset_type == 1) { + dists = new float[npts * dim]; + reader.read((char *)dists, npts * dim * sizeof(float)); + } +} - gt_num = (uint64_t)nptsuint32_t; - uint64_t total_res = (uint64_t)totaluint32_t; +inline void prune_truthset_for_range( + const std::string &bin_file, float range, + std::vector> &groundtruth, size_t &npts) { + size_t read_blk_size = 64 * 1024 * 1024; + cached_ifstream reader(bin_file, read_blk_size); + diskann::cout << "Reading truthset file " << bin_file.c_str() << "... " + << std::endl; + size_t actual_file_size = reader.get_file_size(); + + int npts_i32, dim_i32; + reader.read((char *)&npts_i32, sizeof(int)); + reader.read((char *)&dim_i32, sizeof(int)); + npts = (unsigned)npts_i32; + uint64_t dim = (unsigned)dim_i32; + uint32_t *ids; + float *dists; + + diskann::cout << "Metadata: #pts = " << npts << ", #dims = " << dim << "... " + << std::endl; + + int truthset_type = -1; // 1 means truthset has ids and distances, 2 means + // only ids, -1 is error + size_t expected_file_size_with_dists = + 2 * npts * dim * sizeof(uint32_t) + 2 * sizeof(uint32_t); + + if (actual_file_size == expected_file_size_with_dists) truthset_type = 1; + + if (truthset_type == -1) { + std::stringstream stream; + stream << "Error. File size mismatch. File should have bin format, with " + "npts followed by ngt followed by npts*ngt ids and optionally " + "followed by npts*ngt distance values; actual size: " + << actual_file_size + << ", expected: " << expected_file_size_with_dists; + diskann::cout << stream.str(); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } + + ids = new uint32_t[npts * dim]; + reader.read((char *)ids, npts * dim * sizeof(uint32_t)); + + if (truthset_type == 1) { + dists = new float[npts * dim]; + reader.read((char *)dists, npts * dim * sizeof(float)); + } + float min_dist = std::numeric_limits::max(); + float max_dist = 0; + groundtruth.resize(npts); + for (uint32_t i = 0; i < npts; i++) { + groundtruth[i].clear(); + for (uint32_t j = 0; j < dim; j++) { + if (dists[i * dim + j] <= range) { + groundtruth[i].emplace_back(ids[i * dim + j]); + } + min_dist = min_dist > dists[i * dim + j] ? dists[i * dim + j] : min_dist; + max_dist = max_dist < dists[i * dim + j] ? dists[i * dim + j] : max_dist; + } + // std::cout<> &groundtruth, + uint64_t >_num) { + size_t read_blk_size = 64 * 1024 * 1024; + cached_ifstream reader(bin_file, read_blk_size); + diskann::cout << "Reading truthset file " << bin_file.c_str() << "... " + << std::flush; + size_t actual_file_size = reader.get_file_size(); - size_t expected_file_size = 2 * sizeof(uint32_t) + gt_num * sizeof(uint32_t) + total_res * sizeof(uint32_t); + int nptsuint32_t, totaluint32_t; + reader.read((char *)&nptsuint32_t, sizeof(int)); + reader.read((char *)&totaluint32_t, sizeof(int)); - if (actual_file_size != expected_file_size) - { - std::stringstream stream; - stream << "Error. File size mismatch in range truthset. actual size: " << actual_file_size - << ", expected: " << expected_file_size; - diskann::cout << stream.str(); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - groundtruth.clear(); - groundtruth.resize(gt_num); - std::vector gt_count(gt_num); + gt_num = (uint64_t)nptsuint32_t; + uint64_t total_res = (uint64_t)totaluint32_t; - reader.read((char *)gt_count.data(), sizeof(uint32_t) * gt_num); + diskann::cout << "Metadata: #pts = " << gt_num + << ", #total_results = " << total_res << "..." << std::endl; - std::vector gt_stats(gt_count); - std::sort(gt_stats.begin(), gt_stats.end()); + size_t expected_file_size = 2 * sizeof(uint32_t) + gt_num * sizeof(uint32_t) + + total_res * sizeof(uint32_t); - std::cout << "GT count percentiles:" << std::endl; - for (uint32_t p = 0; p < 100; p += 5) - std::cout << "percentile " << p << ": " << gt_stats[static_cast(std::floor((p / 100.0) * gt_num))] - << std::endl; - std::cout << "percentile 100" - << ": " << gt_stats[gt_num - 1] << std::endl; - - for (uint32_t i = 0; i < gt_num; i++) - { - groundtruth[i].clear(); - groundtruth[i].resize(gt_count[i]); - if (gt_count[i] != 0) - reader.read((char *)groundtruth[i].data(), sizeof(uint32_t) * gt_count[i]); - } + if (actual_file_size != expected_file_size) { + std::stringstream stream; + stream << "Error. File size mismatch in range truthset. actual size: " + << actual_file_size << ", expected: " << expected_file_size; + diskann::cout << stream.str(); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } + groundtruth.clear(); + groundtruth.resize(gt_num); + std::vector gt_count(gt_num); + + reader.read((char *)gt_count.data(), sizeof(uint32_t) * gt_num); + + std::vector gt_stats(gt_count); + std::sort(gt_stats.begin(), gt_stats.end()); + + std::cout << "GT count percentiles:" << std::endl; + for (uint32_t p = 0; p < 100; p += 5) + std::cout << "percentile " << p << ": " + << gt_stats[static_cast(std::floor((p / 100.0) * gt_num))] + << std::endl; + std::cout << "percentile 100" + << ": " << gt_stats[gt_num - 1] << std::endl; + + for (uint32_t i = 0; i < gt_num; i++) { + groundtruth[i].clear(); + groundtruth[i].resize(gt_count[i]); + if (gt_count[i] != 0) + reader.read((char *)groundtruth[i].data(), + sizeof(uint32_t) * gt_count[i]); + } } #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, size_t offset = 0) -{ - T *ptr; - load_bin(files, bin_file, ptr, npts, dim, offset); - data.reset(ptr); +inline void load_bin(MemoryMappedFiles &files, const std::string &bin_file, + std::unique_ptr &data, size_t &npts, size_t &dim, + size_t offset = 0) { + T *ptr; + load_bin(files, bin_file, ptr, npts, dim, offset); + data.reset(ptr); } #endif -inline void copy_file(std::string in_file, std::string out_file) -{ - std::ifstream source(in_file, std::ios::binary); - std::ofstream dest(out_file, std::ios::binary); +inline void copy_file(std::string in_file, std::string out_file) { + std::ifstream source(in_file, std::ios::binary); + std::ofstream dest(out_file, std::ios::binary); - std::istreambuf_iterator begin_source(source); - std::istreambuf_iterator end_source; - std::ostreambuf_iterator begin_dest(dest); - std::copy(begin_source, end_source, begin_dest); + std::istreambuf_iterator begin_source(source); + std::istreambuf_iterator end_source; + std::ostreambuf_iterator begin_dest(dest); + std::copy(begin_source, end_source, begin_dest); - source.close(); - dest.close(); + source.close(); + dest.close(); } -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); -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 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 double calculate_range_search_recall(unsigned num_queries, - std::vector> &groundtruth, - std::vector> &our_results); +DISKANN_DLLEXPORT double calculate_range_search_recall( + unsigned num_queries, std::vector> &groundtruth, + std::vector> &our_results); template -inline void load_bin(const std::string &bin_file, std::unique_ptr &data, size_t &npts, size_t &dim, - size_t offset = 0) -{ - T *ptr; - load_bin(bin_file, ptr, npts, dim, offset); - data.reset(ptr); +inline void load_bin(const std::string &bin_file, std::unique_ptr &data, + size_t &npts, size_t &dim, size_t offset = 0) { + T *ptr; + load_bin(bin_file, ptr, npts, dim, offset); + data.reset(ptr); } -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()) - { - char buff[1024]; +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()) { + char buff[1024]; #ifdef _WINDOWS - auto ret = std::to_string(strerror_s(buff, 1024, errno)); + auto ret = std::to_string(strerror_s(buff, 1024, errno)); #else - auto ret = std::string(strerror_r(errno, buff, 1024)); + auto ret = std::string(strerror_r(errno, buff, 1024)); #endif - std::string error_message = - std::string("Failed to open file") + filename + " for write because " + buff + ", ret=" + ret; - diskann::cerr << error_message << std::endl; - throw diskann::ANNException(error_message, -1); - } + std::string error_message = std::string("Failed to open file") + filename + + " for write because " + buff + ", ret=" + ret; + diskann::cerr << error_message << std::endl; + throw diskann::ANNException(error_message, -1); + } } template -inline size_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; - 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 = " << bytes_written << "B" - << std::endl; - - writer.write((char *)data, npts * ndims * sizeof(T)); - writer.close(); - diskann::cout << "Finished writing bin." << std::endl; - return bytes_written; +inline size_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; + 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 = " << bytes_written << "B" << std::endl; + + writer.write((char *)data, npts * ndims * sizeof(T)); + writer.close(); + diskann::cout << "Finished writing bin." << std::endl; + return bytes_written; } -inline void print_progress(double percentage) -{ - int val = (int)(percentage * 100); - int lpad = (int)(percentage * PBWIDTH); - int rpad = PBWIDTH - lpad; - printf("\r%3d%% [%.*s%*s]", val, lpad, PBSTR, rpad, ""); - fflush(stdout); +inline void print_progress(double percentage) { + int val = (int)(percentage * 100); + int lpad = (int)(percentage * PBWIDTH); + int rpad = PBWIDTH - lpad; + printf("\r%3d%% [%.*s%*s]", val, lpad, PBSTR, rpad, ""); + fflush(stdout); } // load_aligned_bin functions START template -inline void load_aligned_bin_impl(std::basic_istream &reader, size_t actual_file_size, T *&data, size_t &npts, - size_t &dim, size_t &rounded_dim) -{ - 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; - - 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 of " << allocSize << " bytes... " << std::flush; - alloc_aligned(((void **)&data), allocSize, 8 * sizeof(T)); - diskann::cout << "done. Copying data to mem_aligned buffer..." << 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; +inline void load_aligned_bin_impl(std::basic_istream &reader, + size_t actual_file_size, T *&data, + size_t &npts, size_t &dim, + size_t &rounded_dim) { + 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; + + 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 of " << allocSize << " bytes... " + << std::flush; + alloc_aligned(((void **)&data), allocSize, 8 * sizeof(T)); + diskann::cout << "done. Copying data to mem_aligned buffer..." << 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; } #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) -{ - try - { - diskann::cout << "Opening 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 actual_file_size = fc._size; - load_aligned_bin_impl(reader, actual_file_size, data, npts, dim, rounded_dim); - } - catch (std::system_error &e) - { - throw FileException(bin_file, e, __FUNCSIG__, __FILE__, __LINE__); - } +inline void load_aligned_bin(MemoryMappedFiles &files, + const std::string &bin_file, T *&data, + size_t &npts, size_t &dim, size_t &rounded_dim) { + try { + diskann::cout << "Opening 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 actual_file_size = fc._size; + load_aligned_bin_impl(reader, actual_file_size, data, npts, dim, + rounded_dim); + } catch (std::system_error &e) { + throw FileException(bin_file, e, __FUNCSIG__, __FILE__, __LINE__); + } } #endif template -inline void load_aligned_bin(const std::string &bin_file, T *&data, size_t &npts, size_t &dim, size_t &rounded_dim) -{ - std::ifstream reader; - reader.exceptions(std::ifstream::failbit | std::ifstream::badbit); - - try - { - diskann::cout << "Reading (with alignment) bin file " << bin_file << " ..." << std::flush; - reader.open(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); - } - catch (std::system_error &e) - { - throw FileException(bin_file, e, __FUNCSIG__, __FILE__, __LINE__); - } +inline void load_aligned_bin(const std::string &bin_file, T *&data, + size_t &npts, size_t &dim, size_t &rounded_dim) { + std::ifstream reader; + reader.exceptions(std::ifstream::failbit | std::ifstream::badbit); + + try { + diskann::cout << "Reading (with alignment) bin file " << bin_file << " ..." + << std::flush; + reader.open(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); + } catch (std::system_error &e) { + throw FileException(bin_file, e, __FUNCSIG__, __FILE__, __LINE__); + } } template -void convert_types(const InType *srcmat, OutType *destmat, size_t npts, size_t dim) -{ +void convert_types(const InType *srcmat, OutType *destmat, size_t npts, + size_t dim) { #pragma omp parallel for schedule(static, 65536) - for (int64_t i = 0; i < (int64_t)npts; i++) - { - for (uint64_t j = 0; j < dim; j++) - { - destmat[i * dim + j] = (OutType)srcmat[i * dim + j]; - } + for (int64_t i = 0; i < (int64_t)npts; i++) { + for (uint64_t j = 0; j < dim; j++) { + destmat[i * dim + j] = (OutType)srcmat[i * dim + j]; } + } } // this function will take in_file of n*d dimensions and save the output as a @@ -801,295 +782,282 @@ void convert_types(const InType *srcmat, OutType *destmat, size_t npts, size_t d // from MIPS to L2 search from "On Symmetric and Asymmetric LSHs for Inner // Product Search" by Neyshabur and Srebro -template float prepare_base_for_inner_products(const std::string in_file, const std::string out_file) -{ - std::cout << "Pre-processing base file by adding extra coordinate" << std::endl; - std::ifstream in_reader(in_file.c_str(), std::ios::binary); - std::ofstream out_writer(out_file.c_str(), std::ios::binary); - uint64_t npts, in_dims, out_dims; - float max_norm = 0; - - uint32_t npts32, dims32; - in_reader.read((char *)&npts32, sizeof(uint32_t)); - in_reader.read((char *)&dims32, sizeof(uint32_t)); - - npts = npts32; - in_dims = dims32; - out_dims = in_dims + 1; - uint32_t outdims32 = (uint32_t)out_dims; - - out_writer.write((char *)&npts32, sizeof(uint32_t)); - out_writer.write((char *)&outdims32, sizeof(uint32_t)); - - size_t BLOCK_SIZE = 100000; - size_t block_size = npts <= BLOCK_SIZE ? npts : BLOCK_SIZE; - std::unique_ptr in_block_data = std::make_unique(block_size * in_dims); - std::unique_ptr out_block_data = std::make_unique(block_size * out_dims); - - std::memset(out_block_data.get(), 0, sizeof(float) * block_size * out_dims); - uint64_t num_blocks = DIV_ROUND_UP(npts, block_size); - - std::vector norms(npts, 0); - - for (uint64_t b = 0; b < num_blocks; b++) - { - uint64_t start_id = b * block_size; - uint64_t end_id = (b + 1) * block_size < npts ? (b + 1) * block_size : npts; - uint64_t block_pts = end_id - start_id; - in_reader.read((char *)in_block_data.get(), block_pts * in_dims * sizeof(T)); - for (uint64_t p = 0; p < block_pts; p++) - { - for (uint64_t j = 0; j < in_dims; j++) - { - norms[start_id + p] += in_block_data[p * in_dims + j] * in_block_data[p * in_dims + j]; - } - max_norm = max_norm > norms[start_id + p] ? max_norm : norms[start_id + p]; - } +template +float prepare_base_for_inner_products(const std::string in_file, + const std::string out_file) { + std::cout << "Pre-processing base file by adding extra coordinate" + << std::endl; + std::ifstream in_reader(in_file.c_str(), std::ios::binary); + std::ofstream out_writer(out_file.c_str(), std::ios::binary); + uint64_t npts, in_dims, out_dims; + float max_norm = 0; + + uint32_t npts32, dims32; + in_reader.read((char *)&npts32, sizeof(uint32_t)); + in_reader.read((char *)&dims32, sizeof(uint32_t)); + + npts = npts32; + in_dims = dims32; + out_dims = in_dims + 1; + uint32_t outdims32 = (uint32_t)out_dims; + + out_writer.write((char *)&npts32, sizeof(uint32_t)); + out_writer.write((char *)&outdims32, sizeof(uint32_t)); + + size_t BLOCK_SIZE = 100000; + size_t block_size = npts <= BLOCK_SIZE ? npts : BLOCK_SIZE; + std::unique_ptr in_block_data = + std::make_unique(block_size * in_dims); + std::unique_ptr out_block_data = + std::make_unique(block_size * out_dims); + + std::memset(out_block_data.get(), 0, sizeof(float) * block_size * out_dims); + uint64_t num_blocks = DIV_ROUND_UP(npts, block_size); + + std::vector norms(npts, 0); + + for (uint64_t b = 0; b < num_blocks; b++) { + uint64_t start_id = b * block_size; + uint64_t end_id = (b + 1) * block_size < npts ? (b + 1) * block_size : npts; + uint64_t block_pts = end_id - start_id; + in_reader.read((char *)in_block_data.get(), + block_pts * in_dims * sizeof(T)); + for (uint64_t p = 0; p < block_pts; p++) { + for (uint64_t j = 0; j < in_dims; j++) { + norms[start_id + p] += + in_block_data[p * in_dims + j] * in_block_data[p * in_dims + j]; + } + max_norm = + max_norm > norms[start_id + p] ? max_norm : norms[start_id + p]; } - - max_norm = std::sqrt(max_norm); - - in_reader.seekg(2 * sizeof(uint32_t), std::ios::beg); - for (uint64_t b = 0; b < num_blocks; b++) - { - uint64_t start_id = b * block_size; - uint64_t end_id = (b + 1) * block_size < npts ? (b + 1) * block_size : npts; - uint64_t block_pts = end_id - start_id; - in_reader.read((char *)in_block_data.get(), block_pts * in_dims * sizeof(T)); - for (uint64_t p = 0; p < block_pts; p++) - { - for (uint64_t j = 0; j < in_dims; j++) - { - out_block_data[p * out_dims + j] = in_block_data[p * in_dims + j] / max_norm; - } - float res = 1 - (norms[start_id + p] / (max_norm * max_norm)); - res = res <= 0 ? 0 : std::sqrt(res); - out_block_data[p * out_dims + out_dims - 1] = res; - } - out_writer.write((char *)out_block_data.get(), block_pts * out_dims * sizeof(float)); + } + + max_norm = std::sqrt(max_norm); + + in_reader.seekg(2 * sizeof(uint32_t), std::ios::beg); + for (uint64_t b = 0; b < num_blocks; b++) { + uint64_t start_id = b * block_size; + uint64_t end_id = (b + 1) * block_size < npts ? (b + 1) * block_size : npts; + uint64_t block_pts = end_id - start_id; + in_reader.read((char *)in_block_data.get(), + block_pts * in_dims * sizeof(T)); + for (uint64_t p = 0; p < block_pts; p++) { + for (uint64_t j = 0; j < in_dims; j++) { + out_block_data[p * out_dims + j] = + in_block_data[p * in_dims + j] / max_norm; + } + float res = 1 - (norms[start_id + p] / (max_norm * max_norm)); + res = res <= 0 ? 0 : std::sqrt(res); + out_block_data[p * out_dims + out_dims - 1] = res; } - out_writer.close(); - return max_norm; + out_writer.write((char *)out_block_data.get(), + block_pts * out_dims * sizeof(float)); + } + out_writer.close(); + return max_norm; } // 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); +template +void save_Tvecs(const char *filename, T *data, size_t npts, size_t ndims) { + std::string fname(filename); - // create cached ofstream with 64MB cache - cached_ofstream writer(fname, 64 * 1048576); + // create cached ofstream with 64MB cache + cached_ofstream writer(fname, 64 * 1048576); - unsigned dims_u32 = (unsigned)ndims; + unsigned dims_u32 = (unsigned)ndims; - // start writing - for (size_t i = 0; i < npts; i++) - { - // write dims in u32 - writer.write((char *)&dims_u32, sizeof(unsigned)); + // start writing + for (size_t i = 0; i < npts; i++) { + // write dims in u32 + writer.write((char *)&dims_u32, sizeof(unsigned)); - // get cur point in data - T *cur_pt = data + i * ndims; - writer.write((char *)cur_pt, ndims * sizeof(T)); - } + // get cur point in data + T *cur_pt = data + i * ndims; + writer.write((char *)cur_pt, ndims * sizeof(T)); + } } template -inline size_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; - size_t 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; +inline size_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; + size_t 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; } template -inline void copy_aligned_data_from_file(const char *bin_file, T *&data, size_t &npts, size_t &dim, - const size_t &rounded_dim, size_t offset = 0) -{ - if (data == nullptr) - { - diskann::cerr << "Memory was not allocated for " << data << " before calling the load function. Exiting..." - << std::endl; - throw diskann::ANNException("Null pointer passed to copy_aligned_data_from_file function", -1, __FUNCSIG__, - __FILE__, __LINE__); - } - std::ifstream reader; - reader.exceptions(std::ios::badbit | std::ios::failbit); - reader.open(bin_file, std::ios::binary); - reader.seekg(offset, reader.beg); - - 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; - - 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)); - } +inline void copy_aligned_data_from_file(const char *bin_file, T *&data, + size_t &npts, size_t &dim, + const size_t &rounded_dim, + size_t offset = 0) { + if (data == nullptr) { + diskann::cerr << "Memory was not allocated for " << data + << " before calling the load function. Exiting..." + << std::endl; + throw diskann::ANNException( + "Null pointer passed to copy_aligned_data_from_file function", -1, + __FUNCSIG__, __FILE__, __LINE__); + } + std::ifstream reader; + reader.exceptions(std::ios::badbit | std::ios::failbit); + reader.open(bin_file, std::ios::binary); + reader.seekg(offset, reader.beg); + + 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; + + 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)); + } } // NOTE :: good efficiency when total_vec_size is integral multiple of 64 -inline void prefetch_vector(const char *vec, size_t vecsize) -{ - size_t max_prefetch_size = (vecsize / 64) * 64; - for (size_t d = 0; d < max_prefetch_size; d += 64) - _mm_prefetch((const char *)vec + d, _MM_HINT_T0); +inline void prefetch_vector(const char *vec, size_t vecsize) { + size_t max_prefetch_size = (vecsize / 64) * 64; + for (size_t d = 0; d < max_prefetch_size; d += 64) + _mm_prefetch((const char *)vec + d, _MM_HINT_T0); } // NOTE :: good efficiency when total_vec_size is integral multiple of 64 -inline void prefetch_vector_l2(const char *vec, size_t vecsize) -{ - size_t max_prefetch_size = (vecsize / 64) * 64; - for (size_t d = 0; d < max_prefetch_size; d += 64) - _mm_prefetch((const char *)vec + d, _MM_HINT_T1); +inline void prefetch_vector_l2(const char *vec, size_t vecsize) { + size_t max_prefetch_size = (vecsize / 64) * 64; + for (size_t d = 0; d < max_prefetch_size; d += 64) + _mm_prefetch((const char *)vec + d, _MM_HINT_T1); } // NOTE: Implementation in utils.cpp. -void block_convert(std::ofstream &writr, std::ifstream &readr, float *read_buf, uint64_t npts, uint64_t ndims); +void block_convert(std::ofstream &writr, std::ifstream &readr, float *read_buf, + uint64_t npts, uint64_t ndims); -DISKANN_DLLEXPORT void normalize_data_file(const std::string &inFileName, const std::string &outFileName); +DISKANN_DLLEXPORT void normalize_data_file(const std::string &inFileName, + const std::string &outFileName); -}; // namespace diskann +}; // namespace diskann -struct PivotContainer -{ - PivotContainer() = default; +struct PivotContainer { + PivotContainer() = default; - PivotContainer(size_t pivo_id, float pivo_dist) : piv_id{pivo_id}, piv_dist{pivo_dist} - { - } + PivotContainer(size_t pivo_id, float pivo_dist) + : piv_id{pivo_id}, piv_dist{pivo_dist} {} - bool operator<(const PivotContainer &p) const - { - return p.piv_dist < piv_dist; - } + bool operator<(const PivotContainer &p) const { + return p.piv_dist < piv_dist; + } - bool operator>(const PivotContainer &p) const - { - return p.piv_dist > piv_dist; - } + bool operator>(const PivotContainer &p) const { + return p.piv_dist > piv_dist; + } - size_t piv_id; - float piv_dist; + size_t piv_id; + float piv_dist; }; -inline bool validate_index_file_size(std::ifstream &in) -{ - if (!in.is_open()) - throw diskann::ANNException("Index file size check called on unopened file stream", -1, __FUNCSIG__, __FILE__, - __LINE__); - in.seekg(0, in.end); - size_t actual_file_size = in.tellg(); - in.seekg(0, in.beg); - size_t expected_file_size; - in.read((char *)&expected_file_size, sizeof(uint64_t)); - in.seekg(0, in.beg); - if (actual_file_size != expected_file_size) - { - diskann::cerr << "Index file size error. Expected size (metadata): " << expected_file_size - << ", actual file size : " << actual_file_size << "." << std::endl; - return false; - } - return true; +inline bool validate_index_file_size(std::ifstream &in) { + if (!in.is_open()) + throw diskann::ANNException( + "Index file size check called on unopened file stream", -1, __FUNCSIG__, + __FILE__, __LINE__); + in.seekg(0, in.end); + size_t actual_file_size = in.tellg(); + in.seekg(0, in.beg); + size_t expected_file_size; + in.read((char *)&expected_file_size, sizeof(uint64_t)); + in.seekg(0, in.beg); + if (actual_file_size != expected_file_size) { + diskann::cerr << "Index file size error. Expected size (metadata): " + << expected_file_size + << ", actual file size : " << actual_file_size << "." + << std::endl; + return false; + } + return true; } -template inline float get_norm(T *arr, const size_t dim) -{ - float sum = 0.0f; - for (uint32_t i = 0; i < dim; i++) - { - sum += arr[i] * arr[i]; - } - return sqrt(sum); +template +inline float get_norm(T *arr, const size_t dim) { + float sum = 0.0f; + for (uint32_t i = 0; i < dim; i++) { + sum += arr[i] * arr[i]; + } + return sqrt(sum); } // This function is valid only for float data type. -template inline void normalize(T *arr, const size_t dim) -{ - float norm = get_norm(arr, dim); - for (uint32_t i = 0; i < dim; i++) - { - arr[i] = (T)(arr[i] / norm); - } +template +inline void normalize(T *arr, const size_t dim) { + float norm = get_norm(arr, dim); + for (uint32_t i = 0; i < dim; i++) { + arr[i] = (T)(arr[i] / norm); + } } -inline std::vector read_file_to_vector_of_strings(const std::string &filename, bool unique = false) -{ - std::vector result; - std::set elementSet; - if (filename != "") - { - std::ifstream file(filename); - if (file.fail()) - { - throw diskann::ANNException(std::string("Failed to open file ") + filename, -1); - } - std::string line; - while (std::getline(file, line)) - { - if (line.empty()) - { - break; - } - if (line.find(',') != std::string::npos) - { - std::cerr << "Every query must have exactly one filter" << std::endl; - exit(-1); - } - if (!line.empty() && (line.back() == '\r' || line.back() == '\n')) - { - line.erase(line.size() - 1); - } - if (!elementSet.count(line)) - { - result.push_back(line); - } - if (unique) - { - elementSet.insert(line); - } - } - file.close(); +inline std::vector read_file_to_vector_of_strings( + const std::string &filename, bool unique = false) { + std::vector result; + std::set elementSet; + if (filename != "") { + std::ifstream file(filename); + if (file.fail()) { + throw diskann::ANNException( + std::string("Failed to open file ") + filename, -1); } - else - { - throw diskann::ANNException(std::string("Failed to open file. filename can not be blank"), -1); + std::string line; + while (std::getline(file, line)) { + if (line.empty()) { + break; + } + if (line.find(',') != std::string::npos) { + std::cerr << "Every query must have exactly one filter" << std::endl; + exit(-1); + } + if (!line.empty() && (line.back() == '\r' || line.back() == '\n')) { + line.erase(line.size() - 1); + } + if (!elementSet.count(line)) { + result.push_back(line); + } + if (unique) { + elementSet.insert(line); + } } - return result; + file.close(); + } else { + throw diskann::ANNException( + std::string("Failed to open file. filename can not be blank"), -1); + } + return result; } -inline void clean_up_artifacts(tsl::robin_set paths_to_clean, tsl::robin_set path_suffixes) -{ - try - { - for (const auto &path : paths_to_clean) - { - for (const auto &suffix : path_suffixes) - { - std::string curr_path_to_clean(path + "_" + suffix); - if (std::remove(curr_path_to_clean.c_str()) != 0) - diskann::cout << "Warning: Unable to remove file :" << curr_path_to_clean << std::endl; - } - } - diskann::cout << "Cleaned all artifacts" << std::endl; - } - catch (const std::exception &e) - { - diskann::cout << "Warning: Unable to clean all artifacts" << std::endl; +inline void clean_up_artifacts(tsl::robin_set paths_to_clean, + tsl::robin_set path_suffixes) { + try { + for (const auto &path : paths_to_clean) { + for (const auto &suffix : path_suffixes) { + std::string curr_path_to_clean(path + "_" + suffix); + if (std::remove(curr_path_to_clean.c_str()) != 0) + diskann::cout << "Warning: Unable to remove file :" + << curr_path_to_clean << std::endl; + } } + diskann::cout << "Cleaned all artifacts" << std::endl; + } catch (const std::exception &e) { + diskann::cout << "Warning: Unable to clean all artifacts" << std::endl; + } } #ifdef _WINDOWS @@ -1099,57 +1067,53 @@ inline void clean_up_artifacts(tsl::robin_set paths_to_clean, tsl:: extern bool AvxSupportedCPU; extern bool Avx2SupportedCPU; -inline size_t getMemoryUsage() -{ - PROCESS_MEMORY_COUNTERS_EX pmc; - GetProcessMemoryInfo(GetCurrentProcess(), (PROCESS_MEMORY_COUNTERS *)&pmc, sizeof(pmc)); - return pmc.PrivateUsage; +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 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.0 * 1024 * 1024) - << "GB Working set size: " << counters.WorkingSetSize * 1.0 / (1024.0 * 1024 * 1024) - << "GB Private bytes " << counters.PagefileUsage * 1.0 / (1024 * 1024 * 1024) << "GB]" << std::endl; +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.0 * 1024 * 1024) + << "GB Working set size: " + << counters.WorkingSetSize * 1.0 / (1024.0 * 1024 * 1024) + << "GB Private bytes " + << counters.PagefileUsage * 1.0 / (1024 * 1024 * 1024) << "GB]" + << std::endl; } #else // need to check and change this -inline bool avx2Supported() -{ - return true; -} -inline void printProcessMemory(const char *) -{ -} +inline bool avx2Supported() { return true; } +inline void printProcessMemory(const char *) {} -inline size_t getMemoryUsage() -{ // for non-windows, we have not implemented this function - return 0; +inline size_t +getMemoryUsage() { // for non-windows, we have not implemented this function + return 0; } #endif diff --git a/include/windows_aligned_file_reader.h b/include/windows_aligned_file_reader.h index 0d9a3173c..e7092a973 100644 --- a/include/windows_aligned_file_reader.h +++ b/include/windows_aligned_file_reader.h @@ -17,41 +17,39 @@ #include "utils.h" #include "windows_customizations.h" -class WindowsAlignedFileReader : public AlignedFileReader -{ - private: +class WindowsAlignedFileReader : public AlignedFileReader { + private: #ifdef UNICODE - std::wstring m_filename; + std::wstring m_filename; #else - std::string m_filename; + std::string m_filename; #endif - protected: - // virtual IOContext createContext(); - - public: - DISKANN_DLLEXPORT WindowsAlignedFileReader(){}; - DISKANN_DLLEXPORT virtual ~WindowsAlignedFileReader(){}; - - // Open & close ops - // Blocking calls - DISKANN_DLLEXPORT virtual void open(const std::string &fname) override; - DISKANN_DLLEXPORT virtual void close() override; - - DISKANN_DLLEXPORT virtual void register_thread() override; - DISKANN_DLLEXPORT virtual void deregister_thread() override - { - // TODO: Needs implementation. - } - DISKANN_DLLEXPORT virtual void deregister_all_threads() override - { - // TODO: Needs implementation. - } - DISKANN_DLLEXPORT virtual IOContext &get_ctx() override; - - // 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) override; + protected: + // virtual IOContext createContext(); + + public: + DISKANN_DLLEXPORT WindowsAlignedFileReader(){}; + DISKANN_DLLEXPORT virtual ~WindowsAlignedFileReader(){}; + + // Open & close ops + // Blocking calls + DISKANN_DLLEXPORT virtual void open(const std::string &fname) override; + DISKANN_DLLEXPORT virtual void close() override; + + DISKANN_DLLEXPORT virtual void register_thread() override; + DISKANN_DLLEXPORT virtual void deregister_thread() override { + // TODO: Needs implementation. + } + DISKANN_DLLEXPORT virtual void deregister_all_threads() override { + // TODO: Needs implementation. + } + DISKANN_DLLEXPORT virtual IOContext &get_ctx() override; + + // 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) override; }; -#endif // USE_BING_INFRA -#endif //_WINDOWS +#endif // USE_BING_INFRA +#endif //_WINDOWS diff --git a/include/windows_slim_lock.h b/include/windows_slim_lock.h index 5d0d65508..8f37951a2 100644 --- a/include/windows_slim_lock.h +++ b/include/windows_slim_lock.h @@ -6,8 +6,7 @@ #endif #include "Windows.h" -namespace diskann -{ +namespace diskann { // A thin C++ wrapper around Windows exclusive functionality of Windows // SlimReaderWriterLock. // @@ -18,55 +17,42 @@ namespace diskann // // Full documentation can be found at. // https://msdn.microsoft.com/en-us/library/windows/desktop/aa904937(v=vs.85).aspx -class windows_exclusive_slim_lock -{ - public: - windows_exclusive_slim_lock() : _lock(SRWLOCK_INIT) - { - } +class windows_exclusive_slim_lock { + public: + windows_exclusive_slim_lock() : _lock(SRWLOCK_INIT) {} - // The lock is non-copyable. This also disables move constructor/operator=. - windows_exclusive_slim_lock(const windows_exclusive_slim_lock &) = delete; - windows_exclusive_slim_lock &operator=(const windows_exclusive_slim_lock &) = delete; + // The lock is non-copyable. This also disables move constructor/operator=. + windows_exclusive_slim_lock(const windows_exclusive_slim_lock &) = delete; + windows_exclusive_slim_lock &operator=(const windows_exclusive_slim_lock &) = + delete; - void lock() - { - return AcquireSRWLockExclusive(&_lock); - } + void lock() { return AcquireSRWLockExclusive(&_lock); } - bool try_lock() - { - return TryAcquireSRWLockExclusive(&_lock) != FALSE; - } + bool try_lock() { return TryAcquireSRWLockExclusive(&_lock) != FALSE; } - void unlock() - { - return ReleaseSRWLockExclusive(&_lock); - } + void unlock() { return ReleaseSRWLockExclusive(&_lock); } - private: - SRWLOCK _lock; + private: + SRWLOCK _lock; }; // An exclusive lock over a SlimReaderWriterLock. -class windows_exclusive_slim_lock_guard -{ - public: - windows_exclusive_slim_lock_guard(windows_exclusive_slim_lock &p_lock) : _lock(p_lock) - { - _lock.lock(); - } - - // The lock is non-copyable. This also disables move constructor/operator=. - windows_exclusive_slim_lock_guard(const windows_exclusive_slim_lock_guard &) = delete; - windows_exclusive_slim_lock_guard &operator=(const windows_exclusive_slim_lock_guard &) = delete; - - ~windows_exclusive_slim_lock_guard() - { - _lock.unlock(); - } - - private: - windows_exclusive_slim_lock &_lock; +class windows_exclusive_slim_lock_guard { + public: + windows_exclusive_slim_lock_guard(windows_exclusive_slim_lock &p_lock) + : _lock(p_lock) { + _lock.lock(); + } + + // The lock is non-copyable. This also disables move constructor/operator=. + windows_exclusive_slim_lock_guard(const windows_exclusive_slim_lock_guard &) = + delete; + windows_exclusive_slim_lock_guard &operator=( + const windows_exclusive_slim_lock_guard &) = delete; + + ~windows_exclusive_slim_lock_guard() { _lock.unlock(); } + + private: + windows_exclusive_slim_lock &_lock; }; -} // namespace diskann +} // namespace diskann diff --git a/src/abstract_data_store.cpp b/src/abstract_data_store.cpp index a980bd545..e154a9cf4 100644 --- a/src/abstract_data_store.cpp +++ b/src/abstract_data_store.cpp @@ -5,42 +5,35 @@ #include "abstract_data_store.h" -namespace diskann -{ +namespace diskann { template -AbstractDataStore::AbstractDataStore(const location_t capacity, const size_t dim) - : _capacity(capacity), _dim(dim) -{ -} +AbstractDataStore::AbstractDataStore(const location_t capacity, + const size_t dim) + : _capacity(capacity), _dim(dim) {} -template location_t AbstractDataStore::capacity() const -{ - return _capacity; +template +location_t AbstractDataStore::capacity() const { + return _capacity; } -template size_t AbstractDataStore::get_dims() const -{ - return _dim; +template +size_t AbstractDataStore::get_dims() const { + return _dim; } -template location_t AbstractDataStore::resize(const location_t new_num_points) -{ - if (new_num_points > _capacity) - { - return expand(new_num_points); - } - else if (new_num_points < _capacity) - { - return shrink(new_num_points); - } - else - { - return _capacity; - } +template +location_t AbstractDataStore::resize(const location_t new_num_points) { + if (new_num_points > _capacity) { + return expand(new_num_points); + } else if (new_num_points < _capacity) { + return shrink(new_num_points); + } else { + return _capacity; + } } template DISKANN_DLLEXPORT class AbstractDataStore; template DISKANN_DLLEXPORT class AbstractDataStore; template DISKANN_DLLEXPORT class AbstractDataStore; -} // namespace diskann +} // namespace diskann diff --git a/src/ann_exception.cpp b/src/ann_exception.cpp index ba55e3655..36753305a 100644 --- a/src/ann_exception.cpp +++ b/src/ann_exception.cpp @@ -5,32 +5,33 @@ #include #include -namespace diskann -{ +namespace diskann { ANNException::ANNException(const std::string &message, int errorCode) - : std::runtime_error(message), _errorCode(errorCode) -{ -} + : std::runtime_error(message), _errorCode(errorCode) {} -std::string package_string(const std::string &item_name, const std::string &item_val) -{ - return std::string("[") + item_name + ": " + std::string(item_val) + std::string("]"); +std::string package_string(const std::string &item_name, + const std::string &item_val) { + return std::string("[") + item_name + ": " + std::string(item_val) + + std::string("]"); } -ANNException::ANNException(const std::string &message, int errorCode, const std::string &funcSig, +ANNException::ANNException(const std::string &message, int errorCode, + const std::string &funcSig, const std::string &fileName, uint32_t lineNum) - : ANNException(package_string(std::string("FUNC"), funcSig) + package_string(std::string("FILE"), fileName) + - package_string(std::string("LINE"), std::to_string(lineNum)) + " " + message, - errorCode) -{ -} + : ANNException( + package_string(std::string("FUNC"), funcSig) + + package_string(std::string("FILE"), fileName) + + package_string(std::string("LINE"), std::to_string(lineNum)) + + " " + message, + errorCode) {} -FileException::FileException(const std::string &filename, std::system_error &e, const std::string &funcSig, +FileException::FileException(const std::string &filename, std::system_error &e, + const std::string &funcSig, const std::string &fileName, uint32_t lineNum) - : ANNException(std::string(" While opening file \'") + filename + std::string("\', error code: ") + - std::to_string(e.code().value()) + " " + e.code().message(), - e.code().value(), funcSig, fileName, lineNum) -{ -} + : ANNException(std::string(" While opening file \'") + filename + + std::string("\', error code: ") + + std::to_string(e.code().value()) + " " + + e.code().message(), + e.code().value(), funcSig, fileName, lineNum) {} -} // namespace diskann \ No newline at end of file +} // namespace diskann \ No newline at end of file diff --git a/src/disk_utils.cpp b/src/disk_utils.cpp index 0720473bf..a22ebe1fb 100644 --- a/src/disk_utils.cpp +++ b/src/disk_utils.cpp @@ -3,7 +3,8 @@ #include "common_includes.h" -#if defined(RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS) && defined(DISKANN_BUILD) +#if defined(RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS) && \ + defined(DISKANN_BUILD) #include "gperftools/malloc_extension.h" #endif @@ -19,464 +20,456 @@ #include "timer.h" #include "tsl/robin_set.h" -namespace diskann -{ - -void add_new_file_to_single_index(std::string index_file, std::string new_file) -{ - std::unique_ptr metadata; - uint64_t nr, nc; - diskann::load_bin(index_file, metadata, nr, nc); - 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]; - size_t read_blk_size = 64 * 1024 * 1024; - cached_ofstream writer(index_file, read_blk_size); - size_t 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 (uint64_t 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 new_meta; - for (uint64_t i = 0; i < nr; i++) - new_meta.push_back(metadata[i]); - new_meta.push_back(metadata[nr - 1] + fsize); - - diskann::save_bin(index_file, new_meta.data(), new_meta.size(), 1); +namespace diskann { + +void add_new_file_to_single_index(std::string index_file, + std::string new_file) { + std::unique_ptr metadata; + uint64_t nr, nc; + diskann::load_bin(index_file, metadata, nr, nc); + 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]; + size_t read_blk_size = 64 * 1024 * 1024; + cached_ofstream writer(index_file, read_blk_size); + size_t 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 (uint64_t 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 new_meta; + for (uint64_t i = 0; i < nr; i++) new_meta.push_back(metadata[i]); + new_meta.push_back(metadata[nr - 1] + fsize); + + diskann::save_bin(index_file, new_meta.data(), new_meta.size(), 1); } -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 = search_ram_budget - SPACE_FOR_CACHED_NODES_IN_GB; - } - return final_index_ram_limit * 1024 * 1024 * 1024; +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 = 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); +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, - const std::vector ¶m_list) -{ - size_t num_pq_chunks = (size_t)(std::floor)(uint64_t(final_index_ram_limit / (double)points_num)); - diskann::cout << "Calculated num_pq_chunks :" << num_pq_chunks << std::endl; - if (param_list.size() >= 6) - { - float compress_ratio = (float)atof(param_list[5].c_str()); - if (compress_ratio > 0 && compress_ratio <= 1) - { - size_t chunks_by_cr = (size_t)(std::floor)(compress_ratio * dim); - - if (chunks_by_cr > 0 && chunks_by_cr < num_pq_chunks) - { - diskann::cout << "Compress ratio:" << compress_ratio << " new #pq_chunks:" << chunks_by_cr << std::endl; - num_pq_chunks = chunks_by_cr; - } - else - { - diskann::cout << "Compress ratio: " << compress_ratio << " #new pq_chunks: " << chunks_by_cr - << " is either zero or greater than num_pq_chunks: " << num_pq_chunks - << ". num_pq_chunks is unchanged. " << std::endl; - } - } - else - { - diskann::cerr << "Compression ratio: " << compress_ratio << " should be in (0,1]" << std::endl; - } +size_t calculate_num_pq_chunks(double final_index_ram_limit, size_t points_num, + uint32_t dim, + const std::vector ¶m_list) { + size_t num_pq_chunks = (size_t)(std::floor)( + uint64_t(final_index_ram_limit / (double)points_num)); + diskann::cout << "Calculated num_pq_chunks :" << num_pq_chunks << std::endl; + if (param_list.size() >= 6) { + float compress_ratio = (float)atof(param_list[5].c_str()); + if (compress_ratio > 0 && compress_ratio <= 1) { + size_t chunks_by_cr = (size_t)(std::floor)(compress_ratio * dim); + + if (chunks_by_cr > 0 && chunks_by_cr < num_pq_chunks) { + diskann::cout << "Compress ratio:" << compress_ratio + << " new #pq_chunks:" << chunks_by_cr << std::endl; + num_pq_chunks = chunks_by_cr; + } else { + diskann::cout << "Compress ratio: " << compress_ratio + << " #new pq_chunks: " << chunks_by_cr + << " is either zero or greater than num_pq_chunks: " + << num_pq_chunks << ". num_pq_chunks is unchanged. " + << std::endl; + } + } else { + diskann::cerr << "Compression ratio: " << compress_ratio + << " should be in (0,1]" << 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; + 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; + diskann::cout << "Compressing " << dim << "-dimensional data into " + << num_pq_chunks << " bytes per vector." << std::endl; + return num_pq_chunks; } -template T *generateRandomWarmup(uint64_t warmup_num, uint64_t warmup_dim, uint64_t warmup_aligned_dim) -{ - T *warmup = nullptr; - warmup_num = 100000; - diskann::cout << "Generating random warmup file with dim " << warmup_dim << " and aligned dim " - << warmup_aligned_dim << std::flush; - diskann::alloc_aligned(((void **)&warmup), warmup_num * warmup_aligned_dim * sizeof(T), 8 * sizeof(T)); - std::memset(warmup, 0, warmup_num * warmup_aligned_dim * sizeof(T)); - std::random_device rd; - std::mt19937 gen(rd()); - std::uniform_int_distribution<> dis(-128, 127); - for (uint32_t i = 0; i < warmup_num; i++) - { - for (uint32_t d = 0; d < warmup_dim; d++) - { - warmup[i * warmup_aligned_dim + d] = (T)dis(gen); - } +template +T *generateRandomWarmup(uint64_t warmup_num, uint64_t warmup_dim, + uint64_t warmup_aligned_dim) { + T *warmup = nullptr; + warmup_num = 100000; + diskann::cout << "Generating random warmup file with dim " << warmup_dim + << " and aligned dim " << warmup_aligned_dim << std::flush; + diskann::alloc_aligned(((void **)&warmup), + warmup_num * warmup_aligned_dim * sizeof(T), + 8 * sizeof(T)); + std::memset(warmup, 0, warmup_num * warmup_aligned_dim * sizeof(T)); + std::random_device rd; + std::mt19937 gen(rd()); + std::uniform_int_distribution<> dis(-128, 127); + for (uint32_t i = 0; i < warmup_num; i++) { + for (uint32_t d = 0; d < warmup_dim; d++) { + warmup[i * warmup_aligned_dim + d] = (T)dis(gen); } - diskann::cout << "..done" << std::endl; - return warmup; + } + diskann::cout << "..done" << std::endl; + return warmup; } #ifdef EXEC_ENV_OLS template -T *load_warmup(MemoryMappedFiles &files, const std::string &cache_warmup_file, uint64_t &warmup_num, - uint64_t warmup_dim, uint64_t warmup_aligned_dim) -{ - T *warmup = nullptr; - uint64_t file_dim, file_aligned_dim; - - 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 = " << file_dim - << " file_aligned_dim: " << file_aligned_dim << " index_dim: " << warmup_dim - << " index_aligned_dim: " << warmup_aligned_dim << std::endl; - diskann::cerr << stream.str(); - throw diskann::ANNException(stream.str(), -1); - } - } - else - { - warmup = generateRandomWarmup(warmup_num, warmup_dim, warmup_aligned_dim); +T *load_warmup(MemoryMappedFiles &files, const std::string &cache_warmup_file, + uint64_t &warmup_num, uint64_t warmup_dim, + uint64_t warmup_aligned_dim) { + T *warmup = nullptr; + uint64_t file_dim, file_aligned_dim; + + 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 = " << file_dim + << " file_aligned_dim: " << file_aligned_dim + << " index_dim: " << warmup_dim + << " index_aligned_dim: " << warmup_aligned_dim << std::endl; + diskann::cerr << stream.str(); + throw diskann::ANNException(stream.str(), -1); } - return warmup; + } else { + warmup = + generateRandomWarmup(warmup_num, warmup_dim, warmup_aligned_dim); + } + return warmup; } #endif template -T *load_warmup(const std::string &cache_warmup_file, uint64_t &warmup_num, uint64_t warmup_dim, - uint64_t warmup_aligned_dim) -{ - T *warmup = nullptr; - uint64_t file_dim, file_aligned_dim; - - if (file_exists(cache_warmup_file)) - { - diskann::load_aligned_bin(cache_warmup_file, warmup, warmup_num, file_dim, file_aligned_dim); - if (file_dim != warmup_dim || file_aligned_dim != warmup_aligned_dim) - { - std::stringstream stream; - stream << "Mismatched dimensions in sample file. file_dim = " << file_dim - << " file_aligned_dim: " << file_aligned_dim << " index_dim: " << warmup_dim - << " index_aligned_dim: " << warmup_aligned_dim << std::endl; - throw diskann::ANNException(stream.str(), -1); - } - } - else - { - warmup = generateRandomWarmup(warmup_num, warmup_dim, warmup_aligned_dim); +T *load_warmup(const std::string &cache_warmup_file, uint64_t &warmup_num, + uint64_t warmup_dim, uint64_t warmup_aligned_dim) { + T *warmup = nullptr; + uint64_t file_dim, file_aligned_dim; + + if (file_exists(cache_warmup_file)) { + diskann::load_aligned_bin(cache_warmup_file, warmup, warmup_num, + file_dim, file_aligned_dim); + if (file_dim != warmup_dim || file_aligned_dim != warmup_aligned_dim) { + std::stringstream stream; + stream << "Mismatched dimensions in sample file. file_dim = " << file_dim + << " file_aligned_dim: " << file_aligned_dim + << " index_dim: " << warmup_dim + << " index_aligned_dim: " << warmup_aligned_dim << std::endl; + throw diskann::ANNException(stream.str(), -1); } - return warmup; + } else { + warmup = + generateRandomWarmup(warmup_num, warmup_dim, warmup_aligned_dim); + } + return warmup; } /*************************************************** Support for Merging Many Vamana Indices ***************************************************/ -void read_idmap(const std::string &fname, std::vector &ivecs) -{ - uint32_t npts32, dim; - size_t actual_file_size = get_file_size(fname); - 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)) - { - std::stringstream stream; - stream << "Error reading idmap file. Check if the file is bin file with " - "1 dimensional data. Actual: " - << actual_file_size << ", expected: " << (size_t)npts32 + 2 * sizeof(uint32_t) << std::endl; - - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - ivecs.resize(npts32); - reader.read((char *)ivecs.data(), ((size_t)npts32) * sizeof(uint32_t)); - reader.close(); +void read_idmap(const std::string &fname, std::vector &ivecs) { + uint32_t npts32, dim; + size_t actual_file_size = get_file_size(fname); + 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)) { + std::stringstream stream; + stream << "Error reading idmap file. Check if the file is bin file with " + "1 dimensional data. Actual: " + << actual_file_size + << ", expected: " << (size_t)npts32 + 2 * sizeof(uint32_t) + << std::endl; + + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } + ivecs.resize(npts32); + reader.read((char *)ivecs.data(), ((size_t)npts32) * sizeof(uint32_t)); + reader.close(); } -int merge_shards(const std::string &vamana_prefix, const std::string &vamana_suffix, const std::string &idmaps_prefix, - const std::string &idmaps_suffix, const uint64_t nshards, uint32_t max_degree, - const std::string &output_vamana, const std::string &medoids_file, bool use_filters, - const std::string &labels_to_medoids_file) -{ - // Read ID maps - std::vector vamana_names(nshards); - std::vector> idmaps(nshards); - for (uint64_t shard = 0; shard < nshards; shard++) - { - vamana_names[shard] = vamana_prefix + std::to_string(shard) + vamana_suffix; - read_idmap(idmaps_prefix + std::to_string(shard) + idmaps_suffix, idmaps[shard]); +int merge_shards(const std::string &vamana_prefix, + const std::string &vamana_suffix, + const std::string &idmaps_prefix, + const std::string &idmaps_suffix, const uint64_t nshards, + uint32_t max_degree, const std::string &output_vamana, + const std::string &medoids_file, bool use_filters, + const std::string &labels_to_medoids_file) { + // Read ID maps + std::vector vamana_names(nshards); + std::vector> idmaps(nshards); + for (uint64_t shard = 0; shard < nshards; shard++) { + vamana_names[shard] = vamana_prefix + std::to_string(shard) + vamana_suffix; + read_idmap(idmaps_prefix + std::to_string(shard) + idmaps_suffix, + idmaps[shard]); + } + + // find max node id + size_t nnodes = 0; + size_t nelems = 0; + for (auto &idmap : idmaps) { + for (auto &id : idmap) { + nnodes = std::max(nnodes, (size_t)id); } - - // find max node id - size_t nnodes = 0; - size_t nelems = 0; - for (auto &idmap : idmaps) - { - for (auto &id : idmap) - { - nnodes = std::max(nnodes, (size_t)id); - } - nelems += idmap.size(); + nelems += idmap.size(); + } + nnodes++; + diskann::cout << "# nodes: " << nnodes << ", max. degree: " << max_degree + << std::endl; + + // compute inverse map: node -> shards + std::vector> node_shard; + node_shard.reserve(nelems); + for (size_t shard = 0; shard < nshards; shard++) { + diskann::cout << "Creating inverse map -- shard #" << shard << std::endl; + for (size_t idx = 0; idx < idmaps[shard].size(); idx++) { + size_t node_id = idmaps[shard][idx]; + node_shard.push_back(std::make_pair((uint32_t)node_id, (uint32_t)shard)); } - nnodes++; - diskann::cout << "# nodes: " << nnodes << ", max. degree: " << max_degree << std::endl; - - // compute inverse map: node -> shards - std::vector> node_shard; - node_shard.reserve(nelems); - for (size_t shard = 0; shard < nshards; shard++) - { - diskann::cout << "Creating inverse map -- shard #" << shard << std::endl; - for (size_t idx = 0; idx < idmaps[shard].size(); idx++) - { - size_t node_id = idmaps[shard][idx]; - node_shard.push_back(std::make_pair((uint32_t)node_id, (uint32_t)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); + }); + diskann::cout << "Finished computing node -> shards map" << std::endl; + + // will merge all the labels to medoids files of each shard into one + // combined file + if (use_filters) { + std::unordered_map> global_label_to_medoids; + + for (size_t i = 0; i < nshards; i++) { + std::ifstream mapping_reader; + std::string map_file = vamana_names[i] + "_labels_to_medoids.txt"; + mapping_reader.open(map_file); + + std::string line, token; + uint32_t line_cnt = 0; + + while (std::getline(mapping_reader, line)) { + std::istringstream iss(line); + uint32_t cnt = 0; + uint32_t medoid = 0; + uint32_t label = 0; + while (std::getline(iss, token, ',')) { + token.erase(std::remove(token.begin(), token.end(), '\n'), + token.end()); + token.erase(std::remove(token.begin(), token.end(), '\r'), + token.end()); + + uint32_t token_as_num = std::stoul(token); + + if (cnt == 0) + label = token_as_num; + else + medoid = token_as_num; + cnt++; } + global_label_to_medoids[label].push_back(idmaps[i][medoid]); + line_cnt++; + } + mapping_reader.close(); } - 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; - - // will merge all the labels to medoids files of each shard into one - // combined file - if (use_filters) - { - std::unordered_map> global_label_to_medoids; - - for (size_t i = 0; i < nshards; i++) - { - std::ifstream mapping_reader; - std::string map_file = vamana_names[i] + "_labels_to_medoids.txt"; - mapping_reader.open(map_file); - - std::string line, token; - uint32_t line_cnt = 0; - - while (std::getline(mapping_reader, line)) - { - std::istringstream iss(line); - uint32_t cnt = 0; - uint32_t medoid = 0; - uint32_t label = 0; - while (std::getline(iss, token, ',')) - { - token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); - token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); - - uint32_t token_as_num = std::stoul(token); - - if (cnt == 0) - label = token_as_num; - else - medoid = token_as_num; - cnt++; - } - global_label_to_medoids[label].push_back(idmaps[i][medoid]); - line_cnt++; - } - mapping_reader.close(); - } - std::ofstream mapping_writer(labels_to_medoids_file); - assert(mapping_writer.is_open()); - for (auto iter : global_label_to_medoids) - { - mapping_writer << iter.first << ", "; - auto &vec = iter.second; - for (uint32_t idx = 0; idx < vec.size() - 1; idx++) - { - mapping_writer << vec[idx] << ", "; - } - mapping_writer << vec[vec.size() - 1] << std::endl; - } - mapping_writer.close(); - } - - // create cached vamana readers - std::vector vamana_readers(nshards); - for (size_t i = 0; i < nshards; i++) - { - vamana_readers[i].open(vamana_names[i], BUFFER_SIZE_FOR_CACHED_IO); - size_t expected_file_size; - vamana_readers[i].read((char *)&expected_file_size, sizeof(uint64_t)); + std::ofstream mapping_writer(labels_to_medoids_file); + assert(mapping_writer.is_open()); + for (auto iter : global_label_to_medoids) { + mapping_writer << iter.first << ", "; + auto &vec = iter.second; + for (uint32_t idx = 0; idx < vec.size() - 1; idx++) { + mapping_writer << vec[idx] << ", "; + } + mapping_writer << vec[vec.size() - 1] << std::endl; } - - size_t vamana_metadata_size = - sizeof(uint64_t) + sizeof(uint32_t) + sizeof(uint32_t) + sizeof(uint64_t); // expected file size + max degree + - // medoid_id + frozen_point info - - // create cached vamana writers - cached_ofstream merged_vamana_writer(output_vamana, BUFFER_SIZE_FOR_CACHED_IO); - - size_t merged_index_size = vamana_metadata_size; // we initialize the size of the merged index to - // the metadata size - size_t merged_index_frozen = 0; - merged_vamana_writer.write((char *)&merged_index_size, - sizeof(uint64_t)); // we will overwrite the index size at the end - - uint32_t output_width = max_degree; - uint32_t max_input_width = 0; - // read width from each vamana to advance buffer by sizeof(uint32_t) bytes - for (auto &reader : vamana_readers) - { - uint32_t input_width; - reader.read((char *)&input_width, sizeof(uint32_t)); - max_input_width = input_width > max_input_width ? input_width : max_input_width; + mapping_writer.close(); + } + + // create cached vamana readers + std::vector vamana_readers(nshards); + for (size_t i = 0; i < nshards; i++) { + vamana_readers[i].open(vamana_names[i], BUFFER_SIZE_FOR_CACHED_IO); + size_t expected_file_size; + vamana_readers[i].read((char *)&expected_file_size, sizeof(uint64_t)); + } + + size_t vamana_metadata_size = + sizeof(uint64_t) + sizeof(uint32_t) + sizeof(uint32_t) + + sizeof(uint64_t); // expected file size + max degree + + // medoid_id + frozen_point info + + // create cached vamana writers + cached_ofstream merged_vamana_writer(output_vamana, + BUFFER_SIZE_FOR_CACHED_IO); + + size_t merged_index_size = + vamana_metadata_size; // we initialize the size of the merged index to + // the metadata size + size_t merged_index_frozen = 0; + merged_vamana_writer.write( + (char *)&merged_index_size, + sizeof(uint64_t)); // we will overwrite the index size at the end + + uint32_t output_width = max_degree; + uint32_t max_input_width = 0; + // read width from each vamana to advance buffer by sizeof(uint32_t) bytes + for (auto &reader : vamana_readers) { + uint32_t input_width; + reader.read((char *)&input_width, sizeof(uint32_t)); + max_input_width = + input_width > max_input_width ? input_width : max_input_width; + } + + diskann::cout << "Max input width: " << max_input_width + << ", output width: " << output_width << std::endl; + + merged_vamana_writer.write((char *)&output_width, sizeof(uint32_t)); + std::ofstream medoid_writer(medoids_file.c_str(), std::ios::binary); + uint32_t nshards_u32 = (uint32_t)nshards; + uint32_t one_val = 1; + medoid_writer.write((char *)&nshards_u32, sizeof(uint32_t)); + medoid_writer.write((char *)&one_val, sizeof(uint32_t)); + + uint64_t vamana_index_frozen = + 0; // as of now the functionality to merge many overlapping vamana + // indices is supported only for bulk indices without frozen point. + // Hence the final index will also not have any frozen points. + for (uint64_t shard = 0; shard < nshards; shard++) { + uint32_t medoid; + // read medoid + vamana_readers[shard].read((char *)&medoid, sizeof(uint32_t)); + vamana_readers[shard].read((char *)&vamana_index_frozen, sizeof(uint64_t)); + assert(vamana_index_frozen == false); + // rename medoid + medoid = idmaps[shard][medoid]; + + medoid_writer.write((char *)&medoid, sizeof(uint32_t)); + // write renamed medoid + if (shard == (nshards - 1)) //--> uncomment if running hierarchical + merged_vamana_writer.write((char *)&medoid, sizeof(uint32_t)); + } + merged_vamana_writer.write((char *)&merged_index_frozen, sizeof(uint64_t)); + medoid_writer.close(); + + diskann::cout << "Starting merge" << std::endl; + + // Gopal. random_shuffle() is deprecated. + std::random_device rng; + std::mt19937 urng(rng()); + + std::vector nhood_set(nnodes, 0); + std::vector final_nhood; + + uint32_t nnbrs = 0, shard_nnbrs = 0; + uint32_t cur_id = 0; + for (const auto &id_shard : node_shard) { + uint32_t node_id = id_shard.first; + uint32_t shard_id = id_shard.second; + if (cur_id < node_id) { + // Gopal. random_shuffle() is deprecated. + std::shuffle(final_nhood.begin(), final_nhood.end(), urng); + nnbrs = (uint32_t)(std::min)(final_nhood.size(), (uint64_t)max_degree); + // write into merged ofstream + merged_vamana_writer.write((char *)&nnbrs, sizeof(uint32_t)); + merged_vamana_writer.write((char *)final_nhood.data(), + nnbrs * sizeof(uint32_t)); + merged_index_size += (sizeof(uint32_t) + nnbrs * sizeof(uint32_t)); + if (cur_id % 499999 == 1) { + diskann::cout << "." << std::flush; + } + cur_id = node_id; + nnbrs = 0; + for (auto &p : final_nhood) nhood_set[p] = 0; + final_nhood.clear(); } + // read from shard_id ifstream + vamana_readers[shard_id].read((char *)&shard_nnbrs, sizeof(uint32_t)); - diskann::cout << "Max input width: " << max_input_width << ", output width: " << output_width << std::endl; - - merged_vamana_writer.write((char *)&output_width, sizeof(uint32_t)); - std::ofstream medoid_writer(medoids_file.c_str(), std::ios::binary); - uint32_t nshards_u32 = (uint32_t)nshards; - uint32_t one_val = 1; - medoid_writer.write((char *)&nshards_u32, sizeof(uint32_t)); - medoid_writer.write((char *)&one_val, sizeof(uint32_t)); - - uint64_t vamana_index_frozen = 0; // as of now the functionality to merge many overlapping vamana - // indices is supported only for bulk indices without frozen point. - // Hence the final index will also not have any frozen points. - for (uint64_t shard = 0; shard < nshards; shard++) - { - uint32_t medoid; - // read medoid - vamana_readers[shard].read((char *)&medoid, sizeof(uint32_t)); - vamana_readers[shard].read((char *)&vamana_index_frozen, sizeof(uint64_t)); - assert(vamana_index_frozen == false); - // rename medoid - medoid = idmaps[shard][medoid]; - - medoid_writer.write((char *)&medoid, sizeof(uint32_t)); - // write renamed medoid - if (shard == (nshards - 1)) //--> uncomment if running hierarchical - merged_vamana_writer.write((char *)&medoid, sizeof(uint32_t)); + if (shard_nnbrs == 0) { + diskann::cout << "WARNING: shard #" << shard_id << ", node_id " << node_id + << " has 0 nbrs" << std::endl; } - merged_vamana_writer.write((char *)&merged_index_frozen, sizeof(uint64_t)); - medoid_writer.close(); - - diskann::cout << "Starting merge" << std::endl; - - // Gopal. random_shuffle() is deprecated. - std::random_device rng; - std::mt19937 urng(rng()); - - std::vector nhood_set(nnodes, 0); - std::vector final_nhood; - - uint32_t nnbrs = 0, shard_nnbrs = 0; - uint32_t cur_id = 0; - for (const auto &id_shard : node_shard) - { - uint32_t node_id = id_shard.first; - uint32_t shard_id = id_shard.second; - if (cur_id < node_id) - { - // Gopal. random_shuffle() is deprecated. - std::shuffle(final_nhood.begin(), final_nhood.end(), urng); - nnbrs = (uint32_t)(std::min)(final_nhood.size(), (uint64_t)max_degree); - // write into merged ofstream - merged_vamana_writer.write((char *)&nnbrs, sizeof(uint32_t)); - merged_vamana_writer.write((char *)final_nhood.data(), nnbrs * sizeof(uint32_t)); - merged_index_size += (sizeof(uint32_t) + nnbrs * sizeof(uint32_t)); - if (cur_id % 499999 == 1) - { - diskann::cout << "." << std::flush; - } - cur_id = node_id; - nnbrs = 0; - for (auto &p : final_nhood) - nhood_set[p] = 0; - final_nhood.clear(); - } - // read from shard_id ifstream - vamana_readers[shard_id].read((char *)&shard_nnbrs, sizeof(uint32_t)); - - if (shard_nnbrs == 0) - { - diskann::cout << "WARNING: shard #" << shard_id << ", node_id " << node_id << " has 0 nbrs" << std::endl; - } - std::vector shard_nhood(shard_nnbrs); - if (shard_nnbrs > 0) - vamana_readers[shard_id].read((char *)shard_nhood.data(), shard_nnbrs * sizeof(uint32_t)); - // rename nodes - for (uint64_t j = 0; j < shard_nnbrs; j++) - { - if (nhood_set[idmaps[shard_id][shard_nhood[j]]] == 0) - { - nhood_set[idmaps[shard_id][shard_nhood[j]]] = 1; - final_nhood.emplace_back(idmaps[shard_id][shard_nhood[j]]); - } - } + std::vector shard_nhood(shard_nnbrs); + if (shard_nnbrs > 0) + vamana_readers[shard_id].read((char *)shard_nhood.data(), + shard_nnbrs * sizeof(uint32_t)); + // rename nodes + for (uint64_t j = 0; j < shard_nnbrs; j++) { + if (nhood_set[idmaps[shard_id][shard_nhood[j]]] == 0) { + nhood_set[idmaps[shard_id][shard_nhood[j]]] = 1; + final_nhood.emplace_back(idmaps[shard_id][shard_nhood[j]]); + } } - - // Gopal. random_shuffle() is deprecated. - std::shuffle(final_nhood.begin(), final_nhood.end(), urng); - nnbrs = (uint32_t)(std::min)(final_nhood.size(), (uint64_t)max_degree); - // write into merged ofstream - merged_vamana_writer.write((char *)&nnbrs, sizeof(uint32_t)); - if (nnbrs > 0) - { - merged_vamana_writer.write((char *)final_nhood.data(), nnbrs * sizeof(uint32_t)); - } - merged_index_size += (sizeof(uint32_t) + nnbrs * sizeof(uint32_t)); - for (auto &p : final_nhood) - nhood_set[p] = 0; - final_nhood.clear(); - - diskann::cout << "Expected size: " << merged_index_size << std::endl; - - merged_vamana_writer.reset(); - merged_vamana_writer.write((char *)&merged_index_size, sizeof(uint64_t)); - - diskann::cout << "Finished merge" << std::endl; - return 0; + } + + // Gopal. random_shuffle() is deprecated. + std::shuffle(final_nhood.begin(), final_nhood.end(), urng); + nnbrs = (uint32_t)(std::min)(final_nhood.size(), (uint64_t)max_degree); + // write into merged ofstream + merged_vamana_writer.write((char *)&nnbrs, sizeof(uint32_t)); + if (nnbrs > 0) { + merged_vamana_writer.write((char *)final_nhood.data(), + nnbrs * sizeof(uint32_t)); + } + merged_index_size += (sizeof(uint32_t) + nnbrs * sizeof(uint32_t)); + for (auto &p : final_nhood) nhood_set[p] = 0; + final_nhood.clear(); + + diskann::cout << "Expected size: " << merged_index_size << std::endl; + + merged_vamana_writer.reset(); + merged_vamana_writer.write((char *)&merged_index_size, sizeof(uint64_t)); + + diskann::cout << "Finished merge" << std::endl; + return 0; } // TODO: Make this a streaming implementation to avoid exceeding the memory @@ -488,277 +481,281 @@ int merge_shards(const std::string &vamana_prefix, const std::string &vamana_suf the new nodes at the end. The dummy map contains the real graph id of the new nodes added to the graph */ template -void breakup_dense_points(const std::string data_file, const std::string labels_file, uint32_t density, - const std::string out_data_file, const std::string out_labels_file, - const std::string out_metadata_file) -{ - std::string token, line; - std::ifstream labels_stream(labels_file); - T *data; - uint64_t npts, ndims; - diskann::load_bin(data_file, data, npts, ndims); - - std::unordered_map dummy_pt_ids; - uint32_t next_dummy_id = (uint32_t)npts; - - uint32_t point_cnt = 0; - - std::vector> labels_per_point; - labels_per_point.resize(npts); - - uint32_t dense_pts = 0; - if (labels_stream.is_open()) - { - while (getline(labels_stream, line)) - { - std::stringstream iss(line); - uint32_t lbl_cnt = 0; - uint32_t label_host = point_cnt; - while (getline(iss, token, ',')) - { - if (lbl_cnt == density) - { - if (label_host == point_cnt) - dense_pts++; - label_host = next_dummy_id; - labels_per_point.resize(next_dummy_id + 1); - dummy_pt_ids[next_dummy_id] = (uint32_t)point_cnt; - next_dummy_id++; - lbl_cnt = 0; - } - token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); - token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); - uint32_t token_as_num = std::stoul(token); - labels_per_point[label_host].push_back(token_as_num); - lbl_cnt++; - } - point_cnt++; +void breakup_dense_points(const std::string data_file, + const std::string labels_file, uint32_t density, + const std::string out_data_file, + const std::string out_labels_file, + const std::string out_metadata_file) { + std::string token, line; + std::ifstream labels_stream(labels_file); + T *data; + uint64_t npts, ndims; + diskann::load_bin(data_file, data, npts, ndims); + + std::unordered_map dummy_pt_ids; + uint32_t next_dummy_id = (uint32_t)npts; + + uint32_t point_cnt = 0; + + std::vector> labels_per_point; + labels_per_point.resize(npts); + + uint32_t dense_pts = 0; + if (labels_stream.is_open()) { + while (getline(labels_stream, line)) { + std::stringstream iss(line); + uint32_t lbl_cnt = 0; + uint32_t label_host = point_cnt; + while (getline(iss, token, ',')) { + if (lbl_cnt == density) { + if (label_host == point_cnt) dense_pts++; + label_host = next_dummy_id; + labels_per_point.resize(next_dummy_id + 1); + dummy_pt_ids[next_dummy_id] = (uint32_t)point_cnt; + next_dummy_id++; + lbl_cnt = 0; } + token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); + token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); + uint32_t token_as_num = std::stoul(token); + labels_per_point[label_host].push_back(token_as_num); + lbl_cnt++; + } + point_cnt++; } - diskann::cout << "fraction of dense points with >= " << density << " labels = " << (float)dense_pts / (float)npts - << std::endl; + } + diskann::cout << "fraction of dense points with >= " << density + << " labels = " << (float)dense_pts / (float)npts << std::endl; - if (labels_per_point.size() != 0) - { - diskann::cout << labels_per_point.size() << " is the new number of points" << std::endl; - std::ofstream label_writer(out_labels_file); - assert(label_writer.is_open()); - for (uint32_t i = 0; i < labels_per_point.size(); i++) - { - for (uint32_t j = 0; j < (labels_per_point[i].size() - 1); j++) - { - label_writer << labels_per_point[i][j] << ","; - } - if (labels_per_point[i].size() != 0) - label_writer << labels_per_point[i][labels_per_point[i].size() - 1]; - label_writer << std::endl; - } - label_writer.close(); + if (labels_per_point.size() != 0) { + diskann::cout << labels_per_point.size() << " is the new number of points" + << std::endl; + std::ofstream label_writer(out_labels_file); + assert(label_writer.is_open()); + for (uint32_t i = 0; i < labels_per_point.size(); i++) { + for (uint32_t j = 0; j < (labels_per_point[i].size() - 1); j++) { + label_writer << labels_per_point[i][j] << ","; + } + if (labels_per_point[i].size() != 0) + label_writer << labels_per_point[i][labels_per_point[i].size() - 1]; + label_writer << std::endl; } - - if (dummy_pt_ids.size() != 0) - { - diskann::cout << dummy_pt_ids.size() << " is the number of dummy points created" << std::endl; - data = (T *)std::realloc((void *)data, labels_per_point.size() * ndims * sizeof(T)); - std::ofstream dummy_writer(out_metadata_file); - assert(dummy_writer.is_open()); - for (auto i = dummy_pt_ids.begin(); i != dummy_pt_ids.end(); i++) - { - dummy_writer << i->first << "," << i->second << std::endl; - std::memcpy(data + i->first * ndims, data + i->second * ndims, ndims * sizeof(T)); - } - dummy_writer.close(); + label_writer.close(); + } + + if (dummy_pt_ids.size() != 0) { + diskann::cout << dummy_pt_ids.size() + << " is the number of dummy points created" << std::endl; + data = (T *)std::realloc((void *)data, + labels_per_point.size() * ndims * sizeof(T)); + std::ofstream dummy_writer(out_metadata_file); + assert(dummy_writer.is_open()); + for (auto i = dummy_pt_ids.begin(); i != dummy_pt_ids.end(); i++) { + dummy_writer << i->first << "," << i->second << std::endl; + std::memcpy(data + i->first * ndims, data + i->second * ndims, + ndims * sizeof(T)); } + dummy_writer.close(); + } - diskann::save_bin(out_data_file, data, labels_per_point.size(), ndims); + diskann::save_bin(out_data_file, data, labels_per_point.size(), ndims); } -void extract_shard_labels(const std::string &in_label_file, const std::string &shard_ids_bin, - const std::string &shard_label_file) -{ // assumes ith row is for ith - // point in labels file - diskann::cout << "Extracting labels for shard" << std::endl; - - uint32_t *ids = nullptr; - uint64_t num_ids, tmp_dim; - diskann::load_bin(shard_ids_bin, ids, num_ids, tmp_dim); - - uint32_t counter = 0, shard_counter = 0; - std::string cur_line; - - std::ifstream label_reader(in_label_file); - std::ofstream label_writer(shard_label_file); - assert(label_reader.is_open()); - assert(label_reader.is_open()); - if (label_reader && label_writer) - { - while (std::getline(label_reader, cur_line)) - { - if (shard_counter >= num_ids) - { - break; - } - if (counter == ids[shard_counter]) - { - label_writer << cur_line << "\n"; - shard_counter++; - } - counter++; - } +void extract_shard_labels( + const std::string &in_label_file, const std::string &shard_ids_bin, + const std::string &shard_label_file) { // assumes ith row is for ith + // point in labels file + diskann::cout << "Extracting labels for shard" << std::endl; + + uint32_t *ids = nullptr; + uint64_t num_ids, tmp_dim; + diskann::load_bin(shard_ids_bin, ids, num_ids, tmp_dim); + + uint32_t counter = 0, shard_counter = 0; + std::string cur_line; + + std::ifstream label_reader(in_label_file); + std::ofstream label_writer(shard_label_file); + assert(label_reader.is_open()); + assert(label_reader.is_open()); + if (label_reader && label_writer) { + while (std::getline(label_reader, cur_line)) { + if (shard_counter >= num_ids) { + break; + } + if (counter == ids[shard_counter]) { + label_writer << cur_line << "\n"; + shard_counter++; + } + counter++; } - if (ids != nullptr) - delete[] ids; + } + if (ids != nullptr) delete[] ids; } template -int build_merged_vamana_index(std::string base_file, diskann::Metric compareMetric, uint32_t L, uint32_t R, - double sampling_rate, double ram_budget, std::string mem_index_path, - std::string medoids_file, std::string centroids_file, size_t build_pq_bytes, bool use_opq, - uint32_t num_threads, bool use_filters, const std::string &label_file, - const std::string &labels_to_medoids_file, const std::string &universal_label, - const uint32_t Lf) -{ - size_t base_num, base_dim; - diskann::get_bin_metadata(base_file, base_num, base_dim); - - double full_index_ram = estimate_ram_usage(base_num, base_dim, sizeof(T), R); - - // TODO: Make this honest when there is filter support - if (full_index_ram < ram_budget * 1024 * 1024 * 1024) - { - diskann::cout << "Full index fits in RAM budget, should consume at most " - << full_index_ram / (1024 * 1024 * 1024) << "GiBs, so building in one shot" << std::endl; - - diskann::IndexWriteParameters paras = diskann::IndexWriteParametersBuilder(L, R) - .with_filter_list_size(Lf) - .with_saturate_graph(!use_filters) - .with_num_threads(num_threads) - .build(); - using TagT = uint32_t; - diskann::Index _index(compareMetric, base_dim, base_num, false, false, false, - build_pq_bytes > 0, build_pq_bytes, use_opq); - if (!use_filters) - _index.build(base_file.c_str(), base_num, paras); - else - { - if (universal_label != "") - { // indicates no universal label - LabelT unv_label_as_num = 0; - _index.set_universal_label(unv_label_as_num); - } - _index.build_filtered_index(base_file.c_str(), label_file, base_num, paras); - } - _index.save(mem_index_path.c_str()); - - if (use_filters) - { - // need to copy the labels_to_medoids file to the specified input - // file - std::remove(labels_to_medoids_file.c_str()); - std::string mem_labels_to_medoid_file = mem_index_path + "_labels_to_medoids.txt"; - copy_file(mem_labels_to_medoid_file, labels_to_medoids_file); - std::remove(mem_labels_to_medoid_file.c_str()); - } - - std::remove(medoids_file.c_str()); - std::remove(centroids_file.c_str()); - return 0; +int build_merged_vamana_index( + std::string base_file, diskann::Metric compareMetric, uint32_t L, + uint32_t R, double sampling_rate, double ram_budget, + std::string mem_index_path, std::string medoids_file, + std::string centroids_file, size_t build_pq_bytes, bool use_opq, + uint32_t num_threads, bool use_filters, const std::string &label_file, + const std::string &labels_to_medoids_file, + const std::string &universal_label, const uint32_t Lf) { + size_t base_num, base_dim; + diskann::get_bin_metadata(base_file, base_num, base_dim); + + double full_index_ram = estimate_ram_usage(base_num, base_dim, sizeof(T), R); + + // TODO: Make this honest when there is filter support + if (full_index_ram < ram_budget * 1024 * 1024 * 1024) { + diskann::cout << "Full index fits in RAM budget, should consume at most " + << full_index_ram / (1024 * 1024 * 1024) + << "GiBs, so building in one shot" << std::endl; + + diskann::IndexWriteParameters paras = + diskann::IndexWriteParametersBuilder(L, R) + .with_filter_list_size(Lf) + .with_saturate_graph(!use_filters) + .with_num_threads(num_threads) + .build(); + using TagT = uint32_t; + diskann::Index _index( + compareMetric, base_dim, base_num, false, false, false, + build_pq_bytes > 0, build_pq_bytes, use_opq); + if (!use_filters) + _index.build(base_file.c_str(), base_num, paras); + else { + if (universal_label != "") { // indicates no universal label + LabelT unv_label_as_num = 0; + _index.set_universal_label(unv_label_as_num); + } + _index.build_filtered_index(base_file.c_str(), label_file, base_num, + paras); + } + _index.save(mem_index_path.c_str()); + + if (use_filters) { + // need to copy the labels_to_medoids file to the specified input + // file + std::remove(labels_to_medoids_file.c_str()); + std::string mem_labels_to_medoid_file = + mem_index_path + "_labels_to_medoids.txt"; + copy_file(mem_labels_to_medoid_file, labels_to_medoids_file); + std::remove(mem_labels_to_medoid_file.c_str()); } - // where the universal label is to be saved in the final graph - std::string final_index_universal_label_file = mem_index_path + "_universal_label.txt"; - - std::string merged_index_prefix = mem_index_path + "_tempFiles"; - - Timer timer; - int num_parts = - partition_with_ram_budget(base_file, sampling_rate, ram_budget, 2 * R / 3, merged_index_prefix, 2); - diskann::cout << timer.elapsed_seconds_for_step("partitioning data") << std::endl; - - std::string cur_centroid_filepath = merged_index_prefix + "_centroids.bin"; - std::rename(cur_centroid_filepath.c_str(), centroids_file.c_str()); - - timer.reset(); - for (int p = 0; p < num_parts; p++) - { - std::string shard_base_file = merged_index_prefix + "_subshard-" + std::to_string(p) + ".bin"; - - std::string shard_ids_file = merged_index_prefix + "_subshard-" + std::to_string(p) + "_ids_uint32.bin"; - - std::string shard_labels_file = merged_index_prefix + "_subshard-" + std::to_string(p) + "_labels.txt"; - - retrieve_shard_data_from_ids(base_file, shard_ids_file, shard_base_file); - - std::string shard_index_file = merged_index_prefix + "_subshard-" + std::to_string(p) + "_mem.index"; - - diskann::IndexWriteParameters paras = - diskann::IndexWriteParametersBuilder(L, (2 * R / 3)).with_filter_list_size(Lf).build(); - - uint64_t shard_base_dim, shard_base_pts; - get_bin_metadata(shard_base_file, shard_base_pts, shard_base_dim); - diskann::Index _index(compareMetric, shard_base_dim, shard_base_pts, false, false, false, build_pq_bytes > 0, - build_pq_bytes, use_opq); - if (!use_filters) - { - _index.build(shard_base_file.c_str(), shard_base_pts, paras); - } - else - { - diskann::extract_shard_labels(label_file, shard_ids_file, shard_labels_file); - if (universal_label != "") - { // indicates no universal label - LabelT unv_label_as_num = 0; - _index.set_universal_label(unv_label_as_num); - } - _index.build_filtered_index(shard_base_file.c_str(), shard_labels_file, shard_base_pts, paras); - } - _index.save(shard_index_file.c_str()); - // copy universal label file from first shard to the final destination - // index, since all shards anyway share the universal label - if (p == 0) - { - std::string shard_universal_label_file = shard_index_file + "_universal_label.txt"; - if (universal_label != "") - { - copy_file(shard_universal_label_file, final_index_universal_label_file); - } - } - - std::remove(shard_base_file.c_str()); + std::remove(medoids_file.c_str()); + std::remove(centroids_file.c_str()); + return 0; + } + + // where the universal label is to be saved in the final graph + std::string final_index_universal_label_file = + mem_index_path + "_universal_label.txt"; + + std::string merged_index_prefix = mem_index_path + "_tempFiles"; + + Timer timer; + int num_parts = partition_with_ram_budget( + base_file, sampling_rate, ram_budget, 2 * R / 3, merged_index_prefix, 2); + diskann::cout << timer.elapsed_seconds_for_step("partitioning data") + << std::endl; + + std::string cur_centroid_filepath = merged_index_prefix + "_centroids.bin"; + std::rename(cur_centroid_filepath.c_str(), centroids_file.c_str()); + + timer.reset(); + for (int p = 0; p < num_parts; p++) { + std::string shard_base_file = + merged_index_prefix + "_subshard-" + std::to_string(p) + ".bin"; + + std::string shard_ids_file = merged_index_prefix + "_subshard-" + + std::to_string(p) + "_ids_uint32.bin"; + + std::string shard_labels_file = + merged_index_prefix + "_subshard-" + std::to_string(p) + "_labels.txt"; + + retrieve_shard_data_from_ids(base_file, shard_ids_file, shard_base_file); + + std::string shard_index_file = + merged_index_prefix + "_subshard-" + std::to_string(p) + "_mem.index"; + + diskann::IndexWriteParameters paras = + diskann::IndexWriteParametersBuilder(L, (2 * R / 3)) + .with_filter_list_size(Lf) + .build(); + + uint64_t shard_base_dim, shard_base_pts; + get_bin_metadata(shard_base_file, shard_base_pts, shard_base_dim); + diskann::Index _index(compareMetric, shard_base_dim, shard_base_pts, + false, false, false, build_pq_bytes > 0, + build_pq_bytes, use_opq); + if (!use_filters) { + _index.build(shard_base_file.c_str(), shard_base_pts, paras); + } else { + diskann::extract_shard_labels(label_file, shard_ids_file, + shard_labels_file); + if (universal_label != "") { // indicates no universal label + LabelT unv_label_as_num = 0; + _index.set_universal_label(unv_label_as_num); + } + _index.build_filtered_index(shard_base_file.c_str(), shard_labels_file, + shard_base_pts, paras); } - diskann::cout << timer.elapsed_seconds_for_step("building indices on shards") << std::endl; - - timer.reset(); - diskann::merge_shards(merged_index_prefix + "_subshard-", "_mem.index", merged_index_prefix + "_subshard-", - "_ids_uint32.bin", num_parts, R, mem_index_path, medoids_file, use_filters, - labels_to_medoids_file); - diskann::cout << timer.elapsed_seconds_for_step("merging indices") << std::endl; - - // delete tempFiles - for (int p = 0; p < num_parts; p++) - { - std::string shard_base_file = merged_index_prefix + "_subshard-" + std::to_string(p) + ".bin"; - std::string shard_id_file = merged_index_prefix + "_subshard-" + std::to_string(p) + "_ids_uint32.bin"; - std::string shard_labels_file = merged_index_prefix + "_subshard-" + std::to_string(p) + "_labels.txt"; - std::string shard_index_file = merged_index_prefix + "_subshard-" + std::to_string(p) + "_mem.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()); - if (use_filters) - { - std::string shard_index_label_file = shard_index_file + "_labels.txt"; - std::string shard_index_univ_label_file = shard_index_file + "_universal_label.txt"; - std::string shard_index_label_map_file = shard_index_file + "_labels_to_medoids.txt"; - std::remove(shard_labels_file.c_str()); - std::remove(shard_index_label_file.c_str()); - std::remove(shard_index_label_map_file.c_str()); - std::remove(shard_index_univ_label_file.c_str()); - } + _index.save(shard_index_file.c_str()); + // copy universal label file from first shard to the final destination + // index, since all shards anyway share the universal label + if (p == 0) { + std::string shard_universal_label_file = + shard_index_file + "_universal_label.txt"; + if (universal_label != "") { + copy_file(shard_universal_label_file, final_index_universal_label_file); + } } - return 0; + + std::remove(shard_base_file.c_str()); + } + diskann::cout << timer.elapsed_seconds_for_step("building indices on shards") + << std::endl; + + timer.reset(); + diskann::merge_shards(merged_index_prefix + "_subshard-", "_mem.index", + merged_index_prefix + "_subshard-", "_ids_uint32.bin", + num_parts, R, mem_index_path, medoids_file, use_filters, + labels_to_medoids_file); + diskann::cout << timer.elapsed_seconds_for_step("merging indices") + << std::endl; + + // delete tempFiles + for (int p = 0; p < num_parts; p++) { + std::string shard_base_file = + merged_index_prefix + "_subshard-" + std::to_string(p) + ".bin"; + std::string shard_id_file = merged_index_prefix + "_subshard-" + + std::to_string(p) + "_ids_uint32.bin"; + std::string shard_labels_file = + merged_index_prefix + "_subshard-" + std::to_string(p) + "_labels.txt"; + std::string shard_index_file = + merged_index_prefix + "_subshard-" + std::to_string(p) + "_mem.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()); + if (use_filters) { + std::string shard_index_label_file = shard_index_file + "_labels.txt"; + std::string shard_index_univ_label_file = + shard_index_file + "_universal_label.txt"; + std::string shard_index_label_map_file = + shard_index_file + "_labels_to_medoids.txt"; + std::remove(shard_labels_file.c_str()); + std::remove(shard_index_label_file.c_str()); + std::remove(shard_index_label_map_file.c_str()); + std::remove(shard_index_univ_label_file.c_str()); + } + } + return 0; } // General purpose support for DiskANN interface @@ -766,644 +763,690 @@ int build_merged_vamana_index(std::string base_file, diskann::Metric compareMetr // 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, - uint64_t tuning_sample_num, uint64_t 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 (int64_t 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, false, stats + i); - } - 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; +uint32_t optimize_beamwidth( + std::unique_ptr> &pFlashIndex, + T *tuning_sample, uint64_t tuning_sample_num, + uint64_t 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]; - delete[] stats; + auto s = std::chrono::high_resolution_clock::now(); +#pragma omp parallel for schedule(dynamic, 1) num_threads(nthreads) + for (int64_t 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, false, + stats + i); + } + 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; } - return best_bw; + 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, - const std::string reorder_data_file) -{ - uint32_t npts, ndims; - - // amount to read or write in one shot - size_t read_blk_size = 64 * 1024 * 1024; - size_t write_blk_size = read_blk_size; - cached_ifstream base_reader(base_file, read_blk_size); - base_reader.read((char *)&npts, sizeof(uint32_t)); - base_reader.read((char *)&ndims, sizeof(uint32_t)); - - size_t npts_64, ndims_64; - npts_64 = npts; - ndims_64 = ndims; - - // Check if we need to append data for re-ordering - bool append_reorder_data = false; - std::ifstream reorder_data_reader; - - uint32_t npts_reorder_file = 0, ndims_reorder_file = 0; - if (reorder_data_file != std::string("")) - { - append_reorder_data = true; - size_t reorder_data_file_size = get_file_size(reorder_data_file); - reorder_data_reader.exceptions(std::ofstream::failbit | std::ofstream::badbit); - - try - { - reorder_data_reader.open(reorder_data_file, std::ios::binary); - reorder_data_reader.read((char *)&npts_reorder_file, sizeof(uint32_t)); - reorder_data_reader.read((char *)&ndims_reorder_file, sizeof(uint32_t)); - if (npts_reorder_file != npts) - throw ANNException("Mismatch in num_points between reorder " - "data file and base file", - -1, __FUNCSIG__, __FILE__, __LINE__); - if (reorder_data_file_size != 8 + sizeof(float) * (size_t)npts_reorder_file * (size_t)ndims_reorder_file) - throw ANNException("Discrepancy in reorder data file size ", -1, __FUNCSIG__, __FILE__, __LINE__); - } - catch (std::system_error &e) - { - throw FileException(reorder_data_file, e, __FUNCSIG__, __FILE__, __LINE__); - } +void create_disk_layout(const std::string base_file, + const std::string mem_index_file, + const std::string output_file, + const std::string reorder_data_file) { + uint32_t npts, ndims; + + // amount to read or write in one shot + size_t read_blk_size = 64 * 1024 * 1024; + size_t write_blk_size = read_blk_size; + cached_ifstream base_reader(base_file, read_blk_size); + base_reader.read((char *)&npts, sizeof(uint32_t)); + base_reader.read((char *)&ndims, sizeof(uint32_t)); + + size_t npts_64, ndims_64; + npts_64 = npts; + ndims_64 = ndims; + + // Check if we need to append data for re-ordering + bool append_reorder_data = false; + std::ifstream reorder_data_reader; + + uint32_t npts_reorder_file = 0, ndims_reorder_file = 0; + if (reorder_data_file != std::string("")) { + append_reorder_data = true; + size_t reorder_data_file_size = get_file_size(reorder_data_file); + reorder_data_reader.exceptions(std::ofstream::failbit | + std::ofstream::badbit); + + try { + reorder_data_reader.open(reorder_data_file, std::ios::binary); + reorder_data_reader.read((char *)&npts_reorder_file, sizeof(uint32_t)); + reorder_data_reader.read((char *)&ndims_reorder_file, sizeof(uint32_t)); + if (npts_reorder_file != npts) + throw ANNException( + "Mismatch in num_points between reorder " + "data file and base file", + -1, __FUNCSIG__, __FILE__, __LINE__); + if (reorder_data_file_size != 8 + sizeof(float) * + (size_t)npts_reorder_file * + (size_t)ndims_reorder_file) + throw ANNException("Discrepancy in reorder data file size ", -1, + __FUNCSIG__, __FILE__, __LINE__); + } catch (std::system_error &e) { + throw FileException(reorder_data_file, e, __FUNCSIG__, __FILE__, + __LINE__); } - - // create cached reader + writer - size_t actual_file_size = get_file_size(mem_index_file); - diskann::cout << "Vamana index file size=" << actual_file_size << std::endl; - std::ifstream vamana_reader(mem_index_file, std::ios::binary); - cached_ofstream diskann_writer(output_file, write_blk_size); - - // metadata: width, medoid - uint32_t 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__); - } - uint64_t vamana_frozen_num = false, vamana_frozen_loc = 0; - - vamana_reader.read((char *)&width_u32, sizeof(uint32_t)); - vamana_reader.read((char *)&medoid_u32, sizeof(uint32_t)); - vamana_reader.read((char *)&vamana_frozen_num, sizeof(uint64_t)); - // compute - uint64_t medoid, max_node_len, nnodes_per_sector; - npts_64 = (uint64_t)npts; - medoid = (uint64_t)medoid_u32; - if (vamana_frozen_num == 1) - vamana_frozen_loc = medoid; - max_node_len = (((uint64_t)width_u32 + 1) * sizeof(uint32_t)) + (ndims_64 * sizeof(T)); - nnodes_per_sector = SECTOR_LEN / max_node_len; - - diskann::cout << "medoid: " << medoid << "B" << std::endl; - diskann::cout << "max_node_len: " << max_node_len << "B" << std::endl; - diskann::cout << "nnodes_per_sector: " << nnodes_per_sector << "B" << 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); - uint32_t &nnbrs = *(uint32_t *)(node_buf.get() + ndims_64 * sizeof(T)); - uint32_t *nhood_buf = (uint32_t *)(node_buf.get() + (ndims_64 * sizeof(T)) + sizeof(uint32_t)); - - // number of sectors (1 for meta data) - uint64_t n_sectors = ROUND_UP(npts_64, nnodes_per_sector) / nnodes_per_sector; - uint64_t n_reorder_sectors = 0; - uint64_t n_data_nodes_per_sector = 0; - - if (append_reorder_data) - { - n_data_nodes_per_sector = SECTOR_LEN / (ndims_reorder_file * sizeof(float)); - n_reorder_sectors = ROUND_UP(npts_64, n_data_nodes_per_sector) / n_data_nodes_per_sector; + } + + // create cached reader + writer + size_t actual_file_size = get_file_size(mem_index_file); + diskann::cout << "Vamana index file size=" << actual_file_size << std::endl; + std::ifstream vamana_reader(mem_index_file, std::ios::binary); + cached_ofstream diskann_writer(output_file, write_blk_size); + + // metadata: width, medoid + uint32_t 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__); + } + uint64_t vamana_frozen_num = false, vamana_frozen_loc = 0; + + vamana_reader.read((char *)&width_u32, sizeof(uint32_t)); + vamana_reader.read((char *)&medoid_u32, sizeof(uint32_t)); + vamana_reader.read((char *)&vamana_frozen_num, sizeof(uint64_t)); + // compute + uint64_t medoid, max_node_len, nnodes_per_sector; + npts_64 = (uint64_t)npts; + medoid = (uint64_t)medoid_u32; + if (vamana_frozen_num == 1) vamana_frozen_loc = medoid; + max_node_len = + (((uint64_t)width_u32 + 1) * sizeof(uint32_t)) + (ndims_64 * sizeof(T)); + nnodes_per_sector = SECTOR_LEN / max_node_len; + + diskann::cout << "medoid: " << medoid << "B" << std::endl; + diskann::cout << "max_node_len: " << max_node_len << "B" << std::endl; + diskann::cout << "nnodes_per_sector: " << nnodes_per_sector << "B" + << 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); + uint32_t &nnbrs = *(uint32_t *)(node_buf.get() + ndims_64 * sizeof(T)); + uint32_t *nhood_buf = + (uint32_t *)(node_buf.get() + (ndims_64 * sizeof(T)) + sizeof(uint32_t)); + + // number of sectors (1 for meta data) + uint64_t n_sectors = ROUND_UP(npts_64, nnodes_per_sector) / nnodes_per_sector; + uint64_t n_reorder_sectors = 0; + uint64_t n_data_nodes_per_sector = 0; + + if (append_reorder_data) { + n_data_nodes_per_sector = SECTOR_LEN / (ndims_reorder_file * sizeof(float)); + n_reorder_sectors = + ROUND_UP(npts_64, n_data_nodes_per_sector) / n_data_nodes_per_sector; + } + uint64_t disk_index_file_size = + (n_sectors + n_reorder_sectors + 1) * SECTOR_LEN; + + std::vector 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((uint64_t)append_reorder_data); + if (append_reorder_data) { + output_file_meta.push_back(n_sectors + 1); + output_file_meta.push_back(ndims_reorder_file); + output_file_meta.push_back(n_data_nodes_per_sector); + } + output_file_meta.push_back(disk_index_file_size); + + diskann_writer.write(sector_buf.get(), SECTOR_LEN); + + std::unique_ptr cur_node_coords = std::make_unique(ndims_64); + diskann::cout << "# sectors: " << n_sectors << std::endl; + uint64_t cur_node_id = 0; + for (uint64_t sector = 0; sector < n_sectors; sector++) { + if (sector % 100000 == 0) { + diskann::cout << "Sector #" << sector << "written" << std::endl; } - uint64_t disk_index_file_size = (n_sectors + n_reorder_sectors + 1) * SECTOR_LEN; - - std::vector 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((uint64_t)append_reorder_data); - if (append_reorder_data) - { - output_file_meta.push_back(n_sectors + 1); - output_file_meta.push_back(ndims_reorder_file); - output_file_meta.push_back(n_data_nodes_per_sector); + memset(sector_buf.get(), 0, SECTOR_LEN); + for (uint64_t sector_node_id = 0; + sector_node_id < nnodes_per_sector && cur_node_id < npts_64; + sector_node_id++) { + memset(node_buf.get(), 0, max_node_len); + // read cur node's nnbrs + vamana_reader.read((char *)&nnbrs, sizeof(uint32_t)); + + // sanity checks on nnbrs + assert(nnbrs > 0); + assert(nnbrs <= width_u32); + + // read node's nhood + vamana_reader.read((char *)nhood_buf, + (std::min)(nnbrs, width_u32) * sizeof(uint32_t)); + if (nnbrs > width_u32) { + vamana_reader.seekg((nnbrs - width_u32) * sizeof(uint32_t), + vamana_reader.cur); + } + + // write coords of node first + // T *node_coords = data + ((uint64_t) ndims_64 * cur_node_id); + 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 + *(uint32_t *)(node_buf.get() + ndims_64 * sizeof(T)) = + (std::min)(nnbrs, width_u32); + + // write nhood next + memcpy(node_buf.get() + ndims_64 * sizeof(T) + sizeof(uint32_t), + nhood_buf, (std::min)(nnbrs, width_u32) * sizeof(uint32_t)); + + // get offset into sector_buf + char *sector_node_buf = + sector_buf.get() + (sector_node_id * max_node_len); + + // copy node buf into sector_node_buf + memcpy(sector_node_buf, node_buf.get(), max_node_len); + cur_node_id++; } - output_file_meta.push_back(disk_index_file_size); - + // flush sector to disk diskann_writer.write(sector_buf.get(), SECTOR_LEN); + } + if (append_reorder_data) { + diskann::cout << "Index written. Appending reorder data..." << std::endl; - std::unique_ptr cur_node_coords = std::make_unique(ndims_64); - diskann::cout << "# sectors: " << n_sectors << std::endl; - uint64_t cur_node_id = 0; - for (uint64_t sector = 0; sector < n_sectors; sector++) - { - if (sector % 100000 == 0) - { - diskann::cout << "Sector #" << sector << "written" << std::endl; - } - memset(sector_buf.get(), 0, SECTOR_LEN); - for (uint64_t sector_node_id = 0; sector_node_id < nnodes_per_sector && cur_node_id < npts_64; sector_node_id++) - { - memset(node_buf.get(), 0, max_node_len); - // read cur node's nnbrs - vamana_reader.read((char *)&nnbrs, sizeof(uint32_t)); - - // sanity checks on nnbrs - assert(nnbrs > 0); - assert(nnbrs <= width_u32); - - // read node's nhood - vamana_reader.read((char *)nhood_buf, (std::min)(nnbrs, width_u32) * sizeof(uint32_t)); - if (nnbrs > width_u32) - { - vamana_reader.seekg((nnbrs - width_u32) * sizeof(uint32_t), vamana_reader.cur); - } - - // write coords of node first - // T *node_coords = data + ((uint64_t) ndims_64 * cur_node_id); - 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 - *(uint32_t *)(node_buf.get() + ndims_64 * sizeof(T)) = (std::min)(nnbrs, width_u32); - - // write nhood next - memcpy(node_buf.get() + ndims_64 * sizeof(T) + sizeof(uint32_t), nhood_buf, - (std::min)(nnbrs, width_u32) * sizeof(uint32_t)); - - // get offset into sector_buf - char *sector_node_buf = sector_buf.get() + (sector_node_id * max_node_len); - - // copy node buf into sector_node_buf - memcpy(sector_node_buf, node_buf.get(), max_node_len); - cur_node_id++; - } - // flush sector to disk - diskann_writer.write(sector_buf.get(), SECTOR_LEN); - } - if (append_reorder_data) - { - diskann::cout << "Index written. Appending reorder data..." << std::endl; - - auto vec_len = ndims_reorder_file * sizeof(float); - std::unique_ptr vec_buf = std::make_unique(vec_len); - - for (uint64_t sector = 0; sector < n_reorder_sectors; sector++) - { - if (sector % 100000 == 0) - { - diskann::cout << "Reorder data Sector #" << sector << "written" << std::endl; - } - - memset(sector_buf.get(), 0, SECTOR_LEN); - - for (uint64_t sector_node_id = 0; sector_node_id < n_data_nodes_per_sector && sector_node_id < npts_64; - sector_node_id++) - { - memset(vec_buf.get(), 0, vec_len); - reorder_data_reader.read(vec_buf.get(), vec_len); - - // copy node buf into sector_node_buf - memcpy(sector_buf.get() + (sector_node_id * vec_len), vec_buf.get(), vec_len); - } - // flush sector to disk - diskann_writer.write(sector_buf.get(), SECTOR_LEN); - } - } - diskann_writer.close(); - diskann::save_bin(output_file, output_file_meta.data(), output_file_meta.size(), 1, 0); - diskann::cout << "Output disk index file written to " << output_file << std::endl; -} + auto vec_len = ndims_reorder_file * sizeof(float); + std::unique_ptr vec_buf = std::make_unique(vec_len); -template -int build_disk_index(const char *dataFilePath, const char *indexFilePath, const char *indexBuildParameters, - diskann::Metric compareMetric, bool use_opq, const std::string &codebook_prefix, bool use_filters, - const std::string &label_file, const std::string &universal_label, const uint32_t filter_threshold, - const uint32_t Lf) -{ - std::stringstream parser; - parser << std::string(indexBuildParameters); - std::string cur_param; - std::vector param_list; - while (parser >> cur_param) - { - param_list.push_back(cur_param); - } - if (param_list.size() < 5 || param_list.size() > 9) - { - diskann::cout << "Correct usage of parameters is R (max degree)\n" - "L (indexing list size, better if >= R)\n" - "B (RAM limit of final index in GB)\n" - "M (memory limit while indexing)\n" - "T (number of threads for indexing)\n" - "B' (PQ bytes for disk index: optional parameter for " - "very large dimensional data)\n" - "reorder (set true to include full precision in data file" - ": optional paramter, use only when using disk PQ\n" - "build_PQ_byte (number of PQ bytes for inde build; set 0 to use " - "full precision vectors)\n" - "QD Quantized Dimension to overwrite the derived dim from B " + for (uint64_t sector = 0; sector < n_reorder_sectors; sector++) { + if (sector % 100000 == 0) { + diskann::cout << "Reorder data Sector #" << sector << "written" << std::endl; - return -1; - } - - if (!std::is_same::value && compareMetric == diskann::Metric::INNER_PRODUCT) - { - std::stringstream stream; - stream << "DiskANN currently only supports floating point data for Max " - "Inner Product Search. " - << std::endl; - throw diskann::ANNException(stream.str(), -1); - } - - size_t disk_pq_dims = 0; - bool use_disk_pq = false; - size_t build_pq_bytes = 0; - - // if there is a 6th parameter, it means we compress the disk index - // vectors also using PQ data (for very large dimensionality data). If the - // provided parameter is 0, it means we store full vectors. - if (param_list.size() > 5) - { - disk_pq_dims = atoi(param_list[5].c_str()); - use_disk_pq = true; - if (disk_pq_dims == 0) - use_disk_pq = false; - } - - bool reorder_data = false; - if (param_list.size() >= 7) - { - if (1 == atoi(param_list[6].c_str())) - { - reorder_data = true; - } - } - - if (param_list.size() >= 8) - { - build_pq_bytes = atoi(param_list[7].c_str()); + } + + memset(sector_buf.get(), 0, SECTOR_LEN); + + for (uint64_t sector_node_id = 0; + sector_node_id < n_data_nodes_per_sector && sector_node_id < npts_64; + sector_node_id++) { + memset(vec_buf.get(), 0, vec_len); + reorder_data_reader.read(vec_buf.get(), vec_len); + + // copy node buf into sector_node_buf + memcpy(sector_buf.get() + (sector_node_id * vec_len), vec_buf.get(), + vec_len); + } + // flush sector to disk + diskann_writer.write(sector_buf.get(), SECTOR_LEN); } + } + diskann_writer.close(); + diskann::save_bin(output_file, output_file_meta.data(), + output_file_meta.size(), 1, 0); + diskann::cout << "Output disk index file written to " << output_file + << std::endl; +} - std::string base_file(dataFilePath); - std::string data_file_to_use = base_file; - std::string labels_file_original = label_file; - std::string index_prefix_path(indexFilePath); - std::string labels_file_to_use = index_prefix_path + "_label_formatted.txt"; - std::string pq_pivots_path_base = codebook_prefix; - std::string pq_pivots_path = file_exists(pq_pivots_path_base) ? pq_pivots_path_base + "_pq_pivots.bin" - : index_prefix_path + "_pq_pivots.bin"; - std::string pq_compressed_vectors_path = index_prefix_path + "_pq_compressed.bin"; - std::string mem_index_path = index_prefix_path + "_mem.index"; - std::string disk_index_path = index_prefix_path + "_disk.index"; - std::string medoids_path = disk_index_path + "_medoids.bin"; - std::string centroids_path = disk_index_path + "_centroids.bin"; - - std::string labels_to_medoids_path = disk_index_path + "_labels_to_medoids.txt"; - std::string mem_labels_file = mem_index_path + "_labels.txt"; - std::string disk_labels_file = disk_index_path + "_labels.txt"; - std::string mem_univ_label_file = mem_index_path + "_universal_label.txt"; - std::string disk_univ_label_file = disk_index_path + "_universal_label.txt"; - std::string disk_labels_int_map_file = disk_index_path + "_labels_map.txt"; - std::string dummy_remap_file = disk_index_path + "_dummy_remap.txt"; // remap will be used if we break-up points of - // high label-density to create copies - - std::string sample_base_prefix = index_prefix_path + "_sample"; - // optional, used if disk index file must store pq data - std::string disk_pq_pivots_path = index_prefix_path + "_disk.index_pq_pivots.bin"; - // optional, used if disk index must store pq data - std::string disk_pq_compressed_vectors_path = index_prefix_path + "_disk.index_pq_compressed.bin"; - - // output a new base file which contains extra dimension with sqrt(1 - - // ||x||^2/M^2) for every x, M is max norm of all points. Extra space on - // disk needed! - if (compareMetric == diskann::Metric::INNER_PRODUCT) - { - Timer timer; - std::cout << "Using Inner Product search, so need to pre-process base " - "data into temp file. Please ensure there is additional " - "(n*(d+1)*4) bytes for storing pre-processed base vectors, " - "apart from the intermin indices and final index." - << std::endl; - std::string prepped_base = index_prefix_path + "_prepped_base.bin"; - data_file_to_use = prepped_base; - float max_norm_of_base = diskann::prepare_base_for_inner_products(base_file, prepped_base); - std::string norm_file = disk_index_path + "_max_base_norm.bin"; - diskann::save_bin(norm_file, &max_norm_of_base, 1, 1); - diskann::cout << timer.elapsed_seconds_for_step("preprocessing data for inner product") << std::endl; - } - - uint32_t R = (uint32_t)atoi(param_list[0].c_str()); - uint32_t L = (uint32_t)atoi(param_list[1].c_str()); - - 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; - return -1; - } - 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; - return -1; - } - uint32_t num_threads = (uint32_t)atoi(param_list[4].c_str()); - - if (num_threads != 0) - { - omp_set_num_threads(num_threads); - mkl_set_num_threads(num_threads); - } - - 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; - - auto s = std::chrono::high_resolution_clock::now(); - - // If there is filter support, we break-up points which have too many labels - // into replica dummy points which evenly distribute the filters. The rest - // of index build happens on the augmented base and labels - std::string augmented_data_file, augmented_labels_file; - if (use_filters) - { - convert_labels_string_to_int(labels_file_original, labels_file_to_use, disk_labels_int_map_file, - universal_label); - augmented_data_file = index_prefix_path + "_augmented_data.bin"; - augmented_labels_file = index_prefix_path + "_augmented_labels.txt"; - if (filter_threshold != 0) - { - dummy_remap_file = index_prefix_path + "_dummy_remap.txt"; - breakup_dense_points(data_file_to_use, labels_file_to_use, filter_threshold, augmented_data_file, - augmented_labels_file, - dummy_remap_file); // RKNOTE: This has large memory footprint, - // need to make this streaming - data_file_to_use = augmented_data_file; - labels_file_to_use = augmented_labels_file; - } +template +int build_disk_index(const char *dataFilePath, const char *indexFilePath, + const char *indexBuildParameters, + diskann::Metric compareMetric, bool use_opq, + const std::string &codebook_prefix, bool use_filters, + const std::string &label_file, + const std::string &universal_label, + const uint32_t filter_threshold, const uint32_t Lf) { + std::stringstream parser; + parser << std::string(indexBuildParameters); + std::string cur_param; + std::vector param_list; + while (parser >> cur_param) { + param_list.push_back(cur_param); + } + if (param_list.size() < 5 || param_list.size() > 9) { + diskann::cout + << "Correct usage of parameters is R (max degree)\n" + "L (indexing list size, better if >= R)\n" + "B (RAM limit of final index in GB)\n" + "M (memory limit while indexing)\n" + "T (number of threads for indexing)\n" + "B' (PQ bytes for disk index: optional parameter for " + "very large dimensional data)\n" + "reorder (set true to include full precision in data file" + ": optional paramter, use only when using disk PQ\n" + "build_PQ_byte (number of PQ bytes for inde build; set 0 to use " + "full precision vectors)\n" + "QD Quantized Dimension to overwrite the derived dim from B " + << std::endl; + return -1; + } + + if (!std::is_same::value && + compareMetric == diskann::Metric::INNER_PRODUCT) { + std::stringstream stream; + stream << "DiskANN currently only supports floating point data for Max " + "Inner Product Search. " + << std::endl; + throw diskann::ANNException(stream.str(), -1); + } + + size_t disk_pq_dims = 0; + bool use_disk_pq = false; + size_t build_pq_bytes = 0; + + // if there is a 6th parameter, it means we compress the disk index + // vectors also using PQ data (for very large dimensionality data). If the + // provided parameter is 0, it means we store full vectors. + if (param_list.size() > 5) { + disk_pq_dims = atoi(param_list[5].c_str()); + use_disk_pq = true; + if (disk_pq_dims == 0) use_disk_pq = false; + } + + bool reorder_data = false; + if (param_list.size() >= 7) { + if (1 == atoi(param_list[6].c_str())) { + reorder_data = true; } - - size_t points_num, dim; - + } + + if (param_list.size() >= 8) { + build_pq_bytes = atoi(param_list[7].c_str()); + } + + std::string base_file(dataFilePath); + std::string data_file_to_use = base_file; + std::string labels_file_original = label_file; + std::string index_prefix_path(indexFilePath); + std::string labels_file_to_use = index_prefix_path + "_label_formatted.txt"; + std::string pq_pivots_path_base = codebook_prefix; + std::string pq_pivots_path = file_exists(pq_pivots_path_base) + ? pq_pivots_path_base + "_pq_pivots.bin" + : index_prefix_path + "_pq_pivots.bin"; + std::string pq_compressed_vectors_path = + index_prefix_path + "_pq_compressed.bin"; + std::string mem_index_path = index_prefix_path + "_mem.index"; + std::string disk_index_path = index_prefix_path + "_disk.index"; + std::string medoids_path = disk_index_path + "_medoids.bin"; + std::string centroids_path = disk_index_path + "_centroids.bin"; + + std::string labels_to_medoids_path = + disk_index_path + "_labels_to_medoids.txt"; + std::string mem_labels_file = mem_index_path + "_labels.txt"; + std::string disk_labels_file = disk_index_path + "_labels.txt"; + std::string mem_univ_label_file = mem_index_path + "_universal_label.txt"; + std::string disk_univ_label_file = disk_index_path + "_universal_label.txt"; + std::string disk_labels_int_map_file = disk_index_path + "_labels_map.txt"; + std::string dummy_remap_file = + disk_index_path + + "_dummy_remap.txt"; // remap will be used if we break-up points of + // high label-density to create copies + + std::string sample_base_prefix = index_prefix_path + "_sample"; + // optional, used if disk index file must store pq data + std::string disk_pq_pivots_path = + index_prefix_path + "_disk.index_pq_pivots.bin"; + // optional, used if disk index must store pq data + std::string disk_pq_compressed_vectors_path = + index_prefix_path + "_disk.index_pq_compressed.bin"; + + // output a new base file which contains extra dimension with sqrt(1 - + // ||x||^2/M^2) for every x, M is max norm of all points. Extra space on + // disk needed! + if (compareMetric == diskann::Metric::INNER_PRODUCT) { Timer timer; - diskann::get_bin_metadata(data_file_to_use.c_str(), points_num, dim); - const double p_val = ((double)MAX_PQ_TRAINING_SET_SIZE / (double)points_num); - - if (use_disk_pq) - { - generate_disk_quantized_data(data_file_to_use, disk_pq_pivots_path, disk_pq_compressed_vectors_path, - compareMetric, p_val, disk_pq_dims); - } - size_t num_pq_chunks = (size_t)(std::floor)(uint64_t(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; - - if (param_list.size() >= 9 && atoi(param_list[8].c_str()) <= MAX_PQ_CHUNKS && atoi(param_list[8].c_str()) > 0) - { - std::cout << "Use quantized dimension (QD) to overwrite derived quantized dimension from search_DRAM_budget (B)" + std::cout << "Using Inner Product search, so need to pre-process base " + "data into temp file. Please ensure there is additional " + "(n*(d+1)*4) bytes for storing pre-processed base vectors, " + "apart from the intermin indices and final index." + << std::endl; + std::string prepped_base = index_prefix_path + "_prepped_base.bin"; + data_file_to_use = prepped_base; + float max_norm_of_base = + diskann::prepare_base_for_inner_products(base_file, prepped_base); + std::string norm_file = disk_index_path + "_max_base_norm.bin"; + diskann::save_bin(norm_file, &max_norm_of_base, 1, 1); + diskann::cout << timer.elapsed_seconds_for_step( + "preprocessing data for inner product") << std::endl; - num_pq_chunks = atoi(param_list[8].c_str()); + } + + uint32_t R = (uint32_t)atoi(param_list[0].c_str()); + uint32_t L = (uint32_t)atoi(param_list[1].c_str()); + + 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; + return -1; + } + 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; + return -1; + } + uint32_t num_threads = (uint32_t)atoi(param_list[4].c_str()); + + if (num_threads != 0) { + omp_set_num_threads(num_threads); + mkl_set_num_threads(num_threads); + } + + 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; + + auto s = std::chrono::high_resolution_clock::now(); + + // If there is filter support, we break-up points which have too many labels + // into replica dummy points which evenly distribute the filters. The rest + // of index build happens on the augmented base and labels + std::string augmented_data_file, augmented_labels_file; + if (use_filters) { + convert_labels_string_to_int(labels_file_original, labels_file_to_use, + disk_labels_int_map_file, universal_label); + augmented_data_file = index_prefix_path + "_augmented_data.bin"; + augmented_labels_file = index_prefix_path + "_augmented_labels.txt"; + if (filter_threshold != 0) { + dummy_remap_file = index_prefix_path + "_dummy_remap.txt"; + breakup_dense_points( + data_file_to_use, labels_file_to_use, filter_threshold, + augmented_data_file, augmented_labels_file, + dummy_remap_file); // RKNOTE: This has large memory footprint, + // need to make this streaming + data_file_to_use = augmented_data_file; + labels_file_to_use = augmented_labels_file; } - - diskann::cout << "Compressing " << dim << "-dimensional data into " << num_pq_chunks << " bytes per vector." - << std::endl; - - generate_quantized_data(data_file_to_use, pq_pivots_path, pq_compressed_vectors_path, compareMetric, p_val, - num_pq_chunks, use_opq, codebook_prefix); - diskann::cout << timer.elapsed_seconds_for_step("generating quantized data") << std::endl; + } + + size_t points_num, dim; + + Timer timer; + diskann::get_bin_metadata(data_file_to_use.c_str(), points_num, dim); + const double p_val = ((double)MAX_PQ_TRAINING_SET_SIZE / (double)points_num); + + if (use_disk_pq) { + generate_disk_quantized_data(data_file_to_use, disk_pq_pivots_path, + disk_pq_compressed_vectors_path, + compareMetric, p_val, disk_pq_dims); + } + size_t num_pq_chunks = + (size_t)(std::floor)(uint64_t(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; + + if (param_list.size() >= 9 && atoi(param_list[8].c_str()) <= MAX_PQ_CHUNKS && + atoi(param_list[8].c_str()) > 0) { + std::cout << "Use quantized dimension (QD) to overwrite derived quantized " + "dimension from search_DRAM_budget (B)" + << std::endl; + num_pq_chunks = atoi(param_list[8].c_str()); + } + + diskann::cout << "Compressing " << dim << "-dimensional data into " + << num_pq_chunks << " bytes per vector." << std::endl; + + generate_quantized_data(data_file_to_use, pq_pivots_path, + pq_compressed_vectors_path, compareMetric, p_val, + num_pq_chunks, use_opq, codebook_prefix); + diskann::cout << timer.elapsed_seconds_for_step("generating quantized data") + << std::endl; // Gopal. Splitting diskann_dll into separate DLLs for search and build. // This code should only be available in the "build" DLL. -#if defined(RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS) && defined(DISKANN_BUILD) - MallocExtension::instance()->ReleaseFreeMemory(); +#if defined(RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS) && \ + defined(DISKANN_BUILD) + MallocExtension::instance()->ReleaseFreeMemory(); #endif - timer.reset(); - diskann::build_merged_vamana_index(data_file_to_use.c_str(), diskann::Metric::L2, L, R, p_val, - indexing_ram_budget, mem_index_path, medoids_path, centroids_path, - build_pq_bytes, use_opq, num_threads, use_filters, labels_file_to_use, - labels_to_medoids_path, universal_label, Lf); - diskann::cout << timer.elapsed_seconds_for_step("building merged vamana index") << std::endl; - - timer.reset(); - if (!use_disk_pq) - { - diskann::create_disk_layout(data_file_to_use.c_str(), mem_index_path, disk_index_path); - } + timer.reset(); + diskann::build_merged_vamana_index( + data_file_to_use.c_str(), diskann::Metric::L2, L, R, p_val, + indexing_ram_budget, mem_index_path, medoids_path, centroids_path, + build_pq_bytes, use_opq, num_threads, use_filters, labels_file_to_use, + labels_to_medoids_path, universal_label, Lf); + diskann::cout << timer.elapsed_seconds_for_step( + "building merged vamana index") + << std::endl; + + timer.reset(); + if (!use_disk_pq) { + diskann::create_disk_layout(data_file_to_use.c_str(), mem_index_path, + disk_index_path); + } else { + if (!reorder_data) + diskann::create_disk_layout(disk_pq_compressed_vectors_path, + mem_index_path, disk_index_path); else - { - if (!reorder_data) - diskann::create_disk_layout(disk_pq_compressed_vectors_path, mem_index_path, disk_index_path); - else - diskann::create_disk_layout(disk_pq_compressed_vectors_path, mem_index_path, disk_index_path, - data_file_to_use.c_str()); - } - diskann::cout << timer.elapsed_seconds_for_step("generating disk layout") << std::endl; - - 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; - gen_random_slice(data_file_to_use.c_str(), sample_base_prefix, sample_sampling_rate); - if (use_filters) - { - copy_file(labels_file_to_use, disk_labels_file); - std::remove(mem_labels_file.c_str()); - if (universal_label != "") - { - copy_file(mem_univ_label_file, disk_univ_label_file); - std::remove(mem_univ_label_file.c_str()); - } - std::remove(augmented_data_file.c_str()); - std::remove(augmented_labels_file.c_str()); - std::remove(labels_file_to_use.c_str()); + diskann::create_disk_layout(disk_pq_compressed_vectors_path, + mem_index_path, disk_index_path, + data_file_to_use.c_str()); + } + diskann::cout << timer.elapsed_seconds_for_step("generating disk layout") + << std::endl; + + 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; + gen_random_slice(data_file_to_use.c_str(), sample_base_prefix, + sample_sampling_rate); + if (use_filters) { + copy_file(labels_file_to_use, disk_labels_file); + std::remove(mem_labels_file.c_str()); + if (universal_label != "") { + copy_file(mem_univ_label_file, disk_univ_label_file); + std::remove(mem_univ_label_file.c_str()); } + std::remove(augmented_data_file.c_str()); + std::remove(augmented_labels_file.c_str()); + std::remove(labels_file_to_use.c_str()); + } - std::remove(mem_index_path.c_str()); - if (use_disk_pq) - std::remove(disk_pq_compressed_vectors_path.c_str()); + std::remove(mem_index_path.c_str()); + if (use_disk_pq) std::remove(disk_pq_compressed_vectors_path.c_str()); - auto e = std::chrono::high_resolution_clock::now(); - std::chrono::duration diff = e - s; - diskann::cout << "Indexing time: " << diff.count() << std::endl; + auto e = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = e - s; + diskann::cout << "Indexing time: " << diff.count() << std::endl; - return 0; + return 0; } -template DISKANN_DLLEXPORT void create_disk_layout(const std::string base_file, - const std::string mem_index_file, - const std::string output_file, - const std::string reorder_data_file); -template DISKANN_DLLEXPORT void create_disk_layout(const std::string base_file, - const std::string mem_index_file, - const std::string output_file, - const std::string reorder_data_file); -template DISKANN_DLLEXPORT void create_disk_layout(const std::string base_file, const std::string mem_index_file, - const std::string output_file, - const std::string reorder_data_file); - -template DISKANN_DLLEXPORT int8_t *load_warmup(const std::string &cache_warmup_file, uint64_t &warmup_num, - uint64_t warmup_dim, uint64_t warmup_aligned_dim); -template DISKANN_DLLEXPORT uint8_t *load_warmup(const std::string &cache_warmup_file, uint64_t &warmup_num, - uint64_t warmup_dim, uint64_t warmup_aligned_dim); -template DISKANN_DLLEXPORT float *load_warmup(const std::string &cache_warmup_file, uint64_t &warmup_num, - uint64_t warmup_dim, uint64_t warmup_aligned_dim); +template DISKANN_DLLEXPORT void create_disk_layout( + const std::string base_file, const std::string mem_index_file, + const std::string output_file, const std::string reorder_data_file); +template DISKANN_DLLEXPORT void create_disk_layout( + const std::string base_file, const std::string mem_index_file, + const std::string output_file, const std::string reorder_data_file); +template DISKANN_DLLEXPORT void create_disk_layout( + const std::string base_file, const std::string mem_index_file, + const std::string output_file, const std::string reorder_data_file); + +template DISKANN_DLLEXPORT int8_t *load_warmup( + const std::string &cache_warmup_file, uint64_t &warmup_num, + uint64_t warmup_dim, uint64_t warmup_aligned_dim); +template DISKANN_DLLEXPORT uint8_t *load_warmup( + const std::string &cache_warmup_file, uint64_t &warmup_num, + uint64_t warmup_dim, uint64_t warmup_aligned_dim); +template DISKANN_DLLEXPORT float *load_warmup( + const std::string &cache_warmup_file, uint64_t &warmup_num, + uint64_t warmup_dim, uint64_t warmup_aligned_dim); #ifdef EXEC_ENV_OLS -template DISKANN_DLLEXPORT int8_t *load_warmup(MemoryMappedFiles &files, const std::string &cache_warmup_file, - uint64_t &warmup_num, uint64_t warmup_dim, - uint64_t warmup_aligned_dim); -template DISKANN_DLLEXPORT uint8_t *load_warmup(MemoryMappedFiles &files, const std::string &cache_warmup_file, - uint64_t &warmup_num, uint64_t warmup_dim, - uint64_t warmup_aligned_dim); -template DISKANN_DLLEXPORT float *load_warmup(MemoryMappedFiles &files, const std::string &cache_warmup_file, - uint64_t &warmup_num, uint64_t warmup_dim, - uint64_t warmup_aligned_dim); +template DISKANN_DLLEXPORT int8_t *load_warmup( + MemoryMappedFiles &files, const std::string &cache_warmup_file, + uint64_t &warmup_num, uint64_t warmup_dim, uint64_t warmup_aligned_dim); +template DISKANN_DLLEXPORT uint8_t *load_warmup( + MemoryMappedFiles &files, const std::string &cache_warmup_file, + uint64_t &warmup_num, uint64_t warmup_dim, uint64_t warmup_aligned_dim); +template DISKANN_DLLEXPORT float *load_warmup( + MemoryMappedFiles &files, const std::string &cache_warmup_file, + 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, int8_t *tuning_sample, - uint64_t tuning_sample_num, uint64_t tuning_sample_aligned_dim, uint32_t L, uint32_t nthreads, uint32_t start_bw); + std::unique_ptr> &pFlashIndex, + int8_t *tuning_sample, uint64_t tuning_sample_num, + uint64_t 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, - uint64_t tuning_sample_num, uint64_t tuning_sample_aligned_dim, uint32_t L, uint32_t nthreads, uint32_t start_bw); + std::unique_ptr> &pFlashIndex, + uint8_t *tuning_sample, uint64_t tuning_sample_num, + uint64_t 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, - uint64_t tuning_sample_num, uint64_t tuning_sample_aligned_dim, uint32_t L, uint32_t nthreads, uint32_t start_bw); + std::unique_ptr> &pFlashIndex, + float *tuning_sample, uint64_t tuning_sample_num, + uint64_t 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, int8_t *tuning_sample, - uint64_t tuning_sample_num, uint64_t tuning_sample_aligned_dim, uint32_t L, uint32_t nthreads, uint32_t start_bw); + std::unique_ptr> &pFlashIndex, + int8_t *tuning_sample, uint64_t tuning_sample_num, + uint64_t 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, - uint64_t tuning_sample_num, uint64_t tuning_sample_aligned_dim, uint32_t L, uint32_t nthreads, uint32_t start_bw); + std::unique_ptr> &pFlashIndex, + uint8_t *tuning_sample, uint64_t tuning_sample_num, + uint64_t 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, - uint64_t tuning_sample_num, uint64_t tuning_sample_aligned_dim, uint32_t L, uint32_t nthreads, uint32_t start_bw); - -template DISKANN_DLLEXPORT int build_disk_index(const char *dataFilePath, const char *indexFilePath, - const char *indexBuildParameters, - diskann::Metric compareMetric, bool use_opq, - const std::string &codebook_prefix, bool use_filters, - const std::string &label_file, - const std::string &universal_label, - const uint32_t filter_threshold, const uint32_t Lf); -template DISKANN_DLLEXPORT int build_disk_index(const char *dataFilePath, const char *indexFilePath, - const char *indexBuildParameters, - diskann::Metric compareMetric, bool use_opq, - const std::string &codebook_prefix, bool use_filters, - const std::string &label_file, - const std::string &universal_label, - const uint32_t filter_threshold, const uint32_t Lf); -template DISKANN_DLLEXPORT int build_disk_index(const char *dataFilePath, const char *indexFilePath, - const char *indexBuildParameters, - diskann::Metric compareMetric, bool use_opq, - const std::string &codebook_prefix, bool use_filters, - const std::string &label_file, - const std::string &universal_label, - const uint32_t filter_threshold, const uint32_t Lf); + std::unique_ptr> &pFlashIndex, + float *tuning_sample, uint64_t tuning_sample_num, + uint64_t tuning_sample_aligned_dim, uint32_t L, uint32_t nthreads, + uint32_t start_bw); + +template DISKANN_DLLEXPORT int build_disk_index( + const char *dataFilePath, const char *indexFilePath, + const char *indexBuildParameters, diskann::Metric compareMetric, + bool use_opq, const std::string &codebook_prefix, bool use_filters, + const std::string &label_file, const std::string &universal_label, + const uint32_t filter_threshold, const uint32_t Lf); +template DISKANN_DLLEXPORT int build_disk_index( + const char *dataFilePath, const char *indexFilePath, + const char *indexBuildParameters, diskann::Metric compareMetric, + bool use_opq, const std::string &codebook_prefix, bool use_filters, + const std::string &label_file, const std::string &universal_label, + const uint32_t filter_threshold, const uint32_t Lf); +template DISKANN_DLLEXPORT int build_disk_index( + const char *dataFilePath, const char *indexFilePath, + const char *indexBuildParameters, diskann::Metric compareMetric, + bool use_opq, const std::string &codebook_prefix, bool use_filters, + const std::string &label_file, const std::string &universal_label, + const uint32_t filter_threshold, const uint32_t Lf); // LabelT = uint16 -template DISKANN_DLLEXPORT int build_disk_index(const char *dataFilePath, const char *indexFilePath, - const char *indexBuildParameters, - diskann::Metric compareMetric, bool use_opq, - const std::string &codebook_prefix, bool use_filters, - const std::string &label_file, - const std::string &universal_label, - const uint32_t filter_threshold, const uint32_t Lf); -template DISKANN_DLLEXPORT int build_disk_index(const char *dataFilePath, const char *indexFilePath, - const char *indexBuildParameters, - diskann::Metric compareMetric, bool use_opq, - const std::string &codebook_prefix, bool use_filters, - const std::string &label_file, - const std::string &universal_label, - const uint32_t filter_threshold, const uint32_t Lf); -template DISKANN_DLLEXPORT int build_disk_index(const char *dataFilePath, const char *indexFilePath, - const char *indexBuildParameters, - diskann::Metric compareMetric, bool use_opq, - const std::string &codebook_prefix, bool use_filters, - const std::string &label_file, - const std::string &universal_label, - const uint32_t filter_threshold, const uint32_t Lf); +template DISKANN_DLLEXPORT int build_disk_index( + const char *dataFilePath, const char *indexFilePath, + const char *indexBuildParameters, diskann::Metric compareMetric, + bool use_opq, const std::string &codebook_prefix, bool use_filters, + const std::string &label_file, const std::string &universal_label, + const uint32_t filter_threshold, const uint32_t Lf); +template DISKANN_DLLEXPORT int build_disk_index( + const char *dataFilePath, const char *indexFilePath, + const char *indexBuildParameters, diskann::Metric compareMetric, + bool use_opq, const std::string &codebook_prefix, bool use_filters, + const std::string &label_file, const std::string &universal_label, + const uint32_t filter_threshold, const uint32_t Lf); +template DISKANN_DLLEXPORT int build_disk_index( + const char *dataFilePath, const char *indexFilePath, + const char *indexBuildParameters, diskann::Metric compareMetric, + bool use_opq, const std::string &codebook_prefix, bool use_filters, + const std::string &label_file, const std::string &universal_label, + const uint32_t filter_threshold, const uint32_t Lf); template DISKANN_DLLEXPORT int build_merged_vamana_index( - std::string base_file, diskann::Metric compareMetric, uint32_t L, uint32_t R, double sampling_rate, - double ram_budget, std::string mem_index_path, std::string medoids_path, std::string centroids_file, - size_t build_pq_bytes, bool use_opq, uint32_t num_threads, bool use_filters, const std::string &label_file, - const std::string &labels_to_medoids_file, const std::string &universal_label, const uint32_t Lf); + std::string base_file, diskann::Metric compareMetric, uint32_t L, + uint32_t R, double sampling_rate, double ram_budget, + std::string mem_index_path, std::string medoids_path, + std::string centroids_file, size_t build_pq_bytes, bool use_opq, + uint32_t num_threads, bool use_filters, const std::string &label_file, + const std::string &labels_to_medoids_file, + const std::string &universal_label, const uint32_t Lf); template DISKANN_DLLEXPORT int build_merged_vamana_index( - std::string base_file, diskann::Metric compareMetric, uint32_t L, uint32_t R, double sampling_rate, - double ram_budget, std::string mem_index_path, std::string medoids_path, std::string centroids_file, - size_t build_pq_bytes, bool use_opq, uint32_t num_threads, bool use_filters, const std::string &label_file, - const std::string &labels_to_medoids_file, const std::string &universal_label, const uint32_t Lf); + std::string base_file, diskann::Metric compareMetric, uint32_t L, + uint32_t R, double sampling_rate, double ram_budget, + std::string mem_index_path, std::string medoids_path, + std::string centroids_file, size_t build_pq_bytes, bool use_opq, + uint32_t num_threads, bool use_filters, const std::string &label_file, + const std::string &labels_to_medoids_file, + const std::string &universal_label, const uint32_t Lf); template DISKANN_DLLEXPORT int build_merged_vamana_index( - std::string base_file, diskann::Metric compareMetric, uint32_t L, uint32_t R, double sampling_rate, - double ram_budget, std::string mem_index_path, std::string medoids_path, std::string centroids_file, - size_t build_pq_bytes, bool use_opq, uint32_t num_threads, bool use_filters, const std::string &label_file, - const std::string &labels_to_medoids_file, const std::string &universal_label, const uint32_t Lf); + std::string base_file, diskann::Metric compareMetric, uint32_t L, + uint32_t R, double sampling_rate, double ram_budget, + std::string mem_index_path, std::string medoids_path, + std::string centroids_file, size_t build_pq_bytes, bool use_opq, + uint32_t num_threads, bool use_filters, const std::string &label_file, + const std::string &labels_to_medoids_file, + const std::string &universal_label, const uint32_t Lf); // Label=16_t template DISKANN_DLLEXPORT int build_merged_vamana_index( - std::string base_file, diskann::Metric compareMetric, uint32_t L, uint32_t R, double sampling_rate, - double ram_budget, std::string mem_index_path, std::string medoids_path, std::string centroids_file, - size_t build_pq_bytes, bool use_opq, uint32_t num_threads, bool use_filters, const std::string &label_file, - const std::string &labels_to_medoids_file, const std::string &universal_label, const uint32_t Lf); + std::string base_file, diskann::Metric compareMetric, uint32_t L, + uint32_t R, double sampling_rate, double ram_budget, + std::string mem_index_path, std::string medoids_path, + std::string centroids_file, size_t build_pq_bytes, bool use_opq, + uint32_t num_threads, bool use_filters, const std::string &label_file, + const std::string &labels_to_medoids_file, + const std::string &universal_label, const uint32_t Lf); template DISKANN_DLLEXPORT int build_merged_vamana_index( - std::string base_file, diskann::Metric compareMetric, uint32_t L, uint32_t R, double sampling_rate, - double ram_budget, std::string mem_index_path, std::string medoids_path, std::string centroids_file, - size_t build_pq_bytes, bool use_opq, uint32_t num_threads, bool use_filters, const std::string &label_file, - const std::string &labels_to_medoids_file, const std::string &universal_label, const uint32_t Lf); + std::string base_file, diskann::Metric compareMetric, uint32_t L, + uint32_t R, double sampling_rate, double ram_budget, + std::string mem_index_path, std::string medoids_path, + std::string centroids_file, size_t build_pq_bytes, bool use_opq, + uint32_t num_threads, bool use_filters, const std::string &label_file, + const std::string &labels_to_medoids_file, + const std::string &universal_label, const uint32_t Lf); template DISKANN_DLLEXPORT int build_merged_vamana_index( - std::string base_file, diskann::Metric compareMetric, uint32_t L, uint32_t R, double sampling_rate, - double ram_budget, std::string mem_index_path, std::string medoids_path, std::string centroids_file, - size_t build_pq_bytes, bool use_opq, uint32_t num_threads, bool use_filters, const std::string &label_file, - const std::string &labels_to_medoids_file, const std::string &universal_label, const uint32_t Lf); -}; // namespace diskann + std::string base_file, diskann::Metric compareMetric, uint32_t L, + uint32_t R, double sampling_rate, double ram_budget, + std::string mem_index_path, std::string medoids_path, + std::string centroids_file, size_t build_pq_bytes, bool use_opq, + uint32_t num_threads, bool use_filters, const std::string &label_file, + const std::string &labels_to_medoids_file, + const std::string &universal_label, const uint32_t Lf); +}; // namespace diskann diff --git a/src/distance.cpp b/src/distance.cpp index 120801568..e6be334e6 100644 --- a/src/distance.cpp +++ b/src/distance.cpp @@ -19,705 +19,669 @@ #include "logger.h" #include "ann_exception.h" -namespace diskann -{ +namespace diskann { // // Base Class Implementatons // template -float Distance::compare(const T *a, const T *b, const float normA, const float normB, uint32_t length) const -{ - throw std::logic_error("This function is not implemented."); +float Distance::compare(const T *a, const T *b, const float normA, + const float normB, uint32_t length) const { + throw std::logic_error("This function is not implemented."); } -template uint32_t Distance::post_normalization_dimension(uint32_t orig_dimension) const -{ - return orig_dimension; +template +uint32_t Distance::post_normalization_dimension( + uint32_t orig_dimension) const { + return orig_dimension; } -template diskann::Metric Distance::get_metric() const -{ - return _distance_metric; +template +diskann::Metric Distance::get_metric() const { + return _distance_metric; } -template bool Distance::preprocessing_required() const -{ - return false; +template +bool Distance::preprocessing_required() const { + return false; } template -void Distance::preprocess_base_points(T *original_data, const size_t orig_dim, const size_t num_points) -{ -} +void Distance::preprocess_base_points(T *original_data, + const size_t orig_dim, + const size_t num_points) {} template -void Distance::preprocess_query(const T *query_vec, const size_t query_dim, T *scratch_query) -{ - std::memcpy(scratch_query, query_vec, query_dim * sizeof(T)); +void Distance::preprocess_query(const T *query_vec, const size_t query_dim, + T *scratch_query) { + std::memcpy(scratch_query, query_vec, query_dim * sizeof(T)); } -template size_t Distance::get_required_alignment() const -{ - return _alignment_factor; +template +size_t Distance::get_required_alignment() const { + return _alignment_factor; } -template Distance::~Distance() -{ -} +template +Distance::~Distance() {} // // Cosine distance functions. // -float DistanceCosineInt8::compare(const int8_t *a, const int8_t *b, uint32_t length) const -{ +float DistanceCosineInt8::compare(const int8_t *a, const int8_t *b, + uint32_t length) const { #ifdef _WINDOWS - return diskann::CosineSimilarity2(a, b, length); + return diskann::CosineSimilarity2(a, b, length); #else - int magA = 0, magB = 0, scalarProduct = 0; - for (uint32_t i = 0; i < length; i++) - { - magA += ((int32_t)a[i]) * ((int32_t)a[i]); - magB += ((int32_t)b[i]) * ((int32_t)b[i]); - scalarProduct += ((int32_t)a[i]) * ((int32_t)b[i]); - } - // similarity == 1-cosine distance - return 1.0f - (float)(scalarProduct / (sqrt(magA) * sqrt(magB))); + int magA = 0, magB = 0, scalarProduct = 0; + for (uint32_t i = 0; i < length; i++) { + magA += ((int32_t)a[i]) * ((int32_t)a[i]); + magB += ((int32_t)b[i]) * ((int32_t)b[i]); + scalarProduct += ((int32_t)a[i]) * ((int32_t)b[i]); + } + // similarity == 1-cosine distance + return 1.0f - (float)(scalarProduct / (sqrt(magA) * sqrt(magB))); #endif } -float DistanceCosineFloat::compare(const float *a, const float *b, uint32_t length) const -{ +float DistanceCosineFloat::compare(const float *a, const float *b, + uint32_t length) const { #ifdef _WINDOWS - return diskann::CosineSimilarity2(a, b, length); + return diskann::CosineSimilarity2(a, b, length); #else - float magA = 0, magB = 0, scalarProduct = 0; - for (uint32_t i = 0; i < length; i++) - { - magA += (a[i]) * (a[i]); - magB += (b[i]) * (b[i]); - scalarProduct += (a[i]) * (b[i]); - } - // similarity == 1-cosine distance - return 1.0f - (scalarProduct / (sqrt(magA) * sqrt(magB))); + float magA = 0, magB = 0, scalarProduct = 0; + for (uint32_t i = 0; i < length; i++) { + magA += (a[i]) * (a[i]); + magB += (b[i]) * (b[i]); + scalarProduct += (a[i]) * (b[i]); + } + // similarity == 1-cosine distance + return 1.0f - (scalarProduct / (sqrt(magA) * sqrt(magB))); #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))); +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 -{ +float DistanceL2Int8::compare(const int8_t *a, const int8_t *b, + uint32_t size) 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]; + __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 - int32_t result = 0; + int32_t result = 0; #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; + 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 #else - int32_t result = 0; + int32_t result = 0; #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; + 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; +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; + 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 DistanceL2Float::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); +float DistanceL2Float::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 DistanceL2Float::compare(const float *a, const float *b, uint32_t size) const -{ +float DistanceL2Float::compare(const float *a, const float *b, + uint32_t size) const { #endif - float result = 0; + 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); - - sum = _mm256_fmadd_ps(tmp_vec, tmp_vec, sum); - } - - // horizontal add sum - result = _mm256_reduce_add_ps(sum); + // 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); + + 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 (int32_t i = 0; i < (int32_t)size; i++) - { - result += (a[i] - b[i]) * (a[i] - b[i]); - } + for (int32_t i = 0; i < (int32_t)size; i++) { + result += (a[i] - b[i]) * (a[i] - b[i]); + } #endif - return result; + return result; } -template float SlowDistanceL2::compare(const T *a, const T *b, uint32_t length) const -{ - float result = 0.0f; - for (uint32_t i = 0; i < length; i++) - { - result += ((float)(a[i] - b[i])) * (a[i] - b[i]); - } - return result; +template +float SlowDistanceL2::compare(const T *a, const T *b, + uint32_t length) const { + float result = 0.0f; + for (uint32_t i = 0; i < length; i++) { + result += ((float)(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 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]; +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 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; +float AVXDistanceL2Float::compare(const float *, const float *, + uint32_t) const { + return 0; } #endif -template float DistanceInnerProduct::inner_product(const T *a, const T *b, uint32_t size) const -{ - if (!std::is_floating_point::value) - { - diskann::cerr << "ERROR: Inner Product only defined for float currently." << std::endl; - throw diskann::ANNException("ERROR: Inner Product only defined for float currently.", -1, __FUNCSIG__, __FILE__, - __LINE__); - } - - float result = 0; +template +float DistanceInnerProduct::inner_product(const T *a, const T *b, + uint32_t size) const { + if (!std::is_floating_point::value) { + diskann::cerr << "ERROR: Inner Product only defined for float currently." + << std::endl; + throw diskann::ANNException( + "ERROR: Inner Product only defined for float currently.", -1, + __FUNCSIG__, __FILE__, __LINE__); + } + + float result = 0; #ifdef __GNUC__ #ifdef USE_AVX2 -#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; - uint32_t D = (size + 7) & ~7U; - uint32_t DR = D % 16; - uint32_t 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 (uint32_t 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]; +#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; + uint32_t D = (size + 7) & ~7U; + uint32_t DR = D % 16; + uint32_t 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 (uint32_t 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; - uint32_t D = (size + 3) & ~3U; - uint32_t DR = D % 16; - uint32_t 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) - { +#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; + uint32_t D = (size + 3) & ~3U; + uint32_t DR = D % 16; + uint32_t 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); + SSE_DOT(e_l + 8, e_r + 8, sum, l2, r2); case 8: - SSE_DOT(e_l + 4, e_r + 4, sum, l1, r1); + SSE_DOT(e_l + 4, e_r + 4, sum, l1, r1); case 4: - SSE_DOT(e_l, e_r, sum, l0, r0); + SSE_DOT(e_l, e_r, sum, l0, r0); default: - break; - } - for (uint32_t 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]; + break; + } + for (uint32_t 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++; - } + 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; + return result; } -template float DistanceFastL2::compare(const T *a, const T *b, float norm, uint32_t size) const -{ - float result = -2 * DistanceInnerProduct::inner_product(a, b, size); - result += norm; - return result; +template +float DistanceFastL2::compare(const T *a, const T *b, float norm, + uint32_t size) const { + float result = -2 * DistanceInnerProduct::inner_product(a, b, size); + result += norm; + return result; } -template float DistanceFastL2::norm(const T *a, uint32_t size) const -{ - if (!std::is_floating_point::value) - { - diskann::cerr << "ERROR: FastL2 only defined for float currently." << std::endl; - throw diskann::ANNException("ERROR: FastL2 only defined for float currently.", -1, __FUNCSIG__, __FILE__, - __LINE__); - } - float result = 0; +template +float DistanceFastL2::norm(const T *a, uint32_t size) const { + if (!std::is_floating_point::value) { + diskann::cerr << "ERROR: FastL2 only defined for float currently." + << std::endl; + throw diskann::ANNException( + "ERROR: FastL2 only defined for float currently.", -1, __FUNCSIG__, + __FILE__, __LINE__); + } + 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; - uint32_t D = (size + 7) & ~7U; - uint32_t DR = D % 16; - uint32_t 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 (uint32_t 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]; +#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; + uint32_t D = (size + 7) & ~7U; + uint32_t DR = D % 16; + uint32_t 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 (uint32_t 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; - uint32_t D = (size + 3) & ~3U; - uint32_t DR = D % 16; - uint32_t 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) - { +#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; + uint32_t D = (size + 3) & ~3U; + uint32_t DR = D % 16; + uint32_t 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); + SSE_L2NORM(e_l + 8, sum, l2); case 8: - SSE_L2NORM(e_l + 4, sum, l1); + SSE_L2NORM(e_l + 4, sum, l1); case 4: - SSE_L2NORM(e_l, sum, l0); + SSE_L2NORM(e_l, sum, l0); default: - break; - } - for (uint32_t 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]; + break; + } + for (uint32_t 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++; - } + 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; + return result; } -float AVXDistanceInnerProductFloat::compare(const float *a, const float *b, uint32_t size) const -{ - float result = 0.0f; -#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; - uint32_t D = (size + 7) & ~7U; - uint32_t DR = D % 16; - uint32_t 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 AVXDistanceInnerProductFloat::compare(const float *a, const float *b, + uint32_t size) const { + float result = 0.0f; +#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; + uint32_t D = (size + 7) & ~7U; + uint32_t DR = D % 16; + uint32_t 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; #ifndef _WINDOWS - float unpack[8] __attribute__((aligned(32))) = {0, 0, 0, 0, 0, 0, 0, 0}; + float unpack[8] __attribute__((aligned(32))) = {0, 0, 0, 0, 0, 0, 0, 0}; #else - __declspec(align(32)) float unpack[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + __declspec(align(32)) float unpack[8] = {0, 0, 0, 0, 0, 0, 0, 0}; #endif - sum = _mm256_loadu_ps(unpack); - if (DR) - { - AVX_DOT(e_l, e_r, sum, l0, r0); - } + sum = _mm256_loadu_ps(unpack); + if (DR) { + AVX_DOT(e_l, e_r, sum, l0, r0); + } - for (uint32_t 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]; + for (uint32_t 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]; - return -result; + return -result; } -uint32_t AVXNormalizedCosineDistanceFloat::post_normalization_dimension(uint32_t orig_dimension) const -{ - return orig_dimension; +uint32_t AVXNormalizedCosineDistanceFloat::post_normalization_dimension( + uint32_t orig_dimension) const { + return orig_dimension; } -bool AVXNormalizedCosineDistanceFloat::preprocessing_required() const -{ - return true; +bool AVXNormalizedCosineDistanceFloat::preprocessing_required() const { + return true; } -void AVXNormalizedCosineDistanceFloat::preprocess_base_points(float *original_data, const size_t orig_dim, - const size_t num_points) -{ - for (uint32_t i = 0; i < num_points; i++) - { - normalize((float *)(original_data + i * orig_dim), orig_dim); - } +void AVXNormalizedCosineDistanceFloat::preprocess_base_points( + float *original_data, const size_t orig_dim, const size_t num_points) { + for (uint32_t i = 0; i < num_points; i++) { + normalize((float *)(original_data + i * orig_dim), orig_dim); + } } -void AVXNormalizedCosineDistanceFloat::preprocess_query(const float *query_vec, const size_t query_dim, - float *query_scratch) -{ - normalize_and_copy(query_vec, query_dim, query_scratch); +void AVXNormalizedCosineDistanceFloat::preprocess_query(const float *query_vec, + const size_t query_dim, + float *query_scratch) { + normalize_and_copy(query_vec, query_dim, query_scratch); } -void AVXNormalizedCosineDistanceFloat::normalize_and_copy(const float *query_vec, const uint32_t query_dim, - float *query_target) const -{ +void AVXNormalizedCosineDistanceFloat::normalize_and_copy( + const float *query_vec, const uint32_t query_dim, + float *query_target) const { + float norm = get_norm(query_vec, query_dim); - float norm = get_norm(query_vec, query_dim); - - for (uint32_t i = 0; i < query_dim; i++) - { - query_target[i] = query_vec[i] / norm; - } + for (uint32_t i = 0; i < query_dim; i++) { + query_target[i] = query_vec[i] / norm; + } } // 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 DistanceL2Float" << std::endl; - return new diskann::DistanceL2Float(); - } - 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::SlowDistanceL2(); - } - } - else if (m == diskann::Metric::COSINE) - { - diskann::cout << "Cosine: Using either AVX or AVX2 implementation" << std::endl; - return new diskann::DistanceCosineFloat(); - } - else if (m == diskann::Metric::INNER_PRODUCT) - { - diskann::cout << "Inner product: Using AVX2 implementation " - "AVXDistanceInnerProductFloat" - << std::endl; - return new diskann::AVXDistanceInnerProductFloat(); - } - else if (m == diskann::Metric::FAST_L2) - { - diskann::cout << "Fast_L2: Using AVX2 implementation with norm " - "memoization DistanceFastL2" - << std::endl; - return new diskann::DistanceFastL2(); - } - else - { - std::stringstream stream; - stream << "Only L2, cosine, and inner product supported for floating " - "point vectors as of now." - << std::endl; - diskann::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 << "L2: Using AVX2 distance computation DistanceL2Float" + << std::endl; + return new diskann::DistanceL2Float(); + } 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::SlowDistanceL2(); + } + } else if (m == diskann::Metric::COSINE) { + diskann::cout << "Cosine: Using either AVX or AVX2 implementation" + << std::endl; + return new diskann::DistanceCosineFloat(); + } else if (m == diskann::Metric::INNER_PRODUCT) { + diskann::cout << "Inner product: Using AVX2 implementation " + "AVXDistanceInnerProductFloat" + << std::endl; + return new diskann::AVXDistanceInnerProductFloat(); + } else if (m == diskann::Metric::FAST_L2) { + diskann::cout << "Fast_L2: Using AVX2 implementation with norm " + "memoization DistanceFastL2" + << std::endl; + return new diskann::DistanceFastL2(); + } else { + std::stringstream stream; + stream << "Only L2, cosine, and inner product supported for floating " + "point vectors as of now." + << std::endl; + diskann::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 DistanceL2Int8." << 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 " - "SlowDistanceL2Int." - << std::endl; - return new diskann::SlowDistanceL2(); - } - } - else if (m == diskann::Metric::COSINE) - { - diskann::cout << "Using either AVX or AVX2 for Cosine similarity " - "DistanceCosineInt8." - << std::endl; - return new diskann::DistanceCosineInt8(); - } - else - { - std::stringstream stream; - stream << "Only L2 and cosine supported for signed byte vectors." << std::endl; - diskann::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 DistanceL2Int8." + << 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 " + "SlowDistanceL2Int." + << std::endl; + return new diskann::SlowDistanceL2(); + } + } else if (m == diskann::Metric::COSINE) { + diskann::cout << "Using either AVX or AVX2 for Cosine similarity " + "DistanceCosineInt8." + << std::endl; + return new diskann::DistanceCosineInt8(); + } else { + std::stringstream stream; + stream << "Only L2 and cosine supported for signed byte vectors." + << std::endl; + diskann::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) - { +template <> +diskann::Distance *get_distance_function(diskann::Metric m) { + if (m == diskann::Metric::L2) { #ifdef _WINDOWS - diskann::cout << "WARNING: AVX/AVX2 distance function not defined for Uint8. " - "Using " - "slow version. " - "Contact gopalsr@microsoft.com if you need AVX/AVX2 support." - << std::endl; + diskann::cout + << "WARNING: AVX/AVX2 distance function not defined for Uint8. " + "Using " + "slow version. " + "Contact gopalsr@microsoft.com if you need AVX/AVX2 support." + << std::endl; #endif - return new diskann::DistanceL2UInt8(); - } - else if (m == diskann::Metric::COSINE) - { - diskann::cout << "AVX/AVX2 distance function not defined for Uint8. Using " - "slow version SlowDistanceCosineUint8() " - "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 supported for uint32_t byte vectors." << std::endl; - diskann::cerr << stream.str() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } + return new diskann::DistanceL2UInt8(); + } else if (m == diskann::Metric::COSINE) { + diskann::cout + << "AVX/AVX2 distance function not defined for Uint8. Using " + "slow version SlowDistanceCosineUint8() " + "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 supported for uint32_t byte vectors." + << std::endl; + diskann::cerr << stream.str() << std::endl; + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } } template DISKANN_DLLEXPORT class DistanceInnerProduct; @@ -732,4 +696,4 @@ template DISKANN_DLLEXPORT class SlowDistanceL2; template DISKANN_DLLEXPORT class SlowDistanceL2; template DISKANN_DLLEXPORT class SlowDistanceL2; -} // namespace diskann +} // namespace diskann diff --git a/src/filter_utils.cpp b/src/filter_utils.cpp index 965762d1f..c0406c352 100644 --- a/src/filter_utils.cpp +++ b/src/filter_utils.cpp @@ -14,8 +14,7 @@ #include "parameters.h" #include "utils.h" -namespace diskann -{ +namespace diskann { /* * Using passed in parameters and files generated from step 3, * builds a vanilla diskANN index for each label. @@ -24,119 +23,124 @@ namespace diskann * final_index_path_prefix + "_" + label */ template -void generate_label_indices(path input_data_path, path final_index_path_prefix, label_set all_labels, uint32_t R, - uint32_t L, float alpha, uint32_t num_threads) -{ - diskann::IndexWriteParameters label_index_build_parameters = diskann::IndexWriteParametersBuilder(L, R) - .with_saturate_graph(false) - .with_alpha(alpha) - .with_num_threads(num_threads) - .build(); - - std::cout << "Generating indices per label..." << std::endl; - // for each label, build an index on resp. points - double total_indexing_time = 0.0, indexing_percentage = 0.0; - std::cout.setstate(std::ios_base::failbit); - diskann::cout.setstate(std::ios_base::failbit); - for (const auto &lbl : all_labels) - { - path curr_label_input_data_path(input_data_path + "_" + lbl); - path curr_label_index_path(final_index_path_prefix + "_" + lbl); - - size_t number_of_label_points, dimension; - diskann::get_bin_metadata(curr_label_input_data_path, number_of_label_points, dimension); - diskann::Index index(diskann::Metric::L2, dimension, number_of_label_points, false, false); - - auto index_build_timer = std::chrono::high_resolution_clock::now(); - index.build(curr_label_input_data_path.c_str(), number_of_label_points, label_index_build_parameters); - std::chrono::duration current_indexing_time = - std::chrono::high_resolution_clock::now() - index_build_timer; - - total_indexing_time += current_indexing_time.count(); - indexing_percentage += (1 / (double)all_labels.size()); - print_progress(indexing_percentage); - - index.save(curr_label_index_path.c_str()); - } - std::cout.clear(); - diskann::cout.clear(); - - std::cout << "\nDone. Generated per-label indices in " << total_indexing_time << " seconds\n" << std::endl; +void generate_label_indices(path input_data_path, path final_index_path_prefix, + label_set all_labels, uint32_t R, uint32_t L, + float alpha, uint32_t num_threads) { + diskann::IndexWriteParameters label_index_build_parameters = + diskann::IndexWriteParametersBuilder(L, R) + .with_saturate_graph(false) + .with_alpha(alpha) + .with_num_threads(num_threads) + .build(); + + std::cout << "Generating indices per label..." << std::endl; + // for each label, build an index on resp. points + double total_indexing_time = 0.0, indexing_percentage = 0.0; + std::cout.setstate(std::ios_base::failbit); + diskann::cout.setstate(std::ios_base::failbit); + for (const auto &lbl : all_labels) { + path curr_label_input_data_path(input_data_path + "_" + lbl); + path curr_label_index_path(final_index_path_prefix + "_" + lbl); + + size_t number_of_label_points, dimension; + diskann::get_bin_metadata(curr_label_input_data_path, + number_of_label_points, dimension); + diskann::Index index(diskann::Metric::L2, dimension, + number_of_label_points, false, false); + + auto index_build_timer = std::chrono::high_resolution_clock::now(); + index.build(curr_label_input_data_path.c_str(), number_of_label_points, + label_index_build_parameters); + std::chrono::duration current_indexing_time = + std::chrono::high_resolution_clock::now() - index_build_timer; + + total_indexing_time += current_indexing_time.count(); + indexing_percentage += (1 / (double)all_labels.size()); + print_progress(indexing_percentage); + + index.save(curr_label_index_path.c_str()); + } + std::cout.clear(); + diskann::cout.clear(); + + std::cout << "\nDone. Generated per-label indices in " << total_indexing_time + << " seconds\n" + << std::endl; } // for use on systems without writev (i.e. Windows) template -tsl::robin_map> generate_label_specific_vector_files_compat( - path input_data_path, tsl::robin_map labels_to_number_of_points, - std::vector point_ids_to_labels, label_set all_labels) -{ - auto file_writing_timer = std::chrono::high_resolution_clock::now(); - std::ifstream input_data_stream(input_data_path); - - uint32_t number_of_points, dimension; - input_data_stream.read((char *)&number_of_points, sizeof(uint32_t)); - input_data_stream.read((char *)&dimension, sizeof(uint32_t)); - const uint32_t VECTOR_SIZE = dimension * sizeof(T); - if (number_of_points != point_ids_to_labels.size()) - { - std::cerr << "Error: number of points in labels file and data file differ." << std::endl; - throw; - } - - tsl::robin_map labels_to_vectors; - tsl::robin_map labels_to_curr_vector; - tsl::robin_map> label_id_to_orig_id; - - for (const auto &lbl : all_labels) - { - uint32_t number_of_label_pts = labels_to_number_of_points[lbl]; - char *vectors = (char *)malloc(number_of_label_pts * VECTOR_SIZE); - if (vectors == nullptr) - { - throw; - } - labels_to_vectors[lbl] = vectors; - labels_to_curr_vector[lbl] = 0; - label_id_to_orig_id[lbl].reserve(number_of_label_pts); - } - - for (uint32_t point_id = 0; point_id < number_of_points; point_id++) - { - char *curr_vector = (char *)malloc(VECTOR_SIZE); - input_data_stream.read(curr_vector, VECTOR_SIZE); - for (const auto &lbl : point_ids_to_labels[point_id]) - { - char *curr_label_vector_ptr = labels_to_vectors[lbl] + (labels_to_curr_vector[lbl] * VECTOR_SIZE); - memcpy(curr_label_vector_ptr, curr_vector, VECTOR_SIZE); - labels_to_curr_vector[lbl]++; - label_id_to_orig_id[lbl].push_back(point_id); - } - free(curr_vector); +tsl::robin_map> +generate_label_specific_vector_files_compat( + path input_data_path, + tsl::robin_map labels_to_number_of_points, + std::vector point_ids_to_labels, label_set all_labels) { + auto file_writing_timer = std::chrono::high_resolution_clock::now(); + std::ifstream input_data_stream(input_data_path); + + uint32_t number_of_points, dimension; + input_data_stream.read((char *)&number_of_points, sizeof(uint32_t)); + input_data_stream.read((char *)&dimension, sizeof(uint32_t)); + const uint32_t VECTOR_SIZE = dimension * sizeof(T); + if (number_of_points != point_ids_to_labels.size()) { + std::cerr << "Error: number of points in labels file and data file differ." + << std::endl; + throw; + } + + tsl::robin_map labels_to_vectors; + tsl::robin_map labels_to_curr_vector; + tsl::robin_map> label_id_to_orig_id; + + for (const auto &lbl : all_labels) { + uint32_t number_of_label_pts = labels_to_number_of_points[lbl]; + char *vectors = (char *)malloc(number_of_label_pts * VECTOR_SIZE); + if (vectors == nullptr) { + throw; } - - for (const auto &lbl : all_labels) - { - path curr_label_input_data_path(input_data_path + "_" + lbl); - uint32_t number_of_label_pts = labels_to_number_of_points[lbl]; - - std::ofstream label_file_stream; - label_file_stream.exceptions(std::ios::badbit | std::ios::failbit); - label_file_stream.open(curr_label_input_data_path, std::ios_base::binary); - label_file_stream.write((char *)&number_of_label_pts, sizeof(uint32_t)); - label_file_stream.write((char *)&dimension, sizeof(uint32_t)); - label_file_stream.write((char *)labels_to_vectors[lbl], number_of_label_pts * VECTOR_SIZE); - - label_file_stream.close(); - free(labels_to_vectors[lbl]); + labels_to_vectors[lbl] = vectors; + labels_to_curr_vector[lbl] = 0; + label_id_to_orig_id[lbl].reserve(number_of_label_pts); + } + + for (uint32_t point_id = 0; point_id < number_of_points; point_id++) { + char *curr_vector = (char *)malloc(VECTOR_SIZE); + input_data_stream.read(curr_vector, VECTOR_SIZE); + for (const auto &lbl : point_ids_to_labels[point_id]) { + char *curr_label_vector_ptr = + labels_to_vectors[lbl] + (labels_to_curr_vector[lbl] * VECTOR_SIZE); + memcpy(curr_label_vector_ptr, curr_vector, VECTOR_SIZE); + labels_to_curr_vector[lbl]++; + label_id_to_orig_id[lbl].push_back(point_id); } - input_data_stream.close(); - - std::chrono::duration file_writing_time = std::chrono::high_resolution_clock::now() - file_writing_timer; - std::cout << "generated " << all_labels.size() << " label-specific vector files for index building in time " - << file_writing_time.count() << "\n" - << std::endl; - - return label_id_to_orig_id; + free(curr_vector); + } + + for (const auto &lbl : all_labels) { + path curr_label_input_data_path(input_data_path + "_" + lbl); + uint32_t number_of_label_pts = labels_to_number_of_points[lbl]; + + std::ofstream label_file_stream; + label_file_stream.exceptions(std::ios::badbit | std::ios::failbit); + label_file_stream.open(curr_label_input_data_path, std::ios_base::binary); + label_file_stream.write((char *)&number_of_label_pts, sizeof(uint32_t)); + label_file_stream.write((char *)&dimension, sizeof(uint32_t)); + label_file_stream.write((char *)labels_to_vectors[lbl], + number_of_label_pts * VECTOR_SIZE); + + label_file_stream.close(); + free(labels_to_vectors[lbl]); + } + input_data_stream.close(); + + std::chrono::duration file_writing_time = + std::chrono::high_resolution_clock::now() - file_writing_timer; + std::cout << "generated " << all_labels.size() + << " label-specific vector files for index building in time " + << file_writing_time.count() << "\n" + << std::endl; + + return label_id_to_orig_id; } /* @@ -144,36 +148,37 @@ tsl::robin_map> generate_label_specific_vecto * * Returns both the graph index and the size of the file in bytes. */ -load_label_index_return_values load_label_index(path label_index_path, uint32_t label_number_of_points) -{ - std::ifstream label_index_stream; - label_index_stream.exceptions(std::ios::badbit | std::ios::failbit); - label_index_stream.open(label_index_path, std::ios::binary); - - uint64_t index_file_size, index_num_frozen_points; - uint32_t index_max_observed_degree, index_entry_point; - const size_t INDEX_METADATA = 2 * sizeof(uint64_t) + 2 * sizeof(uint32_t); - label_index_stream.read((char *)&index_file_size, sizeof(uint64_t)); - label_index_stream.read((char *)&index_max_observed_degree, sizeof(uint32_t)); - label_index_stream.read((char *)&index_entry_point, sizeof(uint32_t)); - label_index_stream.read((char *)&index_num_frozen_points, sizeof(uint64_t)); - size_t bytes_read = INDEX_METADATA; - - std::vector> label_index(label_number_of_points); - uint32_t nodes_read = 0; - while (bytes_read != index_file_size) - { - uint32_t current_node_num_neighbors; - label_index_stream.read((char *)¤t_node_num_neighbors, sizeof(uint32_t)); - nodes_read++; - - std::vector current_node_neighbors(current_node_num_neighbors); - label_index_stream.read((char *)current_node_neighbors.data(), current_node_num_neighbors * sizeof(uint32_t)); - label_index[nodes_read - 1].swap(current_node_neighbors); - bytes_read += sizeof(uint32_t) * (current_node_num_neighbors + 1); - } - - return std::make_tuple(label_index, index_file_size); +load_label_index_return_values load_label_index( + path label_index_path, uint32_t label_number_of_points) { + std::ifstream label_index_stream; + label_index_stream.exceptions(std::ios::badbit | std::ios::failbit); + label_index_stream.open(label_index_path, std::ios::binary); + + uint64_t index_file_size, index_num_frozen_points; + uint32_t index_max_observed_degree, index_entry_point; + const size_t INDEX_METADATA = 2 * sizeof(uint64_t) + 2 * sizeof(uint32_t); + label_index_stream.read((char *)&index_file_size, sizeof(uint64_t)); + label_index_stream.read((char *)&index_max_observed_degree, sizeof(uint32_t)); + label_index_stream.read((char *)&index_entry_point, sizeof(uint32_t)); + label_index_stream.read((char *)&index_num_frozen_points, sizeof(uint64_t)); + size_t bytes_read = INDEX_METADATA; + + std::vector> label_index(label_number_of_points); + uint32_t nodes_read = 0; + while (bytes_read != index_file_size) { + uint32_t current_node_num_neighbors; + label_index_stream.read((char *)¤t_node_num_neighbors, + sizeof(uint32_t)); + nodes_read++; + + std::vector current_node_neighbors(current_node_num_neighbors); + label_index_stream.read((char *)current_node_neighbors.data(), + current_node_num_neighbors * sizeof(uint32_t)); + label_index[nodes_read - 1].swap(current_node_neighbors); + bytes_read += sizeof(uint32_t) * (current_node_num_neighbors + 1); + } + + return std::make_tuple(label_index, index_file_size); } /* @@ -185,100 +190,95 @@ load_label_index_return_values load_label_index(path label_index_path, uint32_t * 2. map: key is label, value is number of points with the label * 3. the label universe as a set */ -parse_label_file_return_values parse_label_file(path label_data_path, std::string universal_label) -{ - std::ifstream label_data_stream(label_data_path); - std::string line, token; - uint32_t line_cnt = 0; - - // allows us to reserve space for the points_to_labels vector - while (std::getline(label_data_stream, line)) - line_cnt++; - label_data_stream.clear(); - label_data_stream.seekg(0, std::ios::beg); - - // values to return - std::vector point_ids_to_labels(line_cnt); - tsl::robin_map labels_to_number_of_points; - label_set all_labels; - - std::vector points_with_universal_label; - line_cnt = 0; - while (std::getline(label_data_stream, line)) - { - std::istringstream current_labels_comma_separated(line); - label_set current_labels; - - // get point id - uint32_t point_id = line_cnt; - - // parse comma separated labels - bool current_universal_label_check = false; - while (getline(current_labels_comma_separated, token, ',')) - { - token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); - token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); - - // if token is empty, there's no labels for the point - if (token == universal_label) - { - points_with_universal_label.push_back(point_id); - current_universal_label_check = true; - } - else - { - all_labels.insert(token); - current_labels.insert(token); - labels_to_number_of_points[token]++; - } - } - - if (current_labels.size() <= 0 && !current_universal_label_check) - { - std::cerr << "Error: " << point_id << " has no labels." << std::endl; - exit(-1); - } - point_ids_to_labels[point_id] = current_labels; - line_cnt++; +parse_label_file_return_values parse_label_file(path label_data_path, + std::string universal_label) { + std::ifstream label_data_stream(label_data_path); + std::string line, token; + uint32_t line_cnt = 0; + + // allows us to reserve space for the points_to_labels vector + while (std::getline(label_data_stream, line)) line_cnt++; + label_data_stream.clear(); + label_data_stream.seekg(0, std::ios::beg); + + // values to return + std::vector point_ids_to_labels(line_cnt); + tsl::robin_map labels_to_number_of_points; + label_set all_labels; + + std::vector points_with_universal_label; + line_cnt = 0; + while (std::getline(label_data_stream, line)) { + std::istringstream current_labels_comma_separated(line); + label_set current_labels; + + // get point id + uint32_t point_id = line_cnt; + + // parse comma separated labels + bool current_universal_label_check = false; + while (getline(current_labels_comma_separated, token, ',')) { + token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); + token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); + + // if token is empty, there's no labels for the point + if (token == universal_label) { + points_with_universal_label.push_back(point_id); + current_universal_label_check = true; + } else { + all_labels.insert(token); + current_labels.insert(token); + labels_to_number_of_points[token]++; + } } - // for every point with universal label, set its label set to all labels - // also, increment the count for number of points a label has - for (const auto &point_id : points_with_universal_label) - { - point_ids_to_labels[point_id] = all_labels; - for (const auto &lbl : all_labels) - labels_to_number_of_points[lbl]++; + if (current_labels.size() <= 0 && !current_universal_label_check) { + std::cerr << "Error: " << point_id << " has no labels." << std::endl; + exit(-1); } - - std::cout << "Identified " << all_labels.size() << " distinct label(s) for " << point_ids_to_labels.size() - << " points\n" - << std::endl; - - return std::make_tuple(point_ids_to_labels, labels_to_number_of_points, all_labels); + point_ids_to_labels[point_id] = current_labels; + line_cnt++; + } + + // for every point with universal label, set its label set to all labels + // also, increment the count for number of points a label has + for (const auto &point_id : points_with_universal_label) { + point_ids_to_labels[point_id] = all_labels; + for (const auto &lbl : all_labels) labels_to_number_of_points[lbl]++; + } + + std::cout << "Identified " << all_labels.size() << " distinct label(s) for " + << point_ids_to_labels.size() << " points\n" + << std::endl; + + return std::make_tuple(point_ids_to_labels, labels_to_number_of_points, + all_labels); } -template DISKANN_DLLEXPORT void generate_label_indices(path input_data_path, path final_index_path_prefix, - label_set all_labels, uint32_t R, uint32_t L, float alpha, - uint32_t num_threads); -template DISKANN_DLLEXPORT void generate_label_indices(path input_data_path, path final_index_path_prefix, - label_set all_labels, uint32_t R, uint32_t L, - float alpha, uint32_t num_threads); -template DISKANN_DLLEXPORT void generate_label_indices(path input_data_path, path final_index_path_prefix, - label_set all_labels, uint32_t R, uint32_t L, - float alpha, uint32_t num_threads); +template DISKANN_DLLEXPORT void generate_label_indices( + path input_data_path, path final_index_path_prefix, label_set all_labels, + uint32_t R, uint32_t L, float alpha, uint32_t num_threads); +template DISKANN_DLLEXPORT void generate_label_indices( + path input_data_path, path final_index_path_prefix, label_set all_labels, + uint32_t R, uint32_t L, float alpha, uint32_t num_threads); +template DISKANN_DLLEXPORT void generate_label_indices( + path input_data_path, path final_index_path_prefix, label_set all_labels, + uint32_t R, uint32_t L, float alpha, uint32_t num_threads); template DISKANN_DLLEXPORT tsl::robin_map> -generate_label_specific_vector_files_compat(path input_data_path, - tsl::robin_map labels_to_number_of_points, - std::vector point_ids_to_labels, label_set all_labels); +generate_label_specific_vector_files_compat( + path input_data_path, + tsl::robin_map labels_to_number_of_points, + std::vector point_ids_to_labels, label_set all_labels); template DISKANN_DLLEXPORT tsl::robin_map> -generate_label_specific_vector_files_compat(path input_data_path, - tsl::robin_map labels_to_number_of_points, - std::vector point_ids_to_labels, label_set all_labels); +generate_label_specific_vector_files_compat( + path input_data_path, + tsl::robin_map labels_to_number_of_points, + std::vector point_ids_to_labels, label_set all_labels); template DISKANN_DLLEXPORT tsl::robin_map> -generate_label_specific_vector_files_compat(path input_data_path, - tsl::robin_map labels_to_number_of_points, - std::vector point_ids_to_labels, label_set all_labels); +generate_label_specific_vector_files_compat( + path input_data_path, + tsl::robin_map labels_to_number_of_points, + std::vector point_ids_to_labels, label_set all_labels); -} // namespace diskann \ No newline at end of file +} // namespace diskann \ No newline at end of file diff --git a/src/in_mem_data_store.cpp b/src/in_mem_data_store.cpp index 470acc342..690ec60b5 100644 --- a/src/in_mem_data_store.cpp +++ b/src/in_mem_data_store.cpp @@ -6,362 +6,364 @@ #include "utils.h" -namespace diskann -{ +namespace diskann { template -InMemDataStore::InMemDataStore(const location_t num_points, const size_t dim, - std::shared_ptr> distance_metric) - : AbstractDataStore(num_points, dim), _distance_fn(distance_metric) -{ - _aligned_dim = ROUND_UP(dim, _distance_fn->get_required_alignment()); - alloc_aligned(((void **)&_data), this->_capacity * _aligned_dim * sizeof(data_t), 8 * sizeof(data_t)); - std::memset(_data, 0, this->_capacity * _aligned_dim * sizeof(data_t)); +InMemDataStore::InMemDataStore( + const location_t num_points, const size_t dim, + std::shared_ptr> distance_metric) + : AbstractDataStore(num_points, dim), + _distance_fn(distance_metric) { + _aligned_dim = ROUND_UP(dim, _distance_fn->get_required_alignment()); + alloc_aligned(((void **)&_data), + this->_capacity * _aligned_dim * sizeof(data_t), + 8 * sizeof(data_t)); + std::memset(_data, 0, this->_capacity * _aligned_dim * sizeof(data_t)); } -template InMemDataStore::~InMemDataStore() -{ - if (_data != nullptr) - { - aligned_free(this->_data); - } +template +InMemDataStore::~InMemDataStore() { + if (_data != nullptr) { + aligned_free(this->_data); + } } -template size_t InMemDataStore::get_aligned_dim() const -{ - return _aligned_dim; +template +size_t InMemDataStore::get_aligned_dim() const { + return _aligned_dim; } -template size_t InMemDataStore::get_alignment_factor() const -{ - return _distance_fn->get_required_alignment(); +template +size_t InMemDataStore::get_alignment_factor() const { + return _distance_fn->get_required_alignment(); } -template location_t InMemDataStore::load(const std::string &filename) -{ - return load_impl(filename); +template +location_t InMemDataStore::load(const std::string &filename) { + return load_impl(filename); } #ifdef EXEC_ENV_OLS -template location_t Index::load_impl(AlignedFileReader &reader) -{ - size_t file_dim, file_num_points; - - diskann::get_bin_metadata(reader, file_num_points, file_dim); - - if (file_dim != _dim) - { - std::stringstream stream; - stream << "ERROR: Driver requests loading " << _dim << " dimension," - << "but file has " << file_dim << " dimension." << std::endl; - diskann::cerr << stream.str() << std::endl; - aligned_free(_data); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } +template +location_t Index::load_impl(AlignedFileReader &reader) { + size_t file_dim, file_num_points; - if (file_num_points > _max_points + _num_frozen_pts) - { - resize(file_num_points - _num_frozen_pts); - } + diskann::get_bin_metadata(reader, file_num_points, file_dim); + + if (file_dim != _dim) { + std::stringstream stream; + stream << "ERROR: Driver requests loading " << _dim << " dimension," + << "but file has " << file_dim << " dimension." << std::endl; + diskann::cerr << stream.str() << std::endl; + aligned_free(_data); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } + + if (file_num_points > _max_points + _num_frozen_pts) { + resize(file_num_points - _num_frozen_pts); + } - return file_num_points; + return file_num_points; } #endif -template location_t InMemDataStore::load_impl(const std::string &filename) -{ - size_t file_dim, file_num_points; - if (!file_exists(filename)) - { - std::stringstream stream; - stream << "ERROR: data file " << filename << " does not exist." << std::endl; - diskann::cerr << stream.str() << std::endl; - aligned_free(_data); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - diskann::get_bin_metadata(filename, file_num_points, file_dim); - - if (file_dim != this->_dim) - { - std::stringstream stream; - stream << "ERROR: Driver requests loading " << this->_dim << " dimension," - << "but file has " << file_dim << " dimension." << std::endl; - diskann::cerr << stream.str() << std::endl; - aligned_free(_data); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } +template +location_t InMemDataStore::load_impl(const std::string &filename) { + size_t file_dim, file_num_points; + if (!file_exists(filename)) { + std::stringstream stream; + stream << "ERROR: data file " << filename << " does not exist." + << std::endl; + diskann::cerr << stream.str() << std::endl; + aligned_free(_data); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } + diskann::get_bin_metadata(filename, file_num_points, file_dim); + + if (file_dim != this->_dim) { + std::stringstream stream; + stream << "ERROR: Driver requests loading " << this->_dim << " dimension," + << "but file has " << file_dim << " dimension." << std::endl; + diskann::cerr << stream.str() << std::endl; + aligned_free(_data); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } - if (file_num_points > this->capacity()) - { - this->resize(file_num_points); - } + if (file_num_points > this->capacity()) { + this->resize(file_num_points); + } - copy_aligned_data_from_file(filename.c_str(), _data, file_num_points, file_dim, _aligned_dim); + copy_aligned_data_from_file(filename.c_str(), _data, file_num_points, + file_dim, _aligned_dim); - return file_num_points; + return file_num_points; } -template size_t InMemDataStore::save(const std::string &filename, const location_t num_points) -{ - return save_data_in_base_dimensions(filename, _data, num_points, this->get_dims(), this->get_aligned_dim(), 0U); +template +size_t InMemDataStore::save(const std::string &filename, + const location_t num_points) { + return save_data_in_base_dimensions(filename, _data, num_points, + this->get_dims(), this->get_aligned_dim(), + 0U); } -template void InMemDataStore::populate_data(const data_t *vectors, const location_t num_pts) -{ - for (location_t i = 0; i < num_pts; i++) - { - memset(_data + i * _aligned_dim, 0, _aligned_dim * sizeof(data_t)); - std::memmove(_data + i * _aligned_dim, vectors + i * this->_dim, this->_dim * sizeof(data_t)); - } - - if (_distance_fn->preprocessing_required()) - { - _distance_fn->preprocess_base_points(_data, this->_dim, num_pts); - } +template +void InMemDataStore::populate_data(const data_t *vectors, + const location_t num_pts) { + for (location_t i = 0; i < num_pts; i++) { + memset(_data + i * _aligned_dim, 0, _aligned_dim * sizeof(data_t)); + std::memmove(_data + i * _aligned_dim, vectors + i * this->_dim, + this->_dim * sizeof(data_t)); + } + + if (_distance_fn->preprocessing_required()) { + _distance_fn->preprocess_base_points(_data, this->_aligned_dim, num_pts); + } } -template void InMemDataStore::populate_data(const std::string &filename, const size_t offset) -{ - size_t npts, ndim; - copy_aligned_data_from_file(filename.c_str(), _data, npts, ndim, _aligned_dim, offset); - - if ((location_t)npts > this->capacity()) - { - std::stringstream ss; - ss << "Number of points in the file: " << filename - << " is greater than the capacity of data store: " << this->capacity() - << ". Must invoke resize before calling populate_data()" << std::endl; - throw diskann::ANNException(ss.str(), -1); - } - - if ((location_t)ndim != this->get_dims()) - { - std::stringstream ss; - ss << "Number of dimensions of a point in the file: " << filename - << " is not equal to dimensions of data store: " << this->capacity() << "." << std::endl; - throw diskann::ANNException(ss.str(), -1); - } - - if (_distance_fn->preprocessing_required()) - { - _distance_fn->preprocess_base_points(_data, this->_dim, this->capacity()); - } +template +void InMemDataStore::populate_data(const std::string &filename, + const size_t offset) { + size_t npts, ndim; + copy_aligned_data_from_file(filename.c_str(), _data, npts, ndim, _aligned_dim, + offset); + + if ((location_t)npts > this->capacity()) { + std::stringstream ss; + ss << "Number of points in the file: " << filename + << " is greater than the capacity of data store: " << this->capacity() + << ". Must invoke resize before calling populate_data()" << std::endl; + throw diskann::ANNException(ss.str(), -1); + } + + if ((location_t)ndim != this->get_dims()) { + std::stringstream ss; + ss << "Number of dimensions of a point in the file: " << filename + << " is not equal to dimensions of data store: " << this->capacity() + << "." << std::endl; + throw diskann::ANNException(ss.str(), -1); + } + + if (_distance_fn->preprocessing_required()) { + _distance_fn->preprocess_base_points(_data, this->_aligned_dim, + this->capacity()); + } } template -void InMemDataStore::save_data_to_bin(const std::string &filename, const location_t num_points) -{ - save_data_in_base_dimensions(filename, _data, num_points, this->get_dims(), this->get_aligned_dim(), 0U); +void InMemDataStore::save_data_to_bin(const std::string &filename, + const location_t num_points) { + save_data_in_base_dimensions(filename, _data, num_points, this->get_dims(), + this->get_aligned_dim(), 0U); } -template void InMemDataStore::get_vector(const location_t i, data_t *dest) const -{ - memcpy(dest, _data + i * _aligned_dim, this->_dim * sizeof(data_t)); +template +void InMemDataStore::get_vector(const location_t i, + data_t *dest) const { + memcpy(dest, _data + i * _aligned_dim, this->_dim * sizeof(data_t)); } -template void InMemDataStore::set_vector(const location_t loc, const data_t *const vector) -{ - size_t offset_in_data = loc * _aligned_dim; - memset(_data + offset_in_data, 0, _aligned_dim * sizeof(data_t)); - memcpy(_data + offset_in_data, vector, this->_dim * sizeof(data_t)); - if (_distance_fn->preprocessing_required()) - { - _distance_fn->preprocess_base_points(_data + offset_in_data, _aligned_dim, 1); - } +template +void InMemDataStore::set_vector(const location_t loc, + const data_t *const vector) { + size_t offset_in_data = loc * _aligned_dim; + memset(_data + offset_in_data, 0, _aligned_dim * sizeof(data_t)); + memcpy(_data + offset_in_data, vector, this->_dim * sizeof(data_t)); + if (_distance_fn->preprocessing_required()) { + _distance_fn->preprocess_base_points(_data + offset_in_data, _aligned_dim, + 1); + } } -template void InMemDataStore::prefetch_vector(const location_t loc) -{ - diskann::prefetch_vector((const char *)_data + _aligned_dim * (size_t)loc, sizeof(data_t) * _aligned_dim); +template +void InMemDataStore::prefetch_vector(const location_t loc) { + diskann::prefetch_vector((const char *)_data + _aligned_dim * (size_t)loc, + sizeof(data_t) * _aligned_dim); } -template float InMemDataStore::get_distance(const data_t *query, const location_t loc) const -{ - return _distance_fn->compare(query, _data + _aligned_dim * loc, _aligned_dim); +template +float InMemDataStore::get_distance(const data_t *query, + const location_t loc) const { + return _distance_fn->compare(query, _data + _aligned_dim * loc, _aligned_dim); } template -void InMemDataStore::get_distance(const data_t *query, const location_t *locations, - const uint32_t location_count, float *distances) const -{ - for (location_t i = 0; i < location_count; i++) - { - distances[i] = _distance_fn->compare(query, _data + locations[i] * _aligned_dim, this->_aligned_dim); - } +void InMemDataStore::get_distance(const data_t *query, + const location_t *locations, + const uint32_t location_count, + float *distances) const { + for (location_t i = 0; i < location_count; i++) { + distances[i] = _distance_fn->compare( + query, _data + locations[i] * _aligned_dim, this->_aligned_dim); + } } template -float InMemDataStore::get_distance(const location_t loc1, const location_t loc2) const -{ - return _distance_fn->compare(_data + loc1 * _aligned_dim, _data + loc2 * _aligned_dim, this->_aligned_dim); +float InMemDataStore::get_distance(const location_t loc1, + const location_t loc2) const { + return _distance_fn->compare(_data + loc1 * _aligned_dim, + _data + loc2 * _aligned_dim, this->_aligned_dim); } -template location_t InMemDataStore::expand(const location_t new_size) -{ - 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); - } +template +location_t InMemDataStore::expand(const location_t new_size) { + 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, _data, this->capacity() * _aligned_dim * sizeof(data_t)); - aligned_free(_data); - _data = new_data; + data_t *new_data; + alloc_aligned((void **)&new_data, new_size * _aligned_dim * sizeof(data_t), + 8 * sizeof(data_t)); + memcpy(new_data, _data, this->capacity() * _aligned_dim * sizeof(data_t)); + aligned_free(_data); + _data = new_data; #else - realloc_aligned((void **)&_data, new_size * _aligned_dim * sizeof(data_t), 8 * sizeof(data_t)); + realloc_aligned((void **)&_data, new_size * _aligned_dim * sizeof(data_t), + 8 * sizeof(data_t)); #endif - this->_capacity = new_size; - return this->_capacity; + this->_capacity = new_size; + return this->_capacity; } -template location_t InMemDataStore::shrink(const location_t new_size) -{ - 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); - } +template +location_t InMemDataStore::shrink(const location_t new_size) { + 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, _data, new_size * _aligned_dim * sizeof(data_t)); - aligned_free(_data); - _data = new_data; + data_t *new_data; + alloc_aligned((void **)&new_data, new_size * _aligned_dim * sizeof(data_t), + 8 * sizeof(data_t)); + memcpy(new_data, _data, new_size * _aligned_dim * sizeof(data_t)); + aligned_free(_data); + _data = new_data; #else - realloc_aligned((void **)&_data, new_size * _aligned_dim * sizeof(data_t), 8 * sizeof(data_t)); + realloc_aligned((void **)&_data, new_size * _aligned_dim * sizeof(data_t), + 8 * sizeof(data_t)); #endif - this->_capacity = new_size; - return this->_capacity; + this->_capacity = new_size; + return this->_capacity; } template -void InMemDataStore::reposition_points(const location_t old_location_start, const location_t new_location_start, - const location_t num_locations) -{ - if (num_locations == 0 || old_location_start == new_location_start) - { - return; +void InMemDataStore::reposition_points( + const location_t old_location_start, const location_t new_location_start, + const location_t num_locations) { + if (num_locations == 0 || old_location_start == new_location_start) { + return; + } + + /* // Update pointers to the moved nodes. Note: the computation is correct + even + // when new_location_start < old_location_start given the C++ uint32_t + // integer arithmetic rules. + const uint32_t location_delta = new_location_start - old_location_start; + */ + // 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; } - - /* // Update pointers to the moved nodes. Note: the computation is correct even - // when new_location_start < old_location_start given the C++ uint32_t - // integer arithmetic rules. - const uint32_t location_delta = new_location_start - old_location_start; - */ - // 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; - } + } 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_points(old_location_start, new_location_start, num_locations); - memset(_data + _aligned_dim * mem_clear_loc_start, 0, - sizeof(data_t) * _aligned_dim * (mem_clear_loc_end_limit - mem_clear_loc_start)); + // Use memmove to handle overlapping ranges. + copy_points(old_location_start, new_location_start, num_locations); + memset(_data + _aligned_dim * mem_clear_loc_start, 0, + sizeof(data_t) * _aligned_dim * + (mem_clear_loc_end_limit - mem_clear_loc_start)); } template -void InMemDataStore::copy_points(const location_t from_loc, const location_t to_loc, - const location_t num_points) -{ - assert(from_loc < this->_capacity); - assert(to_loc < this->_capacity); - assert(num_points < this->_capacity); - memmove(_data + _aligned_dim * to_loc, _data + _aligned_dim * from_loc, num_points * _aligned_dim * sizeof(data_t)); +void InMemDataStore::copy_points(const location_t from_loc, + const location_t to_loc, + const location_t num_points) { + assert(from_loc < this->_capacity); + assert(to_loc < this->_capacity); + assert(num_points < this->_capacity); + memmove(_data + _aligned_dim * to_loc, _data + _aligned_dim * from_loc, + num_points * _aligned_dim * sizeof(data_t)); } -template location_t InMemDataStore::calculate_medoid() const -{ - // allocate and init centroid - float *center = new float[_aligned_dim]; - for (size_t j = 0; j < _aligned_dim; j++) - center[j] = 0; - - for (size_t i = 0; i < this->capacity(); i++) - for (size_t j = 0; j < _aligned_dim; j++) - center[j] += (float)_data[i * _aligned_dim + j]; +template +location_t InMemDataStore::calculate_medoid() const { + // allocate and init centroid + float *center = new float[_aligned_dim]; + for (size_t j = 0; j < _aligned_dim; j++) center[j] = 0; + for (size_t i = 0; i < this->capacity(); i++) for (size_t j = 0; j < _aligned_dim; j++) - center[j] /= (float)this->capacity(); - - // compute all to one distance - float *distances = new float[this->capacity()]; - - // TODO: REFACTOR. Removing pragma might make this slow. Must revisit. - // Problem is that we need to pass num_threads here, it is not clear - // if data store must be aware of threads! - // #pragma omp parallel for schedule(static, 65536) - for (int64_t i = 0; i < (int64_t)this->capacity(); i++) - { - // extract point and distance reference - float &dist = distances[i]; - const data_t *cur_vec = _data + (i * (size_t)_aligned_dim); - dist = 0; - float diff = 0; - for (size_t j = 0; j < _aligned_dim; j++) - { - diff = (center[j] - (float)cur_vec[j]) * (center[j] - (float)cur_vec[j]); - dist += diff; - } + center[j] += (float)_data[i * _aligned_dim + j]; + + for (size_t j = 0; j < _aligned_dim; j++) + center[j] /= (float)this->capacity(); + + // compute all to one distance + float *distances = new float[this->capacity()]; + + // TODO: REFACTOR. Removing pragma might make this slow. Must revisit. + // Problem is that we need to pass num_threads here, it is not clear + // if data store must be aware of threads! + // #pragma omp parallel for schedule(static, 65536) + for (int64_t i = 0; i < (int64_t)this->capacity(); i++) { + // extract point and distance reference + float &dist = distances[i]; + const data_t *cur_vec = _data + (i * (size_t)_aligned_dim); + dist = 0; + float diff = 0; + for (size_t j = 0; j < _aligned_dim; j++) { + diff = (center[j] - (float)cur_vec[j]) * (center[j] - (float)cur_vec[j]); + dist += diff; } - // find imin - uint32_t min_idx = 0; - float min_dist = distances[0]; - for (uint32_t i = 1; i < this->capacity(); i++) - { - if (distances[i] < min_dist) - { - min_idx = i; - min_dist = distances[i]; - } + } + // find imin + uint32_t min_idx = 0; + float min_dist = distances[0]; + for (uint32_t i = 1; i < this->capacity(); i++) { + if (distances[i] < min_dist) { + min_idx = i; + min_dist = distances[i]; } + } - delete[] distances; - delete[] center; - return min_idx; + delete[] distances; + delete[] center; + return min_idx; } -template Distance *InMemDataStore::get_dist_fn() -{ - return this->_distance_fn.get(); +template +Distance *InMemDataStore::get_dist_fn() { + return this->_distance_fn.get(); } template DISKANN_DLLEXPORT class InMemDataStore; template DISKANN_DLLEXPORT class InMemDataStore; template DISKANN_DLLEXPORT class InMemDataStore; -} // namespace diskann \ No newline at end of file +} // namespace diskann \ No newline at end of file diff --git a/src/in_mem_graph_store.cpp b/src/in_mem_graph_store.cpp index cb83481d3..ad9efbbf4 100644 --- a/src/in_mem_graph_store.cpp +++ b/src/in_mem_graph_store.cpp @@ -4,26 +4,18 @@ #include "in_mem_graph_store.h" #include "utils.h" -namespace diskann -{ +namespace diskann { -InMemGraphStore::InMemGraphStore(const size_t max_pts) : AbstractGraphStore(max_pts) -{ -} +InMemGraphStore::InMemGraphStore(const size_t max_pts) + : AbstractGraphStore(max_pts) {} -int InMemGraphStore::load(const std::string &index_path_prefix) -{ -} -int InMemGraphStore::store(const std::string &index_path_prefix) -{ -} +int InMemGraphStore::load(const std::string &index_path_prefix) {} +int InMemGraphStore::store(const std::string &index_path_prefix) {} -void InMemGraphStore::get_adj_list(const location_t i, std::vector &neighbors) -{ -} +void InMemGraphStore::get_adj_list(const location_t i, + std::vector &neighbors) {} -void InMemGraphStore::set_adj_list(const location_t i, std::vector &neighbors) -{ -} +void InMemGraphStore::set_adj_list(const location_t i, + std::vector &neighbors) {} -} // namespace diskann \ No newline at end of file +} // namespace diskann \ No newline at end of file diff --git a/src/linux_aligned_file_reader.cpp b/src/linux_aligned_file_reader.cpp index 47c7cb1fb..5fef8384f 100644 --- a/src/linux_aligned_file_reader.cpp +++ b/src/linux_aligned_file_reader.cpp @@ -10,215 +10,188 @@ #include "utils.h" #define MAX_EVENTS 1024 -namespace -{ +namespace { typedef struct io_event io_event_t; typedef struct iocb iocb_t; -void execute_io(io_context_t ctx, int fd, std::vector &read_reqs, uint64_t n_retries = 0) -{ +void execute_io(io_context_t ctx, int fd, std::vector &read_reqs, + uint64_t n_retries = 0) { #ifdef DEBUG - for (auto &req : read_reqs) - { - assert(IS_ALIGNED(req.len, 512)); - // std::cout << "request:"<= req.len); - } + for (auto &req : read_reqs) { + assert(IS_ALIGNED(req.len, 512)); + // std::cout << "request:"<= req.len); + } #endif - // break-up requests into chunks of size MAX_EVENTS each - uint64_t n_iters = ROUND_UP(read_reqs.size(), MAX_EVENTS) / MAX_EVENTS; - for (uint64_t iter = 0; iter < n_iters; iter++) - { - uint64_t n_ops = std::min((uint64_t)read_reqs.size() - (iter * MAX_EVENTS), (uint64_t)MAX_EVENTS); - std::vector cbs(n_ops, nullptr); - std::vector evts(n_ops); - std::vector cb(n_ops); - for (uint64_t j = 0; j < n_ops; j++) - { - io_prep_pread(cb.data() + j, fd, read_reqs[j + iter * MAX_EVENTS].buf, read_reqs[j + iter * MAX_EVENTS].len, - read_reqs[j + iter * MAX_EVENTS].offset); - } + // break-up requests into chunks of size MAX_EVENTS each + uint64_t n_iters = ROUND_UP(read_reqs.size(), MAX_EVENTS) / MAX_EVENTS; + for (uint64_t iter = 0; iter < n_iters; iter++) { + uint64_t n_ops = std::min((uint64_t)read_reqs.size() - (iter * MAX_EVENTS), + (uint64_t)MAX_EVENTS); + std::vector cbs(n_ops, nullptr); + std::vector evts(n_ops); + std::vector cb(n_ops); + for (uint64_t j = 0; j < n_ops; j++) { + io_prep_pread(cb.data() + j, fd, read_reqs[j + iter * MAX_EVENTS].buf, + read_reqs[j + iter * MAX_EVENTS].len, + read_reqs[j + iter * MAX_EVENTS].offset); + } - // initialize `cbs` using `cb` array - // + // initialize `cbs` using `cb` array + // - for (uint64_t i = 0; i < n_ops; i++) - { - cbs[i] = cb.data() + i; - } + for (uint64_t i = 0; i < n_ops; i++) { + cbs[i] = cb.data() + i; + } - uint64_t n_tries = 0; - while (n_tries <= n_retries) - { - // issue reads - int64_t ret = io_submit(ctx, (int64_t)n_ops, cbs.data()); - // if requests didn't get accepted - 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"; - exit(-1); - } - else - { - // wait on io_getevents - ret = io_getevents(ctx, (int64_t)n_ops, (int64_t)n_ops, evts.data(), nullptr); - // if requests didn't complete - if (ret != (int64_t)n_ops) - { - std::cerr << "io_getevents() failed; returned " << ret << ", expected=" << n_ops - << ", ernno=" << errno << "=" << ::strerror(-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); + uint64_t n_tries = 0; + while (n_tries <= n_retries) { + // issue reads + int64_t ret = io_submit(ctx, (int64_t)n_ops, cbs.data()); + // if requests didn't get accepted + 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"; + exit(-1); + } else { + // wait on io_getevents + ret = io_getevents(ctx, (int64_t)n_ops, (int64_t)n_ops, evts.data(), + nullptr); + // if requests didn't complete + if (ret != (int64_t)n_ops) { + std::cerr << "io_getevents() failed; returned " << ret + << ", expected=" << n_ops << ", ernno=" << errno << "=" + << ::strerror(-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); + } + */ + } } -} // namespace - -LinuxAlignedFileReader::LinuxAlignedFileReader() -{ - this->file_desc = -1; -} - -LinuxAlignedFileReader::~LinuxAlignedFileReader() -{ - int64_t ret; - // check to make sure file_desc is closed - ret = ::fcntl(this->file_desc, F_GETFD); - if (ret == -1) - { - if (errno != EBADF) - { - std::cerr << "close() not called" << std::endl; - // close file desc - ret = ::close(this->file_desc); - // error checks - if (ret == -1) - { - std::cerr << "close() failed; returned " << ret << ", errno=" << errno << ":" << ::strerror(errno) - << std::endl; - } - } +} // namespace + +LinuxAlignedFileReader::LinuxAlignedFileReader() { this->file_desc = -1; } + +LinuxAlignedFileReader::~LinuxAlignedFileReader() { + int64_t ret; + // check to make sure file_desc is closed + ret = ::fcntl(this->file_desc, F_GETFD); + if (ret == -1) { + if (errno != EBADF) { + std::cerr << "close() not called" << std::endl; + // close file desc + ret = ::close(this->file_desc); + // error checks + if (ret == -1) { + std::cerr << "close() failed; returned " << ret << ", errno=" << errno + << ":" << ::strerror(errno) << std::endl; + } } + } } -io_context_t &LinuxAlignedFileReader::get_ctx() -{ - std::unique_lock lk(ctx_mut); - // perform checks only in DEBUG mode - if (ctx_map.find(std::this_thread::get_id()) == ctx_map.end()) - { - std::cerr << "bad thread access; returning -1 as io_context_t" << std::endl; - return this->bad_ctx; - } - else - { - return ctx_map[std::this_thread::get_id()]; - } +io_context_t &LinuxAlignedFileReader::get_ctx() { + std::unique_lock lk(ctx_mut); + // perform checks only in DEBUG mode + if (ctx_map.find(std::this_thread::get_id()) == ctx_map.end()) { + std::cerr << "bad thread access; returning -1 as io_context_t" << std::endl; + return this->bad_ctx; + } else { + return ctx_map[std::this_thread::get_id()]; + } } -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; - int ret = io_setup(MAX_EVENTS, &ctx); - if (ret != 0) - { - lk.unlock(); - assert(errno != EAGAIN); - assert(errno != ENOMEM); - std::cerr << "io_setup() failed; returned " << ret << ", errno=" << errno << ":" << ::strerror(errno) - << std::endl; - } - else - { - diskann::cout << "allocating ctx: " << ctx << " to thread-id:" << my_id << std::endl; - ctx_map[my_id] = ctx; - } +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; + int ret = io_setup(MAX_EVENTS, &ctx); + if (ret != 0) { lk.unlock(); + assert(errno != EAGAIN); + assert(errno != ENOMEM); + std::cerr << "io_setup() failed; returned " << ret << ", errno=" << errno + << ":" << ::strerror(errno) << std::endl; + } else { + diskann::cout << "allocating ctx: " << ctx << " to thread-id:" << my_id + << std::endl; + ctx_map[my_id] = ctx; + } + lk.unlock(); } -void LinuxAlignedFileReader::deregister_thread() -{ - auto my_id = std::this_thread::get_id(); - std::unique_lock lk(ctx_mut); - assert(ctx_map.find(my_id) != ctx_map.end()); +void LinuxAlignedFileReader::deregister_thread() { + auto my_id = std::this_thread::get_id(); + std::unique_lock lk(ctx_mut); + assert(ctx_map.find(my_id) != ctx_map.end()); + + lk.unlock(); + io_context_t ctx = this->get_ctx(); + io_destroy(ctx); + // assert(ret == 0); + lk.lock(); + ctx_map.erase(my_id); + std::cerr << "returned ctx from thread-id:" << my_id << std::endl; + lk.unlock(); +} - lk.unlock(); - io_context_t ctx = this->get_ctx(); +void LinuxAlignedFileReader::deregister_all_threads() { + std::unique_lock lk(ctx_mut); + for (auto x = ctx_map.begin(); x != ctx_map.end(); x++) { + io_context_t ctx = x.value(); io_destroy(ctx); // assert(ret == 0); - lk.lock(); - ctx_map.erase(my_id); - 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 x = ctx_map.begin(); x != ctx_map.end(); x++) - { - io_context_t ctx = x.value(); - io_destroy(ctx); - // assert(ret == 0); - // lk.lock(); - // ctx_map.erase(my_id); - // std::cerr << "returned ctx from thread-id:" << my_id << std::endl; - } - ctx_map.clear(); - // lk.unlock(); + // lk.lock(); + // ctx_map.erase(my_id); + // std::cerr << "returned ctx from thread-id:" << my_id << std::endl; + } + ctx_map.clear(); + // lk.unlock(); } -void LinuxAlignedFileReader::open(const std::string &fname) -{ - int flags = O_DIRECT | O_RDONLY | O_LARGEFILE; - this->file_desc = ::open(fname.c_str(), flags); - // error checks - assert(this->file_desc != -1); - std::cerr << "Opened file : " << fname << std::endl; +void LinuxAlignedFileReader::open(const std::string &fname) { + int flags = O_DIRECT | O_RDONLY | O_LARGEFILE; + this->file_desc = ::open(fname.c_str(), flags); + // error checks + assert(this->file_desc != -1); + std::cerr << "Opened file : " << fname << std::endl; } -void LinuxAlignedFileReader::close() -{ - // int64_t ret; +void LinuxAlignedFileReader::close() { + // int64_t ret; - // check to make sure file_desc is closed - ::fcntl(this->file_desc, F_GETFD); - // assert(ret != -1); + // check to make sure file_desc is closed + ::fcntl(this->file_desc, F_GETFD); + // assert(ret != -1); - ::close(this->file_desc); - // assert(ret != -1); + ::close(this->file_desc); + // assert(ret != -1); } -void LinuxAlignedFileReader::read(std::vector &read_reqs, io_context_t &ctx, bool async) -{ - if (async == true) - { - diskann::cout << "Async currently not supported in linux." << std::endl; - } - assert(this->file_desc != -1); - execute_io(ctx, this->file_desc, read_reqs); +void LinuxAlignedFileReader::read(std::vector &read_reqs, + io_context_t &ctx, bool async) { + if (async == true) { + diskann::cout << "Async currently not supported in linux." << std::endl; + } + assert(this->file_desc != -1); + execute_io(ctx, this->file_desc, read_reqs); } diff --git a/src/logger.cpp b/src/logger.cpp index 1444487f7..5ab15b84e 100644 --- a/src/logger.cpp +++ b/src/logger.cpp @@ -7,8 +7,7 @@ #include "logger_impl.h" #include "windows_customizations.h" -namespace diskann -{ +namespace diskann { DISKANN_DLLEXPORT ANNStreamBuf coutBuff(stdout); DISKANN_DLLEXPORT ANNStreamBuf cerrBuff(stderr); @@ -19,85 +18,76 @@ DISKANN_DLLEXPORT std::basic_ostream cerr(&cerrBuff); #ifdef EXEC_ENV_OLS std::function g_logger; -void SetCustomLogger(std::function logger) -{ - g_logger = logger; +void SetCustomLogger(std::function logger) { + g_logger = logger; } #endif -ANNStreamBuf::ANNStreamBuf(FILE *fp) -{ - if (fp == nullptr) - { - throw diskann::ANNException("File pointer passed to ANNStreamBuf() cannot be null", -1); - } - if (fp != stdout && fp != stderr) - { - throw diskann::ANNException("The custom logger only supports stdout and stderr.", -1); - } - _fp = fp; - _logLevel = (_fp == stdout) ? LogLevel::LL_Info : LogLevel::LL_Error; +ANNStreamBuf::ANNStreamBuf(FILE *fp) { + if (fp == nullptr) { + throw diskann::ANNException( + "File pointer passed to ANNStreamBuf() cannot be null", -1); + } + if (fp != stdout && fp != stderr) { + throw diskann::ANNException( + "The custom logger only supports stdout and stderr.", -1); + } + _fp = fp; + _logLevel = (_fp == stdout) ? LogLevel::LL_Info : LogLevel::LL_Error; #ifdef EXEC_ENV_OLS - _buf = new char[BUFFER_SIZE + 1]; // See comment in the header + _buf = new char[BUFFER_SIZE + 1]; // See comment in the header #else - _buf = new char[BUFFER_SIZE]; // See comment in the header + _buf = new char[BUFFER_SIZE]; // See comment in the header #endif - std::memset(_buf, 0, (BUFFER_SIZE) * sizeof(char)); - setp(_buf, _buf + BUFFER_SIZE - 1); + std::memset(_buf, 0, (BUFFER_SIZE) * sizeof(char)); + setp(_buf, _buf + BUFFER_SIZE - 1); } -ANNStreamBuf::~ANNStreamBuf() -{ - sync(); - _fp = nullptr; // we'll not close because we can't. - delete[] _buf; +ANNStreamBuf::~ANNStreamBuf() { + sync(); + _fp = nullptr; // we'll not close because we can't. + delete[] _buf; } -int ANNStreamBuf::overflow(int c) -{ - std::lock_guard lock(_mutex); - if (c != EOF) - { - *pptr() = (char)c; - pbump(1); - } - flush(); - return c; +int ANNStreamBuf::overflow(int c) { + std::lock_guard lock(_mutex); + if (c != EOF) { + *pptr() = (char)c; + pbump(1); + } + flush(); + return c; } -int ANNStreamBuf::sync() -{ - std::lock_guard lock(_mutex); - flush(); - return 0; +int ANNStreamBuf::sync() { + std::lock_guard lock(_mutex); + flush(); + return 0; } -int ANNStreamBuf::underflow() -{ - throw diskann::ANNException("Attempt to read on streambuf meant only for writing.", -1); +int ANNStreamBuf::underflow() { + throw diskann::ANNException( + "Attempt to read on streambuf meant only for writing.", -1); } -int ANNStreamBuf::flush() -{ - const int num = (int)(pptr() - pbase()); - logImpl(pbase(), num); - pbump(-num); - return num; +int ANNStreamBuf::flush() { + const int num = (int)(pptr() - pbase()); + logImpl(pbase(), num); + pbump(-num); + return num; } -void ANNStreamBuf::logImpl(char *str, int num) -{ +void ANNStreamBuf::logImpl(char *str, int num) { #ifdef EXEC_ENV_OLS - str[num] = '\0'; // Safe. See the c'tor. - // Invoke the OLS custom logging function. - if (g_logger) - { - g_logger(_logLevel, str); - } + str[num] = '\0'; // Safe. See the c'tor. + // Invoke the OLS custom logging function. + if (g_logger) { + g_logger(_logLevel, str); + } #else - fwrite(str, sizeof(char), num, _fp); - fflush(_fp); + fwrite(str, sizeof(char), num, _fp); + fflush(_fp); #endif } -} // namespace diskann +} // namespace diskann diff --git a/src/math_utils.cpp b/src/math_utils.cpp index 7481da848..c08fd66b1 100644 --- a/src/math_utils.cpp +++ b/src/math_utils.cpp @@ -8,47 +8,42 @@ #include "logger.h" #include "utils.h" -namespace math_utils -{ - -float calc_distance(float *vec_1, float *vec_2, size_t dim) -{ - float dist = 0; - for (size_t j = 0; j < dim; j++) - { - dist += (vec_1[j] - vec_2[j]) * (vec_1[j] - vec_2[j]); - } - return dist; +namespace math_utils { + +float calc_distance(float *vec_1, float *vec_2, size_t dim) { + float dist = 0; + for (size_t j = 0; j < dim; j++) { + dist += (vec_1[j] - vec_2[j]) * (vec_1[j] - vec_2[j]); + } + return dist; } // compute l2-squared norms of data stored in row major num_points * dim, // needs // to be pre-allocated -void compute_vecs_l2sq(float *vecs_l2sq, float *data, const size_t num_points, const size_t dim) -{ +void compute_vecs_l2sq(float *vecs_l2sq, float *data, const size_t num_points, + const size_t dim) { #pragma omp parallel for schedule(static, 8192) - for (int64_t n_iter = 0; n_iter < (int64_t)num_points; n_iter++) - { - vecs_l2sq[n_iter] = cblas_snrm2((MKL_INT)dim, (data + (n_iter * dim)), 1); - vecs_l2sq[n_iter] *= vecs_l2sq[n_iter]; - } + for (int64_t n_iter = 0; n_iter < (int64_t)num_points; n_iter++) { + vecs_l2sq[n_iter] = cblas_snrm2((MKL_INT)dim, (data + (n_iter * dim)), 1); + vecs_l2sq[n_iter] *= vecs_l2sq[n_iter]; + } } -void rotate_data_randomly(float *data, size_t num_points, size_t dim, float *rot_mat, float *&new_mat, - bool transpose_rot) -{ - CBLAS_TRANSPOSE transpose = CblasNoTrans; - if (transpose_rot) - { - diskann::cout << "Transposing rotation matrix.." << std::flush; - transpose = CblasTrans; - } - diskann::cout << "done Rotating data with random matrix.." << std::flush; +void rotate_data_randomly(float *data, size_t num_points, size_t dim, + float *rot_mat, float *&new_mat, bool transpose_rot) { + CBLAS_TRANSPOSE transpose = CblasNoTrans; + if (transpose_rot) { + diskann::cout << "Transposing rotation matrix.." << std::flush; + transpose = CblasTrans; + } + diskann::cout << "done Rotating data with random matrix.." << std::flush; - cblas_sgemm(CblasRowMajor, CblasNoTrans, transpose, (MKL_INT)num_points, (MKL_INT)dim, (MKL_INT)dim, 1.0, data, - (MKL_INT)dim, rot_mat, (MKL_INT)dim, 0, new_mat, (MKL_INT)dim); + cblas_sgemm(CblasRowMajor, CblasNoTrans, transpose, (MKL_INT)num_points, + (MKL_INT)dim, (MKL_INT)dim, 1.0, data, (MKL_INT)dim, rot_mat, + (MKL_INT)dim, 0, new_mat, (MKL_INT)dim); - diskann::cout << "done." << std::endl; + diskann::cout << "done." << std::endl; } // calculate k closest centers to data of num_points * dim (row major) @@ -61,77 +56,70 @@ void rotate_data_randomly(float *data, size_t num_points, size_t dim, float *rot // Default value of k is 1 // Ideally used only by compute_closest_centers -void compute_closest_centers_in_block(const float *const data, const size_t num_points, const size_t dim, - const float *const centers, const size_t num_centers, - const float *const docs_l2sq, const float *const centers_l2sq, - uint32_t *center_index, float *const dist_matrix, size_t k) -{ - if (k > num_centers) - { - diskann::cout << "ERROR: k (" << k << ") > num_center(" << num_centers << ")" << std::endl; - return; - } - - float *ones_a = new float[num_centers]; - float *ones_b = new float[num_points]; - - for (size_t i = 0; i < num_centers; i++) - { - ones_a[i] = 1.0; - } - for (size_t i = 0; i < num_points; i++) - { - ones_b[i] = 1.0; - } - - cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, (MKL_INT)num_points, (MKL_INT)num_centers, (MKL_INT)1, 1.0f, - docs_l2sq, (MKL_INT)1, ones_a, (MKL_INT)1, 0.0f, dist_matrix, (MKL_INT)num_centers); - - cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, (MKL_INT)num_points, (MKL_INT)num_centers, (MKL_INT)1, 1.0f, - ones_b, (MKL_INT)1, centers_l2sq, (MKL_INT)1, 1.0f, dist_matrix, (MKL_INT)num_centers); - - cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, (MKL_INT)num_points, (MKL_INT)num_centers, (MKL_INT)dim, -2.0f, - data, (MKL_INT)dim, centers, (MKL_INT)dim, 1.0f, dist_matrix, (MKL_INT)num_centers); - - if (k == 1) - { +void compute_closest_centers_in_block( + const float *const data, const size_t num_points, const size_t dim, + const float *const centers, const size_t num_centers, + const float *const docs_l2sq, const float *const centers_l2sq, + uint32_t *center_index, float *const dist_matrix, size_t k) { + if (k > num_centers) { + diskann::cout << "ERROR: k (" << k << ") > num_center(" << num_centers + << ")" << std::endl; + return; + } + + float *ones_a = new float[num_centers]; + float *ones_b = new float[num_points]; + + for (size_t i = 0; i < num_centers; i++) { + ones_a[i] = 1.0; + } + for (size_t i = 0; i < num_points; i++) { + ones_b[i] = 1.0; + } + + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, (MKL_INT)num_points, + (MKL_INT)num_centers, (MKL_INT)1, 1.0f, docs_l2sq, (MKL_INT)1, + ones_a, (MKL_INT)1, 0.0f, dist_matrix, (MKL_INT)num_centers); + + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, (MKL_INT)num_points, + (MKL_INT)num_centers, (MKL_INT)1, 1.0f, ones_b, (MKL_INT)1, + centers_l2sq, (MKL_INT)1, 1.0f, dist_matrix, + (MKL_INT)num_centers); + + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, (MKL_INT)num_points, + (MKL_INT)num_centers, (MKL_INT)dim, -2.0f, data, (MKL_INT)dim, + centers, (MKL_INT)dim, 1.0f, dist_matrix, (MKL_INT)num_centers); + + if (k == 1) { #pragma omp parallel for schedule(static, 8192) - for (int64_t i = 0; i < (int64_t)num_points; i++) - { - float min = std::numeric_limits::max(); - float *current = dist_matrix + (i * num_centers); - for (size_t j = 0; j < num_centers; j++) - { - if (current[j] < min) - { - center_index[i] = (uint32_t)j; - min = current[j]; - } - } + for (int64_t i = 0; i < (int64_t)num_points; i++) { + float min = std::numeric_limits::max(); + float *current = dist_matrix + (i * num_centers); + for (size_t j = 0; j < num_centers; j++) { + if (current[j] < min) { + center_index[i] = (uint32_t)j; + min = current[j]; } + } } - else - { + } else { #pragma omp parallel for schedule(static, 8192) - for (int64_t i = 0; i < (int64_t)num_points; i++) - { - std::priority_queue top_k_queue; - float *current = dist_matrix + (i * num_centers); - for (size_t j = 0; j < num_centers; j++) - { - PivotContainer this_piv(j, current[j]); - top_k_queue.push(this_piv); - } - for (size_t j = 0; j < k; j++) - { - PivotContainer this_piv = top_k_queue.top(); - center_index[i * k + j] = (uint32_t)this_piv.piv_id; - top_k_queue.pop(); - } - } + for (int64_t i = 0; i < (int64_t)num_points; i++) { + std::priority_queue top_k_queue; + float *current = dist_matrix + (i * num_centers); + for (size_t j = 0; j < num_centers; j++) { + PivotContainer this_piv(j, current[j]); + top_k_queue.push(this_piv); + } + for (size_t j = 0; j < k; j++) { + PivotContainer this_piv = top_k_queue.top(); + center_index[i * k + j] = (uint32_t)this_piv.piv_id; + top_k_queue.pop(); + } } - delete[] ones_a; - delete[] ones_b; + } + delete[] ones_a; + delete[] ones_b; } // Given data in num_points * new_dim row major @@ -144,92 +132,93 @@ void compute_closest_centers_in_block(const float *const data, const size_t num_ // indices is an empty vector. Additionally, if pts_norms_squared is not null, // then it will assume that point norms are pre-computed and use those values -void compute_closest_centers(float *data, size_t num_points, size_t dim, float *pivot_data, size_t num_centers, - size_t k, uint32_t *closest_centers_ivf, std::vector *inverted_index, - float *pts_norms_squared) -{ - if (k > num_centers) - { - diskann::cout << "ERROR: k (" << k << ") > num_center(" << num_centers << ")" << std::endl; - return; - } - - bool is_norm_given_for_pts = (pts_norms_squared != NULL); - - float *pivs_norms_squared = new float[num_centers]; - if (!is_norm_given_for_pts) - pts_norms_squared = new float[num_points]; - - 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; - - if (!is_norm_given_for_pts) - math_utils::compute_vecs_l2sq(pts_norms_squared, data, num_points, dim); - math_utils::compute_vecs_l2sq(pivs_norms_squared, pivot_data, num_centers, dim); - uint32_t *closest_centers = new uint32_t[PAR_BLOCK_SIZE * k]; - float *distance_matrix = new float[num_centers * PAR_BLOCK_SIZE]; - - for (size_t cur_blk = 0; cur_blk < N_BLOCKS; cur_blk++) - { - float *data_cur_blk = data + cur_blk * PAR_BLOCK_SIZE * dim; - size_t num_pts_blk = std::min(PAR_BLOCK_SIZE, num_points - cur_blk * PAR_BLOCK_SIZE); - float *pts_norms_blk = pts_norms_squared + cur_blk * PAR_BLOCK_SIZE; - - math_utils::compute_closest_centers_in_block(data_cur_blk, num_pts_blk, dim, pivot_data, num_centers, - pts_norms_blk, pivs_norms_squared, closest_centers, - distance_matrix, k); +void compute_closest_centers(float *data, size_t num_points, size_t dim, + float *pivot_data, size_t num_centers, size_t k, + uint32_t *closest_centers_ivf, + std::vector *inverted_index, + float *pts_norms_squared) { + if (k > num_centers) { + diskann::cout << "ERROR: k (" << k << ") > num_center(" << num_centers + << ")" << std::endl; + return; + } + + bool is_norm_given_for_pts = (pts_norms_squared != NULL); + + float *pivs_norms_squared = new float[num_centers]; + if (!is_norm_given_for_pts) pts_norms_squared = new float[num_points]; + + 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; + + if (!is_norm_given_for_pts) + math_utils::compute_vecs_l2sq(pts_norms_squared, data, num_points, dim); + math_utils::compute_vecs_l2sq(pivs_norms_squared, pivot_data, num_centers, + dim); + uint32_t *closest_centers = new uint32_t[PAR_BLOCK_SIZE * k]; + float *distance_matrix = new float[num_centers * PAR_BLOCK_SIZE]; + + for (size_t cur_blk = 0; cur_blk < N_BLOCKS; cur_blk++) { + float *data_cur_blk = data + cur_blk * PAR_BLOCK_SIZE * dim; + size_t num_pts_blk = + std::min(PAR_BLOCK_SIZE, num_points - cur_blk * PAR_BLOCK_SIZE); + float *pts_norms_blk = pts_norms_squared + cur_blk * PAR_BLOCK_SIZE; + + math_utils::compute_closest_centers_in_block( + data_cur_blk, num_pts_blk, dim, pivot_data, num_centers, pts_norms_blk, + pivs_norms_squared, closest_centers, distance_matrix, k); #pragma omp parallel for schedule(static, 1) - for (int64_t j = cur_blk * PAR_BLOCK_SIZE; - j < std::min((int64_t)num_points, (int64_t)((cur_blk + 1) * PAR_BLOCK_SIZE)); j++) - { - for (size_t l = 0; l < k; l++) - { - size_t this_center_id = closest_centers[(j - cur_blk * PAR_BLOCK_SIZE) * k + l]; - closest_centers_ivf[j * k + l] = (uint32_t)this_center_id; - if (inverted_index != NULL) - { + for (int64_t j = cur_blk * PAR_BLOCK_SIZE; + j < std::min((int64_t)num_points, + (int64_t)((cur_blk + 1) * PAR_BLOCK_SIZE)); + j++) { + for (size_t l = 0; l < k; l++) { + size_t this_center_id = + closest_centers[(j - cur_blk * PAR_BLOCK_SIZE) * k + l]; + closest_centers_ivf[j * k + l] = (uint32_t)this_center_id; + if (inverted_index != NULL) { #pragma omp critical - inverted_index[this_center_id].push_back(j); - } - } + inverted_index[this_center_id].push_back(j); } + } } - delete[] closest_centers; - delete[] distance_matrix; - delete[] pivs_norms_squared; - if (!is_norm_given_for_pts) - delete[] pts_norms_squared; + } + delete[] closest_centers; + delete[] distance_matrix; + delete[] pivs_norms_squared; + if (!is_norm_given_for_pts) delete[] pts_norms_squared; } // if to_subtract is 1, will subtract nearest center from each row. Else will // add. Output will be in data_load iself. // Nearest centers need to be provided in closst_centers. -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; +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 < (int64_t)num_points; n_iter++) - { - for (size_t d_iter = 0; d_iter < dim; d_iter++) - { - if (to_subtract == 1) - data_load[n_iter * dim + d_iter] = - data_load[n_iter * dim + d_iter] - cur_pivot_data[closest_centers[n_iter] * dim + d_iter]; - else - data_load[n_iter * dim + d_iter] = - data_load[n_iter * dim + d_iter] + cur_pivot_data[closest_centers[n_iter] * dim + d_iter]; - } + for (int64_t n_iter = 0; n_iter < (int64_t)num_points; n_iter++) { + for (size_t d_iter = 0; d_iter < dim; d_iter++) { + if (to_subtract == 1) + data_load[n_iter * dim + d_iter] = + data_load[n_iter * dim + d_iter] - + cur_pivot_data[closest_centers[n_iter] * dim + d_iter]; + else + data_load[n_iter * dim + d_iter] = + data_load[n_iter * dim + d_iter] + + cur_pivot_data[closest_centers[n_iter] * dim + d_iter]; } + } } -} // namespace math_utils +} // namespace math_utils -namespace kmeans -{ +namespace kmeans { // run Lloyds one iteration // Given data in row major num_points * dim, and centers in row major @@ -239,67 +228,64 @@ namespace kmeans // closest_centers == NULL, will allocate memory and return. Similarly, if // closest_docs == NULL, will allocate memory and return. -float lloyds_iter(float *data, size_t num_points, size_t dim, float *centers, size_t num_centers, float *docs_l2sq, - std::vector *closest_docs, uint32_t *&closest_center) -{ - bool compute_residual = true; - // Timer timer; +float lloyds_iter(float *data, size_t num_points, size_t dim, float *centers, + size_t num_centers, float *docs_l2sq, + std::vector *closest_docs, + uint32_t *&closest_center) { + bool compute_residual = true; + // Timer timer; - if (closest_center == NULL) - closest_center = new uint32_t[num_points]; - if (closest_docs == NULL) - closest_docs = new std::vector[num_centers]; - else - for (size_t c = 0; c < num_centers; ++c) - closest_docs[c].clear(); + if (closest_center == NULL) closest_center = new uint32_t[num_points]; + if (closest_docs == NULL) + closest_docs = new std::vector[num_centers]; + else + for (size_t c = 0; c < num_centers; ++c) closest_docs[c].clear(); - math_utils::compute_closest_centers(data, num_points, dim, centers, num_centers, 1, closest_center, closest_docs, - docs_l2sq); + math_utils::compute_closest_centers(data, num_points, dim, centers, + num_centers, 1, closest_center, + closest_docs, docs_l2sq); - memset(centers, 0, sizeof(float) * (size_t)num_centers * (size_t)dim); + memset(centers, 0, sizeof(float) * (size_t)num_centers * (size_t)dim); #pragma omp parallel for schedule(static, 1) - for (int64_t c = 0; c < (int64_t)num_centers; ++c) - { - float *center = centers + (size_t)c * (size_t)dim; - double *cluster_sum = new double[dim]; - for (size_t i = 0; i < dim; i++) - cluster_sum[i] = 0.0; - for (size_t i = 0; i < closest_docs[c].size(); i++) - { - float *current = data + ((closest_docs[c][i]) * dim); - for (size_t j = 0; j < dim; j++) - { - cluster_sum[j] += (double)current[j]; - } - } - if (closest_docs[c].size() > 0) - { - for (size_t i = 0; i < dim; i++) - center[i] = (float)(cluster_sum[i] / ((double)closest_docs[c].size())); - } - delete[] cluster_sum; + for (int64_t c = 0; c < (int64_t)num_centers; ++c) { + float *center = centers + (size_t)c * (size_t)dim; + double *cluster_sum = new double[dim]; + for (size_t i = 0; i < dim; i++) cluster_sum[i] = 0.0; + for (size_t i = 0; i < closest_docs[c].size(); i++) { + float *current = data + ((closest_docs[c][i]) * dim); + for (size_t j = 0; j < dim; j++) { + cluster_sum[j] += (double)current[j]; + } } + if (closest_docs[c].size() > 0) { + for (size_t i = 0; i < dim; i++) + center[i] = (float)(cluster_sum[i] / ((double)closest_docs[c].size())); + } + delete[] cluster_sum; + } - float residual = 0.0; - if (compute_residual) - { - size_t BUF_PAD = 32; - size_t CHUNK_SIZE = 2 * 8192; - size_t nchunks = num_points / CHUNK_SIZE + (num_points % CHUNK_SIZE == 0 ? 0 : 1); - std::vector residuals(nchunks * BUF_PAD, 0.0); + float residual = 0.0; + if (compute_residual) { + size_t BUF_PAD = 32; + size_t CHUNK_SIZE = 2 * 8192; + size_t nchunks = + num_points / CHUNK_SIZE + (num_points % CHUNK_SIZE == 0 ? 0 : 1); + std::vector residuals(nchunks * BUF_PAD, 0.0); #pragma omp parallel for schedule(static, 32) - for (int64_t chunk = 0; chunk < (int64_t)nchunks; ++chunk) - for (size_t d = chunk * CHUNK_SIZE; d < num_points && d < (chunk + 1) * CHUNK_SIZE; ++d) - residuals[chunk * BUF_PAD] += - math_utils::calc_distance(data + (d * dim), centers + (size_t)closest_center[d] * (size_t)dim, dim); - - for (size_t chunk = 0; chunk < nchunks; ++chunk) - residual += residuals[chunk * BUF_PAD]; - } - - return residual; + for (int64_t chunk = 0; chunk < (int64_t)nchunks; ++chunk) + for (size_t d = chunk * CHUNK_SIZE; + d < num_points && d < (chunk + 1) * CHUNK_SIZE; ++d) + residuals[chunk * BUF_PAD] += math_utils::calc_distance( + data + (d * dim), centers + (size_t)closest_center[d] * (size_t)dim, + dim); + + for (size_t chunk = 0; chunk < nchunks; ++chunk) + residual += residuals[chunk * BUF_PAD]; + } + + return residual; } // Run Lloyds until max_reps or stopping criterion @@ -309,150 +295,142 @@ float lloyds_iter(float *data, size_t num_points, size_t dim, float *centers, si // vector [num_centers], and closest_center = new size_t[num_points] // Final centers are output in centers as row major num_centers * dim // -float run_lloyds(float *data, size_t num_points, size_t dim, float *centers, const size_t num_centers, - const size_t max_reps, std::vector *closest_docs, uint32_t *closest_center) -{ - float residual = std::numeric_limits::max(); - bool ret_closest_docs = true; - bool ret_closest_center = true; - if (closest_docs == NULL) - { - closest_docs = new std::vector[num_centers]; - ret_closest_docs = false; - } - if (closest_center == NULL) - { - closest_center = new uint32_t[num_points]; - ret_closest_center = false; - } - - float *docs_l2sq = new float[num_points]; - math_utils::compute_vecs_l2sq(docs_l2sq, data, num_points, dim); - - float old_residual; - // Timer timer; - for (size_t i = 0; i < max_reps; ++i) - { - old_residual = residual; - - residual = lloyds_iter(data, num_points, dim, centers, num_centers, docs_l2sq, closest_docs, closest_center); - - if (((i != 0) && ((old_residual - residual) / residual) < 0.00001) || - (residual < std::numeric_limits::epsilon())) - { - diskann::cout << "Residuals unchanged: " << old_residual << " becomes " << residual - << ". Early termination." << std::endl; - break; - } +float run_lloyds(float *data, size_t num_points, size_t dim, float *centers, + const size_t num_centers, const size_t max_reps, + std::vector *closest_docs, uint32_t *closest_center) { + float residual = std::numeric_limits::max(); + bool ret_closest_docs = true; + bool ret_closest_center = true; + if (closest_docs == NULL) { + closest_docs = new std::vector[num_centers]; + ret_closest_docs = false; + } + if (closest_center == NULL) { + closest_center = new uint32_t[num_points]; + ret_closest_center = false; + } + + float *docs_l2sq = new float[num_points]; + math_utils::compute_vecs_l2sq(docs_l2sq, data, num_points, dim); + + float old_residual; + // Timer timer; + for (size_t i = 0; i < max_reps; ++i) { + old_residual = residual; + + residual = lloyds_iter(data, num_points, dim, centers, num_centers, + docs_l2sq, closest_docs, closest_center); + + if (((i != 0) && ((old_residual - residual) / residual) < 0.00001) || + (residual < std::numeric_limits::epsilon())) { + diskann::cout << "Residuals unchanged: " << old_residual << " becomes " + << residual << ". Early termination." << std::endl; + break; } - delete[] docs_l2sq; - if (!ret_closest_docs) - delete[] closest_docs; - if (!ret_closest_center) - delete[] closest_center; - return residual; + } + delete[] docs_l2sq; + if (!ret_closest_docs) delete[] closest_docs; + if (!ret_closest_center) delete[] closest_center; + return residual; } // assumes memory allocated for pivot_data as new // float[num_centers*dim] // and select randomly num_centers points as pivots -void selecting_pivots(float *data, size_t num_points, size_t dim, float *pivot_data, size_t num_centers) -{ - // pivot_data = new float[num_centers * dim]; - - std::vector picked; - std::random_device rd; - auto x = rd(); - std::mt19937 generator(x); - std::uniform_int_distribution distribution(0, num_points - 1); - - size_t tmp_pivot; - for (size_t j = 0; j < num_centers; j++) - { - tmp_pivot = distribution(generator); - if (std::find(picked.begin(), picked.end(), tmp_pivot) != picked.end()) - continue; - picked.push_back(tmp_pivot); - std::memcpy(pivot_data + j * dim, data + tmp_pivot * dim, dim * sizeof(float)); - } +void selecting_pivots(float *data, size_t num_points, size_t dim, + float *pivot_data, size_t num_centers) { + // pivot_data = new float[num_centers * dim]; + + std::vector picked; + std::random_device rd; + auto x = rd(); + std::mt19937 generator(x); + std::uniform_int_distribution distribution(0, num_points - 1); + + size_t tmp_pivot; + for (size_t j = 0; j < num_centers; j++) { + tmp_pivot = distribution(generator); + if (std::find(picked.begin(), picked.end(), tmp_pivot) != picked.end()) + continue; + picked.push_back(tmp_pivot); + std::memcpy(pivot_data + j * dim, data + tmp_pivot * dim, + dim * sizeof(float)); + } } -void kmeanspp_selecting_pivots(float *data, size_t num_points, size_t dim, float *pivot_data, size_t num_centers) -{ - if (num_points > 1 << 23) - { - diskann::cout << "ERROR: n_pts " << num_points - << " currently not supported for k-means++, maximum is " - "8388608. Falling back to random pivot " - "selection." - << std::endl; - selecting_pivots(data, num_points, dim, pivot_data, num_centers); - return; - } - - 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); - size_t init_id = int_dist(generator); - size_t num_picked = 1; - - picked.push_back(init_id); - std::memcpy(pivot_data, data + init_id * dim, dim * sizeof(float)); - - float *dist = new float[num_points]; +void kmeanspp_selecting_pivots(float *data, size_t num_points, size_t dim, + float *pivot_data, size_t num_centers) { + if (num_points > 1 << 23) { + diskann::cout << "ERROR: n_pts " << num_points + << " currently not supported for k-means++, maximum is " + "8388608. Falling back to random pivot " + "selection." + << std::endl; + selecting_pivots(data, num_points, dim, pivot_data, num_centers); + return; + } + + 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); + size_t init_id = int_dist(generator); + size_t num_picked = 1; + + picked.push_back(init_id); + std::memcpy(pivot_data, data + init_id * dim, dim * sizeof(float)); + + float *dist = new float[num_points]; #pragma omp parallel for schedule(static, 8192) - for (int64_t i = 0; i < (int64_t)num_points; i++) - { - dist[i] = math_utils::calc_distance(data + i * dim, data + init_id * dim, dim); - } + for (int64_t i = 0; i < (int64_t)num_points; i++) { + dist[i] = + math_utils::calc_distance(data + i * dim, data + init_id * dim, dim); + } - double dart_val; - size_t tmp_pivot; - bool sum_flag = false; + double dart_val; + size_t tmp_pivot; + bool sum_flag = false; - while (num_picked < num_centers) - { - dart_val = distribution(generator); + while (num_picked < num_centers) { + dart_val = distribution(generator); - double sum = 0; - for (size_t i = 0; i < num_points; i++) - { - sum = sum + dist[i]; - } - if (sum == 0) - sum_flag = true; + double sum = 0; + for (size_t i = 0; i < num_points; i++) { + sum = sum + dist[i]; + } + if (sum == 0) sum_flag = true; - dart_val *= sum; + dart_val *= sum; - 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]) - { - break; - } + 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]) { + break; + } - prefix_sum += dist[i]; - } + prefix_sum += dist[i]; + } - if (std::find(picked.begin(), picked.end(), tmp_pivot) != picked.end() && (sum_flag == false)) - continue; - picked.push_back(tmp_pivot); - std::memcpy(pivot_data + num_picked * dim, data + tmp_pivot * dim, dim * sizeof(float)); + if (std::find(picked.begin(), picked.end(), tmp_pivot) != picked.end() && + (sum_flag == false)) + continue; + picked.push_back(tmp_pivot); + std::memcpy(pivot_data + num_picked * dim, data + tmp_pivot * dim, + dim * sizeof(float)); #pragma omp parallel for schedule(static, 8192) - for (int64_t i = 0; i < (int64_t)num_points; i++) - { - dist[i] = (std::min)(dist[i], math_utils::calc_distance(data + i * dim, data + tmp_pivot * dim, dim)); - } - num_picked++; + for (int64_t i = 0; i < (int64_t)num_points; i++) { + dist[i] = + (std::min)(dist[i], math_utils::calc_distance( + data + i * dim, data + tmp_pivot * dim, dim)); } - delete[] dist; + num_picked++; + } + delete[] dist; } -} // namespace kmeans +} // namespace kmeans diff --git a/src/memory_mapper.cpp b/src/memory_mapper.cpp index d1c5ef984..2f6d4fd1f 100644 --- a/src/memory_mapper.cpp +++ b/src/memory_mapper.cpp @@ -8,100 +8,86 @@ using namespace diskann; -MemoryMapper::MemoryMapper(const std::string &filename) : MemoryMapper(filename.c_str()) -{ -} +MemoryMapper::MemoryMapper(const std::string &filename) + : MemoryMapper(filename.c_str()) {} -MemoryMapper::MemoryMapper(const char *filename) -{ +MemoryMapper::MemoryMapper(const char *filename) { #ifndef _WINDOWS - _fd = open(filename, O_RDONLY); - if (_fd <= 0) - { - std::cerr << "Inner vertices file not found" << std::endl; - return; - } - struct stat sb; - if (fstat(_fd, &sb) != 0) - { - std::cerr << "Inner vertices file not dound. " << std::endl; - return; - } - _fileSize = sb.st_size; - diskann::cout << "File Size: " << _fileSize << std::endl; - _buf = (char *)mmap(NULL, _fileSize, PROT_READ, MAP_PRIVATE, _fd, 0); + _fd = open(filename, O_RDONLY); + if (_fd <= 0) { + std::cerr << "Inner vertices file not found" << std::endl; + return; + } + struct stat sb; + if (fstat(_fd, &sb) != 0) { + std::cerr << "Inner vertices file not dound. " << std::endl; + return; + } + _fileSize = sb.st_size; + diskann::cout << "File Size: " << _fileSize << std::endl; + _buf = (char *)mmap(NULL, _fileSize, PROT_READ, MAP_PRIVATE, _fd, 0); #else - _bareFile = - CreateFileA(filename, GENERIC_READ | GENERIC_EXECUTE, 0, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL); - if (_bareFile == nullptr) - { - std::ostringstream message; - message << "CreateFileA(" << filename << ") failed with error " << GetLastError() << std::endl; - std::cerr << message.str(); - throw std::exception(message.str().c_str()); - } + _bareFile = CreateFileA(filename, GENERIC_READ | GENERIC_EXECUTE, 0, NULL, + OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL); + if (_bareFile == nullptr) { + std::ostringstream message; + message << "CreateFileA(" << filename << ") failed with error " + << GetLastError() << std::endl; + std::cerr << message.str(); + throw std::exception(message.str().c_str()); + } - _fd = CreateFileMapping(_bareFile, NULL, PAGE_EXECUTE_READ, 0, 0, NULL); - if (_fd == nullptr) - { - std::ostringstream message; - message << "CreateFileMapping(" << filename << ") failed with error " << GetLastError() << std::endl; - std::cerr << message.str() << std::endl; - throw std::exception(message.str().c_str()); - } + _fd = CreateFileMapping(_bareFile, NULL, PAGE_EXECUTE_READ, 0, 0, NULL); + if (_fd == nullptr) { + std::ostringstream message; + message << "CreateFileMapping(" << filename << ") failed with error " + << GetLastError() << std::endl; + std::cerr << message.str() << std::endl; + throw std::exception(message.str().c_str()); + } - _buf = (char *)MapViewOfFile(_fd, FILE_MAP_READ, 0, 0, 0); - if (_buf == nullptr) - { - std::ostringstream message; - message << "MapViewOfFile(" << filename << ") failed with error: " << GetLastError() << std::endl; - std::cerr << message.str() << std::endl; - throw std::exception(message.str().c_str()); - } + _buf = (char *)MapViewOfFile(_fd, FILE_MAP_READ, 0, 0, 0); + if (_buf == nullptr) { + std::ostringstream message; + message << "MapViewOfFile(" << filename + << ") failed with error: " << GetLastError() << std::endl; + std::cerr << message.str() << std::endl; + throw std::exception(message.str().c_str()); + } - LARGE_INTEGER fSize; - if (TRUE == GetFileSizeEx(_bareFile, &fSize)) - { - _fileSize = fSize.QuadPart; // take the 64-bit value - diskann::cout << "File Size: " << _fileSize << std::endl; - } - else - { - std::cerr << "Failed to get size of file " << filename << std::endl; - } + LARGE_INTEGER fSize; + if (TRUE == GetFileSizeEx(_bareFile, &fSize)) { + _fileSize = fSize.QuadPart; // take the 64-bit value + diskann::cout << "File Size: " << _fileSize << std::endl; + } else { + std::cerr << "Failed to get size of file " << filename << std::endl; + } #endif } -char *MemoryMapper::getBuf() -{ - return _buf; -} +char *MemoryMapper::getBuf() { return _buf; } -size_t MemoryMapper::getFileSize() -{ - return _fileSize; -} +size_t MemoryMapper::getFileSize() { return _fileSize; } -MemoryMapper::~MemoryMapper() -{ +MemoryMapper::~MemoryMapper() { #ifndef _WINDOWS - if (munmap(_buf, _fileSize) != 0) - std::cerr << "ERROR unmapping. CHECK!" << std::endl; - close(_fd); + if (munmap(_buf, _fileSize) != 0) + std::cerr << "ERROR unmapping. CHECK!" << std::endl; + close(_fd); #else - if (FALSE == UnmapViewOfFile(_buf)) - { - std::cerr << "Unmap view of file failed. Error: " << GetLastError() << std::endl; - } + if (FALSE == UnmapViewOfFile(_buf)) { + std::cerr << "Unmap view of file failed. Error: " << GetLastError() + << std::endl; + } - if (FALSE == CloseHandle(_fd)) - { - std::cerr << "Failed to close memory mapped file. Error: " << GetLastError() << std::endl; - } + if (FALSE == CloseHandle(_fd)) { + std::cerr << "Failed to close memory mapped file. Error: " << GetLastError() + << std::endl; + } - if (FALSE == CloseHandle(_bareFile)) - { - std::cerr << "Failed to close file: " << _fileName << " Error: " << GetLastError() << std::endl; - } + if (FALSE == CloseHandle(_bareFile)) { + std::cerr << "Failed to close file: " << _fileName + << " Error: " << GetLastError() << std::endl; + } #endif } diff --git a/src/natural_number_map.cpp b/src/natural_number_map.cpp index 9050831a2..59e742d9b 100644 --- a/src/natural_number_map.cpp +++ b/src/natural_number_map.cpp @@ -6,104 +6,98 @@ #include "natural_number_map.h" -namespace diskann -{ +namespace diskann { static constexpr auto invalid_position = boost::dynamic_bitset<>::npos; template natural_number_map::natural_number_map() - : _size(0), _values_bitset(std::make_unique>()) -{ -} + : _size(0), _values_bitset(std::make_unique>()) {} -template void natural_number_map::reserve(size_t count) -{ - _values_vector.reserve(count); - _values_bitset->reserve(count); +template +void natural_number_map::reserve(size_t count) { + _values_vector.reserve(count); + _values_bitset->reserve(count); } -template size_t natural_number_map::size() const -{ - return _size; +template +size_t natural_number_map::size() const { + return _size; } -template void natural_number_map::set(Key key, Value value) -{ - if (key >= _values_bitset->size()) - { - _values_bitset->resize(static_cast(key) + 1); - _values_vector.resize(_values_bitset->size()); - } - - _values_vector[key] = value; - const bool was_present = _values_bitset->test_set(key, true); - - if (!was_present) - { - ++_size; - } +template +void natural_number_map::set(Key key, Value value) { + if (key >= _values_bitset->size()) { + _values_bitset->resize(static_cast(key) + 1); + _values_vector.resize(_values_bitset->size()); + } + + _values_vector[key] = value; + const bool was_present = _values_bitset->test_set(key, true); + + if (!was_present) { + ++_size; + } } -template void natural_number_map::erase(Key key) -{ - if (key < _values_bitset->size()) - { - const bool was_present = _values_bitset->test_set(key, false); +template +void natural_number_map::erase(Key key) { + if (key < _values_bitset->size()) { + const bool was_present = _values_bitset->test_set(key, false); - if (was_present) - { - --_size; - } + if (was_present) { + --_size; } + } } -template bool natural_number_map::contains(Key key) const -{ - return key < _values_bitset->size() && _values_bitset->test(key); +template +bool natural_number_map::contains(Key key) const { + return key < _values_bitset->size() && _values_bitset->test(key); } -template bool natural_number_map::try_get(Key key, Value &value) const -{ - if (!contains(key)) - { - return false; - } +template +bool natural_number_map::try_get(Key key, Value &value) const { + if (!contains(key)) { + return false; + } - value = _values_vector[key]; - return true; + value = _values_vector[key]; + return true; } template -typename natural_number_map::position natural_number_map::find_first() const -{ - return position{_size > 0 ? _values_bitset->find_first() : invalid_position, 0}; +typename natural_number_map::position +natural_number_map::find_first() const { + return position{_size > 0 ? _values_bitset->find_first() : invalid_position, + 0}; } template -typename natural_number_map::position natural_number_map::find_next( - const position &after_position) const -{ - return position{after_position._keys_already_enumerated < _size ? _values_bitset->find_next(after_position._key) - : invalid_position, - after_position._keys_already_enumerated + 1}; +typename natural_number_map::position +natural_number_map::find_next( + const position &after_position) const { + return position{after_position._keys_already_enumerated < _size + ? _values_bitset->find_next(after_position._key) + : invalid_position, + after_position._keys_already_enumerated + 1}; } -template bool natural_number_map::position::is_valid() const -{ - return _key != invalid_position; +template +bool natural_number_map::position::is_valid() const { + return _key != invalid_position; } -template Value natural_number_map::get(const position &pos) const -{ - assert(pos.is_valid()); - return _values_vector[pos._key]; +template +Value natural_number_map::get(const position &pos) const { + assert(pos.is_valid()); + return _values_vector[pos._key]; } -template void natural_number_map::clear() -{ - _size = 0; - _values_vector.clear(); - _values_bitset->clear(); +template +void natural_number_map::clear() { + _size = 0; + _values_vector.clear(); + _values_bitset->clear(); } // Instantiate used templates. @@ -111,4 +105,4 @@ template class natural_number_map; template class natural_number_map; template class natural_number_map; template class natural_number_map; -} // namespace diskann +} // namespace diskann diff --git a/src/natural_number_set.cpp b/src/natural_number_set.cpp index b36cb5298..5cd7d6f71 100644 --- a/src/natural_number_set.cpp +++ b/src/natural_number_set.cpp @@ -6,65 +6,63 @@ #include "ann_exception.h" #include "natural_number_set.h" -namespace diskann -{ +namespace diskann { template -natural_number_set::natural_number_set() : _values_bitset(std::make_unique>()) -{ -} +natural_number_set::natural_number_set() + : _values_bitset(std::make_unique>()) {} -template bool natural_number_set::is_empty() const -{ - return _values_vector.empty(); +template +bool natural_number_set::is_empty() const { + return _values_vector.empty(); } -template void natural_number_set::reserve(size_t count) -{ - _values_vector.reserve(count); - _values_bitset->reserve(count); +template +void natural_number_set::reserve(size_t count) { + _values_vector.reserve(count); + _values_bitset->reserve(count); } -template void natural_number_set::insert(T id) -{ - _values_vector.emplace_back(id); +template +void natural_number_set::insert(T id) { + _values_vector.emplace_back(id); - if (id >= _values_bitset->size()) - _values_bitset->resize(static_cast(id) + 1); + if (id >= _values_bitset->size()) + _values_bitset->resize(static_cast(id) + 1); - _values_bitset->set(id, true); + _values_bitset->set(id, true); } -template T natural_number_set::pop_any() -{ - if (_values_vector.empty()) - { - throw diskann::ANNException("No values available", -1, __FUNCSIG__, __FILE__, __LINE__); - } +template +T natural_number_set::pop_any() { + if (_values_vector.empty()) { + throw diskann::ANNException("No values available", -1, __FUNCSIG__, + __FILE__, __LINE__); + } - const T id = _values_vector.back(); - _values_vector.pop_back(); + const T id = _values_vector.back(); + _values_vector.pop_back(); - _values_bitset->set(id, false); + _values_bitset->set(id, false); - return id; + return id; } -template void natural_number_set::clear() -{ - _values_vector.clear(); - _values_bitset->clear(); +template +void natural_number_set::clear() { + _values_vector.clear(); + _values_bitset->clear(); } -template size_t natural_number_set::size() const -{ - return _values_vector.size(); +template +size_t natural_number_set::size() const { + return _values_vector.size(); } -template bool natural_number_set::is_in_set(T id) const -{ - return _values_bitset->test(id); +template +bool natural_number_set::is_in_set(T id) const { + return _values_bitset->test(id); } // Instantiate used templates. template class natural_number_set; -} // namespace diskann +} // namespace diskann diff --git a/src/partition.cpp b/src/partition.cpp index f35cadb83..386797e10 100644 --- a/src/partition.cpp +++ b/src/partition.cpp @@ -11,7 +11,8 @@ #include "tsl/robin_map.h" #include "tsl/robin_set.h" -#if defined(RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS) && defined(DISKANN_BUILD) +#if defined(RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS) && \ + defined(DISKANN_BUILD) #include "gperftools/malloc_extension.h" #endif @@ -31,56 +32,59 @@ // #define SAVE_INFLATED_PQ true template -void gen_random_slice(const std::string base_file, const std::string output_prefix, double sampling_rate) -{ - size_t read_blk_size = 64 * 1024 * 1024; - cached_ifstream base_reader(base_file.c_str(), read_blk_size); - std::ofstream sample_writer(std::string(output_prefix + "_data.bin").c_str(), std::ios::binary); - std::ofstream sample_id_writer(std::string(output_prefix + "_ids.bin").c_str(), std::ios::binary); - - std::random_device rd; // Will be used to obtain a seed for the random number engine - auto x = rd(); - std::mt19937 generator(x); // Standard mersenne_twister_engine seeded with rd() - std::uniform_real_distribution distribution(0, 1); - - size_t npts, nd; - uint32_t npts_u32, nd_u32; - uint32_t num_sampled_pts_u32 = 0; - uint32_t one_const = 1; - - base_reader.read((char *)&npts_u32, sizeof(uint32_t)); - base_reader.read((char *)&nd_u32, sizeof(uint32_t)); - diskann::cout << "Loading base " << base_file << ". #points: " << npts_u32 << ". #dim: " << nd_u32 << "." - << std::endl; - sample_writer.write((char *)&num_sampled_pts_u32, sizeof(uint32_t)); - sample_writer.write((char *)&nd_u32, sizeof(uint32_t)); - sample_id_writer.write((char *)&num_sampled_pts_u32, sizeof(uint32_t)); - sample_id_writer.write((char *)&one_const, sizeof(uint32_t)); - - npts = npts_u32; - nd = nd_u32; - std::unique_ptr cur_row = std::make_unique(nd); - - 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) - { - sample_writer.write((char *)cur_row.get(), sizeof(T) * nd); - uint32_t cur_i_u32 = (uint32_t)i; - sample_id_writer.write((char *)&cur_i_u32, sizeof(uint32_t)); - num_sampled_pts_u32++; - } +void gen_random_slice(const std::string base_file, + const std::string output_prefix, double sampling_rate) { + size_t read_blk_size = 64 * 1024 * 1024; + cached_ifstream base_reader(base_file.c_str(), read_blk_size); + std::ofstream sample_writer(std::string(output_prefix + "_data.bin").c_str(), + std::ios::binary); + std::ofstream sample_id_writer( + std::string(output_prefix + "_ids.bin").c_str(), std::ios::binary); + + std::random_device + rd; // Will be used to obtain a seed for the random number engine + auto x = rd(); + std::mt19937 generator( + x); // Standard mersenne_twister_engine seeded with rd() + std::uniform_real_distribution distribution(0, 1); + + size_t npts, nd; + uint32_t npts_u32, nd_u32; + uint32_t num_sampled_pts_u32 = 0; + uint32_t one_const = 1; + + base_reader.read((char *)&npts_u32, sizeof(uint32_t)); + base_reader.read((char *)&nd_u32, sizeof(uint32_t)); + diskann::cout << "Loading base " << base_file << ". #points: " << npts_u32 + << ". #dim: " << nd_u32 << "." << std::endl; + sample_writer.write((char *)&num_sampled_pts_u32, sizeof(uint32_t)); + sample_writer.write((char *)&nd_u32, sizeof(uint32_t)); + sample_id_writer.write((char *)&num_sampled_pts_u32, sizeof(uint32_t)); + sample_id_writer.write((char *)&one_const, sizeof(uint32_t)); + + npts = npts_u32; + nd = nd_u32; + std::unique_ptr cur_row = std::make_unique(nd); + + 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) { + sample_writer.write((char *)cur_row.get(), sizeof(T) * nd); + uint32_t cur_i_u32 = (uint32_t)i; + sample_id_writer.write((char *)&cur_i_u32, sizeof(uint32_t)); + num_sampled_pts_u32++; } - 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); - sample_id_writer.write((char *)&num_sampled_pts_u32, sizeof(uint32_t)); - sample_writer.close(); - sample_id_writer.close(); - diskann::cout << "Wrote " << num_sampled_pts_u32 << " points to sample file: " << output_prefix + "_data.bin" - << std::endl; + } + 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); + sample_id_writer.write((char *)&num_sampled_pts_u32, sizeof(uint32_t)); + sample_writer.close(); + sample_id_writer.close(); + diskann::cout << "Wrote " << num_sampled_pts_u32 + << " points to sample file: " << output_prefix + "_data.bin" + << std::endl; } // streams data from the file, and samples each vector with probability p_val @@ -92,384 +96,388 @@ void gen_random_slice(const std::string base_file, const std::string output_pref ************************************/ template -void gen_random_slice(const std::string data_file, double p_val, float *&sampled_data, size_t &slice_size, - size_t &ndims) -{ - size_t npts; - uint32_t npts32, ndims32; - std::vector> sampled_vectors; - - // amount to read in one shot - size_t read_blk_size = 64 * 1024 * 1024; - // create cached reader + writer - cached_ifstream base_reader(data_file.c_str(), read_blk_size); - - // metadata: npts, ndims - base_reader.read((char *)&npts32, sizeof(uint32_t)); - base_reader.read((char *)&ndims32, sizeof(uint32_t)); - npts = npts32; - ndims = ndims32; - - std::unique_ptr cur_vector_T = std::make_unique(ndims); - p_val = p_val < 1 ? p_val : 1; - - std::random_device rd; // Will be used to obtain a seed for the random number - size_t x = rd(); - std::mt19937 generator((uint32_t)x); - std::uniform_real_distribution distribution(0, 1); - - 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) - { - 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); - } +void gen_random_slice(const std::string data_file, double p_val, + float *&sampled_data, size_t &slice_size, size_t &ndims) { + size_t npts; + uint32_t npts32, ndims32; + std::vector> sampled_vectors; + + // amount to read in one shot + size_t read_blk_size = 64 * 1024 * 1024; + // create cached reader + writer + cached_ifstream base_reader(data_file.c_str(), read_blk_size); + + // metadata: npts, ndims + base_reader.read((char *)&npts32, sizeof(uint32_t)); + base_reader.read((char *)&ndims32, sizeof(uint32_t)); + npts = npts32; + ndims = ndims32; + + std::unique_ptr cur_vector_T = std::make_unique(ndims); + p_val = p_val < 1 ? p_val : 1; + + std::random_device rd; // Will be used to obtain a seed for the random number + size_t x = rd(); + std::mt19937 generator((uint32_t)x); + std::uniform_real_distribution distribution(0, 1); + + 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) { + 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); } - slice_size = sampled_vectors.size(); - 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]; - } + } + slice_size = sampled_vectors.size(); + 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]; } + } } // same as above, but samples from the matrix inputdata instead of a file of // npts*ndims to return sampled_data of size slice_size*ndims. template -void gen_random_slice(const T *inputdata, size_t npts, size_t ndims, double p_val, float *&sampled_data, - size_t &slice_size) -{ - std::vector> sampled_vectors; - const T *cur_vector_T; - - p_val = p_val < 1 ? p_val : 1; - - std::random_device rd; // Will be used to obtain a seed for the random number engine - size_t x = rd(); - std::mt19937 generator((uint32_t)x); // Standard mersenne_twister_engine seeded with rd() - std::uniform_real_distribution distribution(0, 1); - - for (size_t i = 0; i < npts; i++) - { - cur_vector_T = inputdata + ndims * i; - float rnd_val = distribution(generator); - if (rnd_val < 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); - } +void gen_random_slice(const T *inputdata, size_t npts, size_t ndims, + double p_val, float *&sampled_data, size_t &slice_size) { + std::vector> sampled_vectors; + const T *cur_vector_T; + + p_val = p_val < 1 ? p_val : 1; + + std::random_device + rd; // Will be used to obtain a seed for the random number engine + size_t x = rd(); + std::mt19937 generator( + (uint32_t)x); // Standard mersenne_twister_engine seeded with rd() + std::uniform_real_distribution distribution(0, 1); + + for (size_t i = 0; i < npts; i++) { + cur_vector_T = inputdata + ndims * i; + float rnd_val = distribution(generator); + if (rnd_val < 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); } - slice_size = sampled_vectors.size(); - 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]; - } + } + slice_size = sampled_vectors.size(); + 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]; } + } } -int estimate_cluster_sizes(float *test_data_float, size_t num_test, float *pivots, const size_t num_centers, - const size_t test_dim, const size_t k_base, std::vector &cluster_sizes) -{ - cluster_sizes.clear(); - - size_t *shard_counts = new size_t[num_centers]; +int estimate_cluster_sizes(float *test_data_float, size_t num_test, + float *pivots, const size_t num_centers, + const size_t test_dim, const size_t k_base, + std::vector &cluster_sizes) { + cluster_sizes.clear(); - for (size_t i = 0; i < num_centers; i++) - { - shard_counts[i] = 0; - } + size_t *shard_counts = new size_t[num_centers]; - size_t block_size = num_test <= BLOCK_SIZE ? num_test : BLOCK_SIZE; - uint32_t *block_closest_centers = new uint32_t[block_size * k_base]; - float *block_data_float; + for (size_t i = 0; i < num_centers; i++) { + shard_counts[i] = 0; + } - size_t num_blocks = DIV_ROUND_UP(num_test, block_size); + size_t block_size = num_test <= BLOCK_SIZE ? num_test : BLOCK_SIZE; + uint32_t *block_closest_centers = new uint32_t[block_size * k_base]; + float *block_data_float; - for (size_t block = 0; block < num_blocks; block++) - { - size_t start_id = block * block_size; - size_t end_id = (std::min)((block + 1) * block_size, num_test); - size_t cur_blk_size = end_id - start_id; + size_t num_blocks = DIV_ROUND_UP(num_test, block_size); - block_data_float = test_data_float + start_id * test_dim; + for (size_t block = 0; block < num_blocks; block++) { + size_t start_id = block * block_size; + size_t end_id = (std::min)((block + 1) * block_size, num_test); + size_t cur_blk_size = end_id - start_id; - math_utils::compute_closest_centers(block_data_float, cur_blk_size, test_dim, pivots, num_centers, k_base, - block_closest_centers); + block_data_float = test_data_float + start_id * test_dim; - for (size_t p = 0; p < cur_blk_size; p++) - { - for (size_t p1 = 0; p1 < k_base; p1++) - { - size_t shard_id = block_closest_centers[p * k_base + p1]; - shard_counts[shard_id]++; - } - } - } + math_utils::compute_closest_centers(block_data_float, cur_blk_size, + test_dim, pivots, num_centers, k_base, + block_closest_centers); - diskann::cout << "Estimated cluster sizes: "; - for (size_t i = 0; i < num_centers; i++) - { - uint32_t cur_shard_count = (uint32_t)shard_counts[i]; - cluster_sizes.push_back((size_t)cur_shard_count); - diskann::cout << cur_shard_count << " "; + for (size_t p = 0; p < cur_blk_size; p++) { + for (size_t p1 = 0; p1 < k_base; p1++) { + size_t shard_id = block_closest_centers[p * k_base + p1]; + shard_counts[shard_id]++; + } } - diskann::cout << std::endl; - delete[] shard_counts; - delete[] block_closest_centers; - return 0; + } + + diskann::cout << "Estimated cluster sizes: "; + for (size_t i = 0; i < num_centers; i++) { + uint32_t cur_shard_count = (uint32_t)shard_counts[i]; + cluster_sizes.push_back((size_t)cur_shard_count); + diskann::cout << cur_shard_count << " "; + } + diskann::cout << std::endl; + delete[] shard_counts; + delete[] block_closest_centers; + return 0; } template -int shard_data_into_clusters(const std::string data_file, float *pivots, const size_t num_centers, const size_t dim, - const size_t k_base, std::string prefix_path) -{ - size_t read_blk_size = 64 * 1024 * 1024; - // uint64_t write_blk_size = 64 * 1024 * 1024; - // create cached reader + writer - cached_ifstream base_reader(data_file, read_blk_size); - uint32_t npts32; - uint32_t basedim32; - base_reader.read((char *)&npts32, sizeof(uint32_t)); - base_reader.read((char *)&basedim32, sizeof(uint32_t)); - size_t num_points = npts32; - if (basedim32 != dim) - { - diskann::cout << "Error. dimensions dont match for train set and base set" << std::endl; - return -1; - } - - std::unique_ptr shard_counts = std::make_unique(num_centers); - std::vector shard_data_writer(num_centers); - std::vector shard_idmap_writer(num_centers); - uint32_t dummy_size = 0; - uint32_t const_one = 1; - - for (size_t i = 0; i < num_centers; i++) - { - std::string data_filename = prefix_path + "_subshard-" + std::to_string(i) + ".bin"; - std::string idmap_filename = prefix_path + "_subshard-" + std::to_string(i) + "_ids_uint32.bin"; - shard_data_writer[i] = std::ofstream(data_filename.c_str(), std::ios::binary); - shard_idmap_writer[i] = std::ofstream(idmap_filename.c_str(), std::ios::binary); - shard_data_writer[i].write((char *)&dummy_size, sizeof(uint32_t)); - shard_data_writer[i].write((char *)&basedim32, sizeof(uint32_t)); - shard_idmap_writer[i].write((char *)&dummy_size, sizeof(uint32_t)); - shard_idmap_writer[i].write((char *)&const_one, sizeof(uint32_t)); - shard_counts[i] = 0; - } - - size_t block_size = num_points <= BLOCK_SIZE ? num_points : BLOCK_SIZE; - std::unique_ptr block_closest_centers = std::make_unique(block_size * k_base); - std::unique_ptr block_data_T = std::make_unique(block_size * dim); - std::unique_ptr block_data_float = std::make_unique(block_size * dim); - - size_t num_blocks = DIV_ROUND_UP(num_points, block_size); - - for (size_t block = 0; block < num_blocks; block++) - { - size_t start_id = block * block_size; - size_t end_id = (std::min)((block + 1) * block_size, num_points); - size_t cur_blk_size = end_id - start_id; - - base_reader.read((char *)block_data_T.get(), sizeof(T) * (cur_blk_size * dim)); - diskann::convert_types(block_data_T.get(), block_data_float.get(), cur_blk_size, dim); - - math_utils::compute_closest_centers(block_data_float.get(), cur_blk_size, dim, pivots, num_centers, k_base, - block_closest_centers.get()); - - for (size_t p = 0; p < cur_blk_size; p++) - { - for (size_t p1 = 0; p1 < k_base; p1++) - { - size_t shard_id = block_closest_centers[p * k_base + p1]; - uint32_t original_point_map_id = (uint32_t)(start_id + p); - shard_data_writer[shard_id].write((char *)(block_data_T.get() + p * dim), sizeof(T) * dim); - shard_idmap_writer[shard_id].write((char *)&original_point_map_id, sizeof(uint32_t)); - shard_counts[shard_id]++; - } - } - } - - size_t total_count = 0; - diskann::cout << "Actual shard sizes: " << std::flush; - for (size_t i = 0; i < num_centers; i++) - { - uint32_t cur_shard_count = (uint32_t)shard_counts[i]; - total_count += cur_shard_count; - diskann::cout << cur_shard_count << " "; - shard_data_writer[i].seekp(0); - shard_data_writer[i].write((char *)&cur_shard_count, sizeof(uint32_t)); - shard_data_writer[i].close(); - shard_idmap_writer[i].seekp(0); - shard_idmap_writer[i].write((char *)&cur_shard_count, sizeof(uint32_t)); - shard_idmap_writer[i].close(); +int shard_data_into_clusters(const std::string data_file, float *pivots, + const size_t num_centers, const size_t dim, + const size_t k_base, std::string prefix_path) { + size_t read_blk_size = 64 * 1024 * 1024; + // uint64_t write_blk_size = 64 * 1024 * 1024; + // create cached reader + writer + cached_ifstream base_reader(data_file, read_blk_size); + uint32_t npts32; + uint32_t basedim32; + base_reader.read((char *)&npts32, sizeof(uint32_t)); + base_reader.read((char *)&basedim32, sizeof(uint32_t)); + size_t num_points = npts32; + if (basedim32 != dim) { + diskann::cout << "Error. dimensions dont match for train set and base set" + << std::endl; + return -1; + } + + std::unique_ptr shard_counts = + std::make_unique(num_centers); + std::vector shard_data_writer(num_centers); + std::vector shard_idmap_writer(num_centers); + uint32_t dummy_size = 0; + uint32_t const_one = 1; + + for (size_t i = 0; i < num_centers; i++) { + std::string data_filename = + prefix_path + "_subshard-" + std::to_string(i) + ".bin"; + std::string idmap_filename = + prefix_path + "_subshard-" + std::to_string(i) + "_ids_uint32.bin"; + shard_data_writer[i] = + std::ofstream(data_filename.c_str(), std::ios::binary); + shard_idmap_writer[i] = + std::ofstream(idmap_filename.c_str(), std::ios::binary); + shard_data_writer[i].write((char *)&dummy_size, sizeof(uint32_t)); + shard_data_writer[i].write((char *)&basedim32, sizeof(uint32_t)); + shard_idmap_writer[i].write((char *)&dummy_size, sizeof(uint32_t)); + shard_idmap_writer[i].write((char *)&const_one, sizeof(uint32_t)); + shard_counts[i] = 0; + } + + size_t block_size = num_points <= BLOCK_SIZE ? num_points : BLOCK_SIZE; + std::unique_ptr block_closest_centers = + std::make_unique(block_size * k_base); + std::unique_ptr block_data_T = std::make_unique(block_size * dim); + std::unique_ptr block_data_float = + std::make_unique(block_size * dim); + + size_t num_blocks = DIV_ROUND_UP(num_points, block_size); + + for (size_t block = 0; block < num_blocks; block++) { + size_t start_id = block * block_size; + size_t end_id = (std::min)((block + 1) * block_size, num_points); + size_t cur_blk_size = end_id - start_id; + + base_reader.read((char *)block_data_T.get(), + sizeof(T) * (cur_blk_size * dim)); + diskann::convert_types(block_data_T.get(), block_data_float.get(), + cur_blk_size, dim); + + math_utils::compute_closest_centers(block_data_float.get(), cur_blk_size, + dim, pivots, num_centers, k_base, + block_closest_centers.get()); + + for (size_t p = 0; p < cur_blk_size; p++) { + for (size_t p1 = 0; p1 < k_base; p1++) { + size_t shard_id = block_closest_centers[p * k_base + p1]; + uint32_t original_point_map_id = (uint32_t)(start_id + p); + shard_data_writer[shard_id].write( + (char *)(block_data_T.get() + p * dim), sizeof(T) * dim); + shard_idmap_writer[shard_id].write((char *)&original_point_map_id, + sizeof(uint32_t)); + shard_counts[shard_id]++; + } } - - diskann::cout << "\n Partitioned " << num_points << " with replication factor " << k_base << " to get " - << total_count << " points across " << num_centers << " shards " << std::endl; - return 0; + } + + size_t total_count = 0; + diskann::cout << "Actual shard sizes: " << std::flush; + for (size_t i = 0; i < num_centers; i++) { + uint32_t cur_shard_count = (uint32_t)shard_counts[i]; + total_count += cur_shard_count; + diskann::cout << cur_shard_count << " "; + shard_data_writer[i].seekp(0); + shard_data_writer[i].write((char *)&cur_shard_count, sizeof(uint32_t)); + shard_data_writer[i].close(); + shard_idmap_writer[i].seekp(0); + shard_idmap_writer[i].write((char *)&cur_shard_count, sizeof(uint32_t)); + shard_idmap_writer[i].close(); + } + + diskann::cout << "\n Partitioned " << num_points + << " with replication factor " << k_base << " to get " + << total_count << " points across " << num_centers << " shards " + << std::endl; + return 0; } // useful for partitioning large dataset. we first generate only the IDS for // each shard, and retrieve the actual vectors on demand. template -int shard_data_into_clusters_only_ids(const std::string data_file, float *pivots, const size_t num_centers, - const size_t dim, const size_t k_base, std::string prefix_path) -{ - size_t read_blk_size = 64 * 1024 * 1024; - // uint64_t write_blk_size = 64 * 1024 * 1024; - // create cached reader + writer - cached_ifstream base_reader(data_file, read_blk_size); - uint32_t npts32; - uint32_t basedim32; - base_reader.read((char *)&npts32, sizeof(uint32_t)); - base_reader.read((char *)&basedim32, sizeof(uint32_t)); - size_t num_points = npts32; - if (basedim32 != dim) - { - diskann::cout << "Error. dimensions dont match for train set and base set" << std::endl; - return -1; - } - - std::unique_ptr shard_counts = std::make_unique(num_centers); - - std::vector shard_idmap_writer(num_centers); - uint32_t dummy_size = 0; - uint32_t const_one = 1; - - for (size_t i = 0; i < num_centers; i++) - { - std::string idmap_filename = prefix_path + "_subshard-" + std::to_string(i) + "_ids_uint32.bin"; - shard_idmap_writer[i] = std::ofstream(idmap_filename.c_str(), std::ios::binary); - shard_idmap_writer[i].write((char *)&dummy_size, sizeof(uint32_t)); - shard_idmap_writer[i].write((char *)&const_one, sizeof(uint32_t)); - shard_counts[i] = 0; - } - - size_t block_size = num_points <= BLOCK_SIZE ? num_points : BLOCK_SIZE; - std::unique_ptr block_closest_centers = std::make_unique(block_size * k_base); - std::unique_ptr block_data_T = std::make_unique(block_size * dim); - std::unique_ptr block_data_float = std::make_unique(block_size * dim); - - size_t num_blocks = DIV_ROUND_UP(num_points, block_size); - - for (size_t block = 0; block < num_blocks; block++) - { - size_t start_id = block * block_size; - size_t end_id = (std::min)((block + 1) * block_size, num_points); - size_t cur_blk_size = end_id - start_id; - - base_reader.read((char *)block_data_T.get(), sizeof(T) * (cur_blk_size * dim)); - diskann::convert_types(block_data_T.get(), block_data_float.get(), cur_blk_size, dim); - - math_utils::compute_closest_centers(block_data_float.get(), cur_blk_size, dim, pivots, num_centers, k_base, - block_closest_centers.get()); - - for (size_t p = 0; p < cur_blk_size; p++) - { - for (size_t p1 = 0; p1 < k_base; p1++) - { - size_t shard_id = block_closest_centers[p * k_base + p1]; - uint32_t original_point_map_id = (uint32_t)(start_id + p); - shard_idmap_writer[shard_id].write((char *)&original_point_map_id, sizeof(uint32_t)); - shard_counts[shard_id]++; - } - } - } - - size_t total_count = 0; - diskann::cout << "Actual shard sizes: " << std::flush; - for (size_t i = 0; i < num_centers; i++) - { - uint32_t cur_shard_count = (uint32_t)shard_counts[i]; - total_count += cur_shard_count; - diskann::cout << cur_shard_count << " "; - shard_idmap_writer[i].seekp(0); - shard_idmap_writer[i].write((char *)&cur_shard_count, sizeof(uint32_t)); - shard_idmap_writer[i].close(); +int shard_data_into_clusters_only_ids(const std::string data_file, + float *pivots, const size_t num_centers, + const size_t dim, const size_t k_base, + std::string prefix_path) { + size_t read_blk_size = 64 * 1024 * 1024; + // uint64_t write_blk_size = 64 * 1024 * 1024; + // create cached reader + writer + cached_ifstream base_reader(data_file, read_blk_size); + uint32_t npts32; + uint32_t basedim32; + base_reader.read((char *)&npts32, sizeof(uint32_t)); + base_reader.read((char *)&basedim32, sizeof(uint32_t)); + size_t num_points = npts32; + if (basedim32 != dim) { + diskann::cout << "Error. dimensions dont match for train set and base set" + << std::endl; + return -1; + } + + std::unique_ptr shard_counts = + std::make_unique(num_centers); + + std::vector shard_idmap_writer(num_centers); + uint32_t dummy_size = 0; + uint32_t const_one = 1; + + for (size_t i = 0; i < num_centers; i++) { + std::string idmap_filename = + prefix_path + "_subshard-" + std::to_string(i) + "_ids_uint32.bin"; + shard_idmap_writer[i] = + std::ofstream(idmap_filename.c_str(), std::ios::binary); + shard_idmap_writer[i].write((char *)&dummy_size, sizeof(uint32_t)); + shard_idmap_writer[i].write((char *)&const_one, sizeof(uint32_t)); + shard_counts[i] = 0; + } + + size_t block_size = num_points <= BLOCK_SIZE ? num_points : BLOCK_SIZE; + std::unique_ptr block_closest_centers = + std::make_unique(block_size * k_base); + std::unique_ptr block_data_T = std::make_unique(block_size * dim); + std::unique_ptr block_data_float = + std::make_unique(block_size * dim); + + size_t num_blocks = DIV_ROUND_UP(num_points, block_size); + + for (size_t block = 0; block < num_blocks; block++) { + size_t start_id = block * block_size; + size_t end_id = (std::min)((block + 1) * block_size, num_points); + size_t cur_blk_size = end_id - start_id; + + base_reader.read((char *)block_data_T.get(), + sizeof(T) * (cur_blk_size * dim)); + diskann::convert_types(block_data_T.get(), block_data_float.get(), + cur_blk_size, dim); + + math_utils::compute_closest_centers(block_data_float.get(), cur_blk_size, + dim, pivots, num_centers, k_base, + block_closest_centers.get()); + + for (size_t p = 0; p < cur_blk_size; p++) { + for (size_t p1 = 0; p1 < k_base; p1++) { + size_t shard_id = block_closest_centers[p * k_base + p1]; + uint32_t original_point_map_id = (uint32_t)(start_id + p); + shard_idmap_writer[shard_id].write((char *)&original_point_map_id, + sizeof(uint32_t)); + shard_counts[shard_id]++; + } } - - diskann::cout << "\n Partitioned " << num_points << " with replication factor " << k_base << " to get " - << total_count << " points across " << num_centers << " shards " << std::endl; - return 0; + } + + size_t total_count = 0; + diskann::cout << "Actual shard sizes: " << std::flush; + for (size_t i = 0; i < num_centers; i++) { + uint32_t cur_shard_count = (uint32_t)shard_counts[i]; + total_count += cur_shard_count; + diskann::cout << cur_shard_count << " "; + shard_idmap_writer[i].seekp(0); + shard_idmap_writer[i].write((char *)&cur_shard_count, sizeof(uint32_t)); + shard_idmap_writer[i].close(); + } + + diskann::cout << "\n Partitioned " << num_points + << " with replication factor " << k_base << " to get " + << total_count << " points across " << num_centers << " shards " + << std::endl; + return 0; } template -int retrieve_shard_data_from_ids(const std::string data_file, std::string idmap_filename, std::string data_filename) -{ - size_t read_blk_size = 64 * 1024 * 1024; - // uint64_t write_blk_size = 64 * 1024 * 1024; - // create cached reader + writer - cached_ifstream base_reader(data_file, read_blk_size); - uint32_t npts32; - uint32_t basedim32; - base_reader.read((char *)&npts32, sizeof(uint32_t)); - base_reader.read((char *)&basedim32, sizeof(uint32_t)); - size_t num_points = npts32; - size_t dim = basedim32; - - uint32_t dummy_size = 0; - - std::ofstream shard_data_writer(data_filename.c_str(), std::ios::binary); - shard_data_writer.write((char *)&dummy_size, sizeof(uint32_t)); - shard_data_writer.write((char *)&basedim32, sizeof(uint32_t)); - - uint32_t *shard_ids; - uint64_t shard_size, tmp; - diskann::load_bin(idmap_filename, shard_ids, shard_size, tmp); - - uint32_t cur_pos = 0; - uint32_t num_written = 0; - std::cout << "Shard has " << shard_size << " points" << std::endl; - - size_t block_size = num_points <= BLOCK_SIZE ? num_points : BLOCK_SIZE; - std::unique_ptr block_data_T = std::make_unique(block_size * dim); - - size_t num_blocks = DIV_ROUND_UP(num_points, block_size); - - for (size_t block = 0; block < num_blocks; block++) - { - size_t start_id = block * block_size; - size_t end_id = (std::min)((block + 1) * block_size, num_points); - size_t cur_blk_size = end_id - start_id; - - base_reader.read((char *)block_data_T.get(), sizeof(T) * (cur_blk_size * dim)); - - for (size_t p = 0; p < cur_blk_size; p++) - { - uint32_t original_point_map_id = (uint32_t)(start_id + p); - if (cur_pos == shard_size) - break; - if (original_point_map_id == shard_ids[cur_pos]) - { - cur_pos++; - shard_data_writer.write((char *)(block_data_T.get() + p * dim), sizeof(T) * dim); - num_written++; - } - } - if (cur_pos == shard_size) - break; +int retrieve_shard_data_from_ids(const std::string data_file, + std::string idmap_filename, + std::string data_filename) { + size_t read_blk_size = 64 * 1024 * 1024; + // uint64_t write_blk_size = 64 * 1024 * 1024; + // create cached reader + writer + cached_ifstream base_reader(data_file, read_blk_size); + uint32_t npts32; + uint32_t basedim32; + base_reader.read((char *)&npts32, sizeof(uint32_t)); + base_reader.read((char *)&basedim32, sizeof(uint32_t)); + size_t num_points = npts32; + size_t dim = basedim32; + + uint32_t dummy_size = 0; + + std::ofstream shard_data_writer(data_filename.c_str(), std::ios::binary); + shard_data_writer.write((char *)&dummy_size, sizeof(uint32_t)); + shard_data_writer.write((char *)&basedim32, sizeof(uint32_t)); + + uint32_t *shard_ids; + uint64_t shard_size, tmp; + diskann::load_bin(idmap_filename, shard_ids, shard_size, tmp); + + uint32_t cur_pos = 0; + uint32_t num_written = 0; + std::cout << "Shard has " << shard_size << " points" << std::endl; + + size_t block_size = num_points <= BLOCK_SIZE ? num_points : BLOCK_SIZE; + std::unique_ptr block_data_T = std::make_unique(block_size * dim); + + size_t num_blocks = DIV_ROUND_UP(num_points, block_size); + + for (size_t block = 0; block < num_blocks; block++) { + size_t start_id = block * block_size; + size_t end_id = (std::min)((block + 1) * block_size, num_points); + size_t cur_blk_size = end_id - start_id; + + base_reader.read((char *)block_data_T.get(), + sizeof(T) * (cur_blk_size * dim)); + + for (size_t p = 0; p < cur_blk_size; p++) { + uint32_t original_point_map_id = (uint32_t)(start_id + p); + if (cur_pos == shard_size) break; + if (original_point_map_id == shard_ids[cur_pos]) { + cur_pos++; + shard_data_writer.write((char *)(block_data_T.get() + p * dim), + sizeof(T) * dim); + num_written++; + } } + if (cur_pos == shard_size) break; + } - diskann::cout << "Written file with " << num_written << " points" << std::endl; + diskann::cout << "Written file with " << num_written << " points" + << std::endl; - shard_data_writer.seekp(0); - shard_data_writer.write((char *)&num_written, sizeof(uint32_t)); - shard_data_writer.close(); - delete[] shard_ids; - return 0; + shard_data_writer.seekp(0); + shard_data_writer.write((char *)&num_written, sizeof(uint32_t)); + shard_data_writer.close(); + delete[] shard_ids; + return 0; } // partitions a large base file into many shards using k-means hueristic @@ -479,178 +487,197 @@ int retrieve_shard_data_from_ids(const std::string data_file, std::string idmap_ // The total number of points across all shards will be k_base * num_points. template -int partition(const std::string data_file, const float sampling_rate, size_t num_parts, size_t max_k_means_reps, - const std::string prefix_path, size_t k_base) -{ - size_t train_dim; - size_t num_train; - float *train_data_float; +int partition(const std::string data_file, const float sampling_rate, + size_t num_parts, size_t max_k_means_reps, + const std::string prefix_path, size_t k_base) { + size_t train_dim; + size_t num_train; + float *train_data_float; - gen_random_slice(data_file, sampling_rate, train_data_float, num_train, train_dim); + gen_random_slice(data_file, sampling_rate, train_data_float, num_train, + train_dim); - float *pivot_data; + float *pivot_data; - std::string cur_file = std::string(prefix_path); - std::string output_file; + std::string cur_file = std::string(prefix_path); + std::string output_file; - // kmeans_partitioning on training data + // kmeans_partitioning on training data - // cur_file = cur_file + "_kmeans_partitioning-" + - // std::to_string(num_parts); - output_file = cur_file + "_centroids.bin"; + // cur_file = cur_file + "_kmeans_partitioning-" + + // std::to_string(num_parts); + output_file = cur_file + "_centroids.bin"; - pivot_data = new float[num_parts * train_dim]; + pivot_data = new float[num_parts * train_dim]; - // Process Global k-means for kmeans_partitioning Step - diskann::cout << "Processing global k-means (kmeans_partitioning Step)" << std::endl; - kmeans::kmeanspp_selecting_pivots(train_data_float, num_train, train_dim, pivot_data, num_parts); + // Process Global k-means for kmeans_partitioning Step + diskann::cout << "Processing global k-means (kmeans_partitioning Step)" + << std::endl; + kmeans::kmeanspp_selecting_pivots(train_data_float, num_train, train_dim, + pivot_data, num_parts); - kmeans::run_lloyds(train_data_float, num_train, train_dim, pivot_data, num_parts, max_k_means_reps, NULL, NULL); + kmeans::run_lloyds(train_data_float, num_train, train_dim, pivot_data, + num_parts, max_k_means_reps, NULL, NULL); - diskann::cout << "Saving global k-center pivots" << std::endl; - diskann::save_bin(output_file.c_str(), pivot_data, (size_t)num_parts, train_dim); + diskann::cout << "Saving global k-center pivots" << std::endl; + diskann::save_bin(output_file.c_str(), pivot_data, (size_t)num_parts, + train_dim); - // now pivots are ready. need to stream base points and assign them to - // closest clusters. + // now pivots are ready. need to stream base points and assign them to + // closest clusters. - shard_data_into_clusters(data_file, pivot_data, num_parts, train_dim, k_base, prefix_path); - delete[] pivot_data; - delete[] train_data_float; - return 0; + shard_data_into_clusters(data_file, pivot_data, num_parts, train_dim, + k_base, prefix_path); + delete[] pivot_data; + delete[] train_data_float; + return 0; } template -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) -{ - size_t train_dim; - size_t num_train; - float *train_data_float; - size_t max_k_means_reps = 10; - - int num_parts = 3; - bool fit_in_ram = false; - - gen_random_slice(data_file, sampling_rate, train_data_float, num_train, train_dim); - - size_t test_dim; - size_t num_test; - float *test_data_float; - gen_random_slice(data_file, sampling_rate, test_data_float, num_test, test_dim); - - float *pivot_data = nullptr; - - std::string cur_file = std::string(prefix_path); - std::string output_file; - - // kmeans_partitioning on training data - - // cur_file = cur_file + "_kmeans_partitioning-" + - // std::to_string(num_parts); - output_file = cur_file + "_centroids.bin"; - - while (!fit_in_ram) - { - fit_in_ram = true; - - double max_ram_usage = 0; - if (pivot_data != nullptr) - delete[] pivot_data; - - pivot_data = new float[num_parts * train_dim]; - // Process Global k-means for kmeans_partitioning Step - diskann::cout << "Processing global k-means (kmeans_partitioning Step)" << std::endl; - kmeans::kmeanspp_selecting_pivots(train_data_float, num_train, train_dim, pivot_data, num_parts); - - kmeans::run_lloyds(train_data_float, num_train, train_dim, pivot_data, num_parts, max_k_means_reps, NULL, NULL); - - // now pivots are ready. need to stream base points and assign them to - // closest clusters. - - std::vector cluster_sizes; - estimate_cluster_sizes(test_data_float, num_test, pivot_data, num_parts, train_dim, k_base, cluster_sizes); - - for (auto &p : cluster_sizes) - { - // to account for the fact that p is the size of the shard over the - // testing sample. - p = (uint64_t)(p / sampling_rate); - double cur_shard_ram_estimate = diskann::estimate_ram_usage(p, train_dim, sizeof(T), graph_degree); - - if (cur_shard_ram_estimate > max_ram_usage) - max_ram_usage = cur_shard_ram_estimate; - } - diskann::cout << "With " << num_parts - << " parts, max estimated RAM usage: " << max_ram_usage / (1024 * 1024 * 1024) - << "GB, budget given is " << ram_budget << std::endl; - if (max_ram_usage > 1024 * 1024 * 1024 * ram_budget) - { - fit_in_ram = false; - num_parts += 2; - } - } +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) { + size_t train_dim; + size_t num_train; + float *train_data_float; + size_t max_k_means_reps = 10; + + int num_parts = 3; + bool fit_in_ram = false; + + gen_random_slice(data_file, sampling_rate, train_data_float, num_train, + train_dim); + + size_t test_dim; + size_t num_test; + float *test_data_float; + gen_random_slice(data_file, sampling_rate, test_data_float, num_test, + test_dim); + + float *pivot_data = nullptr; + + std::string cur_file = std::string(prefix_path); + std::string output_file; + + // kmeans_partitioning on training data + + // cur_file = cur_file + "_kmeans_partitioning-" + + // std::to_string(num_parts); + output_file = cur_file + "_centroids.bin"; + + while (!fit_in_ram) { + fit_in_ram = true; - diskann::cout << "Saving global k-center pivots" << std::endl; - diskann::save_bin(output_file.c_str(), pivot_data, (size_t)num_parts, train_dim); + double max_ram_usage = 0; + if (pivot_data != nullptr) delete[] pivot_data; - shard_data_into_clusters_only_ids(data_file, pivot_data, num_parts, train_dim, k_base, prefix_path); - delete[] pivot_data; - delete[] train_data_float; - delete[] test_data_float; - return num_parts; + pivot_data = new float[num_parts * train_dim]; + // Process Global k-means for kmeans_partitioning Step + diskann::cout << "Processing global k-means (kmeans_partitioning Step)" + << std::endl; + kmeans::kmeanspp_selecting_pivots(train_data_float, num_train, train_dim, + pivot_data, num_parts); + + kmeans::run_lloyds(train_data_float, num_train, train_dim, pivot_data, + num_parts, max_k_means_reps, NULL, NULL); + + // now pivots are ready. need to stream base points and assign them to + // closest clusters. + + std::vector cluster_sizes; + estimate_cluster_sizes(test_data_float, num_test, pivot_data, num_parts, + train_dim, k_base, cluster_sizes); + + for (auto &p : cluster_sizes) { + // to account for the fact that p is the size of the shard over the + // testing sample. + p = (uint64_t)(p / sampling_rate); + double cur_shard_ram_estimate = + diskann::estimate_ram_usage(p, train_dim, sizeof(T), graph_degree); + + if (cur_shard_ram_estimate > max_ram_usage) + max_ram_usage = cur_shard_ram_estimate; + } + diskann::cout << "With " << num_parts << " parts, max estimated RAM usage: " + << max_ram_usage / (1024 * 1024 * 1024) + << "GB, budget given is " << ram_budget << std::endl; + if (max_ram_usage > 1024 * 1024 * 1024 * ram_budget) { + fit_in_ram = false; + num_parts += 2; + } + } + + diskann::cout << "Saving global k-center pivots" << std::endl; + diskann::save_bin(output_file.c_str(), pivot_data, (size_t)num_parts, + train_dim); + + shard_data_into_clusters_only_ids(data_file, pivot_data, num_parts, + train_dim, k_base, prefix_path); + delete[] pivot_data; + delete[] train_data_float; + delete[] test_data_float; + return num_parts; } // Instantations of supported templates -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); -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 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, float *&sampled_data, size_t &slice_size); -template void DISKANN_DLLEXPORT gen_random_slice(const int8_t *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 std::string data_file, double p_val, float *&sampled_data, - size_t &slice_size, size_t &ndims); -template void DISKANN_DLLEXPORT gen_random_slice(const std::string data_file, double p_val, - float *&sampled_data, size_t &slice_size, size_t &ndims); -template void DISKANN_DLLEXPORT gen_random_slice(const std::string data_file, double p_val, - float *&sampled_data, size_t &slice_size, size_t &ndims); - -template DISKANN_DLLEXPORT int partition(const std::string data_file, const float sampling_rate, - size_t num_centers, size_t max_k_means_reps, - const std::string prefix_path, size_t k_base); -template DISKANN_DLLEXPORT int partition(const std::string data_file, const float sampling_rate, - size_t num_centers, size_t max_k_means_reps, - const std::string prefix_path, size_t k_base); -template DISKANN_DLLEXPORT int partition(const std::string data_file, const float sampling_rate, - size_t num_centers, size_t max_k_means_reps, - const std::string prefix_path, size_t k_base); - -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 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 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 retrieve_shard_data_from_ids(const std::string data_file, - std::string idmap_filename, - std::string data_filename); -template DISKANN_DLLEXPORT int retrieve_shard_data_from_ids(const std::string data_file, - std::string idmap_filename, - std::string data_filename); -template DISKANN_DLLEXPORT int retrieve_shard_data_from_ids(const std::string data_file, - std::string idmap_filename, - std::string data_filename); \ No newline at end of file +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); +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 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, + float *&sampled_data, size_t &slice_size); +template void DISKANN_DLLEXPORT gen_random_slice( + const int8_t *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 std::string data_file, double p_val, float *&sampled_data, + size_t &slice_size, size_t &ndims); +template void DISKANN_DLLEXPORT gen_random_slice( + const std::string data_file, double p_val, float *&sampled_data, + size_t &slice_size, size_t &ndims); +template void DISKANN_DLLEXPORT gen_random_slice( + const std::string data_file, double p_val, float *&sampled_data, + size_t &slice_size, size_t &ndims); + +template DISKANN_DLLEXPORT int partition( + const std::string data_file, const float sampling_rate, size_t num_centers, + size_t max_k_means_reps, const std::string prefix_path, size_t k_base); +template DISKANN_DLLEXPORT int partition( + const std::string data_file, const float sampling_rate, size_t num_centers, + size_t max_k_means_reps, const std::string prefix_path, size_t k_base); +template DISKANN_DLLEXPORT int partition( + const std::string data_file, const float sampling_rate, size_t num_centers, + size_t max_k_means_reps, const std::string prefix_path, size_t k_base); + +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 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 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 retrieve_shard_data_from_ids( + const std::string data_file, std::string idmap_filename, + std::string data_filename); +template DISKANN_DLLEXPORT int retrieve_shard_data_from_ids( + const std::string data_file, std::string idmap_filename, + std::string data_filename); +template DISKANN_DLLEXPORT int retrieve_shard_data_from_ids( + const std::string data_file, std::string idmap_filename, + std::string data_filename); \ No newline at end of file diff --git a/src/pq.cpp b/src/pq.cpp index 3f976d66d..3d453c046 100644 --- a/src/pq.cpp +++ b/src/pq.cpp @@ -11,333 +11,307 @@ // block size for reading/processing large files and matrices in blocks #define BLOCK_SIZE 5000000 -namespace diskann -{ -FixedChunkPQTable::FixedChunkPQTable() -{ -} +namespace diskann { +FixedChunkPQTable::FixedChunkPQTable() {} -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; + 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) -{ +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) -{ +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"; + 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); + 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); + 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__); - } + 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]); + 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]); + 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__); - } + 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; + this->ndims = nc; #ifdef EXEC_ENV_OLS - diskann::load_bin(files, pq_table_file, centroid, nr, nc, file_offset_data[1]); + 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]); + 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; - } + 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]); + 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]); + 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; - - if (file_exists(rotmat_file)) - { + 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; + + if (file_exists(rotmat_file)) { #ifdef EXEC_ENV_OLS - diskann::load_bin(files, rotmat_file, (float *&)rotmat_tr, nr, nc); + diskann::load_bin(files, rotmat_file, (float *&)rotmat_tr, nr, nc); #else - diskann::load_bin(rotmat_file, rotmat_tr, nr, nc); + 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; + 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]; - } + 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); +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)); +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); - } - } +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; - } +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; + } + 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; - } +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) + } + 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::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 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++) - { - memcpy(out + i * ndims, all_coords + ids[i] * ndims, ndims * sizeof(uint8_t)); - } +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++) { + memcpy(out + i * ndims, all_coords + ids[i] * ndims, + ndims * sizeof(uint8_t)); + } } -void pq_dist_lookup(const uint8_t *pq_ids, const size_t n_pts, const size_t pq_nchunks, const float *pq_dists, - std::vector &dists_out) -{ - //_mm_prefetch((char*) dists_out, _MM_HINT_T0); - _mm_prefetch((char *)pq_ids, _MM_HINT_T0); - _mm_prefetch((char *)(pq_ids + 64), _MM_HINT_T0); - _mm_prefetch((char *)(pq_ids + 128), _MM_HINT_T0); - dists_out.clear(); - dists_out.resize(n_pts, 0); - for (size_t chunk = 0; chunk < pq_nchunks; chunk++) - { - const float *chunk_dists = pq_dists + 256 * chunk; - if (chunk < pq_nchunks - 1) - { - _mm_prefetch((char *)(chunk_dists + 256), _MM_HINT_T0); - } - for (size_t idx = 0; idx < n_pts; idx++) - { - uint8_t pq_centerid = pq_ids[pq_nchunks * idx + chunk]; - dists_out[idx] += chunk_dists[pq_centerid]; - } +void pq_dist_lookup(const uint8_t *pq_ids, const size_t n_pts, + const size_t pq_nchunks, const float *pq_dists, + std::vector &dists_out) { + //_mm_prefetch((char*) dists_out, _MM_HINT_T0); + _mm_prefetch((char *)pq_ids, _MM_HINT_T0); + _mm_prefetch((char *)(pq_ids + 64), _MM_HINT_T0); + _mm_prefetch((char *)(pq_ids + 128), _MM_HINT_T0); + dists_out.clear(); + dists_out.resize(n_pts, 0); + for (size_t chunk = 0; chunk < pq_nchunks; chunk++) { + const float *chunk_dists = pq_dists + 256 * chunk; + if (chunk < pq_nchunks - 1) { + _mm_prefetch((char *)(chunk_dists + 256), _MM_HINT_T0); + } + for (size_t idx = 0; idx < n_pts; idx++) { + uint8_t pq_centerid = pq_ids[pq_nchunks * idx + chunk]; + dists_out[idx] += chunk_dists[pq_centerid]; } + } } // Need to replace calls to these functions with calls to vector& based // functions above -void aggregate_coords(const uint32_t *ids, const size_t n_ids, const uint8_t *all_coords, const size_t ndims, - uint8_t *out) -{ - for (size_t i = 0; i < n_ids; i++) - { - memcpy(out + i * ndims, all_coords + ids[i] * ndims, ndims * sizeof(uint8_t)); - } +void aggregate_coords(const uint32_t *ids, const size_t n_ids, + const uint8_t *all_coords, const size_t ndims, + uint8_t *out) { + for (size_t i = 0; i < n_ids; i++) { + memcpy(out + i * ndims, all_coords + ids[i] * ndims, + ndims * sizeof(uint8_t)); + } } -void pq_dist_lookup(const uint8_t *pq_ids, const size_t n_pts, const size_t pq_nchunks, const float *pq_dists, - float *dists_out) -{ - _mm_prefetch((char *)dists_out, _MM_HINT_T0); - _mm_prefetch((char *)pq_ids, _MM_HINT_T0); - _mm_prefetch((char *)(pq_ids + 64), _MM_HINT_T0); - _mm_prefetch((char *)(pq_ids + 128), _MM_HINT_T0); - memset(dists_out, 0, n_pts * sizeof(float)); - for (size_t chunk = 0; chunk < pq_nchunks; chunk++) - { - const float *chunk_dists = pq_dists + 256 * chunk; - if (chunk < pq_nchunks - 1) - { - _mm_prefetch((char *)(chunk_dists + 256), _MM_HINT_T0); - } - for (size_t idx = 0; idx < n_pts; idx++) - { - uint8_t pq_centerid = pq_ids[pq_nchunks * idx + chunk]; - dists_out[idx] += chunk_dists[pq_centerid]; - } +void pq_dist_lookup(const uint8_t *pq_ids, const size_t n_pts, + const size_t pq_nchunks, const float *pq_dists, + float *dists_out) { + _mm_prefetch((char *)dists_out, _MM_HINT_T0); + _mm_prefetch((char *)pq_ids, _MM_HINT_T0); + _mm_prefetch((char *)(pq_ids + 64), _MM_HINT_T0); + _mm_prefetch((char *)(pq_ids + 128), _MM_HINT_T0); + memset(dists_out, 0, n_pts * sizeof(float)); + for (size_t chunk = 0; chunk < pq_nchunks; chunk++) { + const float *chunk_dists = pq_dists + 256 * chunk; + if (chunk < pq_nchunks - 1) { + _mm_prefetch((char *)(chunk_dists + 256), _MM_HINT_T0); + } + for (size_t idx = 0; idx < n_pts; idx++) { + uint8_t pq_centerid = pq_ids[pq_nchunks * idx + chunk]; + dists_out[idx] += chunk_dists[pq_centerid]; } + } } // given training data in train_data of dimensions num_train * dim, generate @@ -345,367 +319,372 @@ void pq_dist_lookup(const uint8_t *pq_ids, const size_t n_pts, const size_t pq_n // 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 -int generate_pq_pivots(const float *const passed_train_data, size_t num_train, uint32_t dim, uint32_t num_centers, - uint32_t num_pq_chunks, uint32_t max_k_means_reps, std::string pq_pivots_path, - bool make_zero_mean) -{ - if (num_pq_chunks > dim) - { - diskann::cout << " Error: number of chunks more than dimension" << std::endl; - return -1; +int generate_pq_pivots(const float *const passed_train_data, size_t num_train, + uint32_t dim, uint32_t num_centers, + uint32_t num_pq_chunks, uint32_t max_k_means_reps, + std::string pq_pivots_path, bool make_zero_mean) { + if (num_pq_chunks > dim) { + diskann::cout << " Error: number of chunks more than dimension" + << std::endl; + return -1; + } + + std::unique_ptr train_data = + std::make_unique(num_train * dim); + std::memcpy(train_data.get(), passed_train_data, + num_train * dim * sizeof(float)); + + 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, METADATA_SIZE); + if (file_dim == dim && file_num_centers == num_centers) { + diskann::cout << "PQ pivot file exists. Not generating again" + << std::endl; + return -1; } - - std::unique_ptr train_data = std::make_unique(num_train * dim); - std::memcpy(train_data.get(), passed_train_data, num_train * dim * sizeof(float)); - - 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, METADATA_SIZE); - 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++) { + centroid[d] = 0; + } + if (make_zero_mean) { // If we use L2 distance, there is an option to + // translate all vectors to make them centered and + // then compute PQ. This needs to be set to false + // when using PQ for MIPS as such translations dont + // preserve inner products. + for (uint64_t d = 0; d < dim; d++) { + for (uint64_t p = 0; p < num_train; p++) { + centroid[d] += train_data[p * dim + d]; + } + centroid[d] /= num_train; } - // Calculate centroid and center the training data - std::unique_ptr centroid = std::make_unique(dim); - for (uint64_t d = 0; d < dim; d++) - { - centroid[d] = 0; + for (uint64_t d = 0; d < dim; d++) { + for (uint64_t p = 0; p < num_train; p++) { + train_data[p * dim + d] -= centroid[d]; + } } - if (make_zero_mean) - { // If we use L2 distance, there is an option to - // translate all vectors to make them centered and - // then compute PQ. This needs to be set to false - // when using PQ for MIPS as such translations dont - // preserve inner products. - for (uint64_t d = 0; d < dim; d++) - { - for (uint64_t p = 0; p < num_train; p++) - { - centroid[d] += train_data[p * dim + d]; - } - centroid[d] /= num_train; - } - - for (uint64_t d = 0; d < dim; d++) - { - for (uint64_t p = 0; p < num_train; p++) - { - train_data[p * dim + d] -= centroid[d]; - } - } + } + + std::vector chunk_offsets; + + size_t low_val = (size_t)std::floor((double)dim / (double)num_pq_chunks); + size_t high_val = (size_t)std::ceil((double)dim / (double)num_pq_chunks); + size_t max_num_high = dim - (low_val * num_pq_chunks); + size_t cur_num_high = 0; + size_t cur_bin_threshold = high_val; + + std::vector> bin_to_dims(num_pq_chunks); + tsl::robin_map dim_to_bin; + std::vector bin_loads(num_pq_chunks, 0); + + // Process dimensions not inserted by previous loop + for (uint32_t d = 0; d < dim; d++) { + if (dim_to_bin.find(d) != dim_to_bin.end()) continue; + auto cur_best = num_pq_chunks + 1; + float cur_best_load = std::numeric_limits::max(); + for (uint32_t b = 0; b < num_pq_chunks; b++) { + if (bin_loads[b] < cur_best_load && + bin_to_dims[b].size() < cur_bin_threshold) { + cur_best = b; + cur_best_load = bin_loads[b]; + } } - - std::vector chunk_offsets; - - size_t low_val = (size_t)std::floor((double)dim / (double)num_pq_chunks); - size_t high_val = (size_t)std::ceil((double)dim / (double)num_pq_chunks); - size_t max_num_high = dim - (low_val * num_pq_chunks); - size_t cur_num_high = 0; - size_t cur_bin_threshold = high_val; - - std::vector> bin_to_dims(num_pq_chunks); - tsl::robin_map dim_to_bin; - std::vector bin_loads(num_pq_chunks, 0); - - // Process dimensions not inserted by previous loop - for (uint32_t d = 0; d < dim; d++) - { - if (dim_to_bin.find(d) != dim_to_bin.end()) - continue; - auto cur_best = num_pq_chunks + 1; - float cur_best_load = std::numeric_limits::max(); - for (uint32_t b = 0; b < num_pq_chunks; b++) - { - if (bin_loads[b] < cur_best_load && bin_to_dims[b].size() < cur_bin_threshold) - { - cur_best = b; - cur_best_load = bin_loads[b]; - } - } - bin_to_dims[cur_best].push_back(d); - if (bin_to_dims[cur_best].size() == high_val) - { - cur_num_high++; - if (cur_num_high == max_num_high) - cur_bin_threshold = low_val; - } + bin_to_dims[cur_best].push_back(d); + if (bin_to_dims[cur_best].size() == high_val) { + cur_num_high++; + if (cur_num_high == max_num_high) cur_bin_threshold = low_val; } + } - chunk_offsets.clear(); - chunk_offsets.push_back(0); + chunk_offsets.clear(); + chunk_offsets.push_back(0); - for (uint32_t b = 0; b < num_pq_chunks; b++) - { - if (b > 0) - chunk_offsets.push_back(chunk_offsets[b - 1] + (uint32_t)bin_to_dims[b - 1].size()); - } - chunk_offsets.push_back(dim); + for (uint32_t b = 0; b < num_pq_chunks; b++) { + if (b > 0) + chunk_offsets.push_back(chunk_offsets[b - 1] + + (uint32_t)bin_to_dims[b - 1].size()); + } + chunk_offsets.push_back(dim); - full_pivot_data.reset(new float[num_centers * dim]); + full_pivot_data.reset(new float[num_centers * dim]); - for (size_t i = 0; i < num_pq_chunks; i++) - { - size_t cur_chunk_size = chunk_offsets[i + 1] - chunk_offsets[i]; + for (size_t i = 0; i < num_pq_chunks; i++) { + size_t cur_chunk_size = chunk_offsets[i + 1] - chunk_offsets[i]; - if (cur_chunk_size == 0) - continue; - std::unique_ptr cur_pivot_data = std::make_unique(num_centers * cur_chunk_size); - std::unique_ptr cur_data = std::make_unique(num_train * cur_chunk_size); - std::unique_ptr closest_center = std::make_unique(num_train); + if (cur_chunk_size == 0) continue; + std::unique_ptr cur_pivot_data = + std::make_unique(num_centers * cur_chunk_size); + std::unique_ptr cur_data = + std::make_unique(num_train * cur_chunk_size); + 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; + diskann::cout << "Processing chunk " << i << " with dimensions [" + << chunk_offsets[i] << ", " << chunk_offsets[i + 1] << ")" + << std::endl; #pragma omp parallel for schedule(static, 65536) - for (int64_t j = 0; j < (int64_t)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)); - } - - kmeans::kmeanspp_selecting_pivots(cur_data.get(), num_train, cur_chunk_size, cur_pivot_data.get(), 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()); - - 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)); - } + for (int64_t j = 0; j < (int64_t)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)); } - std::vector cumul_bytes(4, 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(), chunk_offsets.data(), - chunk_offsets.size(), 1, cumul_bytes[2]); - diskann::save_bin(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; + kmeans::kmeanspp_selecting_pivots(cur_data.get(), num_train, cur_chunk_size, + cur_pivot_data.get(), num_centers); - return 0; -} + 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()); -int generate_opq_pivots(const float *passed_train_data, size_t num_train, uint32_t dim, uint32_t num_centers, - uint32_t num_pq_chunks, std::string opq_pivots_path, bool make_zero_mean) -{ - if (num_pq_chunks > dim) - { - diskann::cout << " Error: number of chunks more than dimension" << std::endl; - return -1; + 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)); } + } + + std::vector cumul_bytes(4, 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(), chunk_offsets.data(), + chunk_offsets.size(), 1, cumul_bytes[2]); + diskann::save_bin(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; + + return 0; +} - std::unique_ptr train_data = std::make_unique(num_train * dim); - std::memcpy(train_data.get(), passed_train_data, num_train * dim * sizeof(float)); - - std::unique_ptr rotated_train_data = std::make_unique(num_train * dim); - std::unique_ptr rotated_and_quantized_train_data = std::make_unique(num_train * dim); - - std::unique_ptr full_pivot_data; - - // rotation matrix for OPQ - std::unique_ptr rotmat_tr; - - // matrices for SVD - std::unique_ptr Umat = std::make_unique(dim * dim); - std::unique_ptr Vmat_T = std::make_unique(dim * dim); - std::unique_ptr singular_values = std::make_unique(dim); - std::unique_ptr correlation_matrix = std::make_unique(dim * dim); - - // Calculate centroid and center the training data - std::unique_ptr centroid = std::make_unique(dim); - for (uint64_t d = 0; d < dim; d++) - { - centroid[d] = 0; +int generate_opq_pivots(const float *passed_train_data, size_t num_train, + uint32_t dim, uint32_t num_centers, + uint32_t num_pq_chunks, std::string opq_pivots_path, + bool make_zero_mean) { + if (num_pq_chunks > dim) { + diskann::cout << " Error: number of chunks more than dimension" + << std::endl; + return -1; + } + + std::unique_ptr train_data = + std::make_unique(num_train * dim); + std::memcpy(train_data.get(), passed_train_data, + num_train * dim * sizeof(float)); + + std::unique_ptr rotated_train_data = + std::make_unique(num_train * dim); + std::unique_ptr rotated_and_quantized_train_data = + std::make_unique(num_train * dim); + + std::unique_ptr full_pivot_data; + + // rotation matrix for OPQ + std::unique_ptr rotmat_tr; + + // matrices for SVD + std::unique_ptr Umat = std::make_unique(dim * dim); + std::unique_ptr Vmat_T = std::make_unique(dim * dim); + std::unique_ptr singular_values = std::make_unique(dim); + std::unique_ptr correlation_matrix = + std::make_unique(dim * dim); + + // Calculate centroid and center the training data + std::unique_ptr centroid = std::make_unique(dim); + for (uint64_t d = 0; d < dim; d++) { + centroid[d] = 0; + } + if (make_zero_mean) { // If we use L2 distance, there is an option to + // translate all vectors to make them centered and + // then compute PQ. This needs to be set to false + // when using PQ for MIPS as such translations dont + // preserve inner products. + for (uint64_t d = 0; d < dim; d++) { + for (uint64_t p = 0; p < num_train; p++) { + centroid[d] += train_data[p * dim + d]; + } + centroid[d] /= num_train; } - if (make_zero_mean) - { // If we use L2 distance, there is an option to - // translate all vectors to make them centered and - // then compute PQ. This needs to be set to false - // when using PQ for MIPS as such translations dont - // preserve inner products. - for (uint64_t d = 0; d < dim; d++) - { - for (uint64_t p = 0; p < num_train; p++) - { - centroid[d] += train_data[p * dim + d]; - } - centroid[d] /= num_train; - } - for (uint64_t d = 0; d < dim; d++) - { - for (uint64_t p = 0; p < num_train; p++) - { - train_data[p * dim + d] -= centroid[d]; - } - } + for (uint64_t d = 0; d < dim; d++) { + for (uint64_t p = 0; p < num_train; p++) { + train_data[p * dim + d] -= centroid[d]; + } } - - std::vector chunk_offsets; - - size_t low_val = (size_t)std::floor((double)dim / (double)num_pq_chunks); - size_t high_val = (size_t)std::ceil((double)dim / (double)num_pq_chunks); - size_t max_num_high = dim - (low_val * num_pq_chunks); - size_t cur_num_high = 0; - size_t cur_bin_threshold = high_val; - - std::vector> bin_to_dims(num_pq_chunks); - tsl::robin_map dim_to_bin; - std::vector bin_loads(num_pq_chunks, 0); - - // Process dimensions not inserted by previous loop - for (uint32_t d = 0; d < dim; d++) - { - if (dim_to_bin.find(d) != dim_to_bin.end()) - continue; - auto cur_best = num_pq_chunks + 1; - float cur_best_load = std::numeric_limits::max(); - for (uint32_t b = 0; b < num_pq_chunks; b++) - { - if (bin_loads[b] < cur_best_load && bin_to_dims[b].size() < cur_bin_threshold) - { - cur_best = b; - cur_best_load = bin_loads[b]; - } - } - bin_to_dims[cur_best].push_back(d); - if (bin_to_dims[cur_best].size() == high_val) - { - cur_num_high++; - if (cur_num_high == max_num_high) - cur_bin_threshold = low_val; - } + } + + std::vector chunk_offsets; + + size_t low_val = (size_t)std::floor((double)dim / (double)num_pq_chunks); + size_t high_val = (size_t)std::ceil((double)dim / (double)num_pq_chunks); + size_t max_num_high = dim - (low_val * num_pq_chunks); + size_t cur_num_high = 0; + size_t cur_bin_threshold = high_val; + + std::vector> bin_to_dims(num_pq_chunks); + tsl::robin_map dim_to_bin; + std::vector bin_loads(num_pq_chunks, 0); + + // Process dimensions not inserted by previous loop + for (uint32_t d = 0; d < dim; d++) { + if (dim_to_bin.find(d) != dim_to_bin.end()) continue; + auto cur_best = num_pq_chunks + 1; + float cur_best_load = std::numeric_limits::max(); + for (uint32_t b = 0; b < num_pq_chunks; b++) { + if (bin_loads[b] < cur_best_load && + bin_to_dims[b].size() < cur_bin_threshold) { + cur_best = b; + cur_best_load = bin_loads[b]; + } } - - chunk_offsets.clear(); - chunk_offsets.push_back(0); - - for (uint32_t b = 0; b < num_pq_chunks; b++) - { - if (b > 0) - chunk_offsets.push_back(chunk_offsets[b - 1] + (uint32_t)bin_to_dims[b - 1].size()); + bin_to_dims[cur_best].push_back(d); + if (bin_to_dims[cur_best].size() == high_val) { + cur_num_high++; + if (cur_num_high == max_num_high) cur_bin_threshold = low_val; } - chunk_offsets.push_back(dim); - - full_pivot_data.reset(new float[num_centers * dim]); - rotmat_tr.reset(new float[dim * dim]); - - std::memset(rotmat_tr.get(), 0, dim * dim * sizeof(float)); - for (uint32_t d1 = 0; d1 < dim; d1++) - *(rotmat_tr.get() + d1 * dim + d1) = 1; - - for (uint32_t rnd = 0; rnd < MAX_OPQ_ITERS; rnd++) - { - // rotate the training data using the current rotation matrix - cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, (MKL_INT)num_train, (MKL_INT)dim, (MKL_INT)dim, 1.0f, - train_data.get(), (MKL_INT)dim, rotmat_tr.get(), (MKL_INT)dim, 0.0f, rotated_train_data.get(), - (MKL_INT)dim); - - // compute the PQ pivots on the rotated space - for (size_t i = 0; i < num_pq_chunks; i++) - { - size_t cur_chunk_size = chunk_offsets[i + 1] - chunk_offsets[i]; - - if (cur_chunk_size == 0) - continue; - std::unique_ptr cur_pivot_data = std::make_unique(num_centers * cur_chunk_size); - std::unique_ptr cur_data = std::make_unique(num_train * cur_chunk_size); - 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; + } + + chunk_offsets.clear(); + chunk_offsets.push_back(0); + + for (uint32_t b = 0; b < num_pq_chunks; b++) { + if (b > 0) + chunk_offsets.push_back(chunk_offsets[b - 1] + + (uint32_t)bin_to_dims[b - 1].size()); + } + chunk_offsets.push_back(dim); + + full_pivot_data.reset(new float[num_centers * dim]); + rotmat_tr.reset(new float[dim * dim]); + + std::memset(rotmat_tr.get(), 0, dim * dim * sizeof(float)); + for (uint32_t d1 = 0; d1 < dim; d1++) *(rotmat_tr.get() + d1 * dim + d1) = 1; + + for (uint32_t rnd = 0; rnd < MAX_OPQ_ITERS; rnd++) { + // rotate the training data using the current rotation matrix + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, (MKL_INT)num_train, + (MKL_INT)dim, (MKL_INT)dim, 1.0f, train_data.get(), + (MKL_INT)dim, rotmat_tr.get(), (MKL_INT)dim, 0.0f, + rotated_train_data.get(), (MKL_INT)dim); + + // compute the PQ pivots on the rotated space + for (size_t i = 0; i < num_pq_chunks; i++) { + size_t cur_chunk_size = chunk_offsets[i + 1] - chunk_offsets[i]; + + if (cur_chunk_size == 0) continue; + std::unique_ptr cur_pivot_data = + std::make_unique(num_centers * cur_chunk_size); + std::unique_ptr cur_data = + std::make_unique(num_train * cur_chunk_size); + 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; #pragma omp parallel for schedule(static, 65536) - for (int64_t j = 0; j < (int64_t)num_train; j++) - { - std::memcpy(cur_data.get() + j * cur_chunk_size, rotated_train_data.get() + j * dim + chunk_offsets[i], - cur_chunk_size * sizeof(float)); - } - - if (rnd == 0) - { - kmeans::kmeanspp_selecting_pivots(cur_data.get(), num_train, cur_chunk_size, cur_pivot_data.get(), - num_centers); - } - else - { - for (uint64_t j = 0; j < num_centers; j++) - { - std::memcpy(cur_pivot_data.get() + j * cur_chunk_size, - full_pivot_data.get() + j * dim + chunk_offsets[i], cur_chunk_size * sizeof(float)); - } - } - - uint32_t num_lloyds_iters = 8; - kmeans::run_lloyds(cur_data.get(), num_train, cur_chunk_size, cur_pivot_data.get(), num_centers, - num_lloyds_iters, NULL, closest_center.get()); - - 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)); - } - - for (size_t j = 0; j < num_train; j++) - { - std::memcpy(rotated_and_quantized_train_data.get() + j * dim + chunk_offsets[i], - cur_pivot_data.get() + (size_t)closest_center[j] * cur_chunk_size, - cur_chunk_size * sizeof(float)); - } - } - - // compute the correlation matrix between the original data and the - // quantized data to compute the new rotation - cblas_sgemm(CblasRowMajor, CblasTrans, CblasNoTrans, (MKL_INT)dim, (MKL_INT)dim, (MKL_INT)num_train, 1.0f, - train_data.get(), (MKL_INT)dim, rotated_and_quantized_train_data.get(), (MKL_INT)dim, 0.0f, - correlation_matrix.get(), (MKL_INT)dim); - - // compute the SVD of the correlation matrix to help determine the new - // rotation matrix - uint32_t errcode = - LAPACKE_sgesdd(LAPACK_ROW_MAJOR, 'A', (MKL_INT)dim, (MKL_INT)dim, correlation_matrix.get(), (MKL_INT)dim, - singular_values.get(), Umat.get(), (MKL_INT)dim, Vmat_T.get(), (MKL_INT)dim); - - if (errcode > 0) - { - std::cout << "SVD failed to converge." << std::endl; - exit(-1); + for (int64_t j = 0; j < (int64_t)num_train; j++) { + std::memcpy(cur_data.get() + j * cur_chunk_size, + rotated_train_data.get() + j * dim + chunk_offsets[i], + cur_chunk_size * sizeof(float)); + } + + if (rnd == 0) { + kmeans::kmeanspp_selecting_pivots(cur_data.get(), num_train, + cur_chunk_size, cur_pivot_data.get(), + num_centers); + } else { + for (uint64_t j = 0; j < num_centers; j++) { + std::memcpy(cur_pivot_data.get() + j * cur_chunk_size, + full_pivot_data.get() + j * dim + chunk_offsets[i], + cur_chunk_size * sizeof(float)); } - - // compute the new rotation matrix from the singular vectors as R^T = U - // V^T - cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, (MKL_INT)dim, (MKL_INT)dim, (MKL_INT)dim, 1.0f, - Umat.get(), (MKL_INT)dim, Vmat_T.get(), (MKL_INT)dim, 0.0f, rotmat_tr.get(), (MKL_INT)dim); + } + + uint32_t num_lloyds_iters = 8; + kmeans::run_lloyds(cur_data.get(), num_train, cur_chunk_size, + cur_pivot_data.get(), num_centers, num_lloyds_iters, + NULL, closest_center.get()); + + 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)); + } + + for (size_t j = 0; j < num_train; j++) { + std::memcpy( + rotated_and_quantized_train_data.get() + j * dim + chunk_offsets[i], + cur_pivot_data.get() + (size_t)closest_center[j] * cur_chunk_size, + cur_chunk_size * sizeof(float)); + } } - std::vector cumul_bytes(4, 0); - cumul_bytes[0] = METADATA_SIZE; - cumul_bytes[1] = cumul_bytes[0] + diskann::save_bin(opq_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(opq_pivots_path.c_str(), centroid.get(), (size_t)dim, 1, cumul_bytes[1]); - cumul_bytes[3] = cumul_bytes[2] + diskann::save_bin(opq_pivots_path.c_str(), chunk_offsets.data(), - chunk_offsets.size(), 1, cumul_bytes[2]); - diskann::save_bin(opq_pivots_path.c_str(), cumul_bytes.data(), cumul_bytes.size(), 1, 0); - - diskann::cout << "Saved opq pivot data to " << opq_pivots_path << " of size " << cumul_bytes[cumul_bytes.size() - 1] - << "B." << std::endl; - - std::string rotmat_path = opq_pivots_path + "_rotation_matrix.bin"; - diskann::save_bin(rotmat_path.c_str(), rotmat_tr.get(), dim, dim); + // compute the correlation matrix between the original data and the + // quantized data to compute the new rotation + cblas_sgemm(CblasRowMajor, CblasTrans, CblasNoTrans, (MKL_INT)dim, + (MKL_INT)dim, (MKL_INT)num_train, 1.0f, train_data.get(), + (MKL_INT)dim, rotated_and_quantized_train_data.get(), + (MKL_INT)dim, 0.0f, correlation_matrix.get(), (MKL_INT)dim); + + // compute the SVD of the correlation matrix to help determine the new + // rotation matrix + uint32_t errcode = LAPACKE_sgesdd( + LAPACK_ROW_MAJOR, 'A', (MKL_INT)dim, (MKL_INT)dim, + correlation_matrix.get(), (MKL_INT)dim, singular_values.get(), + Umat.get(), (MKL_INT)dim, Vmat_T.get(), (MKL_INT)dim); + + if (errcode > 0) { + std::cout << "SVD failed to converge." << std::endl; + exit(-1); + } - return 0; + // compute the new rotation matrix from the singular vectors as R^T = U + // V^T + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, (MKL_INT)dim, + (MKL_INT)dim, (MKL_INT)dim, 1.0f, Umat.get(), (MKL_INT)dim, + Vmat_T.get(), (MKL_INT)dim, 0.0f, rotmat_tr.get(), + (MKL_INT)dim); + } + + std::vector cumul_bytes(4, 0); + cumul_bytes[0] = METADATA_SIZE; + cumul_bytes[1] = + cumul_bytes[0] + + diskann::save_bin(opq_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( + opq_pivots_path.c_str(), centroid.get(), + (size_t)dim, 1, cumul_bytes[1]); + cumul_bytes[3] = + cumul_bytes[2] + + diskann::save_bin(opq_pivots_path.c_str(), chunk_offsets.data(), + chunk_offsets.size(), 1, cumul_bytes[2]); + diskann::save_bin(opq_pivots_path.c_str(), cumul_bytes.data(), + cumul_bytes.size(), 1, 0); + + diskann::cout << "Saved opq pivot data to " << opq_pivots_path << " of size " + << cumul_bytes[cumul_bytes.size() - 1] << "B." << std::endl; + + std::string rotmat_path = opq_pivots_path + "_rotation_matrix.bin"; + diskann::save_bin(rotmat_path.c_str(), rotmat_tr.get(), dim, dim); + + return 0; } // streams the base file (data_file), and computes the closest centers in each @@ -714,344 +693,358 @@ int generate_opq_pivots(const float *passed_train_data, size_t num_train, uint32 // If the numbber of centers is < 256, it stores as byte vector, else as // 4-byte vector in binary format. template -int generate_pq_data_from_pivots(const std::string &data_file, uint32_t num_centers, uint32_t num_pq_chunks, - const std::string &pq_pivots_path, const std::string &pq_compressed_vectors_path, - bool use_opq) -{ - size_t read_blk_size = 64 * 1024 * 1024; - cached_ifstream base_reader(data_file, read_blk_size); - uint32_t npts32; - uint32_t basedim32; - base_reader.read((char *)&npts32, sizeof(uint32_t)); - base_reader.read((char *)&basedim32, sizeof(uint32_t)); - size_t num_points = npts32; - size_t dim = basedim32; - - std::unique_ptr full_pivot_data; - std::unique_ptr rotmat_tr; - std::unique_ptr centroid; - std::unique_ptr chunk_offsets; - - std::string inflated_pq_file = pq_compressed_vectors_path + "_inflated.bin"; - - if (!file_exists(pq_pivots_path)) - { - std::cout << "ERROR: PQ k-means pivot file not found" << std::endl; - throw diskann::ANNException("PQ k-means pivot file not found", -1); - } - else - { - size_t nr, nc; - std::unique_ptr file_offset_data; - - diskann::load_bin(pq_pivots_path.c_str(), file_offset_data, nr, nc, 0); - - if (nr != 4) - { - diskann::cout << "Error reading pq_pivots file " << pq_pivots_path - << ". Offsets dont contain correct metadata, # offsets = " << nr << ", but expecting 4."; - throw diskann::ANNException("Error reading pq_pivots file at offsets data.", -1, __FUNCSIG__, __FILE__, - __LINE__); - } - - diskann::load_bin(pq_pivots_path.c_str(), full_pivot_data, nr, nc, file_offset_data[0]); +int generate_pq_data_from_pivots(const std::string &data_file, + uint32_t num_centers, uint32_t num_pq_chunks, + const std::string &pq_pivots_path, + const std::string &pq_compressed_vectors_path, + bool use_opq) { + size_t read_blk_size = 64 * 1024 * 1024; + cached_ifstream base_reader(data_file, read_blk_size); + uint32_t npts32; + uint32_t basedim32; + base_reader.read((char *)&npts32, sizeof(uint32_t)); + base_reader.read((char *)&basedim32, sizeof(uint32_t)); + size_t num_points = npts32; + size_t dim = basedim32; + + std::unique_ptr full_pivot_data; + std::unique_ptr rotmat_tr; + std::unique_ptr centroid; + std::unique_ptr chunk_offsets; + + std::string inflated_pq_file = pq_compressed_vectors_path + "_inflated.bin"; + + if (!file_exists(pq_pivots_path)) { + std::cout << "ERROR: PQ k-means pivot file not found" << std::endl; + throw diskann::ANNException("PQ k-means pivot file not found", -1); + } else { + size_t nr, nc; + std::unique_ptr file_offset_data; - 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__); - } + diskann::load_bin(pq_pivots_path.c_str(), file_offset_data, nr, nc, + 0); - diskann::load_bin(pq_pivots_path.c_str(), centroid, nr, nc, file_offset_data[1]); + if (nr != 4) { + diskann::cout << "Error reading pq_pivots file " << pq_pivots_path + << ". Offsets dont contain correct metadata, # offsets = " + << nr << ", but expecting 4."; + throw diskann::ANNException( + "Error reading pq_pivots file at offsets data.", -1, __FUNCSIG__, + __FILE__, __LINE__); + } - 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__); - } + 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__); + } - diskann::load_bin(pq_pivots_path.c_str(), chunk_offsets, nr, nc, file_offset_data[2]); + diskann::load_bin(pq_pivots_path.c_str(), centroid, nr, nc, + file_offset_data[1]); - 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__); - } + 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__); + } - if (use_opq) - { - std::string rotmat_path = pq_pivots_path + "_rotation_matrix.bin"; - diskann::load_bin(rotmat_path.c_str(), rotmat_tr, nr, nc); - if (nr != (uint64_t)dim || nc != dim) - { - diskann::cout << "Error reading rotation matrix file." << std::endl; - throw diskann::ANNException("Error reading rotation matrix file.", -1, __FUNCSIG__, __FILE__, __LINE__); - } - } + diskann::load_bin(pq_pivots_path.c_str(), chunk_offsets, nr, nc, + file_offset_data[2]); + + 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; + if (use_opq) { + std::string rotmat_path = pq_pivots_path + "_rotation_matrix.bin"; + diskann::load_bin(rotmat_path.c_str(), rotmat_tr, nr, nc); + if (nr != (uint64_t)dim || nc != dim) { + diskann::cout << "Error reading rotation matrix file." << std::endl; + throw diskann::ANNException("Error reading rotation matrix file.", -1, + __FUNCSIG__, __FILE__, __LINE__); + } } - std::ofstream compressed_file_writer(pq_compressed_vectors_path, std::ios::binary); - uint32_t num_pq_chunks_u32 = num_pq_chunks; + diskann::cout << "Loaded PQ pivot information" << std::endl; + } - compressed_file_writer.write((char *)&num_points, sizeof(uint32_t)); - compressed_file_writer.write((char *)&num_pq_chunks_u32, sizeof(uint32_t)); + std::ofstream compressed_file_writer(pq_compressed_vectors_path, + std::ios::binary); + uint32_t num_pq_chunks_u32 = num_pq_chunks; - size_t block_size = num_points <= BLOCK_SIZE ? num_points : BLOCK_SIZE; + compressed_file_writer.write((char *)&num_points, sizeof(uint32_t)); + compressed_file_writer.write((char *)&num_pq_chunks_u32, sizeof(uint32_t)); + + size_t block_size = num_points <= BLOCK_SIZE ? num_points : BLOCK_SIZE; #ifdef SAVE_INFLATED_PQ - std::ofstream inflated_file_writer(inflated_pq_file, std::ios::binary); - inflated_file_writer.write((char *)&num_points, sizeof(uint32_t)); - inflated_file_writer.write((char *)&basedim32, sizeof(uint32_t)); + std::ofstream inflated_file_writer(inflated_pq_file, std::ios::binary); + inflated_file_writer.write((char *)&num_points, sizeof(uint32_t)); + inflated_file_writer.write((char *)&basedim32, sizeof(uint32_t)); - std::unique_ptr block_inflated_base = std::make_unique(block_size * dim); - std::memset(block_inflated_base.get(), 0, block_size * dim * sizeof(float)); + std::unique_ptr block_inflated_base = + std::make_unique(block_size * dim); + std::memset(block_inflated_base.get(), 0, block_size * dim * sizeof(float)); #endif - std::unique_ptr block_compressed_base = - std::make_unique(block_size * (size_t)num_pq_chunks); - std::memset(block_compressed_base.get(), 0, block_size * (size_t)num_pq_chunks * sizeof(uint32_t)); + std::unique_ptr block_compressed_base = + std::make_unique(block_size * (size_t)num_pq_chunks); + std::memset(block_compressed_base.get(), 0, + block_size * (size_t)num_pq_chunks * sizeof(uint32_t)); - std::unique_ptr block_data_T = std::make_unique(block_size * dim); - std::unique_ptr block_data_float = std::make_unique(block_size * dim); - std::unique_ptr block_data_tmp = std::make_unique(block_size * dim); + std::unique_ptr block_data_T = std::make_unique(block_size * dim); + std::unique_ptr block_data_float = + std::make_unique(block_size * dim); + std::unique_ptr block_data_tmp = + std::make_unique(block_size * dim); - size_t num_blocks = DIV_ROUND_UP(num_points, block_size); + size_t num_blocks = DIV_ROUND_UP(num_points, block_size); - for (size_t block = 0; block < num_blocks; block++) - { - size_t start_id = block * block_size; - size_t end_id = (std::min)((block + 1) * block_size, num_points); - size_t cur_blk_size = end_id - start_id; + for (size_t block = 0; block < num_blocks; block++) { + size_t start_id = block * block_size; + size_t end_id = (std::min)((block + 1) * block_size, num_points); + size_t cur_blk_size = end_id - start_id; - base_reader.read((char *)(block_data_T.get()), sizeof(T) * (cur_blk_size * dim)); - diskann::convert_types(block_data_T.get(), block_data_tmp.get(), cur_blk_size, dim); + base_reader.read((char *)(block_data_T.get()), + sizeof(T) * (cur_blk_size * dim)); + 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 + << ").." << std::flush; - for (size_t p = 0; p < cur_blk_size; p++) - { - for (uint64_t d = 0; d < dim; d++) - { - block_data_tmp[p * dim + d] -= centroid[d]; - } - } + for (size_t p = 0; p < cur_blk_size; p++) { + for (uint64_t d = 0; d < dim; d++) { + block_data_tmp[p * dim + d] -= centroid[d]; + } + } - for (size_t p = 0; p < cur_blk_size; p++) - { - for (uint64_t d = 0; d < dim; d++) - { - block_data_float[p * dim + d] = block_data_tmp[p * dim + d]; - } - } + for (size_t p = 0; p < cur_blk_size; p++) { + for (uint64_t d = 0; d < dim; d++) { + block_data_float[p * dim + d] = block_data_tmp[p * dim + d]; + } + } - if (use_opq) - { - // rotate the current block with the trained rotation matrix before - // PQ - cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, (MKL_INT)cur_blk_size, (MKL_INT)dim, (MKL_INT)dim, - 1.0f, block_data_float.get(), (MKL_INT)dim, rotmat_tr.get(), (MKL_INT)dim, 0.0f, - block_data_tmp.get(), (MKL_INT)dim); - std::memcpy(block_data_float.get(), block_data_tmp.get(), cur_blk_size * dim * sizeof(float)); - } + if (use_opq) { + // rotate the current block with the trained rotation matrix before + // PQ + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, + (MKL_INT)cur_blk_size, (MKL_INT)dim, (MKL_INT)dim, 1.0f, + block_data_float.get(), (MKL_INT)dim, rotmat_tr.get(), + (MKL_INT)dim, 0.0f, block_data_tmp.get(), (MKL_INT)dim); + std::memcpy(block_data_float.get(), block_data_tmp.get(), + cur_blk_size * dim * sizeof(float)); + } - for (size_t i = 0; i < num_pq_chunks; i++) - { - size_t cur_chunk_size = chunk_offsets[i + 1] - chunk_offsets[i]; - if (cur_chunk_size == 0) - continue; + for (size_t i = 0; i < num_pq_chunks; i++) { + size_t cur_chunk_size = chunk_offsets[i + 1] - chunk_offsets[i]; + if (cur_chunk_size == 0) continue; - std::unique_ptr cur_pivot_data = std::make_unique(num_centers * cur_chunk_size); - std::unique_ptr cur_data = std::make_unique(cur_blk_size * cur_chunk_size); - std::unique_ptr closest_center = std::make_unique(cur_blk_size); + std::unique_ptr cur_pivot_data = + std::make_unique(num_centers * cur_chunk_size); + std::unique_ptr cur_data = + std::make_unique(cur_blk_size * cur_chunk_size); + std::unique_ptr closest_center = + std::make_unique(cur_blk_size); #pragma omp parallel for schedule(static, 8192) - for (int64_t j = 0; j < (int64_t)cur_blk_size; j++) - { - for (size_t k = 0; k < cur_chunk_size; k++) - cur_data[j * cur_chunk_size + k] = block_data_float[j * dim + chunk_offsets[i] + k]; - } + for (int64_t j = 0; j < (int64_t)cur_blk_size; j++) { + for (size_t k = 0; k < cur_chunk_size; k++) + cur_data[j * cur_chunk_size + k] = + block_data_float[j * dim + chunk_offsets[i] + k]; + } #pragma omp parallel for schedule(static, 1) - for (int64_t j = 0; j < (int64_t)num_centers; j++) - { - std::memcpy(cur_pivot_data.get() + j * cur_chunk_size, - full_pivot_data.get() + j * dim + chunk_offsets[i], cur_chunk_size * sizeof(float)); - } + for (int64_t j = 0; j < (int64_t)num_centers; j++) { + std::memcpy(cur_pivot_data.get() + j * cur_chunk_size, + full_pivot_data.get() + j * dim + chunk_offsets[i], + cur_chunk_size * sizeof(float)); + } - math_utils::compute_closest_centers(cur_data.get(), cur_blk_size, cur_chunk_size, cur_pivot_data.get(), - num_centers, 1, closest_center.get()); + 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 < (int64_t)cur_blk_size; j++) - { - block_compressed_base[j * num_pq_chunks + i] = closest_center[j]; + for (int64_t j = 0; j < (int64_t)cur_blk_size; j++) { + block_compressed_base[j * num_pq_chunks + i] = closest_center[j]; #ifdef SAVE_INFLATED_PQ - for (size_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]; + for (size_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 - } - } + } + } - if (num_centers > 256) - { - compressed_file_writer.write((char *)(block_compressed_base.get()), - cur_blk_size * num_pq_chunks * sizeof(uint32_t)); - } - else - { - std::unique_ptr pVec = std::make_unique(cur_blk_size * num_pq_chunks); - diskann::convert_types(block_compressed_base.get(), pVec.get(), cur_blk_size, - num_pq_chunks); - compressed_file_writer.write((char *)(pVec.get()), cur_blk_size * num_pq_chunks * sizeof(uint8_t)); - } + if (num_centers > 256) { + compressed_file_writer.write( + (char *)(block_compressed_base.get()), + cur_blk_size * num_pq_chunks * sizeof(uint32_t)); + } else { + std::unique_ptr pVec = + std::make_unique(cur_blk_size * num_pq_chunks); + diskann::convert_types( + block_compressed_base.get(), pVec.get(), cur_blk_size, num_pq_chunks); + compressed_file_writer.write( + (char *)(pVec.get()), cur_blk_size * num_pq_chunks * sizeof(uint8_t)); + } #ifdef SAVE_INFLATED_PQ - inflated_file_writer.write((char *)(block_inflated_base.get()), cur_blk_size * dim * sizeof(float)); + inflated_file_writer.write((char *)(block_inflated_base.get()), + cur_blk_size * dim * sizeof(float)); #endif - diskann::cout << ".done." << std::endl; - } + diskann::cout << ".done." << std::endl; + } // Gopal. Splitting diskann_dll into separate DLLs for search and build. // This code should only be available in the "build" DLL. -#if defined(RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS) && defined(DISKANN_BUILD) - MallocExtension::instance()->ReleaseFreeMemory(); +#if defined(RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS) && \ + defined(DISKANN_BUILD) + MallocExtension::instance()->ReleaseFreeMemory(); #endif - compressed_file_writer.close(); + compressed_file_writer.close(); #ifdef SAVE_INFLATED_PQ - inflated_file_writer.close(); + inflated_file_writer.close(); #endif - return 0; + return 0; } template -void generate_disk_quantized_data(const std::string &data_file_to_use, const std::string &disk_pq_pivots_path, - const std::string &disk_pq_compressed_vectors_path, diskann::Metric compareMetric, - const double p_val, size_t &disk_pq_dims) -{ - size_t train_size, train_dim; - float *train_data; - - // instantiates train_data with random sample updates train_size - gen_random_slice(data_file_to_use.c_str(), p_val, train_data, train_size, train_dim); - diskann::cout << "Training data with " << train_size << " samples loaded." << std::endl; - - if (disk_pq_dims > train_dim) - disk_pq_dims = train_dim; - - std::cout << "Compressing base for disk-PQ into " << disk_pq_dims << " chunks " << std::endl; - generate_pq_pivots(train_data, train_size, (uint32_t)train_dim, 256, (uint32_t)disk_pq_dims, NUM_KMEANS_REPS_PQ, - disk_pq_pivots_path, false); - if (compareMetric == diskann::Metric::INNER_PRODUCT) - generate_pq_data_from_pivots(data_file_to_use, 256, (uint32_t)disk_pq_dims, disk_pq_pivots_path, - disk_pq_compressed_vectors_path); - else - generate_pq_data_from_pivots(data_file_to_use, 256, (uint32_t)disk_pq_dims, disk_pq_pivots_path, - disk_pq_compressed_vectors_path); - - delete[] train_data; +void generate_disk_quantized_data( + const std::string &data_file_to_use, const std::string &disk_pq_pivots_path, + const std::string &disk_pq_compressed_vectors_path, + diskann::Metric compareMetric, const double p_val, size_t &disk_pq_dims) { + size_t train_size, train_dim; + float *train_data; + + // instantiates train_data with random sample updates train_size + gen_random_slice(data_file_to_use.c_str(), p_val, train_data, train_size, + train_dim); + diskann::cout << "Training data with " << train_size << " samples loaded." + << std::endl; + + if (disk_pq_dims > train_dim) disk_pq_dims = train_dim; + + std::cout << "Compressing base for disk-PQ into " << disk_pq_dims + << " chunks " << std::endl; + generate_pq_pivots(train_data, train_size, (uint32_t)train_dim, 256, + (uint32_t)disk_pq_dims, NUM_KMEANS_REPS_PQ, + disk_pq_pivots_path, false); + if (compareMetric == diskann::Metric::INNER_PRODUCT) + generate_pq_data_from_pivots( + data_file_to_use, 256, (uint32_t)disk_pq_dims, disk_pq_pivots_path, + disk_pq_compressed_vectors_path); + else + generate_pq_data_from_pivots(data_file_to_use, 256, + (uint32_t)disk_pq_dims, disk_pq_pivots_path, + disk_pq_compressed_vectors_path); + + delete[] train_data; } template -void generate_quantized_data(const std::string &data_file_to_use, const std::string &pq_pivots_path, - const std::string &pq_compressed_vectors_path, diskann::Metric compareMetric, - const double p_val, const size_t num_pq_chunks, const bool use_opq, - const std::string &codebook_prefix) -{ - size_t train_size, train_dim; - float *train_data; - if (!file_exists(codebook_prefix)) - { - // instantiates train_data with random sample updates train_size - gen_random_slice(data_file_to_use.c_str(), p_val, train_data, train_size, train_dim); - diskann::cout << "Training data with " << train_size << " samples loaded." << std::endl; - - bool make_zero_mean = true; - if (compareMetric == diskann::Metric::INNER_PRODUCT) - make_zero_mean = false; - if (use_opq) // we also do not center the data for OPQ - make_zero_mean = false; - - if (!use_opq) - { - generate_pq_pivots(train_data, train_size, (uint32_t)train_dim, NUM_PQ_CENTROIDS, (uint32_t)num_pq_chunks, - NUM_KMEANS_REPS_PQ, pq_pivots_path, make_zero_mean); - } - else - { - generate_opq_pivots(train_data, train_size, (uint32_t)train_dim, NUM_PQ_CENTROIDS, (uint32_t)num_pq_chunks, - pq_pivots_path, make_zero_mean); - } - delete[] train_data; - } - else - { - diskann::cout << "Skip Training with predefined pivots in: " << pq_pivots_path << std::endl; +void generate_quantized_data(const std::string &data_file_to_use, + const std::string &pq_pivots_path, + const std::string &pq_compressed_vectors_path, + diskann::Metric compareMetric, const double p_val, + const size_t num_pq_chunks, const bool use_opq, + const std::string &codebook_prefix) { + size_t train_size, train_dim; + float *train_data; + if (!file_exists(codebook_prefix)) { + // instantiates train_data with random sample updates train_size + gen_random_slice(data_file_to_use.c_str(), p_val, train_data, train_size, + train_dim); + diskann::cout << "Training data with " << train_size << " samples loaded." + << std::endl; + + bool make_zero_mean = true; + if (compareMetric == diskann::Metric::INNER_PRODUCT) make_zero_mean = false; + if (use_opq) // we also do not center the data for OPQ + make_zero_mean = false; + + if (!use_opq) { + generate_pq_pivots(train_data, train_size, (uint32_t)train_dim, + NUM_PQ_CENTROIDS, (uint32_t)num_pq_chunks, + NUM_KMEANS_REPS_PQ, pq_pivots_path, make_zero_mean); + } else { + generate_opq_pivots(train_data, train_size, (uint32_t)train_dim, + NUM_PQ_CENTROIDS, (uint32_t)num_pq_chunks, + pq_pivots_path, make_zero_mean); } - generate_pq_data_from_pivots(data_file_to_use, NUM_PQ_CENTROIDS, (uint32_t)num_pq_chunks, pq_pivots_path, - pq_compressed_vectors_path, use_opq); + delete[] train_data; + } else { + diskann::cout << "Skip Training with predefined pivots in: " + << pq_pivots_path << std::endl; + } + generate_pq_data_from_pivots(data_file_to_use, NUM_PQ_CENTROIDS, + (uint32_t)num_pq_chunks, pq_pivots_path, + pq_compressed_vectors_path, use_opq); } // Instantations of supported templates -template DISKANN_DLLEXPORT int generate_pq_data_from_pivots(const std::string &data_file, uint32_t num_centers, - uint32_t num_pq_chunks, - const std::string &pq_pivots_path, - const std::string &pq_compressed_vectors_path, - bool use_opq); -template DISKANN_DLLEXPORT int generate_pq_data_from_pivots(const std::string &data_file, uint32_t num_centers, - uint32_t num_pq_chunks, - const std::string &pq_pivots_path, - const std::string &pq_compressed_vectors_path, - bool use_opq); -template DISKANN_DLLEXPORT int generate_pq_data_from_pivots(const std::string &data_file, uint32_t num_centers, - uint32_t num_pq_chunks, - const std::string &pq_pivots_path, - const std::string &pq_compressed_vectors_path, - bool use_opq); - -template DISKANN_DLLEXPORT void generate_disk_quantized_data(const std::string &data_file_to_use, - const std::string &disk_pq_pivots_path, - const std::string &disk_pq_compressed_vectors_path, - diskann::Metric compareMetric, const double p_val, - size_t &disk_pq_dims); +template DISKANN_DLLEXPORT int generate_pq_data_from_pivots( + const std::string &data_file, uint32_t num_centers, uint32_t num_pq_chunks, + const std::string &pq_pivots_path, + const std::string &pq_compressed_vectors_path, bool use_opq); +template DISKANN_DLLEXPORT int generate_pq_data_from_pivots( + const std::string &data_file, uint32_t num_centers, uint32_t num_pq_chunks, + const std::string &pq_pivots_path, + const std::string &pq_compressed_vectors_path, bool use_opq); +template DISKANN_DLLEXPORT int generate_pq_data_from_pivots( + const std::string &data_file, uint32_t num_centers, uint32_t num_pq_chunks, + const std::string &pq_pivots_path, + const std::string &pq_compressed_vectors_path, bool use_opq); + +template DISKANN_DLLEXPORT void generate_disk_quantized_data( + const std::string &data_file_to_use, const std::string &disk_pq_pivots_path, + const std::string &disk_pq_compressed_vectors_path, + diskann::Metric compareMetric, const double p_val, size_t &disk_pq_dims); template DISKANN_DLLEXPORT void generate_disk_quantized_data( const std::string &data_file_to_use, const std::string &disk_pq_pivots_path, - const std::string &disk_pq_compressed_vectors_path, diskann::Metric compareMetric, const double p_val, - size_t &disk_pq_dims); - -template DISKANN_DLLEXPORT void generate_disk_quantized_data(const std::string &data_file_to_use, - const std::string &disk_pq_pivots_path, - const std::string &disk_pq_compressed_vectors_path, - diskann::Metric compareMetric, const double p_val, - size_t &disk_pq_dims); - -template DISKANN_DLLEXPORT void generate_quantized_data(const std::string &data_file_to_use, - const std::string &pq_pivots_path, - const std::string &pq_compressed_vectors_path, - diskann::Metric compareMetric, const double p_val, - const size_t num_pq_chunks, const bool use_opq, - const std::string &codebook_prefix); - -template DISKANN_DLLEXPORT void generate_quantized_data(const std::string &data_file_to_use, - const std::string &pq_pivots_path, - const std::string &pq_compressed_vectors_path, - diskann::Metric compareMetric, const double p_val, - const size_t num_pq_chunks, const bool use_opq, - const std::string &codebook_prefix); - -template DISKANN_DLLEXPORT void generate_quantized_data(const std::string &data_file_to_use, - const std::string &pq_pivots_path, - const std::string &pq_compressed_vectors_path, - diskann::Metric compareMetric, const double p_val, - const size_t num_pq_chunks, const bool use_opq, - const std::string &codebook_prefix); -} // namespace diskann + const std::string &disk_pq_compressed_vectors_path, + diskann::Metric compareMetric, const double p_val, size_t &disk_pq_dims); + +template DISKANN_DLLEXPORT void generate_disk_quantized_data( + const std::string &data_file_to_use, const std::string &disk_pq_pivots_path, + const std::string &disk_pq_compressed_vectors_path, + diskann::Metric compareMetric, const double p_val, size_t &disk_pq_dims); + +template DISKANN_DLLEXPORT void generate_quantized_data( + const std::string &data_file_to_use, const std::string &pq_pivots_path, + const std::string &pq_compressed_vectors_path, + diskann::Metric compareMetric, const double p_val, + const size_t num_pq_chunks, const bool use_opq, + const std::string &codebook_prefix); + +template DISKANN_DLLEXPORT void generate_quantized_data( + const std::string &data_file_to_use, const std::string &pq_pivots_path, + const std::string &pq_compressed_vectors_path, + diskann::Metric compareMetric, const double p_val, + const size_t num_pq_chunks, const bool use_opq, + const std::string &codebook_prefix); + +template DISKANN_DLLEXPORT void generate_quantized_data( + const std::string &data_file_to_use, const std::string &pq_pivots_path, + const std::string &pq_compressed_vectors_path, + diskann::Metric compareMetric, const double p_val, + const size_t num_pq_chunks, const bool use_opq, + const std::string &codebook_prefix); +} // namespace diskann diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index 4d137e560..fa78f689b 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -21,1536 +21,1457 @@ #define NODE_SECTOR_NO(node_id) (((uint64_t)(node_id)) / nnodes_per_sector + 1) // obtains region of sector containing node -#define OFFSET_TO_NODE(sector_buf, node_id) \ - ((char *)sector_buf + (((uint64_t)node_id) % nnodes_per_sector) * max_node_len) +#define OFFSET_TO_NODE(sector_buf, node_id) \ + ((char *)sector_buf + \ + (((uint64_t)node_id) % nnodes_per_sector) * max_node_len) // returns region of `node_buf` containing [NNBRS][NBR_ID(uint32_t)] -#define OFFSET_TO_NODE_NHOOD(node_buf) (unsigned *)((char *)node_buf + disk_bytes_per_point) +#define OFFSET_TO_NODE_NHOOD(node_buf) \ + (unsigned *)((char *)node_buf + disk_bytes_per_point) // returns region of `node_buf` containing [COORD(T)] #define OFFSET_TO_NODE_COORDS(node_buf) (T *)(node_buf) // sector # beyond the end of graph where data for id is present for reordering -#define VECTOR_SECTOR_NO(id) (((uint64_t)(id)) / nvecs_per_sector + reorder_data_start_sector) +#define VECTOR_SECTOR_NO(id) \ + (((uint64_t)(id)) / nvecs_per_sector + reorder_data_start_sector) // sector # beyond the end of graph where data for id is present for reordering -#define VECTOR_SECTOR_OFFSET(id) ((((uint64_t)(id)) % nvecs_per_sector) * data_dim * sizeof(float)) +#define VECTOR_SECTOR_OFFSET(id) \ + ((((uint64_t)(id)) % nvecs_per_sector) * data_dim * sizeof(float)) -namespace diskann -{ +namespace diskann { template -PQFlashIndex::PQFlashIndex(std::shared_ptr &fileReader, diskann::Metric m) - : reader(fileReader), metric(m) -{ - if (m == diskann::Metric::COSINE || m == diskann::Metric::INNER_PRODUCT) - { - if (std::is_floating_point::value) - { - diskann::cout << "Cosine metric chosen for (normalized) float data." - "Changing distance to L2 to boost accuracy." - << std::endl; - metric = diskann::Metric::L2; - } - 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; - } +PQFlashIndex::PQFlashIndex( + std::shared_ptr &fileReader, diskann::Metric m) + : reader(fileReader), metric(m) { + if (m == diskann::Metric::COSINE || m == diskann::Metric::INNER_PRODUCT) { + if (std::is_floating_point::value) { + diskann::cout << "Cosine metric chosen for (normalized) float data." + "Changing distance to L2 to boost accuracy." + << std::endl; + metric = diskann::Metric::L2; + } 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(metric)); - this->dist_cmp_float.reset(diskann::get_distance_function(metric)); + this->dist_cmp.reset(diskann::get_distance_function(metric)); + this->dist_cmp_float.reset(diskann::get_distance_function(metric)); } -template PQFlashIndex::~PQFlashIndex() -{ +template +PQFlashIndex::~PQFlashIndex() { #ifndef EXEC_ENV_OLS - if (data != nullptr) - { - delete[] data; - } + if (data != nullptr) { + delete[] data; + } #endif - if (centroid_data != nullptr) - aligned_free(centroid_data); - // delete backing bufs for nhood and coord cache - if (nhood_cache_buf != nullptr) - { - delete[] nhood_cache_buf; - diskann::aligned_free(coord_cache_buf); - } - - if (load_flag) - { - diskann::cout << "Clearing scratch" << std::endl; - ScratchStoreManager> manager(this->thread_data); - manager.destroy(); - this->reader->deregister_all_threads(); - reader->close(); - } - if (_pts_to_label_offsets != nullptr) - { - delete[] _pts_to_label_offsets; - } + if (centroid_data != nullptr) aligned_free(centroid_data); + // delete backing bufs for nhood and coord cache + if (nhood_cache_buf != nullptr) { + delete[] nhood_cache_buf; + diskann::aligned_free(coord_cache_buf); + } - if (_pts_to_labels != nullptr) - { - delete[] _pts_to_labels; - } + if (load_flag) { + diskann::cout << "Clearing scratch" << std::endl; + ScratchStoreManager> manager(this->thread_data); + manager.destroy(); + this->reader->deregister_all_threads(); + reader->close(); + } + if (_pts_to_label_offsets != nullptr) { + delete[] _pts_to_label_offsets; + } + + if (_pts_to_labels != nullptr) { + delete[] _pts_to_labels; + } } template -void PQFlashIndex::setup_thread_data(uint64_t nthreads, uint64_t visited_reserve) -{ - diskann::cout << "Setting up thread-specific contexts for nthreads: " << nthreads << std::endl; +void PQFlashIndex::setup_thread_data(uint64_t nthreads, + uint64_t visited_reserve) { + diskann::cout << "Setting up thread-specific contexts for nthreads: " + << nthreads << std::endl; // omp parallel for to generate unique thread IDs #pragma omp parallel for num_threads((int)nthreads) - for (int64_t thread = 0; thread < (int64_t)nthreads; thread++) - { + for (int64_t thread = 0; thread < (int64_t)nthreads; thread++) { #pragma omp critical - { - SSDThreadData *data = new SSDThreadData(this->aligned_dim, visited_reserve); - this->reader->register_thread(); - data->ctx = this->reader->get_ctx(); - this->thread_data.push(data); - } + { + SSDThreadData *data = + new SSDThreadData(this->aligned_dim, visited_reserve); + this->reader->register_thread(); + data->ctx = this->reader->get_ctx(); + this->thread_data.push(data); } - load_flag = true; + } + load_flag = true; } -template void PQFlashIndex::load_cache_list(std::vector &node_list) -{ - diskann::cout << "Loading the cache list into memory.." << std::flush; - size_t num_cached_nodes = node_list.size(); - - // borrow thread data - ScratchStoreManager> manager(this->thread_data); - auto this_thread_data = manager.scratch_space(); - IOContext &ctx = this_thread_data->ctx; - - nhood_cache_buf = new uint32_t[num_cached_nodes * (max_degree + 1)]; - memset(nhood_cache_buf, 0, num_cached_nodes * (max_degree + 1)); - - size_t coord_cache_buf_len = num_cached_nodes * aligned_dim; - diskann::alloc_aligned((void **)&coord_cache_buf, coord_cache_buf_len * sizeof(T), 8 * sizeof(T)); - memset(coord_cache_buf, 0, coord_cache_buf_len * sizeof(T)); +template +void PQFlashIndex::load_cache_list( + std::vector &node_list) { + diskann::cout << "Loading the cache list into memory.." << std::flush; + size_t num_cached_nodes = node_list.size(); + + // borrow thread data + ScratchStoreManager> manager(this->thread_data); + auto this_thread_data = manager.scratch_space(); + IOContext &ctx = this_thread_data->ctx; + + nhood_cache_buf = new uint32_t[num_cached_nodes * (max_degree + 1)]; + memset(nhood_cache_buf, 0, num_cached_nodes * (max_degree + 1)); + + size_t coord_cache_buf_len = num_cached_nodes * aligned_dim; + diskann::alloc_aligned((void **)&coord_cache_buf, + coord_cache_buf_len * sizeof(T), 8 * sizeof(T)); + memset(coord_cache_buf, 0, coord_cache_buf_len * sizeof(T)); + + size_t BLOCK_SIZE = 8; + size_t num_blocks = DIV_ROUND_UP(num_cached_nodes, BLOCK_SIZE); + + for (size_t block = 0; block < num_blocks; block++) { + size_t start_idx = block * BLOCK_SIZE; + size_t end_idx = (std::min)(num_cached_nodes, (block + 1) * BLOCK_SIZE); + std::vector read_reqs; + std::vector> nhoods; + for (size_t node_idx = start_idx; node_idx < end_idx; node_idx++) { + AlignedRead read; + char *buf = nullptr; + alloc_aligned((void **)&buf, SECTOR_LEN, SECTOR_LEN); + nhoods.push_back(std::make_pair(node_list[node_idx], buf)); + read.len = SECTOR_LEN; + read.buf = buf; + read.offset = NODE_SECTOR_NO(node_list[node_idx]) * SECTOR_LEN; + read_reqs.push_back(read); + } - size_t BLOCK_SIZE = 8; - size_t num_blocks = DIV_ROUND_UP(num_cached_nodes, BLOCK_SIZE); + reader->read(read_reqs, ctx); - for (size_t block = 0; block < num_blocks; block++) - { - size_t start_idx = block * BLOCK_SIZE; - size_t end_idx = (std::min)(num_cached_nodes, (block + 1) * BLOCK_SIZE); - std::vector read_reqs; - std::vector> nhoods; - for (size_t node_idx = start_idx; node_idx < end_idx; node_idx++) - { - AlignedRead read; - char *buf = nullptr; - alloc_aligned((void **)&buf, SECTOR_LEN, SECTOR_LEN); - nhoods.push_back(std::make_pair(node_list[node_idx], buf)); - read.len = SECTOR_LEN; - read.buf = buf; - read.offset = NODE_SECTOR_NO(node_list[node_idx]) * SECTOR_LEN; - read_reqs.push_back(read); - } - - reader->read(read_reqs, ctx); - - size_t node_idx = start_idx; - for (uint32_t i = 0; i < read_reqs.size(); i++) - { -#if defined(_WINDOWS) && defined(USE_BING_INFRA) // this block is to handle failed reads in - // production settings - if ((*ctx.m_pRequestsStatus)[i] != IOContext::READ_SUCCESS) - { - continue; - } + size_t node_idx = start_idx; + for (uint32_t i = 0; i < read_reqs.size(); i++) { +#if defined(_WINDOWS) && \ + defined(USE_BING_INFRA) // this block is to handle failed reads in + // production settings + if ((*ctx.m_pRequestsStatus)[i] != IOContext::READ_SUCCESS) { + continue; + } #endif - auto &nhood = nhoods[i]; - char *node_buf = OFFSET_TO_NODE(nhood.second, nhood.first); - T *node_coords = OFFSET_TO_NODE_COORDS(node_buf); - T *cached_coords = coord_cache_buf + node_idx * aligned_dim; - memcpy(cached_coords, node_coords, disk_bytes_per_point); - coord_cache.insert(std::make_pair(nhood.first, cached_coords)); - - // insert node nhood into nhood_cache - uint32_t *node_nhood = OFFSET_TO_NODE_NHOOD(node_buf); - - auto nnbrs = *node_nhood; - uint32_t *nbrs = node_nhood + 1; - std::pair cnhood; - cnhood.first = nnbrs; - cnhood.second = nhood_cache_buf + node_idx * (max_degree + 1); - memcpy(cnhood.second, nbrs, nnbrs * sizeof(uint32_t)); - nhood_cache.insert(std::make_pair(nhood.first, cnhood)); - aligned_free(nhood.second); - node_idx++; - } + auto &nhood = nhoods[i]; + char *node_buf = OFFSET_TO_NODE(nhood.second, nhood.first); + T *node_coords = OFFSET_TO_NODE_COORDS(node_buf); + T *cached_coords = coord_cache_buf + node_idx * aligned_dim; + memcpy(cached_coords, node_coords, disk_bytes_per_point); + coord_cache.insert(std::make_pair(nhood.first, cached_coords)); + + // insert node nhood into nhood_cache + uint32_t *node_nhood = OFFSET_TO_NODE_NHOOD(node_buf); + + auto nnbrs = *node_nhood; + uint32_t *nbrs = node_nhood + 1; + std::pair cnhood; + cnhood.first = nnbrs; + cnhood.second = nhood_cache_buf + node_idx * (max_degree + 1); + memcpy(cnhood.second, nbrs, nnbrs * sizeof(uint32_t)); + nhood_cache.insert(std::make_pair(nhood.first, cnhood)); + aligned_free(nhood.second); + node_idx++; } - diskann::cout << "..done." << std::endl; + } + diskann::cout << "..done." << std::endl; } #ifdef EXEC_ENV_OLS template -void PQFlashIndex::generate_cache_list_from_sample_queries(MemoryMappedFiles &files, std::string sample_bin, - uint64_t l_search, uint64_t beamwidth, - uint64_t num_nodes_to_cache, uint32_t nthreads, - std::vector &node_list) -{ +void PQFlashIndex::generate_cache_list_from_sample_queries( + MemoryMappedFiles &files, std::string sample_bin, uint64_t l_search, + uint64_t beamwidth, uint64_t num_nodes_to_cache, uint32_t nthreads, + std::vector &node_list) { #else template -void PQFlashIndex::generate_cache_list_from_sample_queries(std::string sample_bin, uint64_t l_search, - uint64_t beamwidth, uint64_t num_nodes_to_cache, - uint32_t nthreads, - std::vector &node_list) -{ +void PQFlashIndex::generate_cache_list_from_sample_queries( + std::string sample_bin, uint64_t l_search, uint64_t beamwidth, + uint64_t num_nodes_to_cache, uint32_t nthreads, + std::vector &node_list) { #endif - this->count_visited_nodes = true; - this->node_visit_counter.clear(); - this->node_visit_counter.resize(this->num_points); - for (uint32_t i = 0; i < node_visit_counter.size(); i++) - { - this->node_visit_counter[i].first = i; - this->node_visit_counter[i].second = 0; - } + this->count_visited_nodes = true; + this->node_visit_counter.clear(); + this->node_visit_counter.resize(this->num_points); + for (uint32_t i = 0; i < node_visit_counter.size(); i++) { + this->node_visit_counter[i].first = i; + this->node_visit_counter[i].second = 0; + } - uint64_t sample_num, sample_dim, sample_aligned_dim; - T *samples; + uint64_t sample_num, sample_dim, sample_aligned_dim; + T *samples; #ifdef EXEC_ENV_OLS - if (files.fileExists(sample_bin)) - { - diskann::load_aligned_bin(files, sample_bin, samples, sample_num, sample_dim, sample_aligned_dim); - } + if (files.fileExists(sample_bin)) { + diskann::load_aligned_bin(files, sample_bin, samples, sample_num, + sample_dim, sample_aligned_dim); + } #else - if (file_exists(sample_bin)) - { - diskann::load_aligned_bin(sample_bin, samples, sample_num, sample_dim, sample_aligned_dim); - } + if (file_exists(sample_bin)) { + diskann::load_aligned_bin(sample_bin, samples, sample_num, sample_dim, + sample_aligned_dim); + } #endif - else - { - diskann::cerr << "Sample bin file not found. Not generating cache." << std::endl; - return; - } + else { + diskann::cerr << "Sample bin file not found. Not generating cache." + << std::endl; + return; + } - std::vector tmp_result_ids_64(sample_num, 0); - std::vector tmp_result_dists(sample_num, 0); + std::vector tmp_result_ids_64(sample_num, 0); + std::vector tmp_result_dists(sample_num, 0); #pragma omp parallel for schedule(dynamic, 1) num_threads(nthreads) - for (int64_t 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_dists.data() + (i * 1), beamwidth); - } - - std::sort(this->node_visit_counter.begin(), node_visit_counter.end(), - [](std::pair &left, std::pair &right) { - return left.second > right.second; - }); - node_list.clear(); - node_list.shrink_to_fit(); - node_list.reserve(num_nodes_to_cache); - for (uint64_t i = 0; i < num_nodes_to_cache; i++) - { - node_list.push_back(this->node_visit_counter[i].first); - } - this->count_visited_nodes = false; - - diskann::aligned_free(samples); + for (int64_t 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_dists.data() + (i * 1), beamwidth); + } + + std::sort(this->node_visit_counter.begin(), node_visit_counter.end(), + [](std::pair &left, + std::pair &right) { + return left.second > right.second; + }); + node_list.clear(); + node_list.shrink_to_fit(); + node_list.reserve(num_nodes_to_cache); + for (uint64_t i = 0; i < num_nodes_to_cache; i++) { + node_list.push_back(this->node_visit_counter[i].first); + } + this->count_visited_nodes = false; + + diskann::aligned_free(samples); } template -void PQFlashIndex::cache_bfs_levels(uint64_t num_nodes_to_cache, std::vector &node_list, - const bool shuffle) -{ - std::random_device rng; - std::mt19937 urng(rng()); - - tsl::robin_set node_set; - - // Do not cache more than 10% of the nodes in the index - uint64_t tenp_nodes = (uint64_t)(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; +void PQFlashIndex::cache_bfs_levels(uint64_t num_nodes_to_cache, + std::vector &node_list, + const bool shuffle) { + std::random_device rng; + std::mt19937 urng(rng()); + + tsl::robin_set node_set; + + // Do not cache more than 10% of the nodes in the index + uint64_t tenp_nodes = (uint64_t)(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 + ScratchStoreManager> manager(this->thread_data); + auto this_thread_data = manager.scratch_space(); + IOContext &ctx = this_thread_data->ctx; + + std::unique_ptr> cur_level, prev_level; + cur_level = std::make_unique>(); + prev_level = std::make_unique>(); + + for (uint64_t miter = 0; miter < num_medoids; miter++) { + cur_level->insert(medoids[miter]); + } + + uint64_t lvl = 1; + uint64_t prev_node_set_size = 0; + while ((node_set.size() + cur_level->size() < num_nodes_to_cache) && + cur_level->size() != 0) { + // swap prev_level and cur_level + std::swap(prev_level, cur_level); + // clear cur_level + cur_level->clear(); + + std::vector nodes_to_expand; + + for (const uint32_t &id : *prev_level) { + if (node_set.find(id) != node_set.end()) { + continue; + } + node_set.insert(id); + nodes_to_expand.push_back(id); } - diskann::cout << "Caching " << num_nodes_to_cache << "..." << std::endl; - - // borrow thread data - ScratchStoreManager> manager(this->thread_data); - auto this_thread_data = manager.scratch_space(); - IOContext &ctx = this_thread_data->ctx; - - std::unique_ptr> cur_level, prev_level; - cur_level = std::make_unique>(); - prev_level = std::make_unique>(); - for (uint64_t miter = 0; miter < num_medoids; miter++) - { - cur_level->insert(medoids[miter]); - } + if (shuffle) + std::shuffle(nodes_to_expand.begin(), nodes_to_expand.end(), urng); + else + std::sort(nodes_to_expand.begin(), nodes_to_expand.end()); - uint64_t lvl = 1; - uint64_t prev_node_set_size = 0; - while ((node_set.size() + cur_level->size() < num_nodes_to_cache) && cur_level->size() != 0) - { - // swap prev_level and cur_level - std::swap(prev_level, cur_level); - // clear cur_level - cur_level->clear(); - - std::vector nodes_to_expand; - - for (const uint32_t &id : *prev_level) - { - if (node_set.find(id) != node_set.end()) - { - continue; - } - node_set.insert(id); - nodes_to_expand.push_back(id); + diskann::cout << "Level: " << lvl << std::flush; + bool finish_flag = false; + + uint64_t BLOCK_SIZE = 1024; + uint64_t nblocks = DIV_ROUND_UP(nodes_to_expand.size(), BLOCK_SIZE); + for (size_t block = 0; block < nblocks && !finish_flag; block++) { + diskann::cout << "." << std::flush; + 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> nhoods; + for (size_t cur_pt = start; cur_pt < end; cur_pt++) { + char *buf = nullptr; + alloc_aligned((void **)&buf, SECTOR_LEN, SECTOR_LEN); + nhoods.emplace_back(nodes_to_expand[cur_pt], buf); + AlignedRead read; + read.len = SECTOR_LEN; + read.buf = buf; + 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 (uint32_t i = 0; i < read_reqs.size(); i++) { +#if defined(_WINDOWS) && \ + defined(USE_BING_INFRA) // this block is to handle read failures in + // production settings + if ((*ctx.m_pRequestsStatus)[i] != IOContext::READ_SUCCESS) { + continue; } - - if (shuffle) - std::shuffle(nodes_to_expand.begin(), nodes_to_expand.end(), urng); - else - std::sort(nodes_to_expand.begin(), nodes_to_expand.end()); - - diskann::cout << "Level: " << lvl << std::flush; - bool finish_flag = false; - - uint64_t BLOCK_SIZE = 1024; - uint64_t nblocks = DIV_ROUND_UP(nodes_to_expand.size(), BLOCK_SIZE); - for (size_t block = 0; block < nblocks && !finish_flag; block++) - { - diskann::cout << "." << std::flush; - 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> nhoods; - for (size_t cur_pt = start; cur_pt < end; cur_pt++) - { - char *buf = nullptr; - alloc_aligned((void **)&buf, SECTOR_LEN, SECTOR_LEN); - nhoods.emplace_back(nodes_to_expand[cur_pt], buf); - AlignedRead read; - read.len = SECTOR_LEN; - read.buf = buf; - 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 (uint32_t i = 0; i < read_reqs.size(); i++) - { -#if defined(_WINDOWS) && defined(USE_BING_INFRA) // this block is to handle read failures in - // production settings - if ((*ctx.m_pRequestsStatus)[i] != IOContext::READ_SUCCESS) - { - continue; - } #endif - auto &nhood = nhoods[i]; - - // insert node coord into coord_cache - char *node_buf = OFFSET_TO_NODE(nhood.second, nhood.first); - uint32_t *node_nhood = OFFSET_TO_NODE_NHOOD(node_buf); - uint64_t nnbrs = (uint64_t)*node_nhood; - uint32_t *nbrs = node_nhood + 1; - // explore next level - for (uint64_t j = 0; j < nnbrs && !finish_flag; j++) - { - if (node_set.find(nbrs[j]) == node_set.end()) - { - cur_level->insert(nbrs[j]); - } - if (cur_level->size() + node_set.size() >= num_nodes_to_cache) - { - finish_flag = true; - } - } - aligned_free(nhood.second); - } + auto &nhood = nhoods[i]; + + // insert node coord into coord_cache + char *node_buf = OFFSET_TO_NODE(nhood.second, nhood.first); + uint32_t *node_nhood = OFFSET_TO_NODE_NHOOD(node_buf); + uint64_t nnbrs = (uint64_t)*node_nhood; + uint32_t *nbrs = node_nhood + 1; + // explore next level + for (uint64_t j = 0; j < nnbrs && !finish_flag; j++) { + if (node_set.find(nbrs[j]) == node_set.end()) { + cur_level->insert(nbrs[j]); + } + if (cur_level->size() + node_set.size() >= num_nodes_to_cache) { + finish_flag = true; + } } - - diskann::cout << ". #nodes: " << node_set.size() - prev_node_set_size - << ", #nodes thus far: " << node_list.size() << std::endl; - prev_node_set_size = node_set.size(); - lvl++; + aligned_free(nhood.second); + } } - assert(node_set.size() + cur_level->size() == num_nodes_to_cache || cur_level->size() == 0); + diskann::cout << ". #nodes: " << node_set.size() - prev_node_set_size + << ", #nodes thus far: " << node_list.size() << std::endl; + prev_node_set_size = node_set.size(); + lvl++; + } - node_list.clear(); - node_list.reserve(node_set.size() + cur_level->size()); - for (auto node : node_set) - node_list.push_back(node); - for (auto node : *cur_level) - node_list.push_back(node); + assert(node_set.size() + cur_level->size() == num_nodes_to_cache || + cur_level->size() == 0); - diskann::cout << "Level: " << lvl << std::flush; - diskann::cout << ". #nodes: " << node_list.size() - prev_node_set_size << ", #nodes thus far: " << node_list.size() - << std::endl; - diskann::cout << "done" << std::endl; -} - -template void PQFlashIndex::use_medoids_data_as_centroids() -{ - if (centroid_data != nullptr) - aligned_free(centroid_data); - alloc_aligned(((void **)¢roid_data), num_medoids * aligned_dim * sizeof(float), 32); - std::memset(centroid_data, 0, num_medoids * aligned_dim * sizeof(float)); + node_list.clear(); + node_list.reserve(node_set.size() + cur_level->size()); + for (auto node : node_set) node_list.push_back(node); + for (auto node : *cur_level) node_list.push_back(node); - // borrow ctx - ScratchStoreManager> manager(this->thread_data); - auto data = manager.scratch_space(); - IOContext &ctx = data->ctx; - diskann::cout << "Loading centroid data from medoids vector data of " << num_medoids << " medoid(s)" << std::endl; - for (uint64_t cur_m = 0; cur_m < num_medoids; cur_m++) - { - auto medoid = medoids[cur_m]; - // read medoid nhood - char *medoid_buf = nullptr; - alloc_aligned((void **)&medoid_buf, SECTOR_LEN, SECTOR_LEN); - std::vector medoid_read(1); - medoid_read[0].len = SECTOR_LEN; - medoid_read[0].buf = medoid_buf; - medoid_read[0].offset = NODE_SECTOR_NO(medoid) * SECTOR_LEN; - reader->read(medoid_read, ctx); - - // all data about medoid - char *medoid_node_buf = OFFSET_TO_NODE(medoid_buf, medoid); - - // add medoid coords to `coord_cache` - T *medoid_coords = new T[data_dim]; - T *medoid_disk_coords = OFFSET_TO_NODE_COORDS(medoid_node_buf); - memcpy(medoid_coords, medoid_disk_coords, disk_bytes_per_point); - - if (!use_disk_index_pq) - { - for (uint32_t i = 0; i < data_dim; i++) - centroid_data[cur_m * aligned_dim + i] = medoid_coords[i]; - } - else - { - disk_pq_table.inflate_vector((uint8_t *)medoid_coords, (centroid_data + cur_m * aligned_dim)); - } + diskann::cout << "Level: " << lvl << std::flush; + diskann::cout << ". #nodes: " << node_list.size() - prev_node_set_size + << ", #nodes thus far: " << node_list.size() << std::endl; + diskann::cout << "done" << std::endl; +} - aligned_free(medoid_buf); - delete[] medoid_coords; +template +void PQFlashIndex::use_medoids_data_as_centroids() { + if (centroid_data != nullptr) aligned_free(centroid_data); + alloc_aligned(((void **)¢roid_data), + num_medoids * aligned_dim * sizeof(float), 32); + std::memset(centroid_data, 0, num_medoids * aligned_dim * sizeof(float)); + + // borrow ctx + ScratchStoreManager> manager(this->thread_data); + auto data = manager.scratch_space(); + IOContext &ctx = data->ctx; + diskann::cout << "Loading centroid data from medoids vector data of " + << num_medoids << " medoid(s)" << std::endl; + for (uint64_t cur_m = 0; cur_m < num_medoids; cur_m++) { + auto medoid = medoids[cur_m]; + // read medoid nhood + char *medoid_buf = nullptr; + alloc_aligned((void **)&medoid_buf, SECTOR_LEN, SECTOR_LEN); + std::vector medoid_read(1); + medoid_read[0].len = SECTOR_LEN; + medoid_read[0].buf = medoid_buf; + medoid_read[0].offset = NODE_SECTOR_NO(medoid) * SECTOR_LEN; + reader->read(medoid_read, ctx); + + // all data about medoid + char *medoid_node_buf = OFFSET_TO_NODE(medoid_buf, medoid); + + // add medoid coords to `coord_cache` + T *medoid_coords = new T[data_dim]; + T *medoid_disk_coords = OFFSET_TO_NODE_COORDS(medoid_node_buf); + memcpy(medoid_coords, medoid_disk_coords, disk_bytes_per_point); + + if (!use_disk_index_pq) { + for (uint32_t i = 0; i < data_dim; i++) + centroid_data[cur_m * aligned_dim + i] = medoid_coords[i]; + } else { + disk_pq_table.inflate_vector((uint8_t *)medoid_coords, + (centroid_data + cur_m * aligned_dim)); } + + aligned_free(medoid_buf); + delete[] medoid_coords; + } } template -inline int32_t PQFlashIndex::get_filter_number(const LabelT &filter_label) -{ - int idx = -1; - for (uint32_t i = 0; i < _filter_list.size(); i++) - { - if (_filter_list[i] == filter_label) - { - idx = i; - break; - } +inline int32_t PQFlashIndex::get_filter_number( + const LabelT &filter_label) { + int idx = -1; + for (uint32_t i = 0; i < _filter_list.size(); i++) { + if (_filter_list[i] == filter_label) { + idx = i; + break; } - return idx; + } + return idx; } template -std::unordered_map PQFlashIndex::load_label_map(const std::string &labels_map_file) -{ - std::unordered_map string_to_int_mp; - std::ifstream map_reader(labels_map_file); - std::string line, token; - LabelT token_as_num; - std::string label_str; - while (std::getline(map_reader, line)) - { - std::istringstream iss(line); - getline(iss, token, '\t'); - label_str = token; - getline(iss, token, '\t'); - token_as_num = std::stoul(token); - string_to_int_mp[label_str] = token_as_num; - } - return string_to_int_mp; +std::unordered_map PQFlashIndex::load_label_map( + const std::string &labels_map_file) { + std::unordered_map string_to_int_mp; + std::ifstream map_reader(labels_map_file); + std::string line, token; + LabelT token_as_num; + std::string label_str; + while (std::getline(map_reader, line)) { + std::istringstream iss(line); + getline(iss, token, '\t'); + label_str = token; + getline(iss, token, '\t'); + token_as_num = std::stoul(token); + string_to_int_mp[label_str] = token_as_num; + } + return string_to_int_mp; } template -LabelT PQFlashIndex::get_converted_label(const std::string &filter_label) -{ - if (_label_map.find(filter_label) != _label_map.end()) - { - return _label_map[filter_label]; - } - std::stringstream stream; - stream << "Unable to find label in the Label Map"; - diskann::cerr << stream.str() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); +LabelT PQFlashIndex::get_converted_label( + const std::string &filter_label) { + if (_label_map.find(filter_label) != _label_map.end()) { + return _label_map[filter_label]; + } + std::stringstream stream; + stream << "Unable to find label in the Label Map"; + diskann::cerr << stream.str() << std::endl; + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); } template -void PQFlashIndex::get_label_file_metadata(std::string map_file, uint32_t &num_pts, - uint32_t &num_total_labels) -{ - std::ifstream infile(map_file); - std::string line, token; - num_pts = 0; - num_total_labels = 0; - - while (std::getline(infile, line)) - { - std::istringstream iss(line); - while (getline(iss, token, ',')) - { - token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); - token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); - num_total_labels++; - } - num_pts++; +void PQFlashIndex::get_label_file_metadata( + std::string map_file, uint32_t &num_pts, uint32_t &num_total_labels) { + std::ifstream infile(map_file); + std::string line, token; + num_pts = 0; + num_total_labels = 0; + + while (std::getline(infile, line)) { + std::istringstream iss(line); + while (getline(iss, token, ',')) { + token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); + token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); + num_total_labels++; } + num_pts++; + } - diskann::cout << "Labels file metadata: num_points: " << num_pts << ", #total_labels: " << num_total_labels - << std::endl; - infile.close(); + diskann::cout << "Labels file metadata: num_points: " << num_pts + << ", #total_labels: " << num_total_labels << std::endl; + infile.close(); } template -inline bool PQFlashIndex::point_has_label(uint32_t point_id, uint32_t label_id) -{ - uint32_t start_vec = _pts_to_label_offsets[point_id]; - uint32_t num_lbls = _pts_to_labels[start_vec]; - bool ret_val = false; - for (uint32_t i = 0; i < num_lbls; i++) - { - if (_pts_to_labels[start_vec + 1 + i] == label_id) - { - ret_val = true; - break; - } +inline bool PQFlashIndex::point_has_label(uint32_t point_id, + uint32_t label_id) { + uint32_t start_vec = _pts_to_label_offsets[point_id]; + uint32_t num_lbls = _pts_to_labels[start_vec]; + bool ret_val = false; + for (uint32_t i = 0; i < num_lbls; i++) { + if (_pts_to_labels[start_vec + 1 + i] == label_id) { + ret_val = true; + break; } - return ret_val; + } + return ret_val; } template -void PQFlashIndex::parse_label_file(const std::string &label_file, size_t &num_points_labels) -{ - std::ifstream infile(label_file); - if (infile.fail()) - { - throw diskann::ANNException(std::string("Failed to open file ") + label_file, -1); +void PQFlashIndex::parse_label_file(const std::string &label_file, + size_t &num_points_labels) { + std::ifstream infile(label_file); + if (infile.fail()) { + throw diskann::ANNException( + std::string("Failed to open file ") + label_file, -1); + } + + std::string line, token; + uint32_t line_cnt = 0; + + uint32_t num_pts_in_label_file; + uint32_t num_total_labels; + get_label_file_metadata(label_file, num_pts_in_label_file, num_total_labels); + + _pts_to_label_offsets = new uint32_t[num_pts_in_label_file]; + _pts_to_labels = new uint32_t[num_pts_in_label_file + num_total_labels]; + uint32_t counter = 0; + + while (std::getline(infile, line)) { + std::istringstream iss(line); + std::vector lbls(0); + + _pts_to_label_offsets[line_cnt] = counter; + uint32_t &num_lbls_in_cur_pt = _pts_to_labels[counter]; + num_lbls_in_cur_pt = 0; + counter++; + getline(iss, token, '\t'); + std::istringstream new_iss(token); + while (getline(new_iss, token, ',')) { + token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); + token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); + LabelT token_as_num = std::stoul(token); + if (_labels.find(token_as_num) == _labels.end()) { + _filter_list.emplace_back(token_as_num); + } + int32_t filter_num = get_filter_number(token_as_num); + if (filter_num == -1) { + diskann::cout << "Error!! " << std::endl; + exit(-1); + } + _pts_to_labels[counter++] = filter_num; + num_lbls_in_cur_pt++; + _labels.insert(token_as_num); } - std::string line, token; - uint32_t line_cnt = 0; - - uint32_t num_pts_in_label_file; - uint32_t num_total_labels; - get_label_file_metadata(label_file, num_pts_in_label_file, num_total_labels); - - _pts_to_label_offsets = new uint32_t[num_pts_in_label_file]; - _pts_to_labels = new uint32_t[num_pts_in_label_file + num_total_labels]; - uint32_t counter = 0; - - while (std::getline(infile, line)) - { - std::istringstream iss(line); - std::vector lbls(0); - - _pts_to_label_offsets[line_cnt] = counter; - uint32_t &num_lbls_in_cur_pt = _pts_to_labels[counter]; - num_lbls_in_cur_pt = 0; - counter++; - getline(iss, token, '\t'); - std::istringstream new_iss(token); - while (getline(new_iss, token, ',')) - { - token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); - token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); - LabelT token_as_num = std::stoul(token); - if (_labels.find(token_as_num) == _labels.end()) - { - _filter_list.emplace_back(token_as_num); - } - int32_t filter_num = get_filter_number(token_as_num); - if (filter_num == -1) - { - diskann::cout << "Error!! " << std::endl; - exit(-1); - } - _pts_to_labels[counter++] = filter_num; - num_lbls_in_cur_pt++; - _labels.insert(token_as_num); - } - - if (num_lbls_in_cur_pt == 0) - { - diskann::cout << "No label found for point " << line_cnt << std::endl; - exit(-1); - } - line_cnt++; + if (num_lbls_in_cur_pt == 0) { + diskann::cout << "No label found for point " << line_cnt << std::endl; + exit(-1); } - infile.close(); - num_points_labels = line_cnt; + line_cnt++; + } + infile.close(); + num_points_labels = line_cnt; } -template void PQFlashIndex::set_universal_label(const LabelT &label) -{ - int32_t temp_filter_num = get_filter_number(label); - if (temp_filter_num == -1) - { - diskann::cout << "Error, could not find universal label." << std::endl; - } - else - { - _use_universal_label = true; - _universal_filter_num = (uint32_t)temp_filter_num; - } +template +void PQFlashIndex::set_universal_label(const LabelT &label) { + int32_t temp_filter_num = get_filter_number(label); + if (temp_filter_num == -1) { + diskann::cout << "Error, could not find universal label." << std::endl; + } else { + _use_universal_label = true; + _universal_filter_num = (uint32_t)temp_filter_num; + } } #ifdef EXEC_ENV_OLS template -int PQFlashIndex::load(MemoryMappedFiles &files, uint32_t num_threads, const char *index_prefix) -{ +int PQFlashIndex::load(MemoryMappedFiles &files, + uint32_t num_threads, + const char *index_prefix) { #else -template int PQFlashIndex::load(uint32_t num_threads, const char *index_prefix) -{ +template +int PQFlashIndex::load(uint32_t num_threads, + const char *index_prefix) { #endif - std::string pq_table_bin = std::string(index_prefix) + "_pq_pivots.bin"; - std::string pq_compressed_vectors = std::string(index_prefix) + "_pq_compressed.bin"; - std::string disk_index_file = std::string(index_prefix) + "_disk.index"; + std::string pq_table_bin = std::string(index_prefix) + "_pq_pivots.bin"; + std::string pq_compressed_vectors = + std::string(index_prefix) + "_pq_compressed.bin"; + std::string disk_index_file = std::string(index_prefix) + "_disk.index"; #ifdef EXEC_ENV_OLS - return load_from_separate_paths(files, num_threads, disk_index_file.c_str(), pq_table_bin.c_str(), - pq_compressed_vectors.c_str()); + return load_from_separate_paths(files, num_threads, disk_index_file.c_str(), + pq_table_bin.c_str(), + pq_compressed_vectors.c_str()); #else - return load_from_separate_paths(num_threads, disk_index_file.c_str(), pq_table_bin.c_str(), - pq_compressed_vectors.c_str()); + return load_from_separate_paths(num_threads, disk_index_file.c_str(), + pq_table_bin.c_str(), + pq_compressed_vectors.c_str()); #endif } #ifdef EXEC_ENV_OLS template -int PQFlashIndex::load_from_separate_paths(diskann::MemoryMappedFiles &files, uint32_t num_threads, - const char *index_filepath, const char *pivots_filepath, - const char *compressed_filepath) -{ +int PQFlashIndex::load_from_separate_paths( + diskann::MemoryMappedFiles &files, uint32_t num_threads, + const char *index_filepath, const char *pivots_filepath, + const char *compressed_filepath) { #else template -int PQFlashIndex::load_from_separate_paths(uint32_t num_threads, const char *index_filepath, - const char *pivots_filepath, const char *compressed_filepath) -{ +int PQFlashIndex::load_from_separate_paths( + uint32_t num_threads, const char *index_filepath, + const char *pivots_filepath, const char *compressed_filepath) { #endif - std::string pq_table_bin = pivots_filepath; - std::string pq_compressed_vectors = compressed_filepath; - std::string disk_index_file = index_filepath; - std::string medoids_file = std::string(disk_index_file) + "_medoids.bin"; - std::string centroids_file = std::string(disk_index_file) + "_centroids.bin"; - - std::string labels_file = std ::string(disk_index_file) + "_labels.txt"; - std::string labels_to_medoids = std ::string(disk_index_file) + "_labels_to_medoids.txt"; - std::string dummy_map_file = std ::string(disk_index_file) + "_dummy_map.txt"; - std::string labels_map_file = std ::string(disk_index_file) + "_labels_map.txt"; - size_t num_pts_in_label_file = 0; - - size_t pq_file_dim, pq_file_num_centroids; + std::string pq_table_bin = pivots_filepath; + std::string pq_compressed_vectors = compressed_filepath; + std::string disk_index_file = index_filepath; + std::string medoids_file = std::string(disk_index_file) + "_medoids.bin"; + std::string centroids_file = std::string(disk_index_file) + "_centroids.bin"; + + std::string labels_file = std ::string(disk_index_file) + "_labels.txt"; + std::string labels_to_medoids = + std ::string(disk_index_file) + "_labels_to_medoids.txt"; + std::string dummy_map_file = std ::string(disk_index_file) + "_dummy_map.txt"; + std::string labels_map_file = + std ::string(disk_index_file) + "_labels_map.txt"; + size_t num_pts_in_label_file = 0; + + size_t pq_file_dim, pq_file_num_centroids; #ifdef EXEC_ENV_OLS - get_bin_metadata(files, pq_table_bin, pq_file_num_centroids, pq_file_dim, METADATA_SIZE); + get_bin_metadata(files, pq_table_bin, pq_file_num_centroids, pq_file_dim, + METADATA_SIZE); #else - get_bin_metadata(pq_table_bin, pq_file_num_centroids, pq_file_dim, METADATA_SIZE); + get_bin_metadata(pq_table_bin, pq_file_num_centroids, pq_file_dim, + METADATA_SIZE); #endif - this->disk_index_file = disk_index_file; + this->disk_index_file = disk_index_file; - if (pq_file_num_centroids != 256) - { - diskann::cout << "Error. Number of PQ centroids is not 256. Exiting." << std::endl; - return -1; - } - - this->data_dim = pq_file_dim; - // will reset later if we use PQ on disk - this->disk_data_dim = this->data_dim; - // will change later if we use PQ on disk or if we are using - // inner product without PQ - this->disk_bytes_per_point = this->data_dim * sizeof(T); - this->aligned_dim = ROUND_UP(pq_file_dim, 8); - - size_t npts_u64, nchunks_u64; + if (pq_file_num_centroids != 256) { + diskann::cout << "Error. Number of PQ centroids is not 256. Exiting." + << std::endl; + return -1; + } + + this->data_dim = pq_file_dim; + // will reset later if we use PQ on disk + this->disk_data_dim = this->data_dim; + // will change later if we use PQ on disk or if we are using + // inner product without PQ + this->disk_bytes_per_point = this->data_dim * sizeof(T); + this->aligned_dim = ROUND_UP(pq_file_dim, 8); + + size_t npts_u64, nchunks_u64; #ifdef EXEC_ENV_OLS - diskann::load_bin(files, pq_compressed_vectors, this->data, npts_u64, nchunks_u64); + diskann::load_bin(files, pq_compressed_vectors, this->data, npts_u64, + nchunks_u64); #else - diskann::load_bin(pq_compressed_vectors, this->data, npts_u64, nchunks_u64); + diskann::load_bin(pq_compressed_vectors, this->data, npts_u64, + nchunks_u64); #endif - this->num_points = npts_u64; - this->n_chunks = nchunks_u64; - if (file_exists(labels_file)) - { - parse_label_file(labels_file, num_pts_in_label_file); - assert(num_pts_in_label_file == this->num_points); - _label_map = load_label_map(labels_map_file); - if (file_exists(labels_to_medoids)) - { - std::ifstream medoid_stream(labels_to_medoids); - assert(medoid_stream.is_open()); - std::string line, token; - - _filter_to_medoid_id.clear(); - try - { - while (std::getline(medoid_stream, line)) - { - std::istringstream iss(line); - uint32_t cnt = 0; - uint32_t medoid = 0; - LabelT label; - while (std::getline(iss, token, ',')) - { - if (cnt == 0) - label = std::stoul(token); - else - medoid = (uint32_t)stoul(token); - cnt++; - } - _filter_to_medoid_id[label] = medoid; - } - } - catch (std::system_error &e) - { - throw FileException(labels_to_medoids, e, __FUNCSIG__, __FILE__, __LINE__); - } - } - std::string univ_label_file = std ::string(disk_index_file) + "_universal_label.txt"; - if (file_exists(univ_label_file)) - { - std::ifstream universal_label_reader(univ_label_file); - assert(universal_label_reader.is_open()); - std::string univ_label; - universal_label_reader >> univ_label; - universal_label_reader.close(); - LabelT label_as_num = std::stoul(univ_label); - set_universal_label(label_as_num); + this->num_points = npts_u64; + this->n_chunks = nchunks_u64; + if (file_exists(labels_file)) { + parse_label_file(labels_file, num_pts_in_label_file); + assert(num_pts_in_label_file == this->num_points); + _label_map = load_label_map(labels_map_file); + if (file_exists(labels_to_medoids)) { + std::ifstream medoid_stream(labels_to_medoids); + assert(medoid_stream.is_open()); + std::string line, token; + + _filter_to_medoid_id.clear(); + try { + while (std::getline(medoid_stream, line)) { + std::istringstream iss(line); + uint32_t cnt = 0; + uint32_t medoid = 0; + LabelT label; + while (std::getline(iss, token, ',')) { + if (cnt == 0) + label = std::stoul(token); + else + medoid = (uint32_t)stoul(token); + cnt++; + } + _filter_to_medoid_id[label] = medoid; } - if (file_exists(dummy_map_file)) - { - std::ifstream dummy_map_stream(dummy_map_file); - assert(dummy_map_stream.is_open()); - std::string line, token; - - while (std::getline(dummy_map_stream, line)) - { - std::istringstream iss(line); - uint32_t cnt = 0; - uint32_t dummy_id; - uint32_t real_id; - while (std::getline(iss, token, ',')) - { - if (cnt == 0) - dummy_id = (uint32_t)stoul(token); - else - real_id = (uint32_t)stoul(token); - cnt++; - } - _dummy_pts.insert(dummy_id); - _has_dummy_pts.insert(real_id); - _dummy_to_real_map[dummy_id] = real_id; - - if (_real_to_dummy_map.find(real_id) == _real_to_dummy_map.end()) - _real_to_dummy_map[real_id] = std::vector(); - - _real_to_dummy_map[real_id].emplace_back(dummy_id); - } - dummy_map_stream.close(); - diskann::cout << "Loaded dummy map" << std::endl; + } catch (std::system_error &e) { + throw FileException(labels_to_medoids, e, __FUNCSIG__, __FILE__, + __LINE__); + } + } + std::string univ_label_file = + std ::string(disk_index_file) + "_universal_label.txt"; + if (file_exists(univ_label_file)) { + std::ifstream universal_label_reader(univ_label_file); + assert(universal_label_reader.is_open()); + std::string univ_label; + universal_label_reader >> univ_label; + universal_label_reader.close(); + LabelT label_as_num = std::stoul(univ_label); + set_universal_label(label_as_num); + } + if (file_exists(dummy_map_file)) { + std::ifstream dummy_map_stream(dummy_map_file); + assert(dummy_map_stream.is_open()); + std::string line, token; + + while (std::getline(dummy_map_stream, line)) { + std::istringstream iss(line); + uint32_t cnt = 0; + uint32_t dummy_id; + uint32_t real_id; + while (std::getline(iss, token, ',')) { + if (cnt == 0) + dummy_id = (uint32_t)stoul(token); + else + real_id = (uint32_t)stoul(token); + cnt++; } + _dummy_pts.insert(dummy_id); + _has_dummy_pts.insert(real_id); + _dummy_to_real_map[dummy_id] = real_id; + + if (_real_to_dummy_map.find(real_id) == _real_to_dummy_map.end()) + _real_to_dummy_map[real_id] = std::vector(); + + _real_to_dummy_map[real_id].emplace_back(dummy_id); + } + dummy_map_stream.close(); + diskann::cout << "Loaded dummy map" << std::endl; } + } #ifdef EXEC_ENV_OLS - pq_table.load_pq_centroid_bin(files, pq_table_bin.c_str(), nchunks_u64); + 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); + 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 + << "Loaded PQ centroids and in-memory compressed vectors. #points: " + << num_points << " #dim: " << data_dim << " #aligned_dim: " << aligned_dim + << " #chunks: " << n_chunks << std::endl; - if (n_chunks > MAX_PQ_CHUNKS) - { - std::stringstream stream; - stream << "Error loading index. Ensure that max PQ bytes for in-memory " - "PQ data does not exceed " - << MAX_PQ_CHUNKS << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - - std::string disk_pq_pivots_path = this->disk_index_file + "_pq_pivots.bin"; - if (file_exists(disk_pq_pivots_path)) - { - use_disk_index_pq = true; + if (n_chunks > MAX_PQ_CHUNKS) { + std::stringstream stream; + stream << "Error loading index. Ensure that max PQ bytes for in-memory " + "PQ data does not exceed " + << MAX_PQ_CHUNKS << std::endl; + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } + + std::string disk_pq_pivots_path = this->disk_index_file + "_pq_pivots.bin"; + if (file_exists(disk_pq_pivots_path)) { + use_disk_index_pq = true; #ifdef EXEC_ENV_OLS - // giving 0 chunks to make the pq_table infer from the - // chunk_offsets file the correct value - disk_pq_table.load_pq_centroid_bin(files, disk_pq_pivots_path.c_str(), 0); + // giving 0 chunks to make the pq_table infer from the + // chunk_offsets file the correct value + disk_pq_table.load_pq_centroid_bin(files, disk_pq_pivots_path.c_str(), 0); #else - // giving 0 chunks to make the pq_table infer from the - // chunk_offsets file the correct value - disk_pq_table.load_pq_centroid_bin(disk_pq_pivots_path.c_str(), 0); + // giving 0 chunks to make the pq_table infer from the + // 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_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." - << std::endl; - } + disk_pq_n_chunks = disk_pq_table.get_num_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." << 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); + // 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); #else - std::ifstream index_metadata(disk_index_file, std::ios::binary); + std::ifstream index_metadata(disk_index_file, std::ios::binary); #endif - uint32_t nr, nc; // metadata itself is stored as bin format (nr is number of - // metadata, nc should be 1) - READ_U32(index_metadata, nr); - READ_U32(index_metadata, nc); - - uint64_t disk_nnodes; - uint64_t disk_ndims; // can be disk PQ dim if disk_PQ is set to true - READ_U64(index_metadata, disk_nnodes); - READ_U64(index_metadata, disk_ndims); - - if (disk_nnodes != num_points) - { - diskann::cout << "Mismatch in #points for compressed data file and disk " - "index file: " - << disk_nnodes << " vs " << num_points << std::endl; - 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 - disk_bytes_per_point) / sizeof(uint32_t)) - 1; - - if (max_degree > MAX_GRAPH_DEGREE) - { - std::stringstream stream; - stream << "Error loading index. Ensure that max graph degree (R) does " - "not exceed " - << MAX_GRAPH_DEGREE << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - - // setting up concept of frozen points in disk index for streaming-DiskANN - READ_U64(index_metadata, this->num_frozen_points); - uint64_t file_frozen_id; - 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, this->reorder_data_exists); - if (this->reorder_data_exists) - { - if (this->use_disk_index_pq == false) - { - throw ANNException("Reordering is designed for used with disk PQ " - "compression option", - -1, __FUNCSIG__, __FILE__, __LINE__); - } - READ_U64(index_metadata, this->reorder_data_start_sector); - READ_U64(index_metadata, this->ndims_reorder_vecs); - READ_U64(index_metadata, this->nvecs_per_sector); + uint32_t nr, nc; // metadata itself is stored as bin format (nr is number of + // metadata, nc should be 1) + READ_U32(index_metadata, nr); + READ_U32(index_metadata, nc); + + uint64_t disk_nnodes; + uint64_t disk_ndims; // can be disk PQ dim if disk_PQ is set to true + READ_U64(index_metadata, disk_nnodes); + READ_U64(index_metadata, disk_ndims); + + if (disk_nnodes != num_points) { + diskann::cout << "Mismatch in #points for compressed data file and disk " + "index file: " + << disk_nnodes << " vs " << num_points << std::endl; + 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 - disk_bytes_per_point) / sizeof(uint32_t)) - 1; + + if (max_degree > MAX_GRAPH_DEGREE) { + std::stringstream stream; + stream << "Error loading index. Ensure that max graph degree (R) does " + "not exceed " + << MAX_GRAPH_DEGREE << std::endl; + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } + + // setting up concept of frozen points in disk index for streaming-DiskANN + READ_U64(index_metadata, this->num_frozen_points); + uint64_t file_frozen_id; + 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, this->reorder_data_exists); + if (this->reorder_data_exists) { + if (this->use_disk_index_pq == false) { + throw ANNException( + "Reordering is designed for used with disk PQ " + "compression option", + -1, __FUNCSIG__, __FILE__, __LINE__); } + READ_U64(index_metadata, this->reorder_data_start_sector); + READ_U64(index_metadata, this->ndims_reorder_vecs); + READ_U64(index_metadata, this->nvecs_per_sector); + } - 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; + 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; + delete[] bytes; #else - index_metadata.close(); + index_metadata.close(); #endif #ifndef EXEC_ENV_OLS - // open AlignedFileReader handle to index_file - std::string index_fname(disk_index_file); - reader->open(index_fname); - this->setup_thread_data(num_threads); - this->max_nthreads = num_threads; + // open AlignedFileReader handle to index_file + std::string index_fname(disk_index_file); + reader->open(index_fname); + 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; - diskann::load_bin(files, medoids_file, medoids, num_medoids, tmp_dim); + if (files.fileExists(medoids_file)) { + size_t tmp_dim; + diskann::load_bin(files, medoids_file, medoids, num_medoids, + tmp_dim); #else - if (file_exists(medoids_file)) - { - size_t tmp_dim; - diskann::load_bin(medoids_file, medoids, num_medoids, tmp_dim); + if (file_exists(medoids_file)) { + size_t tmp_dim; + diskann::load_bin(medoids_file, medoids, num_medoids, tmp_dim); #endif - if (tmp_dim != 1) - { - std::stringstream stream; - stream << "Error loading medoids file. Expected bin format of m times " - "1 vector of uint32_t." - << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } + if (tmp_dim != 1) { + std::stringstream stream; + stream << "Error loading medoids file. Expected bin format of m times " + "1 vector of uint32_t." + << std::endl; + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } #ifdef EXEC_ENV_OLS - if (!files.fileExists(centroids_file)) - { + if (!files.fileExists(centroids_file)) { #else - if (!file_exists(centroids_file)) - { + if (!file_exists(centroids_file)) { #endif - diskann::cout << "Centroid data file not found. Using corresponding vectors " - "for the medoids " - << std::endl; - use_medoids_data_as_centroids(); - } - else - { - size_t num_centroids, aligned_tmp_dim; + diskann::cout + << "Centroid data file not found. Using corresponding vectors " + "for the medoids " + << std::endl; + use_medoids_data_as_centroids(); + } else { + size_t num_centroids, aligned_tmp_dim; #ifdef EXEC_ENV_OLS - diskann::load_aligned_bin(files, centroids_file, centroid_data, num_centroids, tmp_dim, - aligned_tmp_dim); + diskann::load_aligned_bin(files, centroids_file, centroid_data, + num_centroids, tmp_dim, aligned_tmp_dim); #else - diskann::load_aligned_bin(centroids_file, centroid_data, num_centroids, tmp_dim, aligned_tmp_dim); + diskann::load_aligned_bin(centroids_file, centroid_data, + num_centroids, tmp_dim, aligned_tmp_dim); #endif - if (aligned_tmp_dim != aligned_dim || num_centroids != num_medoids) - { - std::stringstream stream; - 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."; - diskann::cerr << stream.str() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - } - } - else - { - num_medoids = 1; - medoids = new uint32_t[1]; - medoids[0] = (uint32_t)(medoid_id_on_file); - use_medoids_data_as_centroids(); - } - - std::string norm_file = std::string(disk_index_file) + "_max_base_norm.bin"; - - if (file_exists(norm_file) && metric == diskann::Metric::INNER_PRODUCT) - { - uint64_t dumr, dumc; - float *norm_val; - diskann::load_bin(norm_file, norm_val, dumr, dumc); - this->max_base_norm = norm_val[0]; - diskann::cout << "Setting re-scaling factor of base vectors to " << this->max_base_norm << std::endl; - delete[] norm_val; + if (aligned_tmp_dim != aligned_dim || num_centroids != num_medoids) { + std::stringstream stream; + 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."; + diskann::cerr << stream.str() << std::endl; + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } } - diskann::cout << "done.." << std::endl; - return 0; + } else { + num_medoids = 1; + medoids = new uint32_t[1]; + medoids[0] = (uint32_t)(medoid_id_on_file); + use_medoids_data_as_centroids(); + } + + std::string norm_file = std::string(disk_index_file) + "_max_base_norm.bin"; + + if (file_exists(norm_file) && metric == diskann::Metric::INNER_PRODUCT) { + uint64_t dumr, dumc; + float *norm_val; + diskann::load_bin(norm_file, norm_val, dumr, dumc); + this->max_base_norm = norm_val[0]; + diskann::cout << "Setting re-scaling factor of base vectors to " + << this->max_base_norm << std::endl; + delete[] norm_val; + } + diskann::cout << "done.." << std::endl; + return 0; } #ifdef USE_BING_INFRA -bool getNextCompletedRequest(const IOContext &ctx, size_t size, int &completedIndex) -{ - bool waitsRemaining = false; - long completeCount = ctx.m_completeCount; - do - { - for (int i = 0; i < size; i++) - { - auto ithStatus = (*ctx.m_pRequestsStatus)[i]; - if (ithStatus == IOContext::Status::READ_SUCCESS) - { - completedIndex = i; - return true; - } - else if (ithStatus == IOContext::Status::READ_WAIT) - { - waitsRemaining = true; - } - } +bool getNextCompletedRequest(const IOContext &ctx, size_t size, + int &completedIndex) { + bool waitsRemaining = false; + long completeCount = ctx.m_completeCount; + do { + for (int i = 0; i < size; i++) { + auto ithStatus = (*ctx.m_pRequestsStatus)[i]; + if (ithStatus == IOContext::Status::READ_SUCCESS) { + completedIndex = i; + return true; + } else if (ithStatus == IOContext::Status::READ_WAIT) { + waitsRemaining = true; + } + } - // if we didn't find one in READ_SUCCESS, wait for one to complete. - if (waitsRemaining) - { - WaitOnAddress(&ctx.m_completeCount, &completeCount, sizeof(completeCount), 100); - // this assumes the knowledge of the reader behavior (implicit - // contract). need better factoring? - } - } while (waitsRemaining); + // if we didn't find one in READ_SUCCESS, wait for one to complete. + if (waitsRemaining) { + WaitOnAddress(&ctx.m_completeCount, &completeCount, sizeof(completeCount), + 100); + // this assumes the knowledge of the reader behavior (implicit + // contract). need better factoring? + } + } while (waitsRemaining); - completedIndex = -1; - return false; + completedIndex = -1; + return false; } #endif template -void PQFlashIndex::cached_beam_search(const T *query1, const uint64_t k_search, const uint64_t l_search, - uint64_t *indices, float *distances, const uint64_t beam_width, - const bool use_reorder_data, QueryStats *stats) -{ - cached_beam_search(query1, k_search, l_search, indices, distances, beam_width, std::numeric_limits::max(), - use_reorder_data, stats); +void PQFlashIndex::cached_beam_search( + const T *query1, const uint64_t k_search, const uint64_t l_search, + uint64_t *indices, float *distances, const uint64_t beam_width, + const bool use_reorder_data, QueryStats *stats) { + cached_beam_search(query1, k_search, l_search, indices, distances, beam_width, + std::numeric_limits::max(), use_reorder_data, + stats); } template -void PQFlashIndex::cached_beam_search(const T *query1, const uint64_t k_search, const uint64_t l_search, - uint64_t *indices, float *distances, const uint64_t beam_width, - const bool use_filter, const LabelT &filter_label, - const bool use_reorder_data, QueryStats *stats) -{ - cached_beam_search(query1, k_search, l_search, indices, distances, beam_width, use_filter, filter_label, - std::numeric_limits::max(), use_reorder_data, stats); +void PQFlashIndex::cached_beam_search( + const T *query1, const uint64_t k_search, const uint64_t l_search, + uint64_t *indices, float *distances, const uint64_t beam_width, + const bool use_filter, const LabelT &filter_label, + const bool use_reorder_data, QueryStats *stats) { + cached_beam_search(query1, k_search, l_search, indices, distances, beam_width, + use_filter, filter_label, + std::numeric_limits::max(), use_reorder_data, + stats); } template -void PQFlashIndex::cached_beam_search(const T *query1, const uint64_t k_search, const uint64_t l_search, - uint64_t *indices, float *distances, const uint64_t beam_width, - const uint32_t io_limit, const bool use_reorder_data, - QueryStats *stats) -{ - LabelT dummy_filter = 0; - cached_beam_search(query1, k_search, l_search, indices, distances, beam_width, false, dummy_filter, - std::numeric_limits::max(), use_reorder_data, stats); +void PQFlashIndex::cached_beam_search( + const T *query1, const uint64_t k_search, const uint64_t l_search, + uint64_t *indices, float *distances, const uint64_t beam_width, + const uint32_t io_limit, const bool use_reorder_data, QueryStats *stats) { + LabelT dummy_filter = 0; + cached_beam_search(query1, k_search, l_search, indices, distances, beam_width, + false, dummy_filter, std::numeric_limits::max(), + use_reorder_data, stats); } template -void PQFlashIndex::cached_beam_search(const T *query1, const uint64_t k_search, const uint64_t l_search, - uint64_t *indices, float *distances, const uint64_t beam_width, - const bool use_filter, const LabelT &filter_label, - const uint32_t io_limit, const bool use_reorder_data, - QueryStats *stats) -{ - int32_t filter_num = 0; - if (use_filter) - { - filter_num = get_filter_number(filter_label); - if (filter_num < 0) - { - if (!_use_universal_label) - { - return; - } - else - { - filter_num = _universal_filter_num; - } - } +void PQFlashIndex::cached_beam_search( + const T *query1, const uint64_t k_search, const uint64_t l_search, + uint64_t *indices, float *distances, const uint64_t beam_width, + const bool use_filter, const LabelT &filter_label, const uint32_t io_limit, + const bool use_reorder_data, QueryStats *stats) { + int32_t filter_num = 0; + if (use_filter) { + filter_num = get_filter_number(filter_label); + if (filter_num < 0) { + if (!_use_universal_label) { + return; + } else { + filter_num = _universal_filter_num; + } } + } + + if (beam_width > MAX_N_SECTOR_READS) + throw ANNException("Beamwidth can not be higher than MAX_N_SECTOR_READS", + -1, __FUNCSIG__, __FILE__, __LINE__); + + ScratchStoreManager> manager(this->thread_data); + auto data = manager.scratch_space(); + IOContext &ctx = data->ctx; + auto query_scratch = &(data->scratch); + auto pq_query_scratch = query_scratch->_pq_scratch; + + // reset query scratch + query_scratch->reset(); + + // copy query to thread specific aligned and allocated memory (for distance + // calculations we need aligned data) + float query_norm = 0; + T *aligned_query_T = query_scratch->aligned_query_T; + float *query_float = pq_query_scratch->aligned_query_float; + float *query_rotated = pq_query_scratch->rotated_query; + + // if inner product, we laso normalize the query and set the last coordinate + // to 0 (this is the extra coordindate used to convert MIPS to L2 search) + if (metric == diskann::Metric::INNER_PRODUCT) { + for (size_t i = 0; i < this->data_dim - 1; i++) { + aligned_query_T[i] = query1[i]; + query_norm += query1[i] * query1[i]; + } + aligned_query_T[this->data_dim - 1] = 0; - if (beam_width > MAX_N_SECTOR_READS) - throw ANNException("Beamwidth can not be higher than MAX_N_SECTOR_READS", -1, __FUNCSIG__, __FILE__, __LINE__); - - ScratchStoreManager> manager(this->thread_data); - auto data = manager.scratch_space(); - IOContext &ctx = data->ctx; - auto query_scratch = &(data->scratch); - auto pq_query_scratch = query_scratch->_pq_scratch; - - // reset query scratch - query_scratch->reset(); - - // copy query to thread specific aligned and allocated memory (for distance - // calculations we need aligned data) - float query_norm = 0; - T *aligned_query_T = query_scratch->aligned_query_T; - float *query_float = pq_query_scratch->aligned_query_float; - float *query_rotated = pq_query_scratch->rotated_query; - - // if inner product, we laso normalize the query and set the last coordinate - // to 0 (this is the extra coordindate used to convert MIPS to L2 search) - if (metric == diskann::Metric::INNER_PRODUCT) - { - for (size_t i = 0; i < this->data_dim - 1; i++) - { - aligned_query_T[i] = query1[i]; - query_norm += query1[i] * query1[i]; - } - aligned_query_T[this->data_dim - 1] = 0; - - query_norm = std::sqrt(query_norm); + query_norm = std::sqrt(query_norm); - for (size_t i = 0; i < this->data_dim - 1; i++) - { - aligned_query_T[i] /= query_norm; - } - pq_query_scratch->set(this->data_dim, aligned_query_T); + for (size_t i = 0; i < this->data_dim - 1; i++) { + aligned_query_T[i] /= query_norm; } - else - { - for (size_t i = 0; i < this->data_dim; i++) - { - aligned_query_T[i] = query1[i]; - } - pq_query_scratch->set(this->data_dim, aligned_query_T); + pq_query_scratch->set(this->data_dim, aligned_query_T); + } else { + for (size_t i = 0; i < this->data_dim; i++) { + aligned_query_T[i] = query1[i]; } - - // pointers to buffers for data - T *data_buf = query_scratch->coord_scratch; - uint64_t &data_buf_idx = query_scratch->coord_idx; - _mm_prefetch((char *)data_buf, _MM_HINT_T1); - - // sector scratch - char *sector_scratch = query_scratch->sector_scratch; - uint64_t §or_scratch_idx = query_scratch->sector_idx; - - // query <-> PQ chunk centers distances - pq_table.preprocess_query(query_rotated); // center the query and rotate if - // we have a rotation matrix - float *pq_dists = pq_query_scratch->aligned_pqtable_dist_scratch; - pq_table.populate_chunk_distances(query_rotated, pq_dists); - - // query <-> neighbor list - float *dist_scratch = pq_query_scratch->aligned_dist_scratch; - uint8_t *pq_coord_scratch = pq_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 uint32_t *ids, const uint64_t n_ids, - float *dists_out) { - diskann::aggregate_coords(ids, n_ids, this->data, this->n_chunks, pq_coord_scratch); - diskann::pq_dist_lookup(pq_coord_scratch, n_ids, this->n_chunks, pq_dists, dists_out); - }; - Timer query_timer, io_timer, cpu_timer; - - tsl::robin_set &visited = query_scratch->visited; - NeighborPriorityQueue &retset = query_scratch->retset; - retset.reserve(l_search); - std::vector &full_retset = query_scratch->full_retset; - - uint32_t best_medoid = 0; - float best_dist = (std::numeric_limits::max)(); - if (!use_filter) - { - for (uint64_t cur_m = 0; cur_m < num_medoids; cur_m++) - { - float cur_expanded_dist = - dist_cmp_float->compare(query_float, centroid_data + aligned_dim * cur_m, (uint32_t)aligned_dim); - if (cur_expanded_dist < best_dist) - { - best_medoid = medoids[cur_m]; - best_dist = cur_expanded_dist; - } - } + pq_query_scratch->set(this->data_dim, aligned_query_T); + } + + // pointers to buffers for data + T *data_buf = query_scratch->coord_scratch; + uint64_t &data_buf_idx = query_scratch->coord_idx; + _mm_prefetch((char *)data_buf, _MM_HINT_T1); + + // sector scratch + char *sector_scratch = query_scratch->sector_scratch; + uint64_t §or_scratch_idx = query_scratch->sector_idx; + + // query <-> PQ chunk centers distances + pq_table.preprocess_query(query_rotated); // center the query and rotate if + // we have a rotation matrix + float *pq_dists = pq_query_scratch->aligned_pqtable_dist_scratch; + pq_table.populate_chunk_distances(query_rotated, pq_dists); + + // query <-> neighbor list + float *dist_scratch = pq_query_scratch->aligned_dist_scratch; + uint8_t *pq_coord_scratch = pq_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 uint32_t *ids, + const uint64_t n_ids, + float *dists_out) { + diskann::aggregate_coords(ids, n_ids, this->data, this->n_chunks, + pq_coord_scratch); + diskann::pq_dist_lookup(pq_coord_scratch, n_ids, this->n_chunks, pq_dists, + dists_out); + }; + Timer query_timer, io_timer, cpu_timer; + + tsl::robin_set &visited = query_scratch->visited; + NeighborPriorityQueue &retset = query_scratch->retset; + retset.reserve(l_search); + std::vector &full_retset = query_scratch->full_retset; + + uint32_t best_medoid = 0; + float best_dist = (std::numeric_limits::max)(); + if (!use_filter) { + for (uint64_t cur_m = 0; cur_m < num_medoids; cur_m++) { + float cur_expanded_dist = dist_cmp_float->compare( + query_float, centroid_data + aligned_dim * cur_m, + (uint32_t)aligned_dim); + if (cur_expanded_dist < best_dist) { + best_medoid = medoids[cur_m]; + best_dist = cur_expanded_dist; + } } - else if (_filter_to_medoid_id.find(filter_label) != _filter_to_medoid_id.end()) - { - best_medoid = _filter_to_medoid_id[filter_label]; - } - else - { - throw ANNException("Cannot find medoid for specified filter.", -1, __FUNCSIG__, __FILE__, __LINE__); + } else if (_filter_to_medoid_id.find(filter_label) != + _filter_to_medoid_id.end()) { + best_medoid = _filter_to_medoid_id[filter_label]; + } else { + throw ANNException("Cannot find medoid for specified filter.", -1, + __FUNCSIG__, __FILE__, __LINE__); + } + + compute_dists(&best_medoid, 1, dist_scratch); + retset.insert(Neighbor(best_medoid, dist_scratch[0])); + visited.insert(best_medoid); + + uint32_t cmps = 0; + uint32_t hops = 0; + uint32_t num_ios = 0; + + // cleared every iteration + std::vector frontier; + frontier.reserve(2 * beam_width); + std::vector> frontier_nhoods; + frontier_nhoods.reserve(2 * beam_width); + std::vector frontier_read_reqs; + frontier_read_reqs.reserve(2 * beam_width); + std::vector>> + cached_nhoods; + cached_nhoods.reserve(2 * beam_width); + + while (retset.has_unexpanded_node() && num_ios < io_limit) { + // clear iteration state + frontier.clear(); + frontier_nhoods.clear(); + frontier_read_reqs.clear(); + cached_nhoods.clear(); + sector_scratch_idx = 0; + // find new beam + uint32_t num_seen = 0; + while (retset.has_unexpanded_node() && frontier.size() < beam_width && + num_seen < beam_width) { + auto nbr = retset.closest_unexpanded(); + num_seen++; + auto iter = nhood_cache.find(nbr.id); + if (iter != nhood_cache.end()) { + cached_nhoods.push_back(std::make_pair(nbr.id, iter->second)); + if (stats != nullptr) { + stats->n_cache_hits++; + } + } else { + frontier.push_back(nbr.id); + } + if (this->count_visited_nodes) { + reinterpret_cast &>( + this->node_visit_counter[nbr.id].second) + .fetch_add(1); + } } - compute_dists(&best_medoid, 1, dist_scratch); - retset.insert(Neighbor(best_medoid, dist_scratch[0])); - visited.insert(best_medoid); - - uint32_t cmps = 0; - uint32_t hops = 0; - uint32_t num_ios = 0; - - // cleared every iteration - std::vector frontier; - frontier.reserve(2 * beam_width); - std::vector> frontier_nhoods; - frontier_nhoods.reserve(2 * beam_width); - std::vector frontier_read_reqs; - frontier_read_reqs.reserve(2 * beam_width); - std::vector>> cached_nhoods; - cached_nhoods.reserve(2 * beam_width); - - while (retset.has_unexpanded_node() && num_ios < io_limit) - { - // clear iteration state - frontier.clear(); - frontier_nhoods.clear(); - frontier_read_reqs.clear(); - cached_nhoods.clear(); - sector_scratch_idx = 0; - // find new beam - uint32_t num_seen = 0; - while (retset.has_unexpanded_node() && frontier.size() < beam_width && num_seen < beam_width) - { - auto nbr = retset.closest_unexpanded(); - num_seen++; - auto iter = nhood_cache.find(nbr.id); - if (iter != nhood_cache.end()) - { - cached_nhoods.push_back(std::make_pair(nbr.id, iter->second)); - if (stats != nullptr) - { - stats->n_cache_hits++; - } - } - else - { - frontier.push_back(nbr.id); - } - if (this->count_visited_nodes) - { - reinterpret_cast &>(this->node_visit_counter[nbr.id].second).fetch_add(1); - } + // read nhoods of frontier ids + if (!frontier.empty()) { + if (stats != nullptr) stats->n_hops++; + for (uint64_t i = 0; i < frontier.size(); i++) { + auto id = frontier[i]; + std::pair fnhood; + fnhood.first = id; + fnhood.second = sector_scratch + sector_scratch_idx * SECTOR_LEN; + sector_scratch_idx++; + frontier_nhoods.push_back(fnhood); + frontier_read_reqs.emplace_back( + NODE_SECTOR_NO(((size_t)id)) * SECTOR_LEN, SECTOR_LEN, + fnhood.second); + if (stats != nullptr) { + stats->n_4k++; + stats->n_ios++; } - - // read nhoods of frontier ids - if (!frontier.empty()) - { - if (stats != nullptr) - stats->n_hops++; - for (uint64_t i = 0; i < frontier.size(); i++) - { - auto id = frontier[i]; - std::pair fnhood; - fnhood.first = id; - fnhood.second = sector_scratch + sector_scratch_idx * SECTOR_LEN; - sector_scratch_idx++; - frontier_nhoods.push_back(fnhood); - frontier_read_reqs.emplace_back(NODE_SECTOR_NO(((size_t)id)) * SECTOR_LEN, SECTOR_LEN, fnhood.second); - if (stats != nullptr) - { - stats->n_4k++; - stats->n_ios++; - } - num_ios++; - } - io_timer.reset(); + num_ios++; + } + io_timer.reset(); #ifdef USE_BING_INFRA - reader->read(frontier_read_reqs, ctx, - true); // async reader windows. + 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); // synchronous IO linux #endif - if (stats != nullptr) - { - stats->io_us += (double)io_timer.elapsed(); - } - } + if (stats != nullptr) { + 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; - float cur_expanded_dist; - if (!use_disk_index_pq) - { - cur_expanded_dist = dist_cmp->compare(aligned_query_T, node_fp_coords_copy, (uint32_t)aligned_dim); - } - else - { - if (metric == diskann::Metric::INNER_PRODUCT) - cur_expanded_dist = disk_pq_table.inner_product(query_float, (uint8_t *)node_fp_coords_copy); - else - cur_expanded_dist = disk_pq_table.l2_distance( // disk_pq does not support OPQ yet - query_float, (uint8_t *)node_fp_coords_copy); - } - full_retset.push_back(Neighbor((uint32_t)cached_nhood.first, cur_expanded_dist)); - - uint64_t nnbrs = cached_nhood.second.first; - uint32_t *node_nbrs = cached_nhood.second.second; - - // compute node_nbrs <-> query dists in PQ space - cpu_timer.reset(); - compute_dists(node_nbrs, nnbrs, dist_scratch); - if (stats != nullptr) - { - stats->n_cmps += (double)nnbrs; - stats->cpu_us += (double)cpu_timer.elapsed(); - } - - // process prefetched nhood - for (uint64_t m = 0; m < nnbrs; ++m) - { - uint32_t id = node_nbrs[m]; - if (visited.insert(id).second) - { - if (!use_filter && _dummy_pts.find(id) != _dummy_pts.end()) - continue; - - if (use_filter && !point_has_label(id, filter_num) && !point_has_label(id, _universal_filter_num)) - continue; - cmps++; - float dist = dist_scratch[m]; - Neighbor nn(id, dist); - retset.insert(nn); - } - } + // 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; + float cur_expanded_dist; + if (!use_disk_index_pq) { + cur_expanded_dist = dist_cmp->compare( + aligned_query_T, node_fp_coords_copy, (uint32_t)aligned_dim); + } else { + if (metric == diskann::Metric::INNER_PRODUCT) + cur_expanded_dist = disk_pq_table.inner_product( + query_float, (uint8_t *)node_fp_coords_copy); + else + cur_expanded_dist = + disk_pq_table.l2_distance( // disk_pq does not support OPQ yet + query_float, (uint8_t *)node_fp_coords_copy); + } + full_retset.push_back( + Neighbor((uint32_t)cached_nhood.first, cur_expanded_dist)); + + uint64_t nnbrs = cached_nhood.second.first; + uint32_t *node_nbrs = cached_nhood.second.second; + + // compute node_nbrs <-> query dists in PQ space + cpu_timer.reset(); + compute_dists(node_nbrs, nnbrs, dist_scratch); + if (stats != nullptr) { + stats->n_cmps += (double)nnbrs; + stats->cpu_us += (double)cpu_timer.elapsed(); + } + + // process prefetched nhood + for (uint64_t m = 0; m < nnbrs; ++m) { + uint32_t id = node_nbrs[m]; + if (visited.insert(id).second) { + if (!use_filter && _dummy_pts.find(id) != _dummy_pts.end()) continue; + + if (use_filter && !point_has_label(id, filter_num) && + !point_has_label(id, _universal_filter_num)) + continue; + cmps++; + float dist = dist_scratch[m]; + Neighbor nn(id, dist); + retset.insert(nn); } + } + } #ifdef USE_BING_INFRA - // process each frontier nhood - compute distances to unvisited nodes - int completedIndex = -1; - long requestCount = static_cast(frontier_read_reqs.size()); - // If we issued read requests and if a read is complete or there are - // reads in wait state, then enter the while loop. - while (requestCount > 0 && getNextCompletedRequest(ctx, requestCount, completedIndex)) - { - assert(completedIndex >= 0); - auto &frontier_nhood = frontier_nhoods[completedIndex]; - (*ctx.m_pRequestsStatus)[completedIndex] = IOContext::PROCESS_COMPLETE; + // process each frontier nhood - compute distances to unvisited nodes + int completedIndex = -1; + long requestCount = static_cast(frontier_read_reqs.size()); + // If we issued read requests and if a read is complete or there are + // reads in wait state, then enter the while loop. + while (requestCount > 0 && + getNextCompletedRequest(ctx, requestCount, completedIndex)) { + assert(completedIndex >= 0); + auto &frontier_nhood = frontier_nhoods[completedIndex]; + (*ctx.m_pRequestsStatus)[completedIndex] = IOContext::PROCESS_COMPLETE; #else - for (auto &frontier_nhood : frontier_nhoods) - { + for (auto &frontier_nhood : frontier_nhoods) { #endif - char *node_disk_buf = OFFSET_TO_NODE(frontier_nhood.second, frontier_nhood.first); - uint32_t *node_buf = OFFSET_TO_NODE_NHOOD(node_disk_buf); - uint64_t nnbrs = (uint64_t)(*node_buf); - T *node_fp_coords = OFFSET_TO_NODE_COORDS(node_disk_buf); - // assert(data_buf_idx < MAX_N_CMPS); - if (data_buf_idx == MAX_N_CMPS) - data_buf_idx = 0; - - T *node_fp_coords_copy = data_buf + (data_buf_idx * aligned_dim); - data_buf_idx++; - memcpy(node_fp_coords_copy, node_fp_coords, disk_bytes_per_point); - float cur_expanded_dist; - if (!use_disk_index_pq) - { - cur_expanded_dist = dist_cmp->compare(aligned_query_T, node_fp_coords_copy, (uint32_t)aligned_dim); - } - else - { - if (metric == diskann::Metric::INNER_PRODUCT) - cur_expanded_dist = disk_pq_table.inner_product(query_float, (uint8_t *)node_fp_coords_copy); - else - cur_expanded_dist = disk_pq_table.l2_distance(query_float, (uint8_t *)node_fp_coords_copy); - } - full_retset.push_back(Neighbor(frontier_nhood.first, cur_expanded_dist)); - uint32_t *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 += (double)nnbrs; - stats->cpu_us += (double)cpu_timer.elapsed(); - } - - cpu_timer.reset(); - // process prefetch-ed nhood - for (uint64_t m = 0; m < nnbrs; ++m) - { - uint32_t id = node_nbrs[m]; - if (visited.insert(id).second) - { - if (!use_filter && _dummy_pts.find(id) != _dummy_pts.end()) - continue; - - if (use_filter && !point_has_label(id, filter_num) && !point_has_label(id, _universal_filter_num)) - continue; - cmps++; - float dist = dist_scratch[m]; - if (stats != nullptr) - { - stats->n_cmps++; - } - - Neighbor nn(id, dist); - retset.insert(nn); - } - } - - if (stats != nullptr) - { - stats->cpu_us += (double)cpu_timer.elapsed(); - } + char *node_disk_buf = + OFFSET_TO_NODE(frontier_nhood.second, frontier_nhood.first); + uint32_t *node_buf = OFFSET_TO_NODE_NHOOD(node_disk_buf); + uint64_t nnbrs = (uint64_t)(*node_buf); + T *node_fp_coords = OFFSET_TO_NODE_COORDS(node_disk_buf); + // assert(data_buf_idx < MAX_N_CMPS); + if (data_buf_idx == MAX_N_CMPS) data_buf_idx = 0; + + T *node_fp_coords_copy = data_buf + (data_buf_idx * aligned_dim); + data_buf_idx++; + memcpy(node_fp_coords_copy, node_fp_coords, disk_bytes_per_point); + float cur_expanded_dist; + if (!use_disk_index_pq) { + cur_expanded_dist = dist_cmp->compare( + aligned_query_T, node_fp_coords_copy, (uint32_t)aligned_dim); + } else { + if (metric == diskann::Metric::INNER_PRODUCT) + cur_expanded_dist = disk_pq_table.inner_product( + query_float, (uint8_t *)node_fp_coords_copy); + else + cur_expanded_dist = disk_pq_table.l2_distance( + query_float, (uint8_t *)node_fp_coords_copy); + } + full_retset.push_back(Neighbor(frontier_nhood.first, cur_expanded_dist)); + uint32_t *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 += (double)nnbrs; + stats->cpu_us += (double)cpu_timer.elapsed(); + } + + cpu_timer.reset(); + // process prefetch-ed nhood + for (uint64_t m = 0; m < nnbrs; ++m) { + uint32_t id = node_nbrs[m]; + if (visited.insert(id).second) { + if (!use_filter && _dummy_pts.find(id) != _dummy_pts.end()) continue; + + if (use_filter && !point_has_label(id, filter_num) && + !point_has_label(id, _universal_filter_num)) + continue; + cmps++; + float dist = dist_scratch[m]; + if (stats != nullptr) { + stats->n_cmps++; + } + + Neighbor nn(id, dist); + retset.insert(nn); } + } - hops++; + if (stats != nullptr) { + stats->cpu_us += (double)cpu_timer.elapsed(); + } } - // re-sort by distance - std::sort(full_retset.begin(), full_retset.end()); + hops++; + } - if (use_reorder_data) - { - if (!(this->reorder_data_exists)) - { - throw ANNException("Requested use of reordering data which does " - "not exist in index " - "file", - -1, __FUNCSIG__, __FILE__, __LINE__); - } + // re-sort by distance + std::sort(full_retset.begin(), full_retset.end()); - std::vector vec_read_reqs; + if (use_reorder_data) { + if (!(this->reorder_data_exists)) { + throw ANNException( + "Requested use of reordering data which does " + "not exist in index " + "file", + -1, __FUNCSIG__, __FILE__, __LINE__); + } - if (full_retset.size() > k_search * FULL_PRECISION_REORDER_MULTIPLIER) - full_retset.erase(full_retset.begin() + k_search * FULL_PRECISION_REORDER_MULTIPLIER, full_retset.end()); + std::vector vec_read_reqs; - for (size_t i = 0; i < full_retset.size(); ++i) - { - vec_read_reqs.emplace_back(VECTOR_SECTOR_NO(((size_t)full_retset[i].id)) * SECTOR_LEN, SECTOR_LEN, - sector_scratch + i * SECTOR_LEN); + if (full_retset.size() > k_search * FULL_PRECISION_REORDER_MULTIPLIER) + full_retset.erase( + full_retset.begin() + k_search * FULL_PRECISION_REORDER_MULTIPLIER, + full_retset.end()); - if (stats != nullptr) - { - stats->n_4k++; - stats->n_ios++; - } - } + for (size_t i = 0; i < full_retset.size(); ++i) { + vec_read_reqs.emplace_back( + VECTOR_SECTOR_NO(((size_t)full_retset[i].id)) * SECTOR_LEN, + SECTOR_LEN, sector_scratch + i * SECTOR_LEN); + + if (stats != nullptr) { + stats->n_4k++; + stats->n_ios++; + } + } - io_timer.reset(); + io_timer.reset(); #ifdef USE_BING_INFRA - reader->read(vec_read_reqs, ctx, false); // sync reader windows. + reader->read(vec_read_reqs, ctx, false); // sync reader windows. #else - reader->read(vec_read_reqs, ctx); // synchronous IO linux + reader->read(vec_read_reqs, ctx); // synchronous IO linux #endif - if (stats != nullptr) - { - stats->io_us += io_timer.elapsed(); - } - - for (size_t i = 0; i < full_retset.size(); ++i) - { - auto id = full_retset[i].id; - auto location = (sector_scratch + i * SECTOR_LEN) + VECTOR_SECTOR_OFFSET(id); - full_retset[i].distance = dist_cmp->compare(aligned_query_T, (T *)location, this->data_dim); - } + if (stats != nullptr) { + stats->io_us += io_timer.elapsed(); + } - std::sort(full_retset.begin(), full_retset.end()); + for (size_t i = 0; i < full_retset.size(); ++i) { + auto id = full_retset[i].id; + auto location = + (sector_scratch + i * SECTOR_LEN) + VECTOR_SECTOR_OFFSET(id); + full_retset[i].distance = + dist_cmp->compare(aligned_query_T, (T *)location, this->data_dim); } - // copy k_search values - for (uint64_t i = 0; i < k_search; i++) - { - indices[i] = full_retset[i].id; + std::sort(full_retset.begin(), full_retset.end()); + } - if (_dummy_pts.find(indices[i]) != _dummy_pts.end()) - { - indices[i] = _dummy_to_real_map[indices[i]]; - } + // copy k_search values + for (uint64_t i = 0; i < k_search; i++) { + indices[i] = full_retset[i].id; - if (distances != nullptr) - { - distances[i] = full_retset[i].distance; - if (metric == diskann::Metric::INNER_PRODUCT) - { - // flip the sign to convert min to max - distances[i] = (-distances[i]); - // rescale to revert back to original norms (cancelling the - // effect of base and query pre-processing) - if (max_base_norm != 0) - distances[i] *= (max_base_norm * query_norm); - } - } + if (_dummy_pts.find(indices[i]) != _dummy_pts.end()) { + indices[i] = _dummy_to_real_map[indices[i]]; + } + + if (distances != nullptr) { + distances[i] = full_retset[i].distance; + if (metric == diskann::Metric::INNER_PRODUCT) { + // flip the sign to convert min to max + distances[i] = (-distances[i]); + // rescale to revert back to original norms (cancelling the + // effect of base and query pre-processing) + if (max_base_norm != 0) distances[i] *= (max_base_norm * query_norm); + } } + } #ifdef USE_BING_INFRA - ctx.m_completeCount = 0; + ctx.m_completeCount = 0; #endif - if (stats != nullptr) - { - stats->total_us = (double)query_timer.elapsed(); - } + if (stats != nullptr) { + stats->total_us = (double)query_timer.elapsed(); + } } // range search returns results of all neighbors within distance of range. // indices and distances need to be pre-allocated of size l_search and the // return value is the number of matching hits. template -uint32_t PQFlashIndex::range_search(const T *query1, const double range, const uint64_t min_l_search, - const uint64_t max_l_search, std::vector &indices, - std::vector &distances, const uint64_t min_beam_width, - QueryStats *stats) -{ - uint32_t res_count = 0; - - bool stop_flag = false; - - uint32_t l_search = min_l_search; // starting size of the candidate list - while (!stop_flag) - { - indices.resize(l_search); - distances.resize(l_search); - uint64_t cur_bw = min_beam_width > (l_search / 5) ? min_beam_width : l_search / 5; - cur_bw = (cur_bw > 100) ? 100 : cur_bw; - for (auto &x : distances) - x = std::numeric_limits::max(); - this->cached_beam_search(query1, l_search, l_search, indices.data(), distances.data(), cur_bw, false, stats); - for (uint32_t i = 0; i < l_search; i++) - { - if (distances[i] > (float)range) - { - res_count = i; - break; - } - else if (i == l_search - 1) - res_count = l_search; - } - if (res_count < (uint32_t)(l_search / 2.0)) - stop_flag = true; - l_search = l_search * 2; - if (l_search > max_l_search) - stop_flag = true; +uint32_t PQFlashIndex::range_search( + const T *query1, const double range, const uint64_t min_l_search, + const uint64_t max_l_search, std::vector &indices, + std::vector &distances, const uint64_t min_beam_width, + QueryStats *stats) { + uint32_t res_count = 0; + + bool stop_flag = false; + + uint32_t l_search = min_l_search; // starting size of the candidate list + while (!stop_flag) { + indices.resize(l_search); + distances.resize(l_search); + uint64_t cur_bw = + min_beam_width > (l_search / 5) ? min_beam_width : l_search / 5; + cur_bw = (cur_bw > 100) ? 100 : cur_bw; + for (auto &x : distances) x = std::numeric_limits::max(); + this->cached_beam_search(query1, l_search, l_search, indices.data(), + distances.data(), cur_bw, false, stats); + for (uint32_t i = 0; i < l_search; i++) { + if (distances[i] > (float)range) { + res_count = i; + break; + } else if (i == l_search - 1) + res_count = l_search; } - indices.resize(res_count); - distances.resize(res_count); - return res_count; + if (res_count < (uint32_t)(l_search / 2.0)) stop_flag = true; + l_search = l_search * 2; + if (l_search > max_l_search) stop_flag = true; + } + indices.resize(res_count); + distances.resize(res_count); + return res_count; } -template uint64_t PQFlashIndex::get_data_dim() -{ - return data_dim; +template +uint64_t PQFlashIndex::get_data_dim() { + return data_dim; } -template diskann::Metric PQFlashIndex::get_metric() -{ - return this->metric; +template +diskann::Metric PQFlashIndex::get_metric() { + return this->metric; } #ifdef EXEC_ENV_OLS -template char *PQFlashIndex::getHeaderBytes() -{ - IOContext &ctx = reader->get_ctx(); - AlignedRead readReq; - readReq.buf = new char[PQFlashIndex::HEADER_SIZE]; - readReq.len = PQFlashIndex::HEADER_SIZE; - readReq.offset = 0; +template +char *PQFlashIndex::getHeaderBytes() { + IOContext &ctx = reader->get_ctx(); + AlignedRead readReq; + readReq.buf = new char[PQFlashIndex::HEADER_SIZE]; + readReq.len = PQFlashIndex::HEADER_SIZE; + readReq.offset = 0; - std::vector readReqs; - readReqs.push_back(readReq); + std::vector readReqs; + readReqs.push_back(readReq); - reader->read(readReqs, ctx, false); + reader->read(readReqs, ctx, false); - return (char *)readReq.buf; + return (char *)readReq.buf; } #endif @@ -1562,4 +1483,4 @@ template class PQFlashIndex; template class PQFlashIndex; template class PQFlashIndex; -} // namespace diskann +} // namespace diskann diff --git a/src/scratch.cpp b/src/scratch.cpp index 3f8db3bec..02cd6d975 100644 --- a/src/scratch.cpp +++ b/src/scratch.cpp @@ -6,132 +6,133 @@ #include "scratch.h" -namespace diskann -{ +namespace diskann { // // Functions to manage scratch space for in-memory index based search // template -InMemQueryScratch::InMemQueryScratch(uint32_t search_l, uint32_t indexing_l, uint32_t r, uint32_t maxc, size_t dim, - size_t aligned_dim, size_t alignment_factor, +InMemQueryScratch::InMemQueryScratch(uint32_t search_l, uint32_t indexing_l, + uint32_t r, uint32_t maxc, size_t dim, + size_t aligned_dim, + size_t alignment_factor, bool init_pq_scratch) - : _L(0), _R(r), _maxc(maxc) -{ - if (search_l == 0 || indexing_l == 0 || r == 0 || dim == 0) - { - std::stringstream ss; - ss << "In InMemQueryScratch, one of search_l = " << search_l << ", indexing_l = " << indexing_l - << ", dim = " << dim << " or r = " << r << " is zero." << std::endl; - throw diskann::ANNException(ss.str(), -1); - } - - //REFACTOR - //auto aligned_dim = ROUND_UP(dim, 8); - //alloc_aligned(((void **)&_aligned_query), aligned_dim * sizeof(T), 8 * sizeof(T)); - //memset(_aligned_query, 0, aligned_dim * sizeof(T)); - - alloc_aligned(((void **)&_aligned_query), aligned_dim * sizeof(T), alignment_factor * sizeof(T)); - memset(_aligned_query, 0, aligned_dim * sizeof(T)); - - if (init_pq_scratch) - _pq_scratch = new PQScratch(MAX_GRAPH_DEGREE, aligned_dim); - else - _pq_scratch = nullptr; - - _occlude_factor.reserve(maxc); - _inserted_into_pool_bs = new boost::dynamic_bitset<>(); - _id_scratch.reserve(std::ceil(1.5 * GRAPH_SLACK_FACTOR * _R)); - _dist_scratch.reserve(std::ceil(1.5 * GRAPH_SLACK_FACTOR * _R)); - - resize_for_new_L(std::max(search_l, indexing_l)); -} + : _L(0), _R(r), _maxc(maxc) { + if (search_l == 0 || indexing_l == 0 || r == 0 || dim == 0) { + std::stringstream ss; + ss << "In InMemQueryScratch, one of search_l = " << search_l + << ", indexing_l = " << indexing_l << ", dim = " << dim + << " or r = " << r << " is zero." << std::endl; + throw diskann::ANNException(ss.str(), -1); + } + + // REFACTOR + // auto aligned_dim = ROUND_UP(dim, 8); + // alloc_aligned(((void **)&_aligned_query), aligned_dim * sizeof(T), 8 * + // sizeof(T)); memset(_aligned_query, 0, aligned_dim * sizeof(T)); + + alloc_aligned(((void **)&_aligned_query), aligned_dim * sizeof(T), + alignment_factor * sizeof(T)); + memset(_aligned_query, 0, aligned_dim * sizeof(T)); + + if (init_pq_scratch) + _pq_scratch = new PQScratch(MAX_GRAPH_DEGREE, aligned_dim); + else + _pq_scratch = nullptr; -template void InMemQueryScratch::clear() -{ - _pool.clear(); - _best_l_nodes.clear(); - _occlude_factor.clear(); + _occlude_factor.reserve(maxc); + _inserted_into_pool_bs = new boost::dynamic_bitset<>(); + _id_scratch.reserve(std::ceil(1.5 * GRAPH_SLACK_FACTOR * _R)); + _dist_scratch.reserve(std::ceil(1.5 * GRAPH_SLACK_FACTOR * _R)); - _inserted_into_pool_rs.clear(); - _inserted_into_pool_bs->reset(); + resize_for_new_L(std::max(search_l, indexing_l)); +} - _id_scratch.clear(); - _dist_scratch.clear(); +template +void InMemQueryScratch::clear() { + _pool.clear(); + _best_l_nodes.clear(); + _occlude_factor.clear(); - _expanded_nodes_set.clear(); - _expanded_nghrs_vec.clear(); - _occlude_list_output.clear(); + _inserted_into_pool_rs.clear(); + _inserted_into_pool_bs->reset(); -} + _id_scratch.clear(); + _dist_scratch.clear(); -template void InMemQueryScratch::resize_for_new_L(uint32_t new_l) -{ - if (new_l > _L) - { - _L = new_l; - _pool.reserve(3 * _L + _R); - _best_l_nodes.reserve(_L); + _expanded_nodes_set.clear(); + _expanded_nghrs_vec.clear(); + _occlude_list_output.clear(); +} - _inserted_into_pool_rs.reserve(20 * _L); - } +template +void InMemQueryScratch::resize_for_new_L(uint32_t new_l) { + if (new_l > _L) { + _L = new_l; + _pool.reserve(3 * _L + _R); + _best_l_nodes.reserve(_L); + + _inserted_into_pool_rs.reserve(20 * _L); + } } -template InMemQueryScratch::~InMemQueryScratch() -{ - if (_aligned_query != nullptr) - { - aligned_free(_aligned_query); - } +template +InMemQueryScratch::~InMemQueryScratch() { + if (_aligned_query != nullptr) { + aligned_free(_aligned_query); + } - delete _pq_scratch; - delete _inserted_into_pool_bs; + delete _pq_scratch; + delete _inserted_into_pool_bs; } // // Functions to manage scratch space for SSD based search // -template void SSDQueryScratch::reset() -{ - coord_idx = 0; - sector_idx = 0; - visited.clear(); - retset.clear(); - full_retset.clear(); +template +void SSDQueryScratch::reset() { + coord_idx = 0; + sector_idx = 0; + visited.clear(); + retset.clear(); + full_retset.clear(); } -template SSDQueryScratch::SSDQueryScratch(size_t aligned_dim, size_t visited_reserve) -{ - size_t coord_alloc_size = ROUND_UP(MAX_N_CMPS * aligned_dim, 256); +template +SSDQueryScratch::SSDQueryScratch(size_t aligned_dim, + size_t visited_reserve) { + size_t coord_alloc_size = ROUND_UP(MAX_N_CMPS * aligned_dim, 256); - diskann::alloc_aligned((void **)&coord_scratch, coord_alloc_size, 256); - diskann::alloc_aligned((void **)§or_scratch, (size_t)MAX_N_SECTOR_READS * (size_t)SECTOR_LEN, SECTOR_LEN); - diskann::alloc_aligned((void **)&aligned_query_T, aligned_dim * sizeof(T), 8 * sizeof(T)); + diskann::alloc_aligned((void **)&coord_scratch, coord_alloc_size, 256); + diskann::alloc_aligned((void **)§or_scratch, + (size_t)MAX_N_SECTOR_READS * (size_t)SECTOR_LEN, + SECTOR_LEN); + diskann::alloc_aligned((void **)&aligned_query_T, aligned_dim * sizeof(T), + 8 * sizeof(T)); - _pq_scratch = new PQScratch(MAX_GRAPH_DEGREE, aligned_dim); + _pq_scratch = new PQScratch(MAX_GRAPH_DEGREE, aligned_dim); - memset(coord_scratch, 0, MAX_N_CMPS * aligned_dim); - memset(aligned_query_T, 0, aligned_dim * sizeof(T)); + memset(coord_scratch, 0, MAX_N_CMPS * aligned_dim); + memset(aligned_query_T, 0, aligned_dim * sizeof(T)); - visited.reserve(visited_reserve); - full_retset.reserve(visited_reserve); + visited.reserve(visited_reserve); + full_retset.reserve(visited_reserve); } -template SSDQueryScratch::~SSDQueryScratch() -{ - diskann::aligned_free((void *)coord_scratch); - diskann::aligned_free((void *)sector_scratch); +template +SSDQueryScratch::~SSDQueryScratch() { + diskann::aligned_free((void *)coord_scratch); + diskann::aligned_free((void *)sector_scratch); - delete[] _pq_scratch; + delete[] _pq_scratch; } template -SSDThreadData::SSDThreadData(size_t aligned_dim, size_t visited_reserve) : scratch(aligned_dim, visited_reserve) -{ -} +SSDThreadData::SSDThreadData(size_t aligned_dim, size_t visited_reserve) + : scratch(aligned_dim, visited_reserve) {} -template void SSDThreadData::clear() -{ - scratch.reset(); +template +void SSDThreadData::clear() { + scratch.reset(); } template DISKANN_DLLEXPORT class InMemQueryScratch; @@ -145,4 +146,4 @@ template DISKANN_DLLEXPORT class SSDQueryScratch; template DISKANN_DLLEXPORT class SSDThreadData; template DISKANN_DLLEXPORT class SSDThreadData; template DISKANN_DLLEXPORT class SSDThreadData; -} // namespace diskann \ No newline at end of file +} // namespace diskann \ No newline at end of file diff --git a/src/utils.cpp b/src/utils.cpp index b675e656d..5f724964b 100644 --- a/src/utils.cpp +++ b/src/utils.cpp @@ -9,7 +9,7 @@ #include "aligned_file_reader.h" #endif -const uint32_t MAX_REQUEST_SIZE = 1024 * 1024 * 1024; // 64MB +const uint32_t MAX_REQUEST_SIZE = 1024 * 1024 * 1024; // 64MB const uint32_t MAX_SIMULTANEOUS_READ_REQUESTS = 128; #ifdef _WINDOWS @@ -17,48 +17,44 @@ const uint32_t MAX_SIMULTANEOUS_READ_REQUESTS = 128; // Taken from: // https://insufficientlycomplicated.wordpress.com/2011/11/07/detecting-intel-advanced-vector-extensions-avx-in-visual-studio/ -bool cpuHasAvxSupport() -{ - bool avxSupported = false; - - // Checking for AVX requires 3 things: - // 1) CPUID indicates that the OS uses XSAVE and XRSTORE - // instructions (allowing saving YMM registers on context - // switch) - // 2) CPUID indicates support for AVX - // 3) XGETBV indicates the AVX registers will be saved and - // restored on context switch - // - // Note that XGETBV is only available on 686 or later CPUs, so - // the instruction needs to be conditionally run. - int cpuInfo[4]; - __cpuid(cpuInfo, 1); - - bool osUsesXSAVE_XRSTORE = cpuInfo[2] & (1 << 27) || false; - bool cpuAVXSuport = cpuInfo[2] & (1 << 28) || false; - - if (osUsesXSAVE_XRSTORE && cpuAVXSuport) - { - // Check if the OS will save the YMM registers - unsigned long long xcrFeatureMask = _xgetbv(_XCR_XFEATURE_ENABLED_MASK); - avxSupported = (xcrFeatureMask & 0x6) || false; - } - - return avxSupported; +bool cpuHasAvxSupport() { + bool avxSupported = false; + + // Checking for AVX requires 3 things: + // 1) CPUID indicates that the OS uses XSAVE and XRSTORE + // instructions (allowing saving YMM registers on context + // switch) + // 2) CPUID indicates support for AVX + // 3) XGETBV indicates the AVX registers will be saved and + // restored on context switch + // + // Note that XGETBV is only available on 686 or later CPUs, so + // the instruction needs to be conditionally run. + int cpuInfo[4]; + __cpuid(cpuInfo, 1); + + bool osUsesXSAVE_XRSTORE = cpuInfo[2] & (1 << 27) || false; + bool cpuAVXSuport = cpuInfo[2] & (1 << 28) || false; + + if (osUsesXSAVE_XRSTORE && cpuAVXSuport) { + // Check if the OS will save the YMM registers + unsigned long long xcrFeatureMask = _xgetbv(_XCR_XFEATURE_ENABLED_MASK); + avxSupported = (xcrFeatureMask & 0x6) || false; + } + + return avxSupported; } -bool cpuHasAvx2Support() -{ - int cpuInfo[4]; - __cpuid(cpuInfo, 0); - int n = cpuInfo[0]; - if (n >= 7) - { - __cpuidex(cpuInfo, 7, 0); - static int avx2Mask = 0x20; - return (cpuInfo[1] & avx2Mask) > 0; - } - return false; +bool cpuHasAvx2Support() { + int cpuInfo[4]; + __cpuid(cpuInfo, 0); + int n = cpuInfo[0]; + if (n >= 7) { + __cpuidex(cpuInfo, 7, 0); + static int avx2Mask = 0x20; + return (cpuInfo[1] & avx2Mask) > 0; + } + return false; } bool AvxSupportedCPU = cpuHasAvxSupport(); @@ -70,407 +66,423 @@ bool Avx2SupportedCPU = true; bool AvxSupportedCPU = false; #endif -namespace diskann -{ +namespace diskann { -void block_convert(std::ofstream &writr, std::ifstream &readr, float *read_buf, size_t npts, size_t ndims) -{ - readr.read((char *)read_buf, npts * ndims * sizeof(float)); - uint32_t ndims_u32 = (uint32_t)ndims; +void block_convert(std::ofstream &writr, std::ifstream &readr, float *read_buf, + size_t npts, size_t ndims) { + readr.read((char *)read_buf, npts * ndims * sizeof(float)); + uint32_t ndims_u32 = (uint32_t)ndims; #pragma omp parallel for - for (int64_t i = 0; i < (int64_t)npts; i++) - { - float norm_pt = std::numeric_limits::epsilon(); - for (uint32_t 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 (uint32_t dim = 0; dim < ndims_u32; dim++) - { - *(read_buf + i * ndims + dim) = *(read_buf + i * ndims + dim) / norm_pt; - } + for (int64_t i = 0; i < (int64_t)npts; i++) { + float norm_pt = std::numeric_limits::epsilon(); + for (uint32_t dim = 0; dim < ndims_u32; dim++) { + norm_pt += *(read_buf + i * ndims + dim) * *(read_buf + i * ndims + dim); } - writr.write((char *)read_buf, npts * ndims * sizeof(float)); + norm_pt = std::sqrt(norm_pt); + for (uint32_t 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(int32_t)); - readr.read((char *)&ndims_s32, sizeof(int32_t)); - - writr.write((char *)&npts_s32, sizeof(int32_t)); - writr.write((char *)&ndims_s32, sizeof(int32_t)); - - size_t npts = (size_t)npts_s32; - size_t ndims = (size_t)ndims_s32; - diskann::cout << "Normalizing FLOAT vectors in file: " << inFileName << std::endl; - diskann::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims << std::endl; - - size_t blk_size = 131072; - size_t nblks = ROUND_UP(npts, blk_size) / blk_size; - diskann::cout << "# blks: " << nblks << std::endl; +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(int32_t)); + readr.read((char *)&ndims_s32, sizeof(int32_t)); + + writr.write((char *)&npts_s32, sizeof(int32_t)); + writr.write((char *)&ndims_s32, sizeof(int32_t)); + + size_t npts = (size_t)npts_s32; + size_t ndims = (size_t)ndims_s32; + diskann::cout << "Normalizing FLOAT vectors in file: " << inFileName + << std::endl; + diskann::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims + << std::endl; + + size_t blk_size = 131072; + size_t nblks = ROUND_UP(npts, blk_size) / blk_size; + diskann::cout << "# blks: " << nblks << std::endl; + + float *read_buf = new float[npts * ndims]; + for (size_t i = 0; i < nblks; i++) { + size_t 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; +} - float *read_buf = new float[npts * ndims]; - for (size_t i = 0; i < nblks; i++) - { - size_t cblk_size = std::min(npts - i * blk_size, blk_size); - block_convert(writr, readr, read_buf, cblk_size, ndims); +double calculate_recall(uint32_t num_queries, uint32_t *gold_std, + float *gs_dist, uint32_t dim_gs, uint32_t *our_results, + uint32_t dim_or, uint32_t recall_at) { + double total_recall = 0; + std::set gt, res; + + for (size_t i = 0; i < num_queries; i++) { + gt.clear(); + res.clear(); + uint32_t *gt_vec = gold_std + dim_gs * i; + uint32_t *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; + while (tie_breaker < dim_gs && + gt_dist_vec[tie_breaker] == gt_dist_vec[recall_at - 1]) + tie_breaker++; } - delete[] read_buf; - diskann::cout << "Wrote normalized points to file: " << outFileName << std::endl; + gt.insert(gt_vec, gt_vec + tie_breaker); + res.insert(res_vec, + res_vec + recall_at); // change to recall_at for recall k@k + // or dim_or for k@dim_or + uint32_t cur_recall = 0; + for (auto &v : gt) { + if (res.find(v) != res.end()) { + cur_recall++; + } + } + total_recall += cur_recall; + } + return total_recall / (num_queries) * (100.0 / recall_at); } -double calculate_recall(uint32_t num_queries, uint32_t *gold_std, float *gs_dist, uint32_t dim_gs, - uint32_t *our_results, uint32_t dim_or, uint32_t recall_at) -{ - double total_recall = 0; - std::set gt, res; - - for (size_t i = 0; i < num_queries; i++) - { - gt.clear(); - res.clear(); - uint32_t *gt_vec = gold_std + dim_gs * i; - uint32_t *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; - while (tie_breaker < dim_gs && gt_dist_vec[tie_breaker] == gt_dist_vec[recall_at - 1]) - tie_breaker++; - } - - gt.insert(gt_vec, gt_vec + tie_breaker); - res.insert(res_vec, - res_vec + recall_at); // change to recall_at for recall k@k - // or dim_or for k@dim_or - uint32_t cur_recall = 0; - for (auto &v : gt) - { - if (res.find(v) != res.end()) - { - cur_recall++; - } - } - total_recall += cur_recall; +double calculate_recall(uint32_t num_queries, uint32_t *gold_std, + float *gs_dist, uint32_t dim_gs, uint32_t *our_results, + uint32_t dim_or, uint32_t 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(); + uint32_t *gt_vec = gold_std + dim_gs * i; + uint32_t *res_vec = our_results + dim_or * i; + size_t tie_breaker = recall_at; + uint32_t active_points_count = 0; + uint32_t 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++; } - return total_recall / (num_queries) * (100.0 / recall_at); -} -double calculate_recall(uint32_t num_queries, uint32_t *gold_std, float *gs_dist, uint32_t dim_gs, - uint32_t *our_results, uint32_t dim_or, uint32_t 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(); - uint32_t *gt_vec = gold_std + dim_gs * i; - uint32_t *res_vec = our_results + dim_or * i; - size_t tie_breaker = recall_at; - uint32_t active_points_count = 0; - uint32_t 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); - uint32_t cur_recall = 0; - for (auto &v : res) - { - if (gt.find(v) != gt.end()) - { - cur_recall++; - } - } - total_recall += cur_recall; + gt.insert(gt_vec, gt_vec + tie_breaker); + res.insert(res_vec, res_vec + recall_at); + uint32_t cur_recall = 0; + for (auto &v : res) { + if (gt.find(v) != gt.end()) { + cur_recall++; + } } - return ((double)(total_recall / (num_queries))) * ((double)(100.0 / recall_at)); + total_recall += cur_recall; + } + return ((double)(total_recall / (num_queries))) * + ((double)(100.0 / recall_at)); } -double calculate_range_search_recall(uint32_t num_queries, std::vector> &groundtruth, - std::vector> &our_results) -{ - double total_recall = 0; - std::set gt, res; - - for (size_t i = 0; i < num_queries; i++) - { - gt.clear(); - res.clear(); - - gt.insert(groundtruth[i].begin(), groundtruth[i].end()); - res.insert(our_results[i].begin(), our_results[i].end()); - uint32_t cur_recall = 0; - for (auto &v : gt) - { - if (res.find(v) != res.end()) - { - cur_recall++; - } - } - if (gt.size() != 0) - total_recall += ((100.0 * cur_recall) / gt.size()); - else - total_recall += 100; +double calculate_range_search_recall( + uint32_t num_queries, std::vector> &groundtruth, + std::vector> &our_results) { + double total_recall = 0; + std::set gt, res; + + for (size_t i = 0; i < num_queries; i++) { + gt.clear(); + res.clear(); + + gt.insert(groundtruth[i].begin(), groundtruth[i].end()); + res.insert(our_results[i].begin(), our_results[i].end()); + uint32_t cur_recall = 0; + for (auto &v : gt) { + if (res.find(v) != res.end()) { + cur_recall++; + } } - return total_recall / (num_queries); + if (gt.size() != 0) + total_recall += ((100.0 * cur_recall) / gt.size()); + else + total_recall += 100; + } + return total_recall / (num_queries); } #ifdef EXEC_ENV_OLS -void get_bin_metadata(AlignedFileReader &reader, size_t &npts, size_t &ndim, size_t offset) -{ - std::vector readReqs; - AlignedRead readReq; - uint32_t buf[2]; // npts/ndim are uint32_ts. - - readReq.buf = buf; - readReq.offset = offset; - readReq.len = 2 * sizeof(uint32_t); - readReqs.push_back(readReq); - - IOContext &ctx = reader.get_ctx(); - reader.read(readReqs, ctx); // synchronous - if ((*(ctx.m_pRequestsStatus))[0] == IOContext::READ_SUCCESS) - { - npts = buf[0]; - ndim = buf[1]; - diskann::cout << "File has: " << npts << " points, " << ndim << " dimensions at offset: " << offset - << std::endl; - } - else - { - std::stringstream str; - str << "Could not read binary metadata from index file at offset: " << offset << std::endl; - throw diskann::ANNException(str.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } +void get_bin_metadata(AlignedFileReader &reader, size_t &npts, size_t &ndim, + size_t offset) { + std::vector readReqs; + AlignedRead readReq; + uint32_t buf[2]; // npts/ndim are uint32_ts. + + readReq.buf = buf; + readReq.offset = offset; + readReq.len = 2 * sizeof(uint32_t); + readReqs.push_back(readReq); + + IOContext &ctx = reader.get_ctx(); + reader.read(readReqs, ctx); // synchronous + if ((*(ctx.m_pRequestsStatus))[0] == IOContext::READ_SUCCESS) { + npts = buf[0]; + ndim = buf[1]; + diskann::cout << "File has: " << npts << " points, " << ndim + << " dimensions at offset: " << offset << std::endl; + } else { + std::stringstream str; + str << "Could not read binary metadata from index file at offset: " + << offset << std::endl; + throw diskann::ANNException(str.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } } -template void load_bin(AlignedFileReader &reader, T *&data, size_t &npts, size_t &ndim, size_t offset) -{ - // Code assumes that the reader is already setup correctly. - get_bin_metadata(reader, npts, ndim, offset); - data = new T[npts * ndim]; - - size_t data_size = npts * ndim * sizeof(T); - size_t write_offset = 0; - size_t read_start = offset + 2 * sizeof(uint32_t); - - // BingAlignedFileReader can only read uint32_t bytes of data. So, - // we limit ourselves even more to reading 1GB at a time. - std::vector readReqs; - while (data_size > 0) - { - AlignedRead readReq; - readReq.buf = data + write_offset; - readReq.offset = read_start + write_offset; - readReq.len = data_size > MAX_REQUEST_SIZE ? MAX_REQUEST_SIZE : data_size; - readReqs.push_back(readReq); - // in the corner case, the loop will not execute - data_size -= readReq.len; - write_offset += readReq.len; - } - IOContext &ctx = reader.get_ctx(); - reader.read(readReqs, ctx); - for (int i = 0; i < readReqs.size(); i++) - { - // Since we are making sync calls, no request will be in the - // READ_WAIT state. - if ((*(ctx.m_pRequestsStatus))[i] != IOContext::READ_SUCCESS) - { - std::stringstream str; - str << "Could not read binary data from index file at offset: " << readReqs[i].offset << std::endl; - throw diskann::ANNException(str.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } +template +void load_bin(AlignedFileReader &reader, T *&data, size_t &npts, size_t &ndim, + size_t offset) { + // Code assumes that the reader is already setup correctly. + get_bin_metadata(reader, npts, ndim, offset); + data = new T[npts * ndim]; + + size_t data_size = npts * ndim * sizeof(T); + size_t write_offset = 0; + size_t read_start = offset + 2 * sizeof(uint32_t); + + // BingAlignedFileReader can only read uint32_t bytes of data. So, + // we limit ourselves even more to reading 1GB at a time. + std::vector readReqs; + while (data_size > 0) { + AlignedRead readReq; + readReq.buf = data + write_offset; + readReq.offset = read_start + write_offset; + readReq.len = data_size > MAX_REQUEST_SIZE ? MAX_REQUEST_SIZE : data_size; + readReqs.push_back(readReq); + // in the corner case, the loop will not execute + data_size -= readReq.len; + write_offset += readReq.len; + } + IOContext &ctx = reader.get_ctx(); + reader.read(readReqs, ctx); + for (int i = 0; i < readReqs.size(); i++) { + // Since we are making sync calls, no request will be in the + // READ_WAIT state. + if ((*(ctx.m_pRequestsStatus))[i] != IOContext::READ_SUCCESS) { + std::stringstream str; + str << "Could not read binary data from index file at offset: " + << readReqs[i].offset << std::endl; + throw diskann::ANNException(str.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); } + } } template -void load_bin(AlignedFileReader &reader, std::unique_ptr &data, size_t &npts, size_t &ndim, size_t offset) -{ - T *ptr = nullptr; - load_bin(reader, ptr, npts, ndim, offset); - data.reset(ptr); +void load_bin(AlignedFileReader &reader, std::unique_ptr &data, + size_t &npts, size_t &ndim, size_t offset) { + T *ptr = nullptr; + load_bin(reader, ptr, npts, ndim, offset); + data.reset(ptr); } template -void copy_aligned_data_from_file(AlignedFileReader &reader, T *&data, size_t &npts, size_t &ndim, - const size_t &rounded_dim, size_t offset) -{ - if (data == nullptr) - { - diskann::cerr << "Memory was not allocated for " << data << " before calling the load function. Exiting..." - << std::endl; - throw diskann::ANNException("Null pointer passed to copy_aligned_data_from_file()", -1, __FUNCSIG__, __FILE__, - __LINE__); - } - - size_t pts, dim; - get_bin_metadata(reader, pts, dim, offset); - - if (ndim != dim || npts != pts) - { - std::stringstream ss; - ss << "Either file dimension: " << dim << " is != passed dimension: " << ndim << " or file #pts: " << pts - << " is != passed #pts: " << npts << std::endl; - throw diskann::ANNException(ss.str(), -1, __FUNCSIG__, __FILE__, __LINE__); +void copy_aligned_data_from_file(AlignedFileReader &reader, T *&data, + size_t &npts, size_t &ndim, + const size_t &rounded_dim, size_t offset) { + if (data == nullptr) { + diskann::cerr << "Memory was not allocated for " << data + << " before calling the load function. Exiting..." + << std::endl; + throw diskann::ANNException( + "Null pointer passed to copy_aligned_data_from_file()", -1, __FUNCSIG__, + __FILE__, __LINE__); + } + + size_t pts, dim; + get_bin_metadata(reader, pts, dim, offset); + + if (ndim != dim || npts != pts) { + std::stringstream ss; + ss << "Either file dimension: " << dim + << " is != passed dimension: " << ndim << " or file #pts: " << pts + << " is != passed #pts: " << npts << std::endl; + throw diskann::ANNException(ss.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } + + // Instead of reading one point of ndim size and setting (rounded_dim - dim) + // values to zero We'll set everything to zero and read in chunks of data at + // the appropriate locations. + size_t read_offset = offset + 2 * sizeof(uint32_t); + memset(data, 0, npts * rounded_dim * sizeof(T)); + int i = 0; + std::vector read_requests; + + while (i < npts) { + int j = 0; + read_requests.clear(); + while (j < MAX_SIMULTANEOUS_READ_REQUESTS && i < npts) { + AlignedRead read_req; + read_req.buf = data + i * rounded_dim; + read_req.len = dim * sizeof(T); + read_req.offset = read_offset + i * dim * sizeof(T); + read_requests.push_back(read_req); + i++; + j++; } - - // Instead of reading one point of ndim size and setting (rounded_dim - dim) - // values to zero We'll set everything to zero and read in chunks of data at - // the appropriate locations. - size_t read_offset = offset + 2 * sizeof(uint32_t); - memset(data, 0, npts * rounded_dim * sizeof(T)); - int i = 0; - std::vector read_requests; - - while (i < npts) - { - int j = 0; - read_requests.clear(); - while (j < MAX_SIMULTANEOUS_READ_REQUESTS && i < npts) - { - AlignedRead read_req; - read_req.buf = data + i * rounded_dim; - read_req.len = dim * sizeof(T); - read_req.offset = read_offset + i * dim * sizeof(T); - read_requests.push_back(read_req); - i++; - j++; - } - IOContext &ctx = reader.get_ctx(); - reader.read(read_requests, ctx); - for (int k = 0; k < read_requests.size(); k++) - { - if ((*ctx.m_pRequestsStatus)[k] != IOContext::READ_SUCCESS) - { - throw diskann::ANNException("Load data from file using AlignedReader failed.", -1, __FUNCSIG__, - __FILE__, __LINE__); - } - } + IOContext &ctx = reader.get_ctx(); + reader.read(read_requests, ctx); + for (int k = 0; k < read_requests.size(); k++) { + if ((*ctx.m_pRequestsStatus)[k] != IOContext::READ_SUCCESS) { + throw diskann::ANNException( + "Load data from file using AlignedReader failed.", -1, __FUNCSIG__, + __FILE__, __LINE__); + } } + } } // Unlike load_bin, assumes that data is already allocated 'size' entries -template void read_array(AlignedFileReader &reader, T *data, size_t size, size_t offset) -{ - if (data == nullptr) - { - throw diskann::ANNException("read_array requires an allocated buffer.", -1); - if (size * sizeof(T) > MAX_REQUEST_SIZE) - { - std::stringstream ss; - ss << "Cannot read more than " << MAX_REQUEST_SIZE - << " bytes. Current request size: " << std::to_string(size) << " sizeof(T): " << sizeof(T) << std::endl; - throw diskann::ANNException(ss.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - std::vector read_requests; - AlignedRead read_req; - read_req.buf = data; - read_req.len = size * sizeof(T); - read_req.offset = offset; - read_requests.push_back(read_req); - IOContext &ctx = reader.get_ctx(); - reader.read(read_requests, ctx); - - if ((*(ctx.m_pRequestsStatus))[0] != IOContext::READ_SUCCESS) - { - std::stringstream ss; - ss << "Failed to read_array() of size: " << size * sizeof(T) << " at offset: " << offset << " from reader. " - << std::endl; - throw diskann::ANNException(ss.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } +template +void read_array(AlignedFileReader &reader, T *data, size_t size, + size_t offset) { + if (data == nullptr) { + throw diskann::ANNException("read_array requires an allocated buffer.", -1); + if (size * sizeof(T) > MAX_REQUEST_SIZE) { + std::stringstream ss; + ss << "Cannot read more than " << MAX_REQUEST_SIZE + << " bytes. Current request size: " << std::to_string(size) + << " sizeof(T): " << sizeof(T) << std::endl; + throw diskann::ANNException(ss.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); } + std::vector read_requests; + AlignedRead read_req; + read_req.buf = data; + read_req.len = size * sizeof(T); + read_req.offset = offset; + read_requests.push_back(read_req); + IOContext &ctx = reader.get_ctx(); + reader.read(read_requests, ctx); + + if ((*(ctx.m_pRequestsStatus))[0] != IOContext::READ_SUCCESS) { + std::stringstream ss; + ss << "Failed to read_array() of size: " << size * sizeof(T) + << " at offset: " << offset << " from reader. " << std::endl; + throw diskann::ANNException(ss.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } + } } -template void read_value(AlignedFileReader &reader, T &value, size_t offset) -{ - read_array(reader, &value, 1, offset); +template +void read_value(AlignedFileReader &reader, T &value, size_t offset) { + read_array(reader, &value, 1, offset); } -template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, std::unique_ptr &data, - size_t &npts, size_t &ndim, size_t offset); -template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, std::unique_ptr &data, - size_t &npts, size_t &ndim, size_t offset); -template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, std::unique_ptr &data, - size_t &npts, size_t &ndim, size_t offset); -template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, std::unique_ptr &data, - size_t &npts, size_t &ndim, size_t offset); -template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, std::unique_ptr &data, - size_t &npts, size_t &ndim, size_t offset); -template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, std::unique_ptr &data, size_t &npts, - size_t &ndim, size_t offset); - -template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, uint8_t *&data, size_t &npts, size_t &ndim, - size_t offset); -template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, int64_t *&data, size_t &npts, size_t &ndim, - size_t offset); -template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, uint64_t *&data, size_t &npts, - size_t &ndim, size_t offset); -template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, uint32_t *&data, size_t &npts, - size_t &ndim, size_t offset); -template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, int32_t *&data, size_t &npts, size_t &ndim, +template DISKANN_DLLEXPORT void load_bin( + AlignedFileReader &reader, std::unique_ptr &data, size_t &npts, + size_t &ndim, size_t offset); +template DISKANN_DLLEXPORT void load_bin( + AlignedFileReader &reader, std::unique_ptr &data, size_t &npts, + size_t &ndim, size_t offset); +template DISKANN_DLLEXPORT void load_bin( + AlignedFileReader &reader, std::unique_ptr &data, size_t &npts, + size_t &ndim, size_t offset); +template DISKANN_DLLEXPORT void load_bin( + AlignedFileReader &reader, std::unique_ptr &data, size_t &npts, + size_t &ndim, size_t offset); +template DISKANN_DLLEXPORT void load_bin( + AlignedFileReader &reader, std::unique_ptr &data, size_t &npts, + size_t &ndim, size_t offset); +template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, + std::unique_ptr &data, + size_t &npts, size_t &ndim, + size_t offset); + +template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, + uint8_t *&data, size_t &npts, + size_t &ndim, size_t offset); +template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, + int64_t *&data, size_t &npts, + size_t &ndim, size_t offset); +template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, + uint64_t *&data, + size_t &npts, size_t &ndim, + size_t offset); +template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, + uint32_t *&data, + size_t &npts, size_t &ndim, + size_t offset); +template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, + int32_t *&data, size_t &npts, + size_t &ndim, size_t offset); + +template DISKANN_DLLEXPORT void copy_aligned_data_from_file( + AlignedFileReader &reader, uint8_t *&data, size_t &npts, size_t &dim, + const size_t &rounded_dim, size_t offset); +template DISKANN_DLLEXPORT void copy_aligned_data_from_file( + AlignedFileReader &reader, int8_t *&data, size_t &npts, size_t &dim, + const size_t &rounded_dim, size_t offset); +template DISKANN_DLLEXPORT void copy_aligned_data_from_file( + AlignedFileReader &reader, float *&data, size_t &npts, size_t &dim, + const size_t &rounded_dim, size_t offset); + +template DISKANN_DLLEXPORT void read_array(AlignedFileReader &reader, + char *data, size_t size, + size_t offset); + +template DISKANN_DLLEXPORT void read_array(AlignedFileReader &reader, + uint8_t *data, size_t size, + size_t offset); +template DISKANN_DLLEXPORT void read_array(AlignedFileReader &reader, + int8_t *data, size_t size, + size_t offset); +template DISKANN_DLLEXPORT void read_array(AlignedFileReader &reader, + uint32_t *data, + size_t size, + size_t offset); +template DISKANN_DLLEXPORT void read_array(AlignedFileReader &reader, + float *data, size_t size, size_t offset); -template DISKANN_DLLEXPORT void copy_aligned_data_from_file(AlignedFileReader &reader, uint8_t *&data, - size_t &npts, size_t &dim, - const size_t &rounded_dim, size_t offset); -template DISKANN_DLLEXPORT void copy_aligned_data_from_file(AlignedFileReader &reader, int8_t *&data, - size_t &npts, size_t &dim, - const size_t &rounded_dim, size_t offset); -template DISKANN_DLLEXPORT void copy_aligned_data_from_file(AlignedFileReader &reader, float *&data, - size_t &npts, size_t &dim, const size_t &rounded_dim, - size_t offset); - -template DISKANN_DLLEXPORT void read_array(AlignedFileReader &reader, char *data, size_t size, size_t offset); - -template DISKANN_DLLEXPORT void read_array(AlignedFileReader &reader, uint8_t *data, size_t size, +template DISKANN_DLLEXPORT void read_value(AlignedFileReader &reader, + uint8_t &value, size_t offset); -template DISKANN_DLLEXPORT void read_array(AlignedFileReader &reader, int8_t *data, size_t size, size_t offset); -template DISKANN_DLLEXPORT void read_array(AlignedFileReader &reader, uint32_t *data, size_t size, +template DISKANN_DLLEXPORT void read_value(AlignedFileReader &reader, + int8_t &value, + size_t offset); +template DISKANN_DLLEXPORT void read_value(AlignedFileReader &reader, + float &value, size_t offset); +template DISKANN_DLLEXPORT void read_value(AlignedFileReader &reader, + uint32_t &value, + size_t offset); +template DISKANN_DLLEXPORT void read_value(AlignedFileReader &reader, + uint64_t &value, size_t offset); -template DISKANN_DLLEXPORT void read_array(AlignedFileReader &reader, float *data, size_t size, size_t offset); - -template DISKANN_DLLEXPORT void read_value(AlignedFileReader &reader, uint8_t &value, size_t offset); -template DISKANN_DLLEXPORT void read_value(AlignedFileReader &reader, int8_t &value, size_t offset); -template DISKANN_DLLEXPORT void read_value(AlignedFileReader &reader, float &value, size_t offset); -template DISKANN_DLLEXPORT void read_value(AlignedFileReader &reader, uint32_t &value, size_t offset); -template DISKANN_DLLEXPORT void read_value(AlignedFileReader &reader, uint64_t &value, size_t offset); #endif -} // namespace diskann +} // namespace diskann diff --git a/src/windows_aligned_file_reader.cpp b/src/windows_aligned_file_reader.cpp index c65656d3e..a1911efd4 100644 --- a/src/windows_aligned_file_reader.cpp +++ b/src/windows_aligned_file_reader.cpp @@ -9,172 +9,158 @@ #define SECTOR_LEN 4096 -void WindowsAlignedFileReader::open(const std::string &fname) -{ - +void WindowsAlignedFileReader::open(const std::string &fname) { #ifdef UNICODE - m_filename = std::wstring(fname.begin(), fname.end()); + m_filename = std::wstring(fname.begin(), fname.end()); #else - m_filename = fname; + m_filename = fname; #endif - this->register_thread(); + this->register_thread(); } -void WindowsAlignedFileReader::close() -{ - for (auto &k_v : ctx_map) - { - IOContext ctx = ctx_map[k_v.first]; - CloseHandle(ctx.fhandle); - } +void WindowsAlignedFileReader::close() { + for (auto &k_v : ctx_map) { + IOContext ctx = ctx_map[k_v.first]; + CloseHandle(ctx.fhandle); + } } -void WindowsAlignedFileReader::register_thread() -{ - std::unique_lock lk(this->ctx_mut); - if (this->ctx_map.find(std::this_thread::get_id()) != ctx_map.end()) - { - diskann::cout << "Warning:: Duplicate registration for thread_id : " << std::this_thread::get_id() << std::endl; - } - - IOContext ctx; - ctx.fhandle = CreateFile( - m_filename.c_str(), GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, - FILE_ATTRIBUTE_READONLY | FILE_FLAG_NO_BUFFERING | FILE_FLAG_OVERLAPPED | FILE_FLAG_RANDOM_ACCESS, NULL); - if (ctx.fhandle == INVALID_HANDLE_VALUE) - { - diskann::cout << "Error opening " << std::string(m_filename.begin(), m_filename.end()) - << " -- error=" << GetLastError() << std::endl; - } - - // create IOCompletionPort - ctx.iocp = CreateIoCompletionPort(ctx.fhandle, ctx.iocp, 0, 0); - - // create MAX_DEPTH # of reqs - for (uint64_t i = 0; i < MAX_IO_DEPTH; i++) - { - OVERLAPPED os; - memset(&os, 0, sizeof(OVERLAPPED)); - // os.hEvent = CreateEventA(NULL, TRUE, FALSE, NULL); - ctx.reqs.push_back(os); - } - this->ctx_map.insert(std::make_pair(std::this_thread::get_id(), ctx)); +void WindowsAlignedFileReader::register_thread() { + std::unique_lock lk(this->ctx_mut); + if (this->ctx_map.find(std::this_thread::get_id()) != ctx_map.end()) { + diskann::cout << "Warning:: Duplicate registration for thread_id : " + << std::this_thread::get_id() << std::endl; + } + + IOContext ctx; + ctx.fhandle = CreateFile(m_filename.c_str(), GENERIC_READ, FILE_SHARE_READ, + NULL, OPEN_EXISTING, + FILE_ATTRIBUTE_READONLY | FILE_FLAG_NO_BUFFERING | + FILE_FLAG_OVERLAPPED | FILE_FLAG_RANDOM_ACCESS, + NULL); + if (ctx.fhandle == INVALID_HANDLE_VALUE) { + diskann::cout << "Error opening " + << std::string(m_filename.begin(), m_filename.end()) + << " -- error=" << GetLastError() << std::endl; + } + + // create IOCompletionPort + ctx.iocp = CreateIoCompletionPort(ctx.fhandle, ctx.iocp, 0, 0); + + // create MAX_DEPTH # of reqs + for (uint64_t i = 0; i < MAX_IO_DEPTH; i++) { + OVERLAPPED os; + memset(&os, 0, sizeof(OVERLAPPED)); + // os.hEvent = CreateEventA(NULL, TRUE, FALSE, NULL); + ctx.reqs.push_back(os); + } + this->ctx_map.insert(std::make_pair(std::this_thread::get_id(), ctx)); } -IOContext &WindowsAlignedFileReader::get_ctx() -{ - std::unique_lock lk(this->ctx_mut); - if (ctx_map.find(std::this_thread::get_id()) == ctx_map.end()) - { - std::stringstream stream; - stream << "unable to find IOContext for thread_id : " << std::this_thread::get_id() << "\n"; - throw diskann::ANNException(stream.str(), -2, __FUNCSIG__, __FILE__, __LINE__); - } - IOContext &ctx = ctx_map[std::this_thread::get_id()]; - lk.unlock(); - return ctx; +IOContext &WindowsAlignedFileReader::get_ctx() { + std::unique_lock lk(this->ctx_mut); + if (ctx_map.find(std::this_thread::get_id()) == ctx_map.end()) { + std::stringstream stream; + stream << "unable to find IOContext for thread_id : " + << std::this_thread::get_id() << "\n"; + throw diskann::ANNException(stream.str(), -2, __FUNCSIG__, __FILE__, + __LINE__); + } + IOContext &ctx = ctx_map[std::this_thread::get_id()]; + lk.unlock(); + return ctx; } -void WindowsAlignedFileReader::read(std::vector &read_reqs, IOContext &ctx, bool async) -{ - using namespace std::chrono_literals; - // execute each request sequentially - size_t n_reqs = read_reqs.size(); - uint64_t n_batches = ROUND_UP(n_reqs, MAX_IO_DEPTH) / MAX_IO_DEPTH; - for (uint64_t i = 0; i < n_batches; i++) - { - // reset all OVERLAPPED objects - for (auto &os : ctx.reqs) - { - // HANDLE evt = os.hEvent; - memset(&os, 0, sizeof(os)); - // os.hEvent = evt; - - /* - if (ResetEvent(os.hEvent) == 0) { - diskann::cerr << "ResetEvent failed" << std::endl; - exit(-3); - } - */ +void WindowsAlignedFileReader::read(std::vector &read_reqs, + IOContext &ctx, bool async) { + using namespace std::chrono_literals; + // execute each request sequentially + size_t n_reqs = read_reqs.size(); + uint64_t n_batches = ROUND_UP(n_reqs, MAX_IO_DEPTH) / MAX_IO_DEPTH; + for (uint64_t i = 0; i < n_batches; i++) { + // reset all OVERLAPPED objects + for (auto &os : ctx.reqs) { + // HANDLE evt = os.hEvent; + memset(&os, 0, sizeof(os)); + // os.hEvent = evt; + + /* + if (ResetEvent(os.hEvent) == 0) { + diskann::cerr << "ResetEvent failed" << std::endl; + exit(-3); } + */ + } - // batch start/end - uint64_t batch_start = MAX_IO_DEPTH * i; - uint64_t batch_size = std::min((uint64_t)(n_reqs - batch_start), (uint64_t)MAX_IO_DEPTH); - - // fill OVERLAPPED and issue them - for (uint64_t j = 0; j < batch_size; j++) - { - AlignedRead &req = read_reqs[batch_start + j]; - OVERLAPPED &os = ctx.reqs[j]; - - uint64_t offset = req.offset; - uint64_t nbytes = req.len; - char *read_buf = (char *)req.buf; - assert(IS_ALIGNED(read_buf, SECTOR_LEN)); - assert(IS_ALIGNED(offset, SECTOR_LEN)); - assert(IS_ALIGNED(nbytes, SECTOR_LEN)); - - // fill in OVERLAPPED struct - os.Offset = offset & 0xffffffff; - os.OffsetHigh = (offset >> 32); - - BOOL ret = ReadFile(ctx.fhandle, read_buf, nbytes, NULL, &os); - if (ret == FALSE) - { - auto error = GetLastError(); - if (error != ERROR_IO_PENDING) - { - diskann::cerr << "Error queuing IO -- " << error << "\n"; - } - } - else - { - diskann::cerr << "Error queueing IO -- ReadFile returned TRUE" << std::endl; - } + // batch start/end + uint64_t batch_start = MAX_IO_DEPTH * i; + uint64_t batch_size = + std::min((uint64_t)(n_reqs - batch_start), (uint64_t)MAX_IO_DEPTH); + + // fill OVERLAPPED and issue them + for (uint64_t j = 0; j < batch_size; j++) { + AlignedRead &req = read_reqs[batch_start + j]; + OVERLAPPED &os = ctx.reqs[j]; + + uint64_t offset = req.offset; + uint64_t nbytes = req.len; + char *read_buf = (char *)req.buf; + assert(IS_ALIGNED(read_buf, SECTOR_LEN)); + assert(IS_ALIGNED(offset, SECTOR_LEN)); + assert(IS_ALIGNED(nbytes, SECTOR_LEN)); + + // fill in OVERLAPPED struct + os.Offset = offset & 0xffffffff; + os.OffsetHigh = (offset >> 32); + + BOOL ret = ReadFile(ctx.fhandle, read_buf, nbytes, NULL, &os); + if (ret == FALSE) { + auto error = GetLastError(); + if (error != ERROR_IO_PENDING) { + diskann::cerr << "Error queuing IO -- " << error << "\n"; } - DWORD n_read = 0; - uint64_t n_complete = 0; - ULONG_PTR completion_key = 0; - OVERLAPPED *lp_os; - while (n_complete < batch_size) - { - if (GetQueuedCompletionStatus(ctx.iocp, &n_read, &completion_key, &lp_os, INFINITE) != 0) - { - // successfully dequeued a completed I/O - n_complete++; - } - else - { - // failed to dequeue OR dequeued failed I/O - if (lp_os == NULL) - { - DWORD error = GetLastError(); - if (error != WAIT_TIMEOUT) - { - diskann::cerr << "GetQueuedCompletionStatus() failed " - "with error = " - << error << std::endl; - throw diskann::ANNException("GetQueuedCompletionStatus failed with error: ", error, __FUNCSIG__, - __FILE__, __LINE__); - } - // no completion packet dequeued ==> sleep for 5us and try - // again - std::this_thread::sleep_for(5us); - } - else - { - // completion packet for failed IO dequeued - auto op_idx = lp_os - ctx.reqs.data(); - std::stringstream stream; - stream << "I/O failed , offset: " << read_reqs[op_idx].offset - << "with error code: " << GetLastError() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - } + } else { + diskann::cerr << "Error queueing IO -- ReadFile returned TRUE" + << std::endl; + } + } + DWORD n_read = 0; + uint64_t n_complete = 0; + ULONG_PTR completion_key = 0; + OVERLAPPED *lp_os; + while (n_complete < batch_size) { + if (GetQueuedCompletionStatus(ctx.iocp, &n_read, &completion_key, &lp_os, + INFINITE) != 0) { + // successfully dequeued a completed I/O + n_complete++; + } else { + // failed to dequeue OR dequeued failed I/O + if (lp_os == NULL) { + DWORD error = GetLastError(); + if (error != WAIT_TIMEOUT) { + diskann::cerr << "GetQueuedCompletionStatus() failed " + "with error = " + << error << std::endl; + throw diskann::ANNException( + "GetQueuedCompletionStatus failed with error: ", error, + __FUNCSIG__, __FILE__, __LINE__); + } + // no completion packet dequeued ==> sleep for 5us and try + // again + std::this_thread::sleep_for(5us); + } else { + // completion packet for failed IO dequeued + auto op_idx = lp_os - ctx.reqs.data(); + std::stringstream stream; + stream << "I/O failed , offset: " << read_reqs[op_idx].offset + << "with error code: " << GetLastError() << std::endl; + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); } + } } + } } #endif #endif From ec90ca6cae8b204a416fe11277221e3b9e565d36 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Wed, 19 Apr 2023 23:21:36 +0530 Subject: [PATCH 56/69] cleaned up and added comments --- include/abstract_data_store.h | 35 ++++++++++++++++++++++++----------- include/in_mem_data_store.h | 14 +++++++------- src/in_mem_data_store.cpp | 18 +++++++++--------- src/index.cpp | 10 +++++----- 4 files changed, 45 insertions(+), 32 deletions(-) diff --git a/include/abstract_data_store.h b/include/abstract_data_store.h index 6cf83e693..4662fc7dd 100644 --- a/include/abstract_data_store.h +++ b/include/abstract_data_store.h @@ -32,20 +32,23 @@ class AbstractDataStore { // Implementers can choose to return _dim if they are not // concerned about memory alignment. - // Returns _dim aligned to a 8-byte value. Used for allocating - // aligned memory for efficiency/simplicity of code. + // Some distance metrics (like l2) need data vectors to be aligned, so we + // align the dimension by padding zeros. virtual size_t get_aligned_dim() const = 0; // populate the store with vectors (either from a pointer or bin file), - // potentially after normalizing the vectors if the metric deems so + // potentially after pre-processing the vectors if the metric deems so + // e.g., normalizing vectors for cosine distance over floating-point vectors + // useful for bulk or static index building. virtual void populate_data(const data_t *vectors, const location_t num_pts) = 0; virtual void populate_data(const std::string &filename, const size_t offset) = 0; - // reverse of populate, save the first num_pts many points back to bin file - virtual void save_data_to_bin(const std::string &filename, - const location_t num_pts) = 0; + // save the first num_pts many vectors back to bin file + // note: cannot undo the pre-processing done in populate data + virtual void extract_data_to_bin(const std::string &filename, + const location_t num_pts) = 0; // Returns the updated capacity of the datastore. Clients should check // if resize actually changed the capacity to new_num_points before @@ -58,16 +61,24 @@ class AbstractDataStore { virtual location_t resize(const location_t new_num_points); // operations on vectors + // like populate_data function, but over one vector at a time useful for + // streaming setting virtual void get_vector(const location_t i, data_t *dest) const = 0; virtual void set_vector(const location_t i, const data_t *const vector) = 0; virtual void prefetch_vector(const location_t loc) = 0; // internal shuffle operations to move around vectors - virtual void reposition_points(const location_t start_loc, - const location_t end_loc, - const location_t num_points) = 0; - virtual void copy_points(const location_t from_loc, const location_t to_loc, - const location_t num_points) = 0; + // will bulk-move all the vectors in [old_start_loc, old_start_loc + + // num_points) to [new_start_loc, new_start_loc + num_points) and set the old + // positions to zero vectors. + virtual void move_vectors(const location_t old_start_loc, + const location_t new_start_loc, + const location_t num_points) = 0; + + // same as above, without resetting the vectors in [from_loc, from_loc + + // num_points) to zero + virtual void copy_vectors(const location_t from_loc, const location_t to_loc, + const location_t num_points) = 0; // metric specific operations @@ -85,6 +96,8 @@ class AbstractDataStore { virtual location_t calculate_medoid() const = 0; // search helpers + // if the base data is aligned per the request of the metric, this will tell + // how to align the query vector in a consistent manner virtual size_t get_alignment_factor() const = 0; protected: diff --git a/include/in_mem_data_store.h b/include/in_mem_data_store.h index 6314b90bc..db8b0bf20 100644 --- a/include/in_mem_data_store.h +++ b/include/in_mem_data_store.h @@ -36,19 +36,19 @@ class InMemDataStore : public AbstractDataStore { virtual void populate_data(const std::string &filename, const size_t offset) override; - virtual void save_data_to_bin(const std::string &filename, - const location_t num_pts) override; + virtual void extract_data_to_bin(const std::string &filename, + const location_t num_pts) override; virtual void get_vector(const location_t i, data_t *target) const override; virtual void set_vector(const location_t i, const data_t *const vector) override; virtual void prefetch_vector(const location_t loc) override; - virtual void reposition_points(const location_t old_location_start, - const location_t new_location_start, - const location_t num_points) override; - virtual void copy_points(const location_t from_loc, const location_t to_loc, - const location_t num_points) override; + virtual void move_vectors(const location_t old_location_start, + const location_t new_location_start, + const location_t num_points) override; + virtual void copy_vectors(const location_t from_loc, const location_t to_loc, + const location_t num_points) override; virtual float get_distance(const data_t *query, const location_t loc) const override; diff --git a/src/in_mem_data_store.cpp b/src/in_mem_data_store.cpp index 690ec60b5..54757b17a 100644 --- a/src/in_mem_data_store.cpp +++ b/src/in_mem_data_store.cpp @@ -154,8 +154,8 @@ void InMemDataStore::populate_data(const std::string &filename, } template -void InMemDataStore::save_data_to_bin(const std::string &filename, - const location_t num_points) { +void InMemDataStore::extract_data_to_bin(const std::string &filename, + const location_t num_points) { save_data_in_base_dimensions(filename, _data, num_points, this->get_dims(), this->get_aligned_dim(), 0U); } @@ -259,9 +259,9 @@ location_t InMemDataStore::shrink(const location_t new_size) { } template -void InMemDataStore::reposition_points( - const location_t old_location_start, const location_t new_location_start, - const location_t num_locations) { +void InMemDataStore::move_vectors(const location_t old_location_start, + const location_t new_location_start, + const location_t num_locations) { if (num_locations == 0 || old_location_start == new_location_start) { return; } @@ -294,16 +294,16 @@ void InMemDataStore::reposition_points( } // Use memmove to handle overlapping ranges. - copy_points(old_location_start, new_location_start, num_locations); + copy_vectors(old_location_start, new_location_start, num_locations); memset(_data + _aligned_dim * mem_clear_loc_start, 0, sizeof(data_t) * _aligned_dim * (mem_clear_loc_end_limit - mem_clear_loc_start)); } template -void InMemDataStore::copy_points(const location_t from_loc, - const location_t to_loc, - const location_t num_points) { +void InMemDataStore::copy_vectors(const location_t from_loc, + const location_t to_loc, + const location_t num_points) { assert(from_loc < this->_capacity); assert(to_loc < this->_capacity); assert(num_points < this->_capacity); diff --git a/src/index.cpp b/src/index.cpp index 4a4959464..da3b3939d 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -458,7 +458,7 @@ size_t Index::load_data(std::string filename) { copy_aligned_data_from_file(reader, _data, file_num_points, file_dim, _aligned_dim); #else - _data_store->populate_data(filename, 0U); // offset == 0. + _data_store->load(filename); // offset == 0. #endif return file_num_points; } @@ -2112,7 +2112,7 @@ void Index::generate_frozen_point() { _pq_data + res * _num_pq_chunks, _num_pq_chunks * DIV_ROUND_UP(NUM_PQ_BITS, 8)); } else { - _data_store->copy_points(res, _max_points, 1); + _data_store->copy_vectors(res, _max_points, 1); } } @@ -2388,7 +2388,7 @@ void Index::compact_data() { assert(new_location[old] < old); _final_graph[new_location[old]].swap(_final_graph[old]); - _data_store->copy_points(old, new_location[old], 1); + _data_store->copy_vectors(old, new_location[old], 1); } } else { _final_graph[old].clear(); @@ -2536,8 +2536,8 @@ void Index::reposition_points(uint32_t old_location_start, mem_clear_loc_end_limit = new_location_start; } } - _data_store->reposition_points(old_location_start, new_location_start, - num_locations); + _data_store->move_vectors(old_location_start, new_location_start, + num_locations); } template From 6025c337b4a72f61a65297bda4f18d6f4f49804b Mon Sep 17 00:00:00 2001 From: ravishankar Date: Wed, 19 Apr 2023 23:48:26 +0530 Subject: [PATCH 57/69] clang-formatted --- src/dll/dllmain.cpp | 13 +- tests/build_disk_index.cpp | 341 ++-- tests/build_stitched_index.cpp | 706 ++++---- tests/range_search_disk_index.cpp | 598 +++---- tests/search_disk_index.cpp | 822 ++++----- tests/search_memory_index.cpp | 754 ++++---- tests/test_insert_deletes_consolidate.cpp | 797 +++++---- tests/test_streaming_scenario.cpp | 672 +++---- tests/utils/bin_to_fvecs.cpp | 104 +- tests/utils/bin_to_tsv.cpp | 103 +- tests/utils/calculate_recall.cpp | 69 +- tests/utils/compute_groundtruth.cpp | 945 +++++----- .../utils/compute_groundtruth_for_filters.cpp | 1563 ++++++++--------- tests/utils/count_bfs_levels.cpp | 95 +- tests/utils/create_disk_layout.cpp | 57 +- tests/utils/float_bin_to_int8.cpp | 105 +- tests/utils/fvecs_to_bin.cpp | 141 +- tests/utils/fvecs_to_bvecs.cpp | 86 +- tests/utils/gen_random_slice.cpp | 56 +- tests/utils/generate_pq.cpp | 107 +- tests/utils/generate_synthetic_labels.cpp | 326 ++-- tests/utils/int8_to_float.cpp | 28 +- tests/utils/int8_to_float_scale.cpp | 105 +- tests/utils/ivecs_to_bin.cpp | 89 +- tests/utils/merge_shards.cpp | 45 +- tests/utils/partition_data.cpp | 53 +- tests/utils/partition_with_ram_budget.cpp | 53 +- tests/utils/rand_data_gen.cpp | 331 ++-- tests/utils/simulate_aggregate_recall.cpp | 119 +- tests/utils/stats_label_data.cpp | 207 +-- tests/utils/tsv_to_bin.cpp | 191 +- tests/utils/uint32_to_uint8.cpp | 28 +- tests/utils/uint8_to_float.cpp | 28 +- tests/utils/vector_analysis.cpp | 239 ++- 34 files changed, 4998 insertions(+), 4978 deletions(-) diff --git a/src/dll/dllmain.cpp b/src/dll/dllmain.cpp index 9f5ce4420..826323fd0 100644 --- a/src/dll/dllmain.cpp +++ b/src/dll/dllmain.cpp @@ -1,15 +1,14 @@ // dllmain.cpp : Defines the entry point for the DLL application. #include -BOOL APIENTRY DllMain(HMODULE hModule, DWORD ul_reason_for_call, LPVOID lpReserved) -{ - switch (ul_reason_for_call) - { +BOOL APIENTRY DllMain(HMODULE hModule, DWORD ul_reason_for_call, + LPVOID lpReserved) { + switch (ul_reason_for_call) { case DLL_PROCESS_ATTACH: case DLL_THREAD_ATTACH: case DLL_THREAD_DETACH: case DLL_PROCESS_DETACH: - break; - } - return TRUE; + break; + } + return TRUE; } diff --git a/tests/build_disk_index.cpp b/tests/build_disk_index.cpp index 3d097458b..b60c71e8b 100644 --- a/tests/build_disk_index.cpp +++ b/tests/build_disk_index.cpp @@ -12,178 +12,193 @@ namespace po = boost::program_options; -int main(int argc, char **argv) -{ - std::string data_type, dist_fn, data_path, index_path_prefix, codebook_prefix, label_file, universal_label, - label_type; - uint32_t num_threads, R, L, disk_PQ, build_PQ, QD, Lf, filter_threshold; - float B, M; - bool append_reorder_data = false; - bool use_opq = false; +int main(int argc, char **argv) { + std::string data_type, dist_fn, data_path, index_path_prefix, codebook_prefix, + label_file, universal_label, label_type; + uint32_t num_threads, R, L, disk_PQ, build_PQ, QD, Lf, filter_threshold; + float B, M; + bool append_reorder_data = false; + bool use_opq = false; - po::options_description desc{"Arguments"}; - try - { - desc.add_options()("help,h", "Print information on arguments"); - desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); - desc.add_options()("dist_fn", po::value(&dist_fn)->required(), "distance function "); - desc.add_options()("data_path", po::value(&data_path)->required(), - "Input data file in bin format"); - desc.add_options()("index_path_prefix", po::value(&index_path_prefix)->required(), - "Path prefix for saving index file components"); - desc.add_options()("max_degree,R", po::value(&R)->default_value(64), "Maximum graph degree"); - desc.add_options()("Lbuild,L", po::value(&L)->default_value(100), - "Build complexity, higher value results in better graphs"); - desc.add_options()("search_DRAM_budget,B", po::value(&B)->required(), - "DRAM budget in GB for searching the index to set the " - "compressed level for data while search happens"); - desc.add_options()("build_DRAM_budget,M", po::value(&M)->required(), - "DRAM budget in GB for building the index"); - desc.add_options()("num_threads,T", po::value(&num_threads)->default_value(omp_get_num_procs()), - "Number of threads used for building index (defaults to " - "omp_get_num_procs())"); - desc.add_options()("QD", po::value(&QD)->default_value(0), " Quantized Dimension for compression"); - desc.add_options()("codebook_prefix", po::value(&codebook_prefix)->default_value(""), - "Path prefix for pre-trained codebook"); - desc.add_options()("PQ_disk_bytes", po::value(&disk_PQ)->default_value(0), - "Number of bytes to which vectors should be compressed " - "on SSD; 0 for no compression"); - desc.add_options()("append_reorder_data", po::bool_switch()->default_value(false), - "Include full precision data in the index. Use only in " - "conjuction with compressed data on SSD."); - desc.add_options()("build_PQ_bytes", po::value(&build_PQ)->default_value(0), - "Number of PQ bytes to build the index; 0 for full " - "precision build"); - desc.add_options()("use_opq", po::bool_switch()->default_value(false), - "Use Optimized Product Quantization (OPQ)."); - desc.add_options()("label_file", po::value(&label_file)->default_value(""), - "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"); - desc.add_options()("universal_label", po::value(&universal_label)->default_value(""), - "Universal label, Use only in conjuction 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 point instead of comma " - "separated filters for that point"); - desc.add_options()("FilteredLbuild,Lf", po::value(&Lf)->default_value(0), - "Build complexity for filtered points, higher value " - "results in better graphs"); - desc.add_options()("filter_threshold,F", po::value(&filter_threshold)->default_value(0), - "Threshold to break up the existing nodes to generate new graph " - "internally where each node has a maximum F labels."); - desc.add_options()("label_type", po::value(&label_type)->default_value("uint"), - "Storage type of Labels , default value is uint which " - "will consume memory 4 bytes per filter"); + po::options_description desc{"Arguments"}; + try { + desc.add_options()("help,h", "Print information on arguments"); + desc.add_options()("data_type", + po::value(&data_type)->required(), + "data type "); + desc.add_options()("dist_fn", po::value(&dist_fn)->required(), + "distance function "); + desc.add_options()("data_path", + po::value(&data_path)->required(), + "Input data file in bin format"); + desc.add_options()("index_path_prefix", + po::value(&index_path_prefix)->required(), + "Path prefix for saving index file components"); + desc.add_options()("max_degree,R", + po::value(&R)->default_value(64), + "Maximum graph degree"); + desc.add_options()( + "Lbuild,L", po::value(&L)->default_value(100), + "Build complexity, higher value results in better graphs"); + desc.add_options()("search_DRAM_budget,B", po::value(&B)->required(), + "DRAM budget in GB for searching the index to set the " + "compressed level for data while search happens"); + desc.add_options()("build_DRAM_budget,M", po::value(&M)->required(), + "DRAM budget in GB for building the index"); + desc.add_options()( + "num_threads,T", + po::value(&num_threads)->default_value(omp_get_num_procs()), + "Number of threads used for building index (defaults to " + "omp_get_num_procs())"); + desc.add_options()("QD", po::value(&QD)->default_value(0), + " Quantized Dimension for compression"); + desc.add_options()( + "codebook_prefix", + po::value(&codebook_prefix)->default_value(""), + "Path prefix for pre-trained codebook"); + desc.add_options()("PQ_disk_bytes", + po::value(&disk_PQ)->default_value(0), + "Number of bytes to which vectors should be compressed " + "on SSD; 0 for no compression"); + desc.add_options()("append_reorder_data", + po::bool_switch()->default_value(false), + "Include full precision data in the index. Use only in " + "conjuction with compressed data on SSD."); + desc.add_options()("build_PQ_bytes", + po::value(&build_PQ)->default_value(0), + "Number of PQ bytes to build the index; 0 for full " + "precision build"); + desc.add_options()("use_opq", po::bool_switch()->default_value(false), + "Use Optimized Product Quantization (OPQ)."); + desc.add_options()( + "label_file", po::value(&label_file)->default_value(""), + "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"); + desc.add_options()( + "universal_label", + po::value(&universal_label)->default_value(""), + "Universal label, Use only in conjuction 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 point instead of comma " + "separated filters for that point"); + desc.add_options()("FilteredLbuild,Lf", + po::value(&Lf)->default_value(0), + "Build complexity for filtered points, higher value " + "results in better graphs"); + desc.add_options()( + "filter_threshold,F", + po::value(&filter_threshold)->default_value(0), + "Threshold to break up the existing nodes to generate new graph " + "internally where each node has a maximum F labels."); + desc.add_options()( + "label_type", + po::value(&label_type)->default_value("uint"), + "Storage type of Labels , default value is uint which " + "will consume memory 4 bytes per filter"); - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) - { - std::cout << desc; - return 0; - } - po::notify(vm); - if (vm["append_reorder_data"].as()) - append_reorder_data = true; - if (vm["use_opq"].as()) - use_opq = true; - } - catch (const std::exception &ex) - { - std::cerr << ex.what() << '\n'; - return -1; + po::variables_map vm; + po::store(po::parse_command_line(argc, argv, desc), vm); + if (vm.count("help")) { + std::cout << desc; + return 0; } + po::notify(vm); + if (vm["append_reorder_data"].as()) append_reorder_data = true; + if (vm["use_opq"].as()) use_opq = true; + } catch (const std::exception &ex) { + std::cerr << ex.what() << '\n'; + return -1; + } - bool use_filters = false; - if (label_file != "") - { - use_filters = true; - } + bool use_filters = false; + if (label_file != "") { + use_filters = true; + } - diskann::Metric metric; - if (dist_fn == std::string("l2")) - metric = diskann::Metric::L2; - else if (dist_fn == std::string("mips")) - metric = diskann::Metric::INNER_PRODUCT; - else - { - std::cout << "Error. Only l2 and mips distance functions are supported" << std::endl; - return -1; - } + diskann::Metric metric; + if (dist_fn == std::string("l2")) + metric = diskann::Metric::L2; + else if (dist_fn == std::string("mips")) + metric = diskann::Metric::INNER_PRODUCT; + else { + std::cout << "Error. Only l2 and mips distance functions are supported" + << std::endl; + return -1; + } - if (append_reorder_data) - { - if (disk_PQ == 0) - { - std::cout << "Error: It is not necessary to append data for reordering " - "when vectors are not compressed on disk." - << std::endl; - return -1; - } - if (data_type != std::string("float")) - { - std::cout << "Error: Appending data for reordering currently only " - "supported for float data type." - << std::endl; - return -1; - } + if (append_reorder_data) { + if (disk_PQ == 0) { + std::cout << "Error: It is not necessary to append data for reordering " + "when vectors are not compressed on disk." + << std::endl; + return -1; + } + if (data_type != std::string("float")) { + std::cout << "Error: Appending data for reordering currently only " + "supported for float data type." + << std::endl; + return -1; } + } - std::string params = std::string(std::to_string(R)) + " " + std::string(std::to_string(L)) + " " + - std::string(std::to_string(B)) + " " + std::string(std::to_string(M)) + " " + - std::string(std::to_string(num_threads)) + " " + std::string(std::to_string(disk_PQ)) + " " + - std::string(std::to_string(append_reorder_data)) + " " + - std::string(std::to_string(build_PQ)) + " " + std::string(std::to_string(QD)); + std::string params = std::string(std::to_string(R)) + " " + + std::string(std::to_string(L)) + " " + + std::string(std::to_string(B)) + " " + + std::string(std::to_string(M)) + " " + + std::string(std::to_string(num_threads)) + " " + + std::string(std::to_string(disk_PQ)) + " " + + std::string(std::to_string(append_reorder_data)) + " " + + std::string(std::to_string(build_PQ)) + " " + + std::string(std::to_string(QD)); - try - { - if (label_file != "" && label_type == "ushort") - { - if (data_type == std::string("int8")) - return diskann::build_disk_index(data_path.c_str(), index_path_prefix.c_str(), params.c_str(), - metric, use_opq, codebook_prefix, use_filters, label_file, - universal_label, filter_threshold, Lf); - else if (data_type == std::string("uint8")) - return diskann::build_disk_index( - data_path.c_str(), index_path_prefix.c_str(), params.c_str(), metric, use_opq, codebook_prefix, - use_filters, label_file, universal_label, filter_threshold, Lf); - else if (data_type == std::string("float")) - return diskann::build_disk_index( - data_path.c_str(), index_path_prefix.c_str(), params.c_str(), metric, use_opq, codebook_prefix, - use_filters, label_file, universal_label, filter_threshold, Lf); - else - { - diskann::cerr << "Error. Unsupported data type" << std::endl; - return -1; - } - } - else - { - if (data_type == std::string("int8")) - return diskann::build_disk_index(data_path.c_str(), index_path_prefix.c_str(), params.c_str(), - metric, use_opq, codebook_prefix, use_filters, label_file, - universal_label, filter_threshold, Lf); - else if (data_type == std::string("uint8")) - return diskann::build_disk_index(data_path.c_str(), index_path_prefix.c_str(), params.c_str(), - metric, use_opq, codebook_prefix, use_filters, label_file, - universal_label, filter_threshold, Lf); - else if (data_type == std::string("float")) - return diskann::build_disk_index(data_path.c_str(), index_path_prefix.c_str(), params.c_str(), - metric, use_opq, codebook_prefix, use_filters, label_file, - universal_label, filter_threshold, Lf); - else - { - diskann::cerr << "Error. Unsupported data type" << std::endl; - return -1; - } - } - } - catch (const std::exception &e) - { - std::cout << std::string(e.what()) << std::endl; - diskann::cerr << "Index build failed." << std::endl; + try { + if (label_file != "" && label_type == "ushort") { + if (data_type == std::string("int8")) + return diskann::build_disk_index( + data_path.c_str(), index_path_prefix.c_str(), params.c_str(), + metric, use_opq, codebook_prefix, use_filters, label_file, + universal_label, filter_threshold, Lf); + else if (data_type == std::string("uint8")) + return diskann::build_disk_index( + data_path.c_str(), index_path_prefix.c_str(), params.c_str(), + metric, use_opq, codebook_prefix, use_filters, label_file, + universal_label, filter_threshold, Lf); + else if (data_type == std::string("float")) + return diskann::build_disk_index( + data_path.c_str(), index_path_prefix.c_str(), params.c_str(), + metric, use_opq, codebook_prefix, use_filters, label_file, + universal_label, filter_threshold, Lf); + else { + diskann::cerr << "Error. Unsupported data type" << std::endl; + return -1; + } + } else { + if (data_type == std::string("int8")) + return diskann::build_disk_index( + data_path.c_str(), index_path_prefix.c_str(), params.c_str(), + metric, use_opq, codebook_prefix, use_filters, label_file, + universal_label, filter_threshold, Lf); + else if (data_type == std::string("uint8")) + return diskann::build_disk_index( + data_path.c_str(), index_path_prefix.c_str(), params.c_str(), + metric, use_opq, codebook_prefix, use_filters, label_file, + universal_label, filter_threshold, Lf); + else if (data_type == std::string("float")) + return diskann::build_disk_index( + data_path.c_str(), index_path_prefix.c_str(), params.c_str(), + metric, use_opq, codebook_prefix, use_filters, label_file, + universal_label, filter_threshold, Lf); + else { + diskann::cerr << "Error. Unsupported data type" << std::endl; return -1; + } } + } catch (const std::exception &e) { + std::cout << std::string(e.what()) << std::endl; + diskann::cerr << "Index build failed." << std::endl; + return -1; + } } diff --git a/tests/build_stitched_index.cpp b/tests/build_stitched_index.cpp index 770a11f83..c19ea6da1 100644 --- a/tests/build_stitched_index.cpp +++ b/tests/build_stitched_index.cpp @@ -20,29 +20,28 @@ #include "utils.h" namespace po = boost::program_options; -typedef std::tuple>, uint64_t> stitch_indices_return_values; +typedef std::tuple>, uint64_t> + stitch_indices_return_values; /* * Inline function to display progress bar. */ -inline void print_progress(double percentage) -{ - int val = (int)(percentage * 100); - int lpad = (int)(percentage * PBWIDTH); - int rpad = PBWIDTH - lpad; - printf("\r%3d%% [%.*s%*s]", val, lpad, PBSTR, rpad, ""); - fflush(stdout); +inline void print_progress(double percentage) { + int val = (int)(percentage * 100); + int lpad = (int)(percentage * PBWIDTH); + int rpad = PBWIDTH - lpad; + printf("\r%3d%% [%.*s%*s]", val, lpad, PBSTR, rpad, ""); + fflush(stdout); } /* * Inline function to generate a random integer in a range. */ -inline size_t random(size_t range_from, size_t range_to) -{ - std::random_device rand_dev; - std::mt19937 generator(rand_dev()); - std::uniform_int_distribution distr(range_from, range_to); - return distr(generator); +inline size_t random(size_t range_from, size_t range_to) { + std::random_device rand_dev; + std::mt19937 generator(rand_dev()); + std::uniform_int_distribution distr(range_from, range_to); + return distr(generator); } /* @@ -50,52 +49,63 @@ inline size_t random(size_t range_from, size_t range_to) * * Arguments are merely the inputs from the command line. */ -void handle_args(int argc, char **argv, std::string &data_type, path &input_data_path, path &final_index_path_prefix, - path &label_data_path, std::string &universal_label, uint32_t &num_threads, uint32_t &R, uint32_t &L, - uint32_t &stitched_R, float &alpha) -{ - po::options_description desc{"Arguments"}; - try - { - desc.add_options()("help,h", "Print information on arguments"); - desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); - desc.add_options()("data_path", po::value(&input_data_path)->required(), "Input data file in bin format"); - desc.add_options()("index_path_prefix", po::value(&final_index_path_prefix)->required(), - "Path prefix for saving index file components"); - desc.add_options()("max_degree,R", po::value(&R)->default_value(64), "Maximum graph degree"); - desc.add_options()("Lbuild,L", po::value(&L)->default_value(100), - "Build complexity, higher value results in better graphs"); - desc.add_options()("stitched_R", po::value(&stitched_R)->default_value(100), - "Degree to prune final graph down to"); - desc.add_options()("alpha", po::value(&alpha)->default_value(1.2f), - "alpha controls density and diameter of graph, set " - "1 for sparse graph, " - "1.2 or 1.4 for denser graphs with lower diameter"); - desc.add_options()("num_threads,T", po::value(&num_threads)->default_value(omp_get_num_procs()), - "Number of threads used for building index (defaults to " - "omp_get_num_procs())"); - desc.add_options()("label_file", po::value(&label_data_path)->default_value(""), - "Input label file in txt format if present"); - desc.add_options()("universal_label", po::value(&universal_label)->default_value(""), - "If a point comes with the specified universal label (and only the " - "univ. " - "label), then the point is considered to have every possible " - "label"); - - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) - { - std::cout << desc; - exit(0); - } - po::notify(vm); - } - catch (const std::exception &ex) - { - std::cerr << ex.what() << '\n'; - throw; +void handle_args(int argc, char **argv, std::string &data_type, + path &input_data_path, path &final_index_path_prefix, + path &label_data_path, std::string &universal_label, + uint32_t &num_threads, uint32_t &R, uint32_t &L, + uint32_t &stitched_R, float &alpha) { + po::options_description desc{"Arguments"}; + try { + desc.add_options()("help,h", "Print information on arguments"); + desc.add_options()("data_type", + po::value(&data_type)->required(), + "data type "); + desc.add_options()("data_path", + po::value(&input_data_path)->required(), + "Input data file in bin format"); + desc.add_options()("index_path_prefix", + po::value(&final_index_path_prefix)->required(), + "Path prefix for saving index file components"); + desc.add_options()("max_degree,R", + po::value(&R)->default_value(64), + "Maximum graph degree"); + desc.add_options()( + "Lbuild,L", po::value(&L)->default_value(100), + "Build complexity, higher value results in better graphs"); + desc.add_options()("stitched_R", + po::value(&stitched_R)->default_value(100), + "Degree to prune final graph down to"); + desc.add_options()("alpha", po::value(&alpha)->default_value(1.2f), + "alpha controls density and diameter of graph, set " + "1 for sparse graph, " + "1.2 or 1.4 for denser graphs with lower diameter"); + desc.add_options()( + "num_threads,T", + po::value(&num_threads)->default_value(omp_get_num_procs()), + "Number of threads used for building index (defaults to " + "omp_get_num_procs())"); + desc.add_options()("label_file", + po::value(&label_data_path)->default_value(""), + "Input label file in txt format if present"); + desc.add_options()( + "universal_label", + po::value(&universal_label)->default_value(""), + "If a point comes with the specified universal label (and only the " + "univ. " + "label), then the point is considered to have every possible " + "label"); + + po::variables_map vm; + po::store(po::parse_command_line(argc, argv, desc), vm); + if (vm.count("help")) { + std::cout << desc; + exit(0); } + po::notify(vm); + } catch (const std::exception &ex) { + std::cerr << ex.what() << '\n'; + throw; + } } /* @@ -106,98 +116,109 @@ void handle_args(int argc, char **argv, std::string &data_type, path &input_data * 3. data (redundant for static indices) * 4. labels (redundant for static indices) */ -void save_full_index(path final_index_path_prefix, path input_data_path, uint64_t final_index_size, +void save_full_index(path final_index_path_prefix, path input_data_path, + uint64_t final_index_size, std::vector> stitched_graph, - tsl::robin_map entry_points, std::string universal_label, - path label_data_path) -{ - // aux. file 1 - auto saving_index_timer = std::chrono::high_resolution_clock::now(); - std::ifstream original_label_data_stream; - original_label_data_stream.exceptions(std::ios::badbit | std::ios::failbit); - original_label_data_stream.open(label_data_path, std::ios::binary); - std::ofstream new_label_data_stream; - new_label_data_stream.exceptions(std::ios::badbit | std::ios::failbit); - new_label_data_stream.open(final_index_path_prefix + "_labels.txt", std::ios::binary); - new_label_data_stream << original_label_data_stream.rdbuf(); - original_label_data_stream.close(); - new_label_data_stream.close(); - - // aux. file 2 - std::ifstream original_input_data_stream; - original_input_data_stream.exceptions(std::ios::badbit | std::ios::failbit); - original_input_data_stream.open(input_data_path, std::ios::binary); - std::ofstream new_input_data_stream; - new_input_data_stream.exceptions(std::ios::badbit | std::ios::failbit); - new_input_data_stream.open(final_index_path_prefix + ".data", std::ios::binary); - new_input_data_stream << original_input_data_stream.rdbuf(); - original_input_data_stream.close(); - new_input_data_stream.close(); - - // aux. file 3 - std::ofstream labels_to_medoids_writer; - labels_to_medoids_writer.exceptions(std::ios::badbit | std::ios::failbit); - labels_to_medoids_writer.open(final_index_path_prefix + "_labels_to_medoids.txt"); - for (auto iter : entry_points) - labels_to_medoids_writer << iter.first << ", " << iter.second << std::endl; - labels_to_medoids_writer.close(); - - // aux. file 4 (only if we're using a universal label) - if (universal_label != "") - { - std::ofstream universal_label_writer; - universal_label_writer.exceptions(std::ios::badbit | std::ios::failbit); - universal_label_writer.open(final_index_path_prefix + "_universal_label.txt"); - universal_label_writer << universal_label << std::endl; - universal_label_writer.close(); - } - - // main index - uint64_t index_num_frozen_points = 0, index_num_edges = 0; - uint32_t index_max_observed_degree = 0, index_entry_point = 0; - const size_t METADATA = 2 * sizeof(uint64_t) + 2 * sizeof(uint32_t); - for (auto &point_neighbors : stitched_graph) - { - index_max_observed_degree = std::max(index_max_observed_degree, (uint32_t)point_neighbors.size()); + tsl::robin_map entry_points, + std::string universal_label, path label_data_path) { + // aux. file 1 + auto saving_index_timer = std::chrono::high_resolution_clock::now(); + std::ifstream original_label_data_stream; + original_label_data_stream.exceptions(std::ios::badbit | std::ios::failbit); + original_label_data_stream.open(label_data_path, std::ios::binary); + std::ofstream new_label_data_stream; + new_label_data_stream.exceptions(std::ios::badbit | std::ios::failbit); + new_label_data_stream.open(final_index_path_prefix + "_labels.txt", + std::ios::binary); + new_label_data_stream << original_label_data_stream.rdbuf(); + original_label_data_stream.close(); + new_label_data_stream.close(); + + // aux. file 2 + std::ifstream original_input_data_stream; + original_input_data_stream.exceptions(std::ios::badbit | std::ios::failbit); + original_input_data_stream.open(input_data_path, std::ios::binary); + std::ofstream new_input_data_stream; + new_input_data_stream.exceptions(std::ios::badbit | std::ios::failbit); + new_input_data_stream.open(final_index_path_prefix + ".data", + std::ios::binary); + new_input_data_stream << original_input_data_stream.rdbuf(); + original_input_data_stream.close(); + new_input_data_stream.close(); + + // aux. file 3 + std::ofstream labels_to_medoids_writer; + labels_to_medoids_writer.exceptions(std::ios::badbit | std::ios::failbit); + labels_to_medoids_writer.open(final_index_path_prefix + + "_labels_to_medoids.txt"); + for (auto iter : entry_points) + labels_to_medoids_writer << iter.first << ", " << iter.second << std::endl; + labels_to_medoids_writer.close(); + + // aux. file 4 (only if we're using a universal label) + if (universal_label != "") { + std::ofstream universal_label_writer; + universal_label_writer.exceptions(std::ios::badbit | std::ios::failbit); + universal_label_writer.open(final_index_path_prefix + + "_universal_label.txt"); + universal_label_writer << universal_label << std::endl; + universal_label_writer.close(); + } + + // main index + uint64_t index_num_frozen_points = 0, index_num_edges = 0; + uint32_t index_max_observed_degree = 0, index_entry_point = 0; + const size_t METADATA = 2 * sizeof(uint64_t) + 2 * sizeof(uint32_t); + for (auto &point_neighbors : stitched_graph) { + index_max_observed_degree = + std::max(index_max_observed_degree, (uint32_t)point_neighbors.size()); + } + + std::ofstream stitched_graph_writer; + stitched_graph_writer.exceptions(std::ios::badbit | std::ios::failbit); + stitched_graph_writer.open(final_index_path_prefix, std::ios_base::binary); + + stitched_graph_writer.write((char *)&final_index_size, sizeof(uint64_t)); + stitched_graph_writer.write((char *)&index_max_observed_degree, + sizeof(uint32_t)); + stitched_graph_writer.write((char *)&index_entry_point, sizeof(uint32_t)); + stitched_graph_writer.write((char *)&index_num_frozen_points, + sizeof(uint64_t)); + + size_t bytes_written = METADATA; + for (uint32_t node_point = 0; node_point < stitched_graph.size(); + node_point++) { + uint32_t current_node_num_neighbors = stitched_graph[node_point].size(); + std::vector current_node_neighbors = stitched_graph[node_point]; + stitched_graph_writer.write((char *)¤t_node_num_neighbors, + sizeof(uint32_t)); + bytes_written += sizeof(uint32_t); + for (const auto ¤t_node_neighbor : current_node_neighbors) { + stitched_graph_writer.write((char *)¤t_node_neighbor, + sizeof(uint32_t)); + bytes_written += sizeof(uint32_t); } + index_num_edges += current_node_num_neighbors; + } - std::ofstream stitched_graph_writer; - stitched_graph_writer.exceptions(std::ios::badbit | std::ios::failbit); - stitched_graph_writer.open(final_index_path_prefix, std::ios_base::binary); - - stitched_graph_writer.write((char *)&final_index_size, sizeof(uint64_t)); - stitched_graph_writer.write((char *)&index_max_observed_degree, sizeof(uint32_t)); - stitched_graph_writer.write((char *)&index_entry_point, sizeof(uint32_t)); - stitched_graph_writer.write((char *)&index_num_frozen_points, sizeof(uint64_t)); - - size_t bytes_written = METADATA; - for (uint32_t node_point = 0; node_point < stitched_graph.size(); node_point++) - { - uint32_t current_node_num_neighbors = stitched_graph[node_point].size(); - std::vector current_node_neighbors = stitched_graph[node_point]; - stitched_graph_writer.write((char *)¤t_node_num_neighbors, sizeof(uint32_t)); - bytes_written += sizeof(uint32_t); - for (const auto ¤t_node_neighbor : current_node_neighbors) - { - stitched_graph_writer.write((char *)¤t_node_neighbor, sizeof(uint32_t)); - bytes_written += sizeof(uint32_t); - } - index_num_edges += current_node_num_neighbors; - } - - if (bytes_written != final_index_size) - { - std::cerr << "Error: written bytes does not match allocated space" << std::endl; - throw; - } - - stitched_graph_writer.close(); - - std::chrono::duration saving_index_time = std::chrono::high_resolution_clock::now() - saving_index_timer; - std::cout << "Stitched graph written in " << saving_index_time.count() << " seconds" << std::endl; - std::cout << "Stitched graph average degree: " << ((float)index_num_edges) / ((float)(stitched_graph.size())) + if (bytes_written != final_index_size) { + std::cerr << "Error: written bytes does not match allocated space" << std::endl; - std::cout << "Stitched graph max degree: " << index_max_observed_degree << std::endl << std::endl; + throw; + } + + stitched_graph_writer.close(); + + std::chrono::duration saving_index_time = + std::chrono::high_resolution_clock::now() - saving_index_timer; + std::cout << "Stitched graph written in " << saving_index_time.count() + << " seconds" << std::endl; + std::cout << "Stitched graph average degree: " + << ((float)index_num_edges) / ((float)(stitched_graph.size())) + << std::endl; + std::cout << "Stitched graph max degree: " << index_max_observed_degree + << std::endl + << std::endl; } /* @@ -208,52 +229,55 @@ void save_full_index(path final_index_path_prefix, path input_data_path, uint64_ */ template stitch_indices_return_values stitch_label_indices( - path final_index_path_prefix, uint32_t total_number_of_points, label_set all_labels, + path final_index_path_prefix, uint32_t total_number_of_points, + label_set all_labels, tsl::robin_map labels_to_number_of_points, tsl::robin_map &label_entry_points, - tsl::robin_map> label_id_to_orig_id_map) -{ - size_t final_index_size = 0; - std::vector> stitched_graph(total_number_of_points); - - auto stitching_index_timer = std::chrono::high_resolution_clock::now(); - for (const auto &lbl : all_labels) - { - path curr_label_index_path(final_index_path_prefix + "_" + lbl); - std::vector> curr_label_index; - uint64_t curr_label_index_size; - uint32_t curr_label_entry_point; - - std::tie(curr_label_index, curr_label_index_size) = - diskann::load_label_index(curr_label_index_path, labels_to_number_of_points[lbl]); - curr_label_entry_point = random(0, curr_label_index.size()); - label_entry_points[lbl] = label_id_to_orig_id_map[lbl][curr_label_entry_point]; - - for (uint32_t node_point = 0; node_point < curr_label_index.size(); node_point++) - { - uint32_t original_point_id = label_id_to_orig_id_map[lbl][node_point]; - for (auto &node_neighbor : curr_label_index[node_point]) - { - uint32_t original_neighbor_id = label_id_to_orig_id_map[lbl][node_neighbor]; - std::vector curr_point_neighbors = stitched_graph[original_point_id]; - if (std::find(curr_point_neighbors.begin(), curr_point_neighbors.end(), original_neighbor_id) == - curr_point_neighbors.end()) - { - stitched_graph[original_point_id].push_back(original_neighbor_id); - final_index_size += sizeof(uint32_t); - } - } + tsl::robin_map> + label_id_to_orig_id_map) { + size_t final_index_size = 0; + std::vector> stitched_graph(total_number_of_points); + + auto stitching_index_timer = std::chrono::high_resolution_clock::now(); + for (const auto &lbl : all_labels) { + path curr_label_index_path(final_index_path_prefix + "_" + lbl); + std::vector> curr_label_index; + uint64_t curr_label_index_size; + uint32_t curr_label_entry_point; + + std::tie(curr_label_index, curr_label_index_size) = + diskann::load_label_index(curr_label_index_path, + labels_to_number_of_points[lbl]); + curr_label_entry_point = random(0, curr_label_index.size()); + label_entry_points[lbl] = + label_id_to_orig_id_map[lbl][curr_label_entry_point]; + + for (uint32_t node_point = 0; node_point < curr_label_index.size(); + node_point++) { + uint32_t original_point_id = label_id_to_orig_id_map[lbl][node_point]; + for (auto &node_neighbor : curr_label_index[node_point]) { + uint32_t original_neighbor_id = + label_id_to_orig_id_map[lbl][node_neighbor]; + std::vector curr_point_neighbors = + stitched_graph[original_point_id]; + if (std::find(curr_point_neighbors.begin(), curr_point_neighbors.end(), + original_neighbor_id) == curr_point_neighbors.end()) { + stitched_graph[original_point_id].push_back(original_neighbor_id); + final_index_size += sizeof(uint32_t); } + } } + } - const size_t METADATA = 2 * sizeof(uint64_t) + 2 * sizeof(uint32_t); - final_index_size += (total_number_of_points * sizeof(uint32_t) + METADATA); + const size_t METADATA = 2 * sizeof(uint64_t) + 2 * sizeof(uint32_t); + final_index_size += (total_number_of_points * sizeof(uint32_t) + METADATA); - std::chrono::duration stitching_index_time = - std::chrono::high_resolution_clock::now() - stitching_index_timer; - std::cout << "stitched graph generated in memory in " << stitching_index_time.count() << " seconds" << std::endl; + std::chrono::duration stitching_index_time = + std::chrono::high_resolution_clock::now() - stitching_index_timer; + std::cout << "stitched graph generated in memory in " + << stitching_index_time.count() << " seconds" << std::endl; - return std::make_tuple(stitched_graph, final_index_size); + return std::make_tuple(stitched_graph, final_index_size); } /* @@ -264,31 +288,37 @@ stitch_indices_return_values stitch_label_indices( * and pruned graph. */ template -void prune_and_save(path final_index_path_prefix, path full_index_path_prefix, path input_data_path, - std::vector> stitched_graph, uint32_t stitched_R, - tsl::robin_map label_entry_points, std::string universal_label, - path label_data_path, uint32_t num_threads) -{ - size_t dimension, number_of_label_points; - auto diskann_cout_buffer = diskann::cout.rdbuf(nullptr); - auto std_cout_buffer = std::cout.rdbuf(nullptr); - auto pruning_index_timer = std::chrono::high_resolution_clock::now(); - - diskann::get_bin_metadata(input_data_path, number_of_label_points, dimension); - diskann::Index index(diskann::Metric::L2, dimension, number_of_label_points, false, false); - - // not searching this index, set search_l to 0 - index.load(full_index_path_prefix.c_str(), num_threads, 1); - - std::cout << "parsing labels" << std::endl; - - index.prune_all_neighbors(stitched_R, 750, 1.2); - index.save((final_index_path_prefix).c_str()); - - diskann::cout.rdbuf(diskann_cout_buffer); - std::cout.rdbuf(std_cout_buffer); - std::chrono::duration pruning_index_time = std::chrono::high_resolution_clock::now() - pruning_index_timer; - std::cout << "pruning performed in " << pruning_index_time.count() << " seconds\n" << std::endl; +void prune_and_save(path final_index_path_prefix, path full_index_path_prefix, + path input_data_path, + std::vector> stitched_graph, + uint32_t stitched_R, + tsl::robin_map label_entry_points, + std::string universal_label, path label_data_path, + uint32_t num_threads) { + size_t dimension, number_of_label_points; + auto diskann_cout_buffer = diskann::cout.rdbuf(nullptr); + auto std_cout_buffer = std::cout.rdbuf(nullptr); + auto pruning_index_timer = std::chrono::high_resolution_clock::now(); + + diskann::get_bin_metadata(input_data_path, number_of_label_points, dimension); + diskann::Index index(diskann::Metric::L2, dimension, + number_of_label_points, false, false); + + // not searching this index, set search_l to 0 + index.load(full_index_path_prefix.c_str(), num_threads, 1); + + std::cout << "parsing labels" << std::endl; + + index.prune_all_neighbors(stitched_R, 750, 1.2); + index.save((final_index_path_prefix).c_str()); + + diskann::cout.rdbuf(diskann_cout_buffer); + std::cout.rdbuf(std_cout_buffer); + std::chrono::duration pruning_index_time = + std::chrono::high_resolution_clock::now() - pruning_index_timer; + std::cout << "pruning performed in " << pruning_index_time.count() + << " seconds\n" + << std::endl; } /* @@ -299,131 +329,157 @@ void prune_and_save(path final_index_path_prefix, path full_index_path_prefix, p * 2. the separate diskANN indices built for each label * 3. the '.data' file created while generating the indices */ -void clean_up_artifacts(path input_data_path, path final_index_path_prefix, label_set all_labels) -{ - for (const auto &lbl : all_labels) - { - path curr_label_input_data_path(input_data_path + "_" + lbl); - path curr_label_index_path(final_index_path_prefix + "_" + lbl); - path curr_label_index_path_data(curr_label_index_path + ".data"); - - if (std::remove(curr_label_index_path.c_str()) != 0) - throw; - if (std::remove(curr_label_input_data_path.c_str()) != 0) - throw; - if (std::remove(curr_label_index_path_data.c_str()) != 0) - throw; - } +void clean_up_artifacts(path input_data_path, path final_index_path_prefix, + label_set all_labels) { + for (const auto &lbl : all_labels) { + path curr_label_input_data_path(input_data_path + "_" + lbl); + path curr_label_index_path(final_index_path_prefix + "_" + lbl); + path curr_label_index_path_data(curr_label_index_path + ".data"); + + if (std::remove(curr_label_index_path.c_str()) != 0) throw; + if (std::remove(curr_label_input_data_path.c_str()) != 0) throw; + if (std::remove(curr_label_index_path_data.c_str()) != 0) throw; + } } -int main(int argc, char **argv) -{ - // 1. handle cmdline inputs - std::string data_type; - path input_data_path, final_index_path_prefix, label_data_path; - std::string universal_label; - uint32_t num_threads, R, L, stitched_R; - float alpha; +int main(int argc, char **argv) { + // 1. handle cmdline inputs + std::string data_type; + path input_data_path, final_index_path_prefix, label_data_path; + std::string universal_label; + uint32_t num_threads, R, L, stitched_R; + float alpha; - auto index_timer = std::chrono::high_resolution_clock::now(); - handle_args(argc, argv, data_type, input_data_path, final_index_path_prefix, label_data_path, universal_label, - num_threads, R, L, stitched_R, alpha); + auto index_timer = std::chrono::high_resolution_clock::now(); + handle_args(argc, argv, data_type, input_data_path, final_index_path_prefix, + label_data_path, universal_label, num_threads, R, L, stitched_R, + alpha); - path labels_file_to_use = final_index_path_prefix + "_label_formatted.txt"; - path labels_map_file = final_index_path_prefix + "_labels_map.txt"; + path labels_file_to_use = final_index_path_prefix + "_label_formatted.txt"; + path labels_map_file = final_index_path_prefix + "_labels_map.txt"; - convert_labels_string_to_int(label_data_path, labels_file_to_use, labels_map_file, universal_label); + convert_labels_string_to_int(label_data_path, labels_file_to_use, + labels_map_file, universal_label); - // 2. parse label file and create necessary data structures - std::vector point_ids_to_labels; - tsl::robin_map labels_to_number_of_points; - label_set all_labels; + // 2. parse label file and create necessary data structures + std::vector point_ids_to_labels; + tsl::robin_map labels_to_number_of_points; + label_set all_labels; - std::tie(point_ids_to_labels, labels_to_number_of_points, all_labels) = - diskann::parse_label_file(labels_file_to_use, universal_label); + std::tie(point_ids_to_labels, labels_to_number_of_points, all_labels) = + diskann::parse_label_file(labels_file_to_use, universal_label); - // 3. for each label, make a separate data file - tsl::robin_map> label_id_to_orig_id_map; - uint32_t total_number_of_points = point_ids_to_labels.size(); + // 3. for each label, make a separate data file + tsl::robin_map> label_id_to_orig_id_map; + uint32_t total_number_of_points = point_ids_to_labels.size(); #ifndef _WINDOWS - if (data_type == "uint8") - label_id_to_orig_id_map = diskann::generate_label_specific_vector_files( - input_data_path, labels_to_number_of_points, point_ids_to_labels, all_labels); - else if (data_type == "int8") - label_id_to_orig_id_map = diskann::generate_label_specific_vector_files( - input_data_path, labels_to_number_of_points, point_ids_to_labels, all_labels); - else if (data_type == "float") - label_id_to_orig_id_map = diskann::generate_label_specific_vector_files( - input_data_path, labels_to_number_of_points, point_ids_to_labels, all_labels); - else - throw; + if (data_type == "uint8") + label_id_to_orig_id_map = + diskann::generate_label_specific_vector_files( + input_data_path, labels_to_number_of_points, point_ids_to_labels, + all_labels); + else if (data_type == "int8") + label_id_to_orig_id_map = + diskann::generate_label_specific_vector_files( + input_data_path, labels_to_number_of_points, point_ids_to_labels, + all_labels); + else if (data_type == "float") + label_id_to_orig_id_map = + diskann::generate_label_specific_vector_files( + input_data_path, labels_to_number_of_points, point_ids_to_labels, + all_labels); + else + throw; #else - if (data_type == "uint8") - label_id_to_orig_id_map = diskann::generate_label_specific_vector_files_compat( - input_data_path, labels_to_number_of_points, point_ids_to_labels, all_labels); - else if (data_type == "int8") - label_id_to_orig_id_map = diskann::generate_label_specific_vector_files_compat( - input_data_path, labels_to_number_of_points, point_ids_to_labels, all_labels); - else if (data_type == "float") - label_id_to_orig_id_map = diskann::generate_label_specific_vector_files_compat( - input_data_path, labels_to_number_of_points, point_ids_to_labels, all_labels); - else - throw; + if (data_type == "uint8") + label_id_to_orig_id_map = + diskann::generate_label_specific_vector_files_compat( + input_data_path, labels_to_number_of_points, point_ids_to_labels, + all_labels); + else if (data_type == "int8") + label_id_to_orig_id_map = + diskann::generate_label_specific_vector_files_compat( + input_data_path, labels_to_number_of_points, point_ids_to_labels, + all_labels); + else if (data_type == "float") + label_id_to_orig_id_map = + diskann::generate_label_specific_vector_files_compat( + input_data_path, labels_to_number_of_points, point_ids_to_labels, + all_labels); + else + throw; #endif - // 4. for each created data file, create a vanilla diskANN index - if (data_type == "uint8") - diskann::generate_label_indices(input_data_path, final_index_path_prefix, all_labels, R, L, alpha, - num_threads); - else if (data_type == "int8") - diskann::generate_label_indices(input_data_path, final_index_path_prefix, all_labels, R, L, alpha, - num_threads); - else if (data_type == "float") - diskann::generate_label_indices(input_data_path, final_index_path_prefix, all_labels, R, L, alpha, - num_threads); - else - throw; - - // 5. "stitch" the indices together - std::vector> stitched_graph; - tsl::robin_map label_entry_points; - uint64_t stitched_graph_size; - - if (data_type == "uint8") - std::tie(stitched_graph, stitched_graph_size) = - stitch_label_indices(final_index_path_prefix, total_number_of_points, all_labels, - labels_to_number_of_points, label_entry_points, label_id_to_orig_id_map); - else if (data_type == "int8") - std::tie(stitched_graph, stitched_graph_size) = - stitch_label_indices(final_index_path_prefix, total_number_of_points, all_labels, - labels_to_number_of_points, label_entry_points, label_id_to_orig_id_map); - else if (data_type == "float") - std::tie(stitched_graph, stitched_graph_size) = - stitch_label_indices(final_index_path_prefix, total_number_of_points, all_labels, - labels_to_number_of_points, label_entry_points, label_id_to_orig_id_map); - else - throw; - path full_index_path_prefix = final_index_path_prefix + "_full"; - // 5a. save the stitched graph to disk - save_full_index(full_index_path_prefix, input_data_path, stitched_graph_size, stitched_graph, label_entry_points, - universal_label, labels_file_to_use); - - // 6. run a prune on the stitched index, and save to disk - if (data_type == "uint8") - prune_and_save(final_index_path_prefix, full_index_path_prefix, input_data_path, stitched_graph, - stitched_R, label_entry_points, universal_label, labels_file_to_use, num_threads); - else if (data_type == "int8") - prune_and_save(final_index_path_prefix, full_index_path_prefix, input_data_path, stitched_graph, - stitched_R, label_entry_points, universal_label, labels_file_to_use, num_threads); - else if (data_type == "float") - prune_and_save(final_index_path_prefix, full_index_path_prefix, input_data_path, stitched_graph, - stitched_R, label_entry_points, universal_label, labels_file_to_use, num_threads); - else - throw; - - std::chrono::duration index_time = std::chrono::high_resolution_clock::now() - index_timer; - std::cout << "pruned/stitched graph generated in " << index_time.count() << " seconds" << std::endl; - - clean_up_artifacts(input_data_path, final_index_path_prefix, all_labels); + // 4. for each created data file, create a vanilla diskANN index + if (data_type == "uint8") + diskann::generate_label_indices( + input_data_path, final_index_path_prefix, all_labels, R, L, alpha, + num_threads); + else if (data_type == "int8") + diskann::generate_label_indices(input_data_path, + final_index_path_prefix, all_labels, + R, L, alpha, num_threads); + else if (data_type == "float") + diskann::generate_label_indices(input_data_path, + final_index_path_prefix, all_labels, + R, L, alpha, num_threads); + else + throw; + + // 5. "stitch" the indices together + std::vector> stitched_graph; + tsl::robin_map label_entry_points; + uint64_t stitched_graph_size; + + if (data_type == "uint8") + std::tie(stitched_graph, stitched_graph_size) = + stitch_label_indices( + final_index_path_prefix, total_number_of_points, all_labels, + labels_to_number_of_points, label_entry_points, + label_id_to_orig_id_map); + else if (data_type == "int8") + std::tie(stitched_graph, stitched_graph_size) = + stitch_label_indices( + final_index_path_prefix, total_number_of_points, all_labels, + labels_to_number_of_points, label_entry_points, + label_id_to_orig_id_map); + else if (data_type == "float") + std::tie(stitched_graph, stitched_graph_size) = stitch_label_indices( + final_index_path_prefix, total_number_of_points, all_labels, + labels_to_number_of_points, label_entry_points, + label_id_to_orig_id_map); + else + throw; + path full_index_path_prefix = final_index_path_prefix + "_full"; + // 5a. save the stitched graph to disk + save_full_index(full_index_path_prefix, input_data_path, stitched_graph_size, + stitched_graph, label_entry_points, universal_label, + labels_file_to_use); + + // 6. run a prune on the stitched index, and save to disk + if (data_type == "uint8") + prune_and_save(final_index_path_prefix, full_index_path_prefix, + input_data_path, stitched_graph, stitched_R, + label_entry_points, universal_label, + labels_file_to_use, num_threads); + else if (data_type == "int8") + prune_and_save(final_index_path_prefix, full_index_path_prefix, + input_data_path, stitched_graph, stitched_R, + label_entry_points, universal_label, + labels_file_to_use, num_threads); + else if (data_type == "float") + prune_and_save(final_index_path_prefix, full_index_path_prefix, + input_data_path, stitched_graph, stitched_R, + label_entry_points, universal_label, + labels_file_to_use, num_threads); + else + throw; + + std::chrono::duration index_time = + std::chrono::high_resolution_clock::now() - index_timer; + std::cout << "pruned/stitched graph generated in " << index_time.count() + << " seconds" << std::endl; + + clean_up_artifacts(input_data_path, final_index_path_prefix, all_labels); } diff --git a/tests/range_search_disk_index.cpp b/tests/range_search_disk_index.cpp index a67dac378..671925d9c 100644 --- a/tests/range_search_disk_index.cpp +++ b/tests/range_search_disk_index.cpp @@ -33,332 +33,334 @@ namespace po = boost::program_options; #define WARMUP false -void print_stats(std::string category, std::vector percentiles, std::vector results) -{ - diskann::cout << std::setw(20) << category << ": " << std::flush; - for (uint32_t s = 0; s < percentiles.size(); s++) - { - diskann::cout << std::setw(8) << percentiles[s] << "%"; - } - diskann::cout << std::endl; - diskann::cout << std::setw(22) << " " << std::flush; - for (uint32_t s = 0; s < percentiles.size(); s++) - { - diskann::cout << std::setw(9) << results[s]; - } - diskann::cout << std::endl; +void print_stats(std::string category, std::vector percentiles, + std::vector results) { + diskann::cout << std::setw(20) << category << ": " << std::flush; + for (uint32_t s = 0; s < percentiles.size(); s++) { + diskann::cout << std::setw(8) << percentiles[s] << "%"; + } + diskann::cout << std::endl; + diskann::cout << std::setw(22) << " " << std::flush; + for (uint32_t s = 0; s < percentiles.size(); s++) { + diskann::cout << std::setw(9) << results[s]; + } + diskann::cout << std::endl; } template -int search_disk_index(diskann::Metric &metric, const std::string &index_path_prefix, const std::string &query_file, - std::string >_file, const uint32_t num_threads, const float search_range, - const uint32_t beamwidth, const uint32_t num_nodes_to_cache, const std::vector &Lvec) -{ - std::string pq_prefix = index_path_prefix + "_pq"; - std::string disk_index_file = index_path_prefix + "_disk.index"; - std::string warmup_query_file = index_path_prefix + "_sample_data.bin"; - - diskann::cout << "Search parameters: #threads: " << num_threads << ", "; - if (beamwidth <= 0) - diskann::cout << "beamwidth to be optimized for each L value" << std::endl; - else - diskann::cout << " beamwidth: " << beamwidth << std::endl; - - // load query bin - T *query = nullptr; - std::vector> groundtruth_ids; - size_t query_num, query_dim, query_aligned_dim, gt_num; - diskann::load_aligned_bin(query_file, query, query_num, query_dim, query_aligned_dim); - - bool calc_recall_flag = false; - if (gt_file != std::string("null") && file_exists(gt_file)) - { - diskann::load_range_truthset(gt_file, groundtruth_ids, - gt_num); // use for range search type of truthset - // diskann::prune_truthset_for_range(gt_file, search_range, - // groundtruth_ids, gt_num); // use for traditional truthset - if (gt_num != query_num) - { - diskann::cout << "Error. Mismatch in number of queries and ground truth data" << std::endl; - return -1; - } - calc_recall_flag = true; +int search_disk_index(diskann::Metric &metric, + const std::string &index_path_prefix, + const std::string &query_file, std::string >_file, + const uint32_t num_threads, const float search_range, + const uint32_t beamwidth, + const uint32_t num_nodes_to_cache, + const std::vector &Lvec) { + std::string pq_prefix = index_path_prefix + "_pq"; + std::string disk_index_file = index_path_prefix + "_disk.index"; + std::string warmup_query_file = index_path_prefix + "_sample_data.bin"; + + diskann::cout << "Search parameters: #threads: " << num_threads << ", "; + if (beamwidth <= 0) + diskann::cout << "beamwidth to be optimized for each L value" << std::endl; + else + diskann::cout << " beamwidth: " << beamwidth << std::endl; + + // load query bin + T *query = nullptr; + std::vector> groundtruth_ids; + size_t query_num, query_dim, query_aligned_dim, gt_num; + diskann::load_aligned_bin(query_file, query, query_num, query_dim, + query_aligned_dim); + + bool calc_recall_flag = false; + if (gt_file != std::string("null") && file_exists(gt_file)) { + diskann::load_range_truthset( + gt_file, groundtruth_ids, + gt_num); // use for range search type of truthset + // diskann::prune_truthset_for_range(gt_file, search_range, + // groundtruth_ids, gt_num); // use for traditional truthset + if (gt_num != query_num) { + diskann::cout + << "Error. Mismatch in number of queries and ground truth data" + << std::endl; + return -1; } + calc_recall_flag = true; + } - std::shared_ptr reader = nullptr; + std::shared_ptr reader = nullptr; #ifdef _WINDOWS #ifndef USE_BING_INFRA - reader.reset(new WindowsAlignedFileReader()); + reader.reset(new WindowsAlignedFileReader()); #else - reader.reset(new diskann::BingAlignedFileReader()); + reader.reset(new diskann::BingAlignedFileReader()); #endif #else - reader.reset(new LinuxAlignedFileReader()); + reader.reset(new LinuxAlignedFileReader()); #endif - std::unique_ptr> _pFlashIndex( - new diskann::PQFlashIndex(reader, metric)); - - int res = _pFlashIndex->load(num_threads, index_path_prefix.c_str()); - - if (res != 0) - { - return res; - } - // 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->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_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, warmup_dim, warmup_aligned_dim); - } - else - { - warmup_num = (std::min)((uint32_t)150000, (uint32_t)15000 * num_threads); - warmup_dim = query_dim; - warmup_aligned_dim = query_aligned_dim; - diskann::alloc_aligned(((void **)&warmup), warmup_num * warmup_aligned_dim * sizeof(T), 8 * sizeof(T)); - std::memset(warmup, 0, warmup_num * warmup_aligned_dim * sizeof(T)); - std::random_device rd; - std::mt19937 gen(rd()); - std::uniform_int_distribution<> dis(-128, 127); - for (uint32_t i = 0; i < warmup_num; i++) - { - for (uint32_t d = 0; d < warmup_dim; d++) - { - warmup[i * warmup_aligned_dim + d] = (T)dis(gen); - } - } + std::unique_ptr> _pFlashIndex( + new diskann::PQFlashIndex(reader, metric)); + + int res = _pFlashIndex->load(num_threads, index_path_prefix.c_str()); + + if (res != 0) { + return res; + } + // 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->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_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, + warmup_dim, warmup_aligned_dim); + } else { + warmup_num = (std::min)((uint32_t)150000, (uint32_t)15000 * num_threads); + warmup_dim = query_dim; + warmup_aligned_dim = query_aligned_dim; + diskann::alloc_aligned(((void **)&warmup), + warmup_num * warmup_aligned_dim * sizeof(T), + 8 * sizeof(T)); + std::memset(warmup, 0, warmup_num * warmup_aligned_dim * sizeof(T)); + std::random_device rd; + std::mt19937 gen(rd()); + std::uniform_int_distribution<> dis(-128, 127); + for (uint32_t i = 0; i < warmup_num; i++) { + for (uint32_t d = 0; d < warmup_dim; d++) { + warmup[i * warmup_aligned_dim + d] = (T)dis(gen); } - diskann::cout << "Warming up index... " << std::flush; - std::vector warmup_result_ids_64(warmup_num, 0); - std::vector warmup_result_dists(warmup_num, 0); - -#pragma omp parallel for schedule(dynamic, 1) - for (int64_t 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); - } - diskann::cout << "..done" << std::endl; + } } + diskann::cout << "Warming up index... " << std::flush; + std::vector warmup_result_ids_64(warmup_num, 0); + std::vector warmup_result_dists(warmup_num, 0); - diskann::cout.setf(std::ios_base::fixed, std::ios_base::floatfield); - diskann::cout.precision(2); - - std::string recall_string = "Recall@rng=" + std::to_string(search_range); - diskann::cout << std::setw(6) << "L" << std::setw(12) << "Beamwidth" << std::setw(16) << "QPS" << std::setw(16) - << "Mean Latency" << std::setw(16) << "99.9 Latency" << std::setw(16) << "Mean IOs" << std::setw(16) - << "CPU (s)"; - if (calc_recall_flag) - { - diskann::cout << std::setw(16) << recall_string << std::endl; +#pragma omp parallel for schedule(dynamic, 1) + for (int64_t 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); } - else - diskann::cout << std::endl; - diskann::cout << "===============================================================" - "===========================================" - << std::endl; + diskann::cout << "..done" << std::endl; + } + + diskann::cout.setf(std::ios_base::fixed, std::ios_base::floatfield); + diskann::cout.precision(2); + + std::string recall_string = "Recall@rng=" + std::to_string(search_range); + diskann::cout << std::setw(6) << "L" << std::setw(12) << "Beamwidth" + << std::setw(16) << "QPS" << std::setw(16) << "Mean Latency" + << std::setw(16) << "99.9 Latency" << std::setw(16) + << "Mean IOs" << std::setw(16) << "CPU (s)"; + if (calc_recall_flag) { + diskann::cout << std::setw(16) << recall_string << std::endl; + } else + diskann::cout << std::endl; + diskann::cout + << "===============================================================" + "===========================================" + << std::endl; - std::vector>> query_result_ids(Lvec.size()); + std::vector>> query_result_ids(Lvec.size()); - uint32_t optimized_beamwidth = 2; - uint32_t max_list_size = 10000; + uint32_t optimized_beamwidth = 2; + uint32_t max_list_size = 10000; - for (uint32_t test_id = 0; test_id < Lvec.size(); test_id++) - { - uint64_t L = Lvec[test_id]; + for (uint32_t test_id = 0; test_id < Lvec.size(); test_id++) { + uint64_t L = Lvec[test_id]; - if (beamwidth <= 0) - { - optimized_beamwidth = - optimize_beamwidth(_pFlashIndex, warmup, warmup_num, warmup_aligned_dim, L, optimized_beamwidth); - } - else - optimized_beamwidth = beamwidth; + if (beamwidth <= 0) { + optimized_beamwidth = + optimize_beamwidth(_pFlashIndex, warmup, warmup_num, + warmup_aligned_dim, L, optimized_beamwidth); + } else + optimized_beamwidth = beamwidth; - query_result_ids[test_id].clear(); - query_result_ids[test_id].resize(query_num); + query_result_ids[test_id].clear(); + query_result_ids[test_id].resize(query_num); - diskann::QueryStats *stats = new diskann::QueryStats[query_num]; + diskann::QueryStats *stats = new diskann::QueryStats[query_num]; - auto s = std::chrono::high_resolution_clock::now(); + auto s = std::chrono::high_resolution_clock::now(); #pragma omp parallel for schedule(dynamic, 1) - for (int64_t i = 0; i < (int64_t)query_num; i++) - { - std::vector indices; - std::vector distances; - uint32_t res_count = - _pFlashIndex->range_search(query + (i * query_aligned_dim), search_range, L, max_list_size, indices, - distances, optimized_beamwidth, stats + i); - query_result_ids[test_id][i].reserve(res_count); - query_result_ids[test_id][i].resize(res_count); - for (uint32_t idx = 0; idx < res_count; idx++) - query_result_ids[test_id][i][idx] = indices[idx]; - } - auto e = std::chrono::high_resolution_clock::now(); - std::chrono::duration diff = e - s; - auto qps = (1.0 * query_num) / (1.0 * diff.count()); - - auto mean_latency = diskann::get_mean_stats( - stats, query_num, [](const diskann::QueryStats &stats) { return stats.total_us; }); - - auto latency_999 = diskann::get_percentile_stats( - stats, query_num, 0.999, [](const diskann::QueryStats &stats) { return stats.total_us; }); - - auto mean_ios = diskann::get_mean_stats(stats, query_num, - [](const diskann::QueryStats &stats) { return stats.n_ios; }); - - float mean_cpuus = diskann::get_mean_stats( - stats, query_num, [](const diskann::QueryStats &stats) { return stats.cpu_us; }); - - float recall = 0; - float ratio_of_sums = 0; - if (calc_recall_flag) - { - recall = diskann::calculate_range_search_recall(query_num, groundtruth_ids, query_result_ids[test_id]); - - uint32_t total_true_positive = 0; - uint32_t total_positive = 0; - for (uint32_t i = 0; i < query_num; i++) - { - total_true_positive += query_result_ids[test_id][i].size(); - total_positive += groundtruth_ids[i].size(); - } - - ratio_of_sums = (1.0 * total_true_positive) / (1.0 * total_positive); - } - - diskann::cout << std::setw(6) << L << std::setw(12) << optimized_beamwidth << std::setw(16) << qps - << std::setw(16) << mean_latency << std::setw(16) << latency_999 << std::setw(16) << mean_ios - << std::setw(16) << mean_cpuus; - if (calc_recall_flag) - { - diskann::cout << std::setw(16) << recall << "," << ratio_of_sums << std::endl; - } - else - diskann::cout << std::endl; - } - - diskann::cout << "Done searching. " << std::endl; - - diskann::aligned_free(query); - if (warmup != nullptr) - diskann::aligned_free(warmup); - return 0; -} - -int main(int argc, char **argv) -{ - std::string data_type, dist_fn, index_path_prefix, result_path_prefix, query_file, gt_file; - uint32_t num_threads, W, num_nodes_to_cache; - std::vector Lvec; - float range; - - po::options_description desc{"Arguments"}; - try - { - desc.add_options()("help,h", "Print information on arguments"); - desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); - desc.add_options()("dist_fn", po::value(&dist_fn)->required(), - "distance function "); - desc.add_options()("index_path_prefix", po::value(&index_path_prefix)->required(), - "Path prefix to the index"); - desc.add_options()("query_file", po::value(&query_file)->required(), - "Query file in binary format"); - desc.add_options()("gt_file", po::value(>_file)->default_value(std::string("null")), - "ground truth file for the queryset"); - desc.add_options()("range_threshold,K", po::value(&range)->required(), - "Number of neighbors to be returned"); - desc.add_options()("search_list,L", po::value>(&Lvec)->multitoken(), - "List of L values of search"); - desc.add_options()("beamwidth,W", po::value(&W)->default_value(2), "Beamwidth for search"); - desc.add_options()("num_nodes_to_cache", po::value(&num_nodes_to_cache)->default_value(100000), - "Beamwidth for search"); - desc.add_options()("num_threads,T", po::value(&num_threads)->default_value(omp_get_num_procs()), - "Number of threads used for building index (defaults to " - "omp_get_num_procs())"); - - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) - { - std::cout << desc; - return 0; - } - po::notify(vm); - } - catch (const std::exception &ex) - { - std::cerr << ex.what() << '\n'; - return -1; - } - - diskann::Metric metric; - if (dist_fn == std::string("mips")) - { - metric = diskann::Metric::INNER_PRODUCT; + for (int64_t i = 0; i < (int64_t)query_num; i++) { + std::vector indices; + std::vector distances; + uint32_t res_count = _pFlashIndex->range_search( + query + (i * query_aligned_dim), search_range, L, max_list_size, + indices, distances, optimized_beamwidth, stats + i); + query_result_ids[test_id][i].reserve(res_count); + query_result_ids[test_id][i].resize(res_count); + for (uint32_t idx = 0; idx < res_count; idx++) + query_result_ids[test_id][i][idx] = indices[idx]; } - else if (dist_fn == std::string("l2")) - { - metric = diskann::Metric::L2; - } - else if (dist_fn == std::string("cosine")) - { - metric = diskann::Metric::COSINE; - } - else - { - std::cout << "Unsupported distance function. Currently only L2/ Inner " - "Product/Cosine are supported." - << std::endl; - return -1; + auto e = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = e - s; + auto qps = (1.0 * query_num) / (1.0 * diff.count()); + + auto mean_latency = diskann::get_mean_stats( + stats, query_num, + [](const diskann::QueryStats &stats) { return stats.total_us; }); + + auto latency_999 = diskann::get_percentile_stats( + stats, query_num, 0.999, + [](const diskann::QueryStats &stats) { return stats.total_us; }); + + auto mean_ios = diskann::get_mean_stats( + stats, query_num, + [](const diskann::QueryStats &stats) { return stats.n_ios; }); + + float mean_cpuus = diskann::get_mean_stats( + stats, query_num, + [](const diskann::QueryStats &stats) { return stats.cpu_us; }); + + float recall = 0; + float ratio_of_sums = 0; + if (calc_recall_flag) { + recall = diskann::calculate_range_search_recall( + query_num, groundtruth_ids, query_result_ids[test_id]); + + uint32_t total_true_positive = 0; + uint32_t total_positive = 0; + for (uint32_t i = 0; i < query_num; i++) { + total_true_positive += query_result_ids[test_id][i].size(); + total_positive += groundtruth_ids[i].size(); + } + + ratio_of_sums = (1.0 * total_true_positive) / (1.0 * total_positive); } - if ((data_type != std::string("float")) && (metric == diskann::Metric::INNER_PRODUCT)) - { - std::cout << "Currently support only floating point data for Inner Product." << std::endl; - return -1; - } + diskann::cout << std::setw(6) << L << std::setw(12) << optimized_beamwidth + << std::setw(16) << qps << std::setw(16) << mean_latency + << std::setw(16) << latency_999 << std::setw(16) << mean_ios + << std::setw(16) << mean_cpuus; + if (calc_recall_flag) { + diskann::cout << std::setw(16) << recall << "," << ratio_of_sums + << std::endl; + } else + diskann::cout << std::endl; + } + + diskann::cout << "Done searching. " << std::endl; + + diskann::aligned_free(query); + if (warmup != nullptr) diskann::aligned_free(warmup); + return 0; +} - try - { - if (data_type == std::string("float")) - return search_disk_index(metric, index_path_prefix, query_file, gt_file, num_threads, range, W, - num_nodes_to_cache, Lvec); - else if (data_type == std::string("int8")) - return search_disk_index(metric, index_path_prefix, query_file, gt_file, num_threads, range, W, - num_nodes_to_cache, Lvec); - else if (data_type == std::string("uint8")) - return search_disk_index(metric, index_path_prefix, query_file, gt_file, num_threads, range, W, - num_nodes_to_cache, Lvec); - else - { - std::cerr << "Unsupported data type. Use float or int8 or uint8" << std::endl; - return -1; - } +int main(int argc, char **argv) { + std::string data_type, dist_fn, index_path_prefix, result_path_prefix, + query_file, gt_file; + uint32_t num_threads, W, num_nodes_to_cache; + std::vector Lvec; + float range; + + po::options_description desc{"Arguments"}; + try { + desc.add_options()("help,h", "Print information on arguments"); + desc.add_options()("data_type", + po::value(&data_type)->required(), + "data type "); + desc.add_options()("dist_fn", po::value(&dist_fn)->required(), + "distance function "); + desc.add_options()("index_path_prefix", + po::value(&index_path_prefix)->required(), + "Path prefix to the index"); + desc.add_options()("query_file", + po::value(&query_file)->required(), + "Query file in binary format"); + desc.add_options()( + "gt_file", + po::value(>_file)->default_value(std::string("null")), + "ground truth file for the queryset"); + desc.add_options()("range_threshold,K", + po::value(&range)->required(), + "Number of neighbors to be returned"); + desc.add_options()("search_list,L", + po::value>(&Lvec)->multitoken(), + "List of L values of search"); + desc.add_options()("beamwidth,W", po::value(&W)->default_value(2), + "Beamwidth for search"); + desc.add_options()( + "num_nodes_to_cache", + po::value(&num_nodes_to_cache)->default_value(100000), + "Beamwidth for search"); + desc.add_options()( + "num_threads,T", + po::value(&num_threads)->default_value(omp_get_num_procs()), + "Number of threads used for building index (defaults to " + "omp_get_num_procs())"); + + po::variables_map vm; + po::store(po::parse_command_line(argc, argv, desc), vm); + if (vm.count("help")) { + std::cout << desc; + return 0; } - catch (const std::exception &e) - { - std::cout << std::string(e.what()) << std::endl; - diskann::cerr << "Index search failed." << std::endl; - return -1; + po::notify(vm); + } catch (const std::exception &ex) { + std::cerr << ex.what() << '\n'; + return -1; + } + + diskann::Metric metric; + if (dist_fn == std::string("mips")) { + metric = diskann::Metric::INNER_PRODUCT; + } else if (dist_fn == std::string("l2")) { + metric = diskann::Metric::L2; + } else if (dist_fn == std::string("cosine")) { + metric = diskann::Metric::COSINE; + } else { + std::cout << "Unsupported distance function. Currently only L2/ Inner " + "Product/Cosine are supported." + << std::endl; + return -1; + } + + if ((data_type != std::string("float")) && + (metric == diskann::Metric::INNER_PRODUCT)) { + std::cout << "Currently support only floating point data for Inner Product." + << std::endl; + return -1; + } + + try { + if (data_type == std::string("float")) + return search_disk_index(metric, index_path_prefix, query_file, + gt_file, num_threads, range, W, + num_nodes_to_cache, Lvec); + else if (data_type == std::string("int8")) + return search_disk_index(metric, index_path_prefix, query_file, + gt_file, num_threads, range, W, + num_nodes_to_cache, Lvec); + else if (data_type == std::string("uint8")) + return search_disk_index(metric, index_path_prefix, query_file, + gt_file, num_threads, range, W, + num_nodes_to_cache, Lvec); + else { + std::cerr << "Unsupported data type. Use float or int8 or uint8" + << std::endl; + return -1; } + } catch (const std::exception &e) { + std::cout << std::string(e.what()) << std::endl; + diskann::cerr << "Index search failed." << std::endl; + return -1; + } } diff --git a/tests/search_disk_index.cpp b/tests/search_disk_index.cpp index a2528c92b..e4a87ac7d 100644 --- a/tests/search_disk_index.cpp +++ b/tests/search_disk_index.cpp @@ -30,450 +30,460 @@ namespace po = boost::program_options; -void print_stats(std::string category, std::vector percentiles, std::vector results) -{ - diskann::cout << std::setw(20) << category << ": " << std::flush; - for (uint32_t s = 0; s < percentiles.size(); s++) - { - diskann::cout << std::setw(8) << percentiles[s] << "%"; - } - diskann::cout << std::endl; - diskann::cout << std::setw(22) << " " << std::flush; - for (uint32_t s = 0; s < percentiles.size(); s++) - { - diskann::cout << std::setw(9) << results[s]; - } - diskann::cout << std::endl; +void print_stats(std::string category, std::vector percentiles, + std::vector results) { + diskann::cout << std::setw(20) << category << ": " << std::flush; + for (uint32_t s = 0; s < percentiles.size(); s++) { + diskann::cout << std::setw(8) << percentiles[s] << "%"; + } + diskann::cout << std::endl; + diskann::cout << std::setw(22) << " " << std::flush; + for (uint32_t s = 0; s < percentiles.size(); s++) { + diskann::cout << std::setw(9) << results[s]; + } + diskann::cout << std::endl; } template -int search_disk_index(diskann::Metric &metric, const std::string &index_path_prefix, - const std::string &result_output_prefix, const std::string &query_file, std::string >_file, - const uint32_t num_threads, const uint32_t recall_at, const uint32_t beamwidth, - const uint32_t num_nodes_to_cache, const uint32_t search_io_limit, - const std::vector &Lvec, const float fail_if_recall_below, - const std::vector &query_filters, const bool use_reorder_data = false) -{ - diskann::cout << "Search parameters: #threads: " << num_threads << ", "; - if (beamwidth <= 0) - diskann::cout << "beamwidth to be optimized for each L value" << std::flush; - else - diskann::cout << " beamwidth: " << beamwidth << std::flush; - if (search_io_limit == std::numeric_limits::max()) - diskann::cout << "." << std::endl; - else - diskann::cout << ", io_limit: " << search_io_limit << "." << std::endl; - - std::string warmup_query_file = index_path_prefix + "_sample_data.bin"; - - // load query bin - T *query = nullptr; - uint32_t *gt_ids = nullptr; - float *gt_dists = nullptr; - size_t query_num, query_dim, query_aligned_dim, gt_num, gt_dim; - diskann::load_aligned_bin(query_file, query, query_num, query_dim, query_aligned_dim); - - bool filtered_search = false; - if (!query_filters.empty()) - { - filtered_search = true; - if (query_filters.size() != 1 && query_filters.size() != query_num) - { - std::cout << "Error. Mismatch in number of queries and size of query filters file" << std::endl; - return -1; // To return -1 or some other error handling? - } +int search_disk_index( + diskann::Metric &metric, const std::string &index_path_prefix, + const std::string &result_output_prefix, const std::string &query_file, + std::string >_file, const uint32_t num_threads, const uint32_t recall_at, + const uint32_t beamwidth, const uint32_t num_nodes_to_cache, + const uint32_t search_io_limit, const std::vector &Lvec, + const float fail_if_recall_below, + const std::vector &query_filters, + const bool use_reorder_data = false) { + diskann::cout << "Search parameters: #threads: " << num_threads << ", "; + if (beamwidth <= 0) + diskann::cout << "beamwidth to be optimized for each L value" << std::flush; + else + diskann::cout << " beamwidth: " << beamwidth << std::flush; + if (search_io_limit == std::numeric_limits::max()) + diskann::cout << "." << std::endl; + else + diskann::cout << ", io_limit: " << search_io_limit << "." << std::endl; + + std::string warmup_query_file = index_path_prefix + "_sample_data.bin"; + + // load query bin + T *query = nullptr; + uint32_t *gt_ids = nullptr; + float *gt_dists = nullptr; + size_t query_num, query_dim, query_aligned_dim, gt_num, gt_dim; + diskann::load_aligned_bin(query_file, query, query_num, query_dim, + query_aligned_dim); + + bool filtered_search = false; + if (!query_filters.empty()) { + filtered_search = true; + if (query_filters.size() != 1 && query_filters.size() != query_num) { + std::cout << "Error. Mismatch in number of queries and size of query " + "filters file" + << std::endl; + return -1; // To return -1 or some other error handling? } - - bool calc_recall_flag = false; - if (gt_file != std::string("null") && gt_file != std::string("NULL") && file_exists(gt_file)) - { - diskann::load_truthset(gt_file, gt_ids, gt_dists, gt_num, gt_dim); - if (gt_num != query_num) - { - diskann::cout << "Error. Mismatch in number of queries and ground truth data" << std::endl; - } - calc_recall_flag = true; + } + + bool calc_recall_flag = false; + if (gt_file != std::string("null") && gt_file != std::string("NULL") && + file_exists(gt_file)) { + diskann::load_truthset(gt_file, gt_ids, gt_dists, gt_num, gt_dim); + if (gt_num != query_num) { + diskann::cout + << "Error. Mismatch in number of queries and ground truth data" + << std::endl; } + calc_recall_flag = true; + } - std::shared_ptr reader = nullptr; + std::shared_ptr reader = nullptr; #ifdef _WINDOWS #ifndef USE_BING_INFRA - reader.reset(new WindowsAlignedFileReader()); + reader.reset(new WindowsAlignedFileReader()); #else - reader.reset(new diskann::BingAlignedFileReader()); + reader.reset(new diskann::BingAlignedFileReader()); #endif #else - reader.reset(new LinuxAlignedFileReader()); + reader.reset(new LinuxAlignedFileReader()); #endif - std::unique_ptr> _pFlashIndex( - new diskann::PQFlashIndex(reader, metric)); - - int res = _pFlashIndex->load(num_threads, index_path_prefix.c_str()); - - if (res != 0) - { - return res; - } - // 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); - if (num_nodes_to_cache > 0) - _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_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, warmup_dim, warmup_aligned_dim); + std::unique_ptr> _pFlashIndex( + new diskann::PQFlashIndex(reader, metric)); + + int res = _pFlashIndex->load(num_threads, index_path_prefix.c_str()); + + if (res != 0) { + return res; + } + // 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); + if (num_nodes_to_cache > 0) + _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_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, + warmup_dim, warmup_aligned_dim); + } else { + warmup_num = (std::min)((uint32_t)150000, (uint32_t)15000 * num_threads); + warmup_dim = query_dim; + warmup_aligned_dim = query_aligned_dim; + diskann::alloc_aligned(((void **)&warmup), + warmup_num * warmup_aligned_dim * sizeof(T), + 8 * sizeof(T)); + std::memset(warmup, 0, warmup_num * warmup_aligned_dim * sizeof(T)); + std::random_device rd; + std::mt19937 gen(rd()); + std::uniform_int_distribution<> dis(-128, 127); + for (uint32_t i = 0; i < warmup_num; i++) { + for (uint32_t d = 0; d < warmup_dim; d++) { + warmup[i * warmup_aligned_dim + d] = (T)dis(gen); } - else - { - warmup_num = (std::min)((uint32_t)150000, (uint32_t)15000 * num_threads); - warmup_dim = query_dim; - warmup_aligned_dim = query_aligned_dim; - diskann::alloc_aligned(((void **)&warmup), warmup_num * warmup_aligned_dim * sizeof(T), 8 * sizeof(T)); - std::memset(warmup, 0, warmup_num * warmup_aligned_dim * sizeof(T)); - std::random_device rd; - std::mt19937 gen(rd()); - std::uniform_int_distribution<> dis(-128, 127); - for (uint32_t i = 0; i < warmup_num; i++) - { - for (uint32_t d = 0; d < warmup_dim; d++) - { - warmup[i * warmup_aligned_dim + d] = (T)dis(gen); - } - } - } - diskann::cout << "Warming up index... " << std::flush; - std::vector warmup_result_ids_64(warmup_num, 0); - std::vector warmup_result_dists(warmup_num, 0); - -#pragma omp parallel for schedule(dynamic, 1) - for (int64_t 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); - } - diskann::cout << "..done" << std::endl; + } } + diskann::cout << "Warming up index... " << std::flush; + std::vector warmup_result_ids_64(warmup_num, 0); + std::vector warmup_result_dists(warmup_num, 0); - diskann::cout.setf(std::ios_base::fixed, std::ios_base::floatfield); - diskann::cout.precision(2); - - std::string recall_string = "Recall@" + std::to_string(recall_at); - diskann::cout << std::setw(6) << "L" << std::setw(12) << "Beamwidth" << std::setw(16) << "QPS" << std::setw(16) - << "Mean Latency" << std::setw(16) << "99.9 Latency" << std::setw(16) << "Mean IOs" << std::setw(16) - << "CPU (s)"; - if (calc_recall_flag) - { - diskann::cout << std::setw(16) << recall_string << std::endl; +#pragma omp parallel for schedule(dynamic, 1) + for (int64_t 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); } - else - diskann::cout << std::endl; - diskann::cout << "===============================================================" - "=======================================================" - << std::endl; + diskann::cout << "..done" << std::endl; + } + + diskann::cout.setf(std::ios_base::fixed, std::ios_base::floatfield); + diskann::cout.precision(2); + + std::string recall_string = "Recall@" + std::to_string(recall_at); + diskann::cout << std::setw(6) << "L" << std::setw(12) << "Beamwidth" + << std::setw(16) << "QPS" << std::setw(16) << "Mean Latency" + << std::setw(16) << "99.9 Latency" << std::setw(16) + << "Mean IOs" << std::setw(16) << "CPU (s)"; + if (calc_recall_flag) { + diskann::cout << std::setw(16) << recall_string << std::endl; + } else + diskann::cout << std::endl; + diskann::cout + << "===============================================================" + "=======================================================" + << std::endl; - std::vector> query_result_ids(Lvec.size()); - std::vector> query_result_dists(Lvec.size()); + std::vector> query_result_ids(Lvec.size()); + std::vector> query_result_dists(Lvec.size()); - uint32_t optimized_beamwidth = 2; + uint32_t optimized_beamwidth = 2; - float best_recall = 0.0; + float best_recall = 0.0; - for (uint32_t test_id = 0; test_id < Lvec.size(); test_id++) - { - uint64_t L = Lvec[test_id]; + for (uint32_t test_id = 0; test_id < Lvec.size(); test_id++) { + uint64_t L = Lvec[test_id]; - if (L < recall_at) - { - diskann::cout << "Ignoring search with L:" << L << " since it's smaller than K:" << recall_at << std::endl; - continue; - } + if (L < recall_at) { + diskann::cout << "Ignoring search with L:" << L + << " since it's smaller than K:" << recall_at << std::endl; + continue; + } - if (beamwidth <= 0) - { - diskann::cout << "Tuning beamwidth.." << std::endl; - optimized_beamwidth = - optimize_beamwidth(_pFlashIndex, warmup, warmup_num, warmup_aligned_dim, L, optimized_beamwidth); - } - else - optimized_beamwidth = beamwidth; + if (beamwidth <= 0) { + diskann::cout << "Tuning beamwidth.." << std::endl; + optimized_beamwidth = + optimize_beamwidth(_pFlashIndex, warmup, warmup_num, + warmup_aligned_dim, 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_ids[test_id].resize(recall_at * query_num); + query_result_dists[test_id].resize(recall_at * query_num); - auto stats = new diskann::QueryStats[query_num]; + auto stats = new diskann::QueryStats[query_num]; - std::vector query_result_ids_64(recall_at * query_num); - auto s = std::chrono::high_resolution_clock::now(); + std::vector query_result_ids_64(recall_at * query_num); + auto s = std::chrono::high_resolution_clock::now(); #pragma omp parallel for schedule(dynamic, 1) - for (int64_t i = 0; i < (int64_t)query_num; i++) - { - if (!filtered_search) - { - _pFlashIndex->cached_beam_search(query + (i * query_aligned_dim), recall_at, L, - query_result_ids_64.data() + (i * recall_at), - query_result_dists[test_id].data() + (i * recall_at), - optimized_beamwidth, use_reorder_data, stats + i); - } - else - { - LabelT label_for_search; - if (query_filters.size() == 1) - { // one label for all queries - label_for_search = _pFlashIndex->get_converted_label(query_filters[0]); - } - else - { // one label for each query - label_for_search = _pFlashIndex->get_converted_label(query_filters[i]); - } - _pFlashIndex->cached_beam_search( - query + (i * query_aligned_dim), recall_at, L, query_result_ids_64.data() + (i * recall_at), - query_result_dists[test_id].data() + (i * recall_at), optimized_beamwidth, true, label_for_search, - use_reorder_data, 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()); - - diskann::convert_types(query_result_ids_64.data(), query_result_ids[test_id].data(), - query_num, recall_at); - - auto mean_latency = diskann::get_mean_stats( - stats, query_num, [](const diskann::QueryStats &stats) { return stats.total_us; }); - - auto latency_999 = diskann::get_percentile_stats( - stats, query_num, 0.999, [](const diskann::QueryStats &stats) { return stats.total_us; }); - - auto mean_ios = diskann::get_mean_stats(stats, query_num, - [](const diskann::QueryStats &stats) { return stats.n_ios; }); - - auto mean_cpuus = diskann::get_mean_stats(stats, query_num, - [](const diskann::QueryStats &stats) { return stats.cpu_us; }); - - 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); - best_recall = std::max(recall, best_recall); + for (int64_t i = 0; i < (int64_t)query_num; i++) { + if (!filtered_search) { + _pFlashIndex->cached_beam_search( + query + (i * query_aligned_dim), recall_at, L, + query_result_ids_64.data() + (i * recall_at), + query_result_dists[test_id].data() + (i * recall_at), + optimized_beamwidth, use_reorder_data, stats + i); + } else { + LabelT label_for_search; + if (query_filters.size() == 1) { // one label for all queries + label_for_search = + _pFlashIndex->get_converted_label(query_filters[0]); + } else { // one label for each query + label_for_search = + _pFlashIndex->get_converted_label(query_filters[i]); } - - diskann::cout << std::setw(6) << L << std::setw(12) << optimized_beamwidth << std::setw(16) << qps - << std::setw(16) << mean_latency << std::setw(16) << latency_999 << std::setw(16) << mean_ios - << std::setw(16) << mean_cpuus; - if (calc_recall_flag) - { - diskann::cout << std::setw(16) << recall << std::endl; - } - else - diskann::cout << std::endl; - delete[] stats; + _pFlashIndex->cached_beam_search( + query + (i * query_aligned_dim), recall_at, L, + query_result_ids_64.data() + (i * recall_at), + query_result_dists[test_id].data() + (i * recall_at), + optimized_beamwidth, true, label_for_search, use_reorder_data, + stats + i); + } } - - diskann::cout << "Done searching. Now saving results " << std::endl; - uint64_t test_id = 0; - for (auto L : Lvec) - { - if (L < recall_at) - continue; - - std::string cur_result_path = result_output_prefix + "_" + std::to_string(L) + "_idx_uint32.bin"; - diskann::save_bin(cur_result_path, query_result_ids[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); + auto e = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = e - s; + float qps = (1.0 * query_num) / (1.0 * diff.count()); + + diskann::convert_types(query_result_ids_64.data(), + query_result_ids[test_id].data(), + query_num, recall_at); + + auto mean_latency = diskann::get_mean_stats( + stats, query_num, + [](const diskann::QueryStats &stats) { return stats.total_us; }); + + auto latency_999 = diskann::get_percentile_stats( + stats, query_num, 0.999, + [](const diskann::QueryStats &stats) { return stats.total_us; }); + + auto mean_ios = diskann::get_mean_stats( + stats, query_num, + [](const diskann::QueryStats &stats) { return stats.n_ios; }); + + auto mean_cpuus = diskann::get_mean_stats( + stats, query_num, + [](const diskann::QueryStats &stats) { return stats.cpu_us; }); + + 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); + best_recall = std::max(recall, best_recall); } - diskann::aligned_free(query); - if (warmup != nullptr) - diskann::aligned_free(warmup); - return best_recall >= fail_if_recall_below ? 0 : -1; + diskann::cout << std::setw(6) << L << std::setw(12) << optimized_beamwidth + << std::setw(16) << qps << std::setw(16) << mean_latency + << std::setw(16) << latency_999 << std::setw(16) << mean_ios + << std::setw(16) << mean_cpuus; + if (calc_recall_flag) { + diskann::cout << std::setw(16) << recall << std::endl; + } else + diskann::cout << std::endl; + delete[] stats; + } + + diskann::cout << "Done searching. Now saving results " << std::endl; + uint64_t test_id = 0; + for (auto L : Lvec) { + if (L < recall_at) continue; + + std::string cur_result_path = + result_output_prefix + "_" + std::to_string(L) + "_idx_uint32.bin"; + diskann::save_bin(cur_result_path, + query_result_ids[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); + return best_recall >= fail_if_recall_below ? 0 : -1; } -int main(int argc, char **argv) -{ - std::string data_type, dist_fn, index_path_prefix, result_path_prefix, query_file, gt_file, filter_label, - label_type, query_filters_file; - uint32_t num_threads, K, W, num_nodes_to_cache, search_io_limit; - std::vector Lvec; - bool use_reorder_data = false; - float fail_if_recall_below = 0.0f; - - po::options_description desc{"Arguments"}; - try - { - desc.add_options()("help,h", "Print information on arguments"); - desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); - desc.add_options()("dist_fn", po::value(&dist_fn)->required(), - "distance function "); - desc.add_options()("index_path_prefix", po::value(&index_path_prefix)->required(), - "Path prefix to the index"); - desc.add_options()("result_path", po::value(&result_path_prefix)->required(), - "Path prefix for saving results of the queries"); - desc.add_options()("query_file", po::value(&query_file)->required(), - "Query file in binary format"); - desc.add_options()("gt_file", po::value(>_file)->default_value(std::string("null")), - "ground truth file for the queryset"); - desc.add_options()("recall_at,K", po::value(&K)->required(), "Number of neighbors to be returned"); - desc.add_options()("search_list,L", po::value>(&Lvec)->multitoken(), - "List of L values of search"); - desc.add_options()("beamwidth,W", po::value(&W)->default_value(2), - "Beamwidth for search. Set 0 to optimize internally."); - desc.add_options()("num_nodes_to_cache", po::value(&num_nodes_to_cache)->default_value(0), - "Beamwidth for search"); - desc.add_options()("search_io_limit", - po::value(&search_io_limit)->default_value(std::numeric_limits::max()), - "Max #IOs for search"); - desc.add_options()("num_threads,T", po::value(&num_threads)->default_value(omp_get_num_procs()), - "Number of threads used for building index (defaults to " - "omp_get_num_procs())"); - desc.add_options()("use_reorder_data", po::bool_switch()->default_value(false), - "Include full precision data in the index. Use only in " - "conjuction with compressed data on SSD."); - desc.add_options()("filter_label", po::value(&filter_label)->default_value(std::string("")), - "Filter Label for Filtered Search"); - desc.add_options()("query_filters_file", - po::value(&query_filters_file)->default_value(std::string("")), - "Filter file for Queries for Filtered Search "); - desc.add_options()("label_type", po::value(&label_type)->default_value("uint"), - "Storage type of Labels , default value is uint which " - "will consume memory 4 bytes per filter"); - desc.add_options()("fail_if_recall_below", po::value(&fail_if_recall_below)->default_value(0.0f), - "If set to a value >0 and <100%, program returns -1 if best recall " - "found is below this threshold. "); - - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) - { - std::cout << desc; - return 0; - } - po::notify(vm); - if (vm["use_reorder_data"].as()) - use_reorder_data = true; - } - catch (const std::exception &ex) - { - std::cerr << ex.what() << '\n'; - return -1; - } - - diskann::Metric metric; - if (dist_fn == std::string("mips")) - { - metric = diskann::Metric::INNER_PRODUCT; - } - else if (dist_fn == std::string("l2")) - { - metric = diskann::Metric::L2; - } - else if (dist_fn == std::string("cosine")) - { - metric = diskann::Metric::COSINE; +int main(int argc, char **argv) { + std::string data_type, dist_fn, index_path_prefix, result_path_prefix, + query_file, gt_file, filter_label, label_type, query_filters_file; + uint32_t num_threads, K, W, num_nodes_to_cache, search_io_limit; + std::vector Lvec; + bool use_reorder_data = false; + float fail_if_recall_below = 0.0f; + + po::options_description desc{"Arguments"}; + try { + desc.add_options()("help,h", "Print information on arguments"); + desc.add_options()("data_type", + po::value(&data_type)->required(), + "data type "); + desc.add_options()("dist_fn", po::value(&dist_fn)->required(), + "distance function "); + desc.add_options()("index_path_prefix", + po::value(&index_path_prefix)->required(), + "Path prefix to the index"); + desc.add_options()("result_path", + po::value(&result_path_prefix)->required(), + "Path prefix for saving results of the queries"); + desc.add_options()("query_file", + po::value(&query_file)->required(), + "Query file in binary format"); + desc.add_options()( + "gt_file", + po::value(>_file)->default_value(std::string("null")), + "ground truth file for the queryset"); + desc.add_options()("recall_at,K", po::value(&K)->required(), + "Number of neighbors to be returned"); + desc.add_options()("search_list,L", + po::value>(&Lvec)->multitoken(), + "List of L values of search"); + desc.add_options()("beamwidth,W", po::value(&W)->default_value(2), + "Beamwidth for search. Set 0 to optimize internally."); + desc.add_options()( + "num_nodes_to_cache", + po::value(&num_nodes_to_cache)->default_value(0), + "Beamwidth for search"); + desc.add_options()( + "search_io_limit", + po::value(&search_io_limit) + ->default_value(std::numeric_limits::max()), + "Max #IOs for search"); + desc.add_options()( + "num_threads,T", + po::value(&num_threads)->default_value(omp_get_num_procs()), + "Number of threads used for building index (defaults to " + "omp_get_num_procs())"); + desc.add_options()("use_reorder_data", + po::bool_switch()->default_value(false), + "Include full precision data in the index. Use only in " + "conjuction with compressed data on SSD."); + desc.add_options()( + "filter_label", + po::value(&filter_label)->default_value(std::string("")), + "Filter Label for Filtered Search"); + desc.add_options()("query_filters_file", + po::value(&query_filters_file) + ->default_value(std::string("")), + "Filter file for Queries for Filtered Search "); + desc.add_options()( + "label_type", + po::value(&label_type)->default_value("uint"), + "Storage type of Labels , default value is uint which " + "will consume memory 4 bytes per filter"); + desc.add_options()( + "fail_if_recall_below", + po::value(&fail_if_recall_below)->default_value(0.0f), + "If set to a value >0 and <100%, program returns -1 if best recall " + "found is below this threshold. "); + + po::variables_map vm; + po::store(po::parse_command_line(argc, argv, desc), vm); + if (vm.count("help")) { + std::cout << desc; + return 0; } - else - { - std::cout << "Unsupported distance function. Currently only L2/ Inner " - "Product/Cosine are supported." + po::notify(vm); + if (vm["use_reorder_data"].as()) use_reorder_data = true; + } catch (const std::exception &ex) { + std::cerr << ex.what() << '\n'; + return -1; + } + + diskann::Metric metric; + if (dist_fn == std::string("mips")) { + metric = diskann::Metric::INNER_PRODUCT; + } else if (dist_fn == std::string("l2")) { + metric = diskann::Metric::L2; + } else if (dist_fn == std::string("cosine")) { + metric = diskann::Metric::COSINE; + } else { + std::cout << "Unsupported distance function. Currently only L2/ Inner " + "Product/Cosine are supported." + << std::endl; + return -1; + } + + if ((data_type != std::string("float")) && + (metric == diskann::Metric::INNER_PRODUCT)) { + std::cout << "Currently support only floating point data for Inner Product." + << std::endl; + return -1; + } + + if (use_reorder_data && data_type != std::string("float")) { + std::cout << "Error: Reorder data for reordering currently only " + "supported for float data type." + << std::endl; + return -1; + } + + if (filter_label != "" && query_filters_file != "") { + std::cerr + << "Only one of filter_label and query_filters_file should be provided" + << std::endl; + return -1; + } + + std::vector query_filters; + if (filter_label != "") { + query_filters.push_back(filter_label); + } else if (query_filters_file != "") { + query_filters = read_file_to_vector_of_strings(query_filters_file); + } + + try { + if (!query_filters.empty() && label_type == "ushort") { + if (data_type == std::string("float")) + return search_disk_index( + metric, index_path_prefix, result_path_prefix, query_file, gt_file, + num_threads, K, W, num_nodes_to_cache, search_io_limit, Lvec, + fail_if_recall_below, query_filters, use_reorder_data); + else if (data_type == std::string("int8")) + return search_disk_index( + metric, index_path_prefix, result_path_prefix, query_file, gt_file, + num_threads, K, W, num_nodes_to_cache, search_io_limit, Lvec, + fail_if_recall_below, query_filters, use_reorder_data); + else if (data_type == std::string("uint8")) + return search_disk_index( + metric, index_path_prefix, result_path_prefix, query_file, gt_file, + num_threads, K, W, num_nodes_to_cache, search_io_limit, Lvec, + fail_if_recall_below, query_filters, use_reorder_data); + else { + std::cerr << "Unsupported data type. Use float or int8 or uint8" << std::endl; return -1; - } - - if ((data_type != std::string("float")) && (metric == diskann::Metric::INNER_PRODUCT)) - { - std::cout << "Currently support only floating point data for Inner Product." << std::endl; - return -1; - } - - if (use_reorder_data && data_type != std::string("float")) - { - std::cout << "Error: Reorder data for reordering currently only " - "supported for float data type." + } + } else { + if (data_type == std::string("float")) + return search_disk_index( + metric, index_path_prefix, result_path_prefix, query_file, gt_file, + num_threads, K, W, num_nodes_to_cache, search_io_limit, Lvec, + fail_if_recall_below, query_filters, use_reorder_data); + else if (data_type == std::string("int8")) + return search_disk_index( + metric, index_path_prefix, result_path_prefix, query_file, gt_file, + num_threads, K, W, num_nodes_to_cache, search_io_limit, Lvec, + fail_if_recall_below, query_filters, use_reorder_data); + else if (data_type == std::string("uint8")) + return search_disk_index( + metric, index_path_prefix, result_path_prefix, query_file, gt_file, + num_threads, K, W, num_nodes_to_cache, search_io_limit, Lvec, + fail_if_recall_below, query_filters, use_reorder_data); + else { + std::cerr << "Unsupported data type. Use float or int8 or uint8" << std::endl; return -1; + } } - - if (filter_label != "" && query_filters_file != "") - { - std::cerr << "Only one of filter_label and query_filters_file should be provided" << std::endl; - return -1; - } - - std::vector query_filters; - if (filter_label != "") - { - query_filters.push_back(filter_label); - } - else if (query_filters_file != "") - { - query_filters = read_file_to_vector_of_strings(query_filters_file); - } - - try - { - if (!query_filters.empty() && label_type == "ushort") - { - if (data_type == std::string("float")) - return search_disk_index( - metric, index_path_prefix, result_path_prefix, query_file, gt_file, num_threads, K, W, - num_nodes_to_cache, search_io_limit, Lvec, fail_if_recall_below, query_filters, use_reorder_data); - else if (data_type == std::string("int8")) - return search_disk_index( - metric, index_path_prefix, result_path_prefix, query_file, gt_file, num_threads, K, W, - num_nodes_to_cache, search_io_limit, Lvec, fail_if_recall_below, query_filters, use_reorder_data); - else if (data_type == std::string("uint8")) - return search_disk_index( - metric, index_path_prefix, result_path_prefix, query_file, gt_file, num_threads, K, W, - num_nodes_to_cache, search_io_limit, Lvec, fail_if_recall_below, query_filters, use_reorder_data); - else - { - std::cerr << "Unsupported data type. Use float or int8 or uint8" << std::endl; - return -1; - } - } - else - { - if (data_type == std::string("float")) - return search_disk_index(metric, index_path_prefix, result_path_prefix, query_file, gt_file, - num_threads, K, W, num_nodes_to_cache, search_io_limit, Lvec, - fail_if_recall_below, query_filters, use_reorder_data); - else if (data_type == std::string("int8")) - return search_disk_index(metric, index_path_prefix, result_path_prefix, query_file, gt_file, - num_threads, K, W, num_nodes_to_cache, search_io_limit, Lvec, - fail_if_recall_below, query_filters, use_reorder_data); - else if (data_type == std::string("uint8")) - return search_disk_index(metric, index_path_prefix, result_path_prefix, query_file, gt_file, - num_threads, K, W, num_nodes_to_cache, search_io_limit, Lvec, - fail_if_recall_below, query_filters, use_reorder_data); - else - { - std::cerr << "Unsupported data type. Use float or int8 or uint8" << std::endl; - return -1; - } - } - } - catch (const std::exception &e) - { - std::cout << std::string(e.what()) << std::endl; - diskann::cerr << "Index search failed." << std::endl; - return -1; - } + } catch (const std::exception &e) { + std::cout << std::string(e.what()) << std::endl; + diskann::cerr << "Index search failed." << std::endl; + return -1; + } } \ No newline at end of file diff --git a/tests/search_memory_index.cpp b/tests/search_memory_index.cpp index 8fa17d4e8..e9ba47c3e 100644 --- a/tests/search_memory_index.cpp +++ b/tests/search_memory_index.cpp @@ -24,409 +24,395 @@ 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 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) -{ - // Load the query file - T *query = nullptr; - uint32_t *gt_ids = nullptr; - float *gt_dists = nullptr; - size_t query_num, query_dim, query_aligned_dim, gt_num, gt_dim; - diskann::load_aligned_bin(query_file, query, query_num, query_dim, query_aligned_dim); - - // Check for ground truth - bool calc_recall_flag = false; - if (truthset_file != std::string("null") && file_exists(truthset_file)) - { - diskann::load_truthset(truthset_file, 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; - } - calc_recall_flag = true; - } - else - { - diskann::cout << " Truthset file " << truthset_file << " not found. Not computing recall." << std::endl; - } - - bool filtered_search = false; - if (!query_filters.empty()) - { - filtered_search = true; - if (query_filters.size() != 1 && query_filters.size() != query_num) - { - std::cout << "Error. Mismatch in number of queries and size of query filters file" << std::endl; - return -1; // To return -1 or some other error handling? - } - } - - using TagT = uint32_t; - const bool concurrent = false, pq_dist_build = false, use_opq = false; - const size_t num_pq_chunks = 0; - using IndexType = diskann::Index; - const size_t num_frozen_pts = IndexType::get_graph_num_frozen_points(index_path); - IndexType index(metric, query_dim, 0, dynamic, tags, concurrent, pq_dist_build, num_pq_chunks, use_opq, - num_frozen_pts); - std::cout << "Index class instantiated" << std::endl; - index.load(index_path.c_str(), num_threads, *(std::max_element(Lvec.begin(), Lvec.end()))); - std::cout << "Index loaded" << std::endl; - if (metric == diskann::FAST_L2) - index.optimize_index_layout(); - - std::cout << "Using " << num_threads << " threads to search" << std::endl; - std::cout.setf(std::ios_base::fixed, std::ios_base::floatfield); - std::cout.precision(2); - const std::string qps_title = show_qps_per_thread ? "QPS/thread" : "QPS"; - uint32_t table_width = 0; - if (tags) - { - std::cout << std::setw(4) << "Ls" << std::setw(12) << qps_title << std::setw(20) << "Mean Latency (mus)" - << std::setw(15) << "99.9 Latency"; - table_width += 4 + 12 + 20 + 15; - } - else - { - std::cout << std::setw(4) << "Ls" << std::setw(12) << qps_title << std::setw(18) << "Avg dist cmps" - << std::setw(20) << "Mean Latency (mus)" << std::setw(15) << "99.9 Latency"; - table_width += 4 + 12 + 18 + 20 + 15; +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 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) { + // Load the query file + T *query = nullptr; + uint32_t *gt_ids = nullptr; + float *gt_dists = nullptr; + size_t query_num, query_dim, query_aligned_dim, gt_num, gt_dim; + diskann::load_aligned_bin(query_file, query, query_num, query_dim, + query_aligned_dim); + + // Check for ground truth + bool calc_recall_flag = false; + if (truthset_file != std::string("null") && file_exists(truthset_file)) { + diskann::load_truthset(truthset_file, 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; } - uint32_t recalls_to_print = 0; - const uint32_t first_recall = print_all_recalls ? 1 : recall_at; - if (calc_recall_flag) - { - for (uint32_t curr_recall = first_recall; curr_recall <= recall_at; curr_recall++) - { - std::cout << std::setw(12) << ("Recall@" + std::to_string(curr_recall)); - } - recalls_to_print = recall_at + 1 - first_recall; - table_width += recalls_to_print * 12; + calc_recall_flag = true; + } else { + diskann::cout << " Truthset file " << truthset_file + << " not found. Not computing recall." << std::endl; + } + + bool filtered_search = false; + if (!query_filters.empty()) { + filtered_search = true; + if (query_filters.size() != 1 && query_filters.size() != query_num) { + std::cout << "Error. Mismatch in number of queries and size of query " + "filters file" + << std::endl; + return -1; // To return -1 or some other error handling? } - std::cout << std::endl; - std::cout << std::string(table_width, '=') << std::endl; - - std::vector> query_result_ids(Lvec.size()); - std::vector> query_result_dists(Lvec.size()); - std::vector latency_stats(query_num, 0); - std::vector cmp_stats; - if (not tags) - { - cmp_stats = std::vector(query_num, 0); + } + + using TagT = uint32_t; + const bool concurrent = false, pq_dist_build = false, use_opq = false; + const size_t num_pq_chunks = 0; + using IndexType = diskann::Index; + const size_t num_frozen_pts = + IndexType::get_graph_num_frozen_points(index_path); + IndexType index(metric, query_dim, 0, dynamic, tags, concurrent, + pq_dist_build, num_pq_chunks, use_opq, num_frozen_pts); + std::cout << "Index class instantiated" << std::endl; + index.load(index_path.c_str(), num_threads, + *(std::max_element(Lvec.begin(), Lvec.end()))); + std::cout << "Index loaded" << std::endl; + if (metric == diskann::FAST_L2) index.optimize_index_layout(); + + std::cout << "Using " << num_threads << " threads to search" << std::endl; + std::cout.setf(std::ios_base::fixed, std::ios_base::floatfield); + std::cout.precision(2); + const std::string qps_title = show_qps_per_thread ? "QPS/thread" : "QPS"; + uint32_t table_width = 0; + if (tags) { + std::cout << std::setw(4) << "Ls" << std::setw(12) << qps_title + << std::setw(20) << "Mean Latency (mus)" << std::setw(15) + << "99.9 Latency"; + table_width += 4 + 12 + 20 + 15; + } else { + std::cout << std::setw(4) << "Ls" << std::setw(12) << qps_title + << std::setw(18) << "Avg dist cmps" << std::setw(20) + << "Mean Latency (mus)" << std::setw(15) << "99.9 Latency"; + table_width += 4 + 12 + 18 + 20 + 15; + } + uint32_t recalls_to_print = 0; + const uint32_t first_recall = print_all_recalls ? 1 : recall_at; + if (calc_recall_flag) { + for (uint32_t curr_recall = first_recall; curr_recall <= recall_at; + curr_recall++) { + std::cout << std::setw(12) << ("Recall@" + std::to_string(curr_recall)); } - - std::vector query_result_tags; - if (tags) - { - query_result_tags.resize(recall_at * query_num); + recalls_to_print = recall_at + 1 - first_recall; + table_width += recalls_to_print * 12; + } + std::cout << std::endl; + std::cout << std::string(table_width, '=') << std::endl; + + std::vector> query_result_ids(Lvec.size()); + std::vector> query_result_dists(Lvec.size()); + std::vector latency_stats(query_num, 0); + std::vector cmp_stats; + if (not tags) { + cmp_stats = std::vector(query_num, 0); + } + + std::vector query_result_tags; + if (tags) { + query_result_tags.resize(recall_at * query_num); + } + + float best_recall = 0.0; + + for (uint32_t test_id = 0; test_id < Lvec.size(); test_id++) { + uint64_t L = Lvec[test_id]; + if (L < recall_at) { + diskann::cout << "Ignoring search with L:" << L + << " since it's smaller than K:" << recall_at << std::endl; + continue; } - float best_recall = 0.0; - - for (uint32_t test_id = 0; test_id < Lvec.size(); test_id++) - { - uint64_t L = Lvec[test_id]; - if (L < recall_at) - { - diskann::cout << "Ignoring search with L:" << L << " since it's smaller than K:" << recall_at << std::endl; - continue; - } - - query_result_ids[test_id].resize(recall_at * query_num); - query_result_dists[test_id].resize(recall_at * query_num); - std::vector res = std::vector(); + query_result_ids[test_id].resize(recall_at * query_num); + query_result_dists[test_id].resize(recall_at * query_num); + std::vector res = std::vector(); - auto s = std::chrono::high_resolution_clock::now(); - omp_set_num_threads(num_threads); + auto s = std::chrono::high_resolution_clock::now(); + omp_set_num_threads(num_threads); #pragma omp parallel for schedule(dynamic, 1) - for (int64_t i = 0; i < (int64_t)query_num; i++) - { - auto qs = std::chrono::high_resolution_clock::now(); - if (filtered_search) - { - LabelT filter_label_as_num; - if (query_filters.size() == 1) - { - filter_label_as_num = index.get_converted_label(query_filters[0]); - } - else - { - filter_label_as_num = index.get_converted_label(query_filters[i]); - } - auto retval = index.search_with_filters(query + i * query_aligned_dim, filter_label_as_num, recall_at, - L, query_result_ids[test_id].data() + i * recall_at, - query_result_dists[test_id].data() + i * recall_at); - cmp_stats[i] = retval.second; - } - else if (metric == diskann::FAST_L2) - { - index.search_with_optimized_layout(query + i * query_aligned_dim, recall_at, L, - query_result_ids[test_id].data() + i * recall_at); - } - else if (tags) - { - index.search_with_tags(query + i * query_aligned_dim, recall_at, L, - query_result_tags.data() + i * recall_at, nullptr, res); - for (int64_t r = 0; r < (int64_t)recall_at; r++) - { - query_result_ids[test_id][recall_at * i + r] = query_result_tags[recall_at * i + r]; - } - } - else - { - cmp_stats[i] = index - .search(query + i * query_aligned_dim, recall_at, L, - query_result_ids[test_id].data() + i * recall_at) - .second; - } - auto qe = std::chrono::high_resolution_clock::now(); - std::chrono::duration diff = qe - qs; - latency_stats[i] = diff.count() * 1000000; - } - std::chrono::duration diff = std::chrono::high_resolution_clock::now() - s; - - float displayed_qps = static_cast(query_num) / diff.count(); - - if (show_qps_per_thread) - displayed_qps /= num_threads; - - std::vector recalls; - if (calc_recall_flag) - { - recalls.reserve(recalls_to_print); - for (uint32_t curr_recall = first_recall; curr_recall <= recall_at; curr_recall++) - { - recalls.push_back(diskann::calculate_recall(query_num, gt_ids, gt_dists, gt_dim, - query_result_ids[test_id].data(), recall_at, curr_recall)); - } - } - - std::sort(latency_stats.begin(), latency_stats.end()); - float mean_latency = - std::accumulate(latency_stats.begin(), latency_stats.end(), 0.0) / static_cast(query_num); - - float avg_cmps = (float)std::accumulate(cmp_stats.begin(), cmp_stats.end(), 0) / (float)query_num; - - if (tags) - { - std::cout << std::setw(4) << L << std::setw(12) << displayed_qps << std::setw(20) << (float)mean_latency - << std::setw(15) << (float)latency_stats[(uint64_t)(0.999 * query_num)]; - } - else - { - std::cout << std::setw(4) << L << std::setw(12) << displayed_qps << std::setw(18) << avg_cmps - << std::setw(20) << (float)mean_latency << std::setw(15) - << (float)latency_stats[(uint64_t)(0.999 * query_num)]; - } - for (float recall : recalls) - { - std::cout << std::setw(12) << recall; - best_recall = std::max(recall, best_recall); - } - std::cout << std::endl; - } - - std::cout << "Done searching. Now saving results " << std::endl; - uint64_t test_id = 0; - for (auto L : Lvec) - { - if (L < recall_at) - { - diskann::cout << "Ignoring search with L:" << L << " since it's smaller than K:" << recall_at << std::endl; - continue; + for (int64_t i = 0; i < (int64_t)query_num; i++) { + auto qs = std::chrono::high_resolution_clock::now(); + if (filtered_search) { + LabelT filter_label_as_num; + if (query_filters.size() == 1) { + filter_label_as_num = index.get_converted_label(query_filters[0]); + } else { + filter_label_as_num = index.get_converted_label(query_filters[i]); } - std::string cur_result_path = result_path_prefix + "_" + std::to_string(L) + "_idx_uint32.bin"; - diskann::save_bin(cur_result_path, query_result_ids[test_id].data(), query_num, recall_at); - test_id++; - } - - diskann::aligned_free(query); - - return best_recall >= fail_if_recall_below ? 0 : -1; -} - -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; - std::vector Lvec; - bool print_all_recalls, dynamic, tags, show_qps_per_thread; - float fail_if_recall_below = 0.0f; - - po::options_description desc{"Arguments"}; - try - { - desc.add_options()("help,h", "Print information on arguments"); - desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); - desc.add_options()("dist_fn", po::value(&dist_fn)->required(), - "distance function "); - desc.add_options()("index_path_prefix", po::value(&index_path_prefix)->required(), - "Path prefix to the index"); - desc.add_options()("result_path", po::value(&result_path)->required(), - "Path prefix for saving results of the queries"); - desc.add_options()("query_file", po::value(&query_file)->required(), - "Query file in binary format"); - desc.add_options()("filter_label", po::value(&filter_label)->default_value(std::string("")), - "Filter Label for Filtered Search"); - desc.add_options()("query_filters_file", - po::value(&query_filters_file)->default_value(std::string("")), - "Filter file for Queries for Filtered Search "); - desc.add_options()("label_type", po::value(&label_type)->default_value("uint"), - "Storage type of Labels , default value is uint which " - "will consume memory 4 bytes per filter"); - desc.add_options()("gt_file", po::value(>_file)->default_value(std::string("null")), - "ground truth file for the queryset"); - desc.add_options()("recall_at,K", po::value(&K)->required(), "Number of neighbors to be returned"); - desc.add_options()("print_all_recalls", po::bool_switch(&print_all_recalls), - "Print recalls at all positions, from 1 up to specified " - "recall_at value"); - desc.add_options()("search_list,L", po::value>(&Lvec)->multitoken(), - "List of L values of search"); - desc.add_options()("num_threads,T", po::value(&num_threads)->default_value(omp_get_num_procs()), - "Number of threads used for building index (defaults to " - "omp_get_num_procs())"); - desc.add_options()("dynamic", po::value(&dynamic)->default_value(false), - "Whether the index is dynamic. Default false."); - desc.add_options()("tags", po::value(&tags)->default_value(false), - "Whether to search with tags. Default false."); - desc.add_options()("qps_per_thread", po::bool_switch(&show_qps_per_thread), - "Print overall QPS divided by the number of threads in " - "the output table"); - desc.add_options()("fail_if_recall_below", po::value(&fail_if_recall_below)->default_value(0.0f), - "If set to a value >0 and <100%, program returns -1 if best recall " - "found is below this threshold. "); - - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) - { - std::cout << desc; - return 0; + auto retval = index.search_with_filters( + query + i * query_aligned_dim, filter_label_as_num, recall_at, L, + query_result_ids[test_id].data() + i * recall_at, + query_result_dists[test_id].data() + i * recall_at); + cmp_stats[i] = retval.second; + } else if (metric == diskann::FAST_L2) { + index.search_with_optimized_layout( + query + i * query_aligned_dim, recall_at, L, + query_result_ids[test_id].data() + i * recall_at); + } else if (tags) { + index.search_with_tags(query + i * query_aligned_dim, recall_at, L, + query_result_tags.data() + i * recall_at, + nullptr, res); + for (int64_t r = 0; r < (int64_t)recall_at; r++) { + query_result_ids[test_id][recall_at * i + r] = + query_result_tags[recall_at * i + r]; } - po::notify(vm); + } else { + cmp_stats[i] = + index + .search(query + i * query_aligned_dim, recall_at, L, + query_result_ids[test_id].data() + i * recall_at) + .second; + } + auto qe = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = qe - qs; + latency_stats[i] = diff.count() * 1000000; } - catch (const std::exception &ex) - { - std::cerr << ex.what() << '\n'; - return -1; + std::chrono::duration diff = + std::chrono::high_resolution_clock::now() - s; + + float displayed_qps = static_cast(query_num) / diff.count(); + + if (show_qps_per_thread) displayed_qps /= num_threads; + + std::vector recalls; + if (calc_recall_flag) { + recalls.reserve(recalls_to_print); + for (uint32_t curr_recall = first_recall; curr_recall <= recall_at; + curr_recall++) { + recalls.push_back(diskann::calculate_recall( + query_num, gt_ids, gt_dists, gt_dim, + query_result_ids[test_id].data(), recall_at, curr_recall)); + } } - diskann::Metric metric; - if ((dist_fn == std::string("mips")) && (data_type == std::string("float"))) - { - metric = diskann::Metric::INNER_PRODUCT; + std::sort(latency_stats.begin(), latency_stats.end()); + float mean_latency = + std::accumulate(latency_stats.begin(), latency_stats.end(), 0.0) / + static_cast(query_num); + + float avg_cmps = + (float)std::accumulate(cmp_stats.begin(), cmp_stats.end(), 0) / + (float)query_num; + + if (tags) { + std::cout << std::setw(4) << L << std::setw(12) << displayed_qps + << std::setw(20) << (float)mean_latency << std::setw(15) + << (float)latency_stats[(uint64_t)(0.999 * query_num)]; + } else { + std::cout << std::setw(4) << L << std::setw(12) << displayed_qps + << std::setw(18) << avg_cmps << std::setw(20) + << (float)mean_latency << std::setw(15) + << (float)latency_stats[(uint64_t)(0.999 * query_num)]; } - else if (dist_fn == std::string("l2")) - { - metric = diskann::Metric::L2; - } - else if (dist_fn == std::string("cosine")) - { - metric = diskann::Metric::COSINE; - } - else if ((dist_fn == std::string("fast_l2")) && (data_type == std::string("float"))) - { - metric = diskann::Metric::FAST_L2; - } - else - { - std::cout << "Unsupported distance function. Currently only l2/ cosine are " - "supported in general, and mips/fast_l2 only for floating " - "point data." - << std::endl; - return -1; - } - - if (dynamic && not tags) - { - std::cerr << "Tags must be enabled while searching dynamically built indices" << std::endl; - return -1; + for (float recall : recalls) { + std::cout << std::setw(12) << recall; + best_recall = std::max(recall, best_recall); } - - if (fail_if_recall_below < 0.0 || fail_if_recall_below >= 100.0) - { - std::cerr << "fail_if_recall_below parameter must be between 0 and 100%" << std::endl; - return -1; + std::cout << std::endl; + } + + std::cout << "Done searching. Now saving results " << std::endl; + uint64_t test_id = 0; + for (auto L : Lvec) { + if (L < recall_at) { + diskann::cout << "Ignoring search with L:" << L + << " since it's smaller than K:" << recall_at << std::endl; + continue; } + std::string cur_result_path = + result_path_prefix + "_" + std::to_string(L) + "_idx_uint32.bin"; + diskann::save_bin(cur_result_path, + query_result_ids[test_id].data(), query_num, + recall_at); + test_id++; + } - if (filter_label != "" && query_filters_file != "") - { - std::cerr << "Only one of filter_label and query_filters_file should be provided" << std::endl; - return -1; - } + diskann::aligned_free(query); - std::vector query_filters; - if (filter_label != "") - { - query_filters.push_back(filter_label); - } - else if (query_filters_file != "") - { - query_filters = read_file_to_vector_of_strings(query_filters_file); - } + return best_recall >= fail_if_recall_below ? 0 : -1; +} - try - { - if (!query_filters.empty() && label_type == "ushort") - { - 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, - 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, - 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, - num_threads, K, print_all_recalls, Lvec, dynamic, tags, - show_qps_per_thread, query_filters, fail_if_recall_below); - } - else - { - std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; - return -1; - } - } - else - { - 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, 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, 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, - num_threads, K, print_all_recalls, Lvec, dynamic, tags, - show_qps_per_thread, query_filters, fail_if_recall_below); - } - else - { - std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; - return -1; - } - } +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; + std::vector Lvec; + bool print_all_recalls, dynamic, tags, show_qps_per_thread; + float fail_if_recall_below = 0.0f; + + po::options_description desc{"Arguments"}; + try { + desc.add_options()("help,h", "Print information on arguments"); + desc.add_options()("data_type", + po::value(&data_type)->required(), + "data type "); + desc.add_options()("dist_fn", po::value(&dist_fn)->required(), + "distance function "); + desc.add_options()("index_path_prefix", + po::value(&index_path_prefix)->required(), + "Path prefix to the index"); + desc.add_options()("result_path", + po::value(&result_path)->required(), + "Path prefix for saving results of the queries"); + desc.add_options()("query_file", + po::value(&query_file)->required(), + "Query file in binary format"); + desc.add_options()( + "filter_label", + po::value(&filter_label)->default_value(std::string("")), + "Filter Label for Filtered Search"); + desc.add_options()("query_filters_file", + po::value(&query_filters_file) + ->default_value(std::string("")), + "Filter file for Queries for Filtered Search "); + desc.add_options()( + "label_type", + po::value(&label_type)->default_value("uint"), + "Storage type of Labels , default value is uint which " + "will consume memory 4 bytes per filter"); + desc.add_options()( + "gt_file", + po::value(>_file)->default_value(std::string("null")), + "ground truth file for the queryset"); + desc.add_options()("recall_at,K", po::value(&K)->required(), + "Number of neighbors to be returned"); + desc.add_options()("print_all_recalls", po::bool_switch(&print_all_recalls), + "Print recalls at all positions, from 1 up to specified " + "recall_at value"); + desc.add_options()("search_list,L", + po::value>(&Lvec)->multitoken(), + "List of L values of search"); + desc.add_options()( + "num_threads,T", + po::value(&num_threads)->default_value(omp_get_num_procs()), + "Number of threads used for building index (defaults to " + "omp_get_num_procs())"); + desc.add_options()("dynamic", + po::value(&dynamic)->default_value(false), + "Whether the index is dynamic. Default false."); + desc.add_options()("tags", po::value(&tags)->default_value(false), + "Whether to search with tags. Default false."); + desc.add_options()("qps_per_thread", po::bool_switch(&show_qps_per_thread), + "Print overall QPS divided by the number of threads in " + "the output table"); + desc.add_options()( + "fail_if_recall_below", + po::value(&fail_if_recall_below)->default_value(0.0f), + "If set to a value >0 and <100%, program returns -1 if best recall " + "found is below this threshold. "); + + po::variables_map vm; + po::store(po::parse_command_line(argc, argv, desc), vm); + if (vm.count("help")) { + std::cout << desc; + return 0; } - catch (std::exception &e) - { - std::cout << std::string(e.what()) << std::endl; - diskann::cerr << "Index search failed." << std::endl; + po::notify(vm); + } catch (const std::exception &ex) { + std::cerr << ex.what() << '\n'; + return -1; + } + + diskann::Metric metric; + if ((dist_fn == std::string("mips")) && (data_type == std::string("float"))) { + metric = diskann::Metric::INNER_PRODUCT; + } else if (dist_fn == std::string("l2")) { + metric = diskann::Metric::L2; + } else if (dist_fn == std::string("cosine")) { + metric = diskann::Metric::COSINE; + } else if ((dist_fn == std::string("fast_l2")) && + (data_type == std::string("float"))) { + metric = diskann::Metric::FAST_L2; + } else { + std::cout << "Unsupported distance function. Currently only l2/ cosine are " + "supported in general, and mips/fast_l2 only for floating " + "point data." + << std::endl; + return -1; + } + + if (dynamic && not tags) { + std::cerr + << "Tags must be enabled while searching dynamically built indices" + << std::endl; + return -1; + } + + if (fail_if_recall_below < 0.0 || fail_if_recall_below >= 100.0) { + std::cerr << "fail_if_recall_below parameter must be between 0 and 100%" + << std::endl; + return -1; + } + + if (filter_label != "" && query_filters_file != "") { + std::cerr + << "Only one of filter_label and query_filters_file should be provided" + << std::endl; + return -1; + } + + std::vector query_filters; + if (filter_label != "") { + query_filters.push_back(filter_label); + } else if (query_filters_file != "") { + query_filters = read_file_to_vector_of_strings(query_filters_file); + } + + try { + if (!query_filters.empty() && label_type == "ushort") { + 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, 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, 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, + num_threads, K, print_all_recalls, Lvec, dynamic, tags, + show_qps_per_thread, query_filters, fail_if_recall_below); + } else { + std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; + return -1; + } + } else { + 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, 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, 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, + num_threads, K, print_all_recalls, Lvec, dynamic, tags, + show_qps_per_thread, query_filters, fail_if_recall_below); + } else { + std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; return -1; + } } + } catch (std::exception &e) { + std::cout << std::string(e.what()) << std::endl; + diskann::cerr << "Index search failed." << std::endl; + return -1; + } } diff --git a/tests/test_insert_deletes_consolidate.cpp b/tests/test_insert_deletes_consolidate.cpp index 40d2c6344..75579fb1f 100644 --- a/tests/test_insert_deletes_consolidate.cpp +++ b/tests/test_insert_deletes_consolidate.cpp @@ -25,407 +25,452 @@ namespace po = boost::program_options; // load_aligned_bin modified to read pieces of the file, but using ifstream // instead of cached_ifstream. template -inline void load_aligned_bin_part(const std::string &bin_file, T *data, size_t offset_points, size_t points_to_read) -{ - diskann::Timer timer; - std::ifstream reader; - reader.exceptions(std::ios::failbit | std::ios::badbit); - reader.open(bin_file, std::ios::binary | std::ios::ate); - size_t actual_file_size = reader.tellg(); - reader.seekg(0, std::ios::beg); - - int npts_i32, dim_i32; - reader.read((char *)&npts_i32, sizeof(int)); - reader.read((char *)&dim_i32, sizeof(int)); - size_t npts = (uint32_t)npts_i32; - size_t dim = (uint32_t)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; - std::cout << stream.str(); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - - if (offset_points + points_to_read > npts) - { - std::stringstream stream; - stream << "Error. Not enough points in file. Requested " << offset_points << " offset and " << points_to_read - << " points, but have only " << npts << " points" << std::endl; - std::cout << stream.str(); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - - reader.seekg(2 * sizeof(uint32_t) + offset_points * dim * sizeof(T)); - - const size_t rounded_dim = ROUND_UP(dim, 8); - - for (size_t i = 0; i < points_to_read; i++) - { - reader.read((char *)(data + i * rounded_dim), dim * sizeof(T)); - memset(data + i * rounded_dim + dim, 0, (rounded_dim - dim) * sizeof(T)); - } - reader.close(); - - const double elapsedSeconds = timer.elapsed() / 1000000.0; - std::cout << "Read " << points_to_read << " points using non-cached reads in " << elapsedSeconds << std::endl; +inline void load_aligned_bin_part(const std::string &bin_file, T *data, + size_t offset_points, size_t points_to_read) { + diskann::Timer timer; + std::ifstream reader; + reader.exceptions(std::ios::failbit | std::ios::badbit); + reader.open(bin_file, std::ios::binary | std::ios::ate); + size_t actual_file_size = reader.tellg(); + reader.seekg(0, std::ios::beg); + + int npts_i32, dim_i32; + reader.read((char *)&npts_i32, sizeof(int)); + reader.read((char *)&dim_i32, sizeof(int)); + size_t npts = (uint32_t)npts_i32; + size_t dim = (uint32_t)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; + std::cout << stream.str(); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } + + if (offset_points + points_to_read > npts) { + std::stringstream stream; + stream << "Error. Not enough points in file. Requested " << offset_points + << " offset and " << points_to_read << " points, but have only " + << npts << " points" << std::endl; + std::cout << stream.str(); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } + + reader.seekg(2 * sizeof(uint32_t) + offset_points * dim * sizeof(T)); + + const size_t rounded_dim = ROUND_UP(dim, 8); + + for (size_t i = 0; i < points_to_read; i++) { + reader.read((char *)(data + i * rounded_dim), dim * sizeof(T)); + memset(data + i * rounded_dim + dim, 0, (rounded_dim - dim) * sizeof(T)); + } + reader.close(); + + const double elapsedSeconds = timer.elapsed() / 1000000.0; + std::cout << "Read " << points_to_read << " points using non-cached reads in " + << elapsedSeconds << std::endl; } -std::string get_save_filename(const std::string &save_path, size_t points_to_skip, size_t points_deleted, - size_t last_point_threshold) -{ - std::string final_path = save_path; - if (points_to_skip > 0) - { - final_path += "skip" + std::to_string(points_to_skip) + "-"; - } - - final_path += "del" + std::to_string(points_deleted) + "-"; - final_path += std::to_string(last_point_threshold); - return final_path; +std::string get_save_filename(const std::string &save_path, + size_t points_to_skip, size_t points_deleted, + size_t last_point_threshold) { + std::string final_path = save_path; + if (points_to_skip > 0) { + final_path += "skip" + std::to_string(points_to_skip) + "-"; + } + + final_path += "del" + std::to_string(points_deleted) + "-"; + final_path += std::to_string(last_point_threshold); + return final_path; } template -void insert_till_next_checkpoint(diskann::Index &index, size_t start, size_t end, size_t thread_count, T *data, - size_t aligned_dim) -{ - diskann::Timer insert_timer; +void insert_till_next_checkpoint(diskann::Index &index, size_t start, + size_t end, size_t thread_count, T *data, + size_t aligned_dim) { + diskann::Timer insert_timer; #pragma omp parallel for num_threads(thread_count) schedule(dynamic) - for (int64_t j = start; j < (int64_t)end; j++) - { - index.insert_point(&data[(j - start) * aligned_dim], 1 + static_cast(j)); - } - const double elapsedSeconds = insert_timer.elapsed() / 1000000.0; - std::cout << "Insertion time " << elapsedSeconds << " seconds (" << (end - start) / elapsedSeconds - << " points/second overall, " << (end - start) / elapsedSeconds / thread_count << " per thread)\n "; + for (int64_t j = start; j < (int64_t)end; j++) { + index.insert_point(&data[(j - start) * aligned_dim], + 1 + static_cast(j)); + } + const double elapsedSeconds = insert_timer.elapsed() / 1000000.0; + std::cout << "Insertion time " << elapsedSeconds << " seconds (" + << (end - start) / elapsedSeconds << " points/second overall, " + << (end - start) / elapsedSeconds / thread_count + << " per thread)\n "; } template -void delete_from_beginning(diskann::Index &index, diskann::IndexWriteParameters &delete_params, - size_t points_to_skip, size_t points_to_delete_from_beginning) -{ - try - { - std::cout << std::endl - << "Lazy deleting points " << points_to_skip << " to " - << points_to_skip + points_to_delete_from_beginning << "... "; - for (size_t i = points_to_skip; i < points_to_skip + points_to_delete_from_beginning; ++i) - index.lazy_delete(i + 1); // Since tags are data location + 1 - std::cout << "done." << std::endl; - - auto report = index.consolidate_deletes(delete_params); - std::cout << "#active points: " << report._active_points << std::endl - << "max points: " << report._max_points << std::endl - << "empty slots: " << report._empty_slots << std::endl - << "deletes processed: " << report._slots_released << std::endl - << "latest delete size: " << report._delete_set_size << std::endl - << "rate: (" << points_to_delete_from_beginning / report._time << " points/second overall, " - << points_to_delete_from_beginning / report._time / delete_params.num_threads << " per thread)" - << std::endl; - } - catch (std::system_error &e) - { - std::cout << "Exception caught in deletion thread: " << e.what() << std::endl; - } +void delete_from_beginning(diskann::Index &index, + diskann::IndexWriteParameters &delete_params, + size_t points_to_skip, + size_t points_to_delete_from_beginning) { + try { + std::cout << std::endl + << "Lazy deleting points " << points_to_skip << " to " + << points_to_skip + points_to_delete_from_beginning << "... "; + for (size_t i = points_to_skip; + i < points_to_skip + points_to_delete_from_beginning; ++i) + index.lazy_delete(i + 1); // Since tags are data location + 1 + std::cout << "done." << std::endl; + + auto report = index.consolidate_deletes(delete_params); + std::cout << "#active points: " << report._active_points << std::endl + << "max points: " << report._max_points << std::endl + << "empty slots: " << report._empty_slots << std::endl + << "deletes processed: " << report._slots_released << std::endl + << "latest delete size: " << report._delete_set_size << std::endl + << "rate: (" << points_to_delete_from_beginning / report._time + << " points/second overall, " + << points_to_delete_from_beginning / report._time / + delete_params.num_threads + << " per thread)" << std::endl; + } catch (std::system_error &e) { + std::cout << "Exception caught in deletion thread: " << e.what() + << std::endl; + } } template -void build_incremental_index(const std::string &data_path, const uint32_t L, const uint32_t R, const float alpha, - const uint32_t thread_count, size_t points_to_skip, size_t max_points_to_insert, - size_t beginning_index_size, float start_point_norm, 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) -{ - - diskann::IndexWriteParameters params = diskann::IndexWriteParametersBuilder(L, R) - .with_max_occlusion_size(500) // C = 500 - .with_alpha(alpha) - .with_num_rounds(1) - .with_num_threads(thread_count) - .with_num_frozen_points(num_start_pts) - .build(); - - size_t dim, aligned_dim; - size_t num_points; - - diskann::get_bin_metadata(data_path, num_points, dim); - aligned_dim = ROUND_UP(dim, 8); - - if (points_to_skip > num_points) - { - throw diskann::ANNException("Asked to skip more points than in data file", -1, __FUNCSIG__, __FILE__, __LINE__); - } - - if (max_points_to_insert == 0) - { - max_points_to_insert = num_points; - } - - if (points_to_skip + max_points_to_insert > num_points) - { - max_points_to_insert = num_points - points_to_skip; - std::cerr << "WARNING: Reducing max_points_to_insert to " << max_points_to_insert - << " points since the data file has only that many" << std::endl; - } - - using TagT = uint32_t; - const bool enable_tags = true; +void build_incremental_index( + const std::string &data_path, const uint32_t L, const uint32_t R, + const float alpha, const uint32_t thread_count, size_t points_to_skip, + size_t max_points_to_insert, size_t beginning_index_size, + float start_point_norm, 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) { + diskann::IndexWriteParameters params = + diskann::IndexWriteParametersBuilder(L, R) + .with_max_occlusion_size(500) // C = 500 + .with_alpha(alpha) + .with_num_rounds(1) + .with_num_threads(thread_count) + .with_num_frozen_points(num_start_pts) + .build(); + + size_t dim, aligned_dim; + size_t num_points; + + diskann::get_bin_metadata(data_path, num_points, dim); + aligned_dim = ROUND_UP(dim, 8); + + if (points_to_skip > num_points) { + throw diskann::ANNException("Asked to skip more points than in data file", + -1, __FUNCSIG__, __FILE__, __LINE__); + } + + if (max_points_to_insert == 0) { + max_points_to_insert = num_points; + } + + if (points_to_skip + max_points_to_insert > num_points) { + max_points_to_insert = num_points - points_to_skip; + std::cerr << "WARNING: Reducing max_points_to_insert to " + << max_points_to_insert + << " points since the data file has only that many" << std::endl; + } + + using TagT = uint32_t; + const bool enable_tags = true; + + diskann::Index index(diskann::L2, dim, max_points_to_insert, true, + params, L, thread_count, enable_tags, + concurrent); + + size_t current_point_offset = points_to_skip; + const size_t last_point_threshold = points_to_skip + max_points_to_insert; + + if (beginning_index_size > max_points_to_insert) { + beginning_index_size = max_points_to_insert; + std::cerr << "WARNING: Reducing beginning index size to " + << beginning_index_size + << " points since the data file has only that many" << std::endl; + } + if (checkpoints_per_snapshot > 0 && + beginning_index_size > points_per_checkpoint) { + beginning_index_size = points_per_checkpoint; + std::cerr << "WARNING: Reducing beginning index size to " + << beginning_index_size << std::endl; + } + + T *data = nullptr; + diskann::alloc_aligned((void **)&data, + std::max(points_per_checkpoint, beginning_index_size) * + aligned_dim * sizeof(T), + 8 * sizeof(T)); + + std::vector tags(beginning_index_size); + std::iota(tags.begin(), tags.end(), + 1 + static_cast(current_point_offset)); + + load_aligned_bin_part(data_path, data, current_point_offset, + beginning_index_size); + std::cout << "load aligned bin succeeded" << std::endl; + diskann::Timer timer; + + if (beginning_index_size > 0) { + index.build(data, beginning_index_size, params, tags); + index.enable_delete(); + } else { + index.set_start_points_at_random(static_cast(start_point_norm)); + index.enable_delete(); + } + + const double elapsedSeconds = timer.elapsed() / 1000000.0; + std::cout << "Initial non-incremental index build time for " + << beginning_index_size << " points took " << elapsedSeconds + << " seconds (" << beginning_index_size / elapsedSeconds + << " points/second)\n "; + + current_point_offset = beginning_index_size; + + if (points_to_delete_from_beginning > max_points_to_insert) { + points_to_delete_from_beginning = + static_cast(max_points_to_insert); + std::cerr << "WARNING: Reducing points to delete from beginning to " + << points_to_delete_from_beginning + << " points since the data file has only that many" << std::endl; + } + + if (concurrent) { + int sub_threads = (thread_count + 1) / 2; + bool delete_launched = false; + std::future delete_task; - diskann::Index index(diskann::L2, dim, max_points_to_insert, true, params, L, thread_count, enable_tags, - concurrent); - - size_t current_point_offset = points_to_skip; - const size_t last_point_threshold = points_to_skip + max_points_to_insert; - - if (beginning_index_size > max_points_to_insert) - { - beginning_index_size = max_points_to_insert; - std::cerr << "WARNING: Reducing beginning index size to " << beginning_index_size - << " points since the data file has only that many" << std::endl; - } - if (checkpoints_per_snapshot > 0 && beginning_index_size > points_per_checkpoint) - { - beginning_index_size = points_per_checkpoint; - std::cerr << "WARNING: Reducing beginning index size to " << beginning_index_size << std::endl; - } - - T *data = nullptr; - diskann::alloc_aligned( - (void **)&data, std::max(points_per_checkpoint, beginning_index_size) * aligned_dim * sizeof(T), 8 * sizeof(T)); - - std::vector tags(beginning_index_size); - std::iota(tags.begin(), tags.end(), 1 + static_cast(current_point_offset)); - - load_aligned_bin_part(data_path, data, current_point_offset, beginning_index_size); - std::cout << "load aligned bin succeeded" << std::endl; diskann::Timer timer; - if (beginning_index_size > 0) - { - index.build(data, beginning_index_size, params, tags); - index.enable_delete(); + for (size_t start = current_point_offset; start < last_point_threshold; + start += points_per_checkpoint, + current_point_offset += points_per_checkpoint) { + const size_t end = + std::min(start + points_per_checkpoint, last_point_threshold); + std::cout << std::endl + << "Inserting from " << start << " to " << end << std::endl; + + auto insert_task = std::async(std::launch::async, [&]() { + load_aligned_bin_part(data_path, data, start, end - start); + insert_till_next_checkpoint(index, start, end, sub_threads, data, + aligned_dim); + }); + insert_task.wait(); + + if (!delete_launched && end >= start_deletes_after && + end >= points_to_skip + points_to_delete_from_beginning) { + delete_launched = true; + diskann::IndexWriteParameters delete_params = + diskann::IndexWriteParametersBuilder(params) + .with_num_threads(sub_threads) + .build(); + + delete_task = std::async(std::launch::async, [&]() { + delete_from_beginning(index, delete_params, points_to_skip, + points_to_delete_from_beginning); + }); + } } - else - { - index.set_start_points_at_random(static_cast(start_point_norm)); - index.enable_delete(); + delete_task.wait(); + + std::cout << "Time Elapsed " << timer.elapsed() / 1000 << "ms\n"; + const auto save_path_inc = get_save_filename( + save_path + ".after-concurrent-delete-", points_to_skip, + points_to_delete_from_beginning, last_point_threshold); + index.save(save_path_inc.c_str(), true); + } else { + size_t last_snapshot_points_threshold = 0; + size_t num_checkpoints_till_snapshot = checkpoints_per_snapshot; + + for (size_t start = current_point_offset; start < last_point_threshold; + start += points_per_checkpoint, + current_point_offset += points_per_checkpoint) { + const size_t end = + std::min(start + points_per_checkpoint, last_point_threshold); + std::cout << std::endl + << "Inserting from " << start << " to " << end << std::endl; + + load_aligned_bin_part(data_path, data, start, end - start); + insert_till_next_checkpoint(index, start, end, thread_count, data, + aligned_dim); + + if (checkpoints_per_snapshot > 0 && + --num_checkpoints_till_snapshot == 0) { + diskann::Timer save_timer; + + const auto save_path_inc = + get_save_filename(save_path + ".inc-", points_to_skip, + points_to_delete_from_beginning, end); + index.save(save_path_inc.c_str(), false); + const double elapsedSeconds = save_timer.elapsed() / 1000000.0; + const size_t points_saved = end - points_to_skip; + + std::cout << "Saved " << points_saved << " points in " << elapsedSeconds + << " seconds (" << points_saved / elapsedSeconds + << " points/second)\n"; + + num_checkpoints_till_snapshot = checkpoints_per_snapshot; + last_snapshot_points_threshold = end; + } + + std::cout << "Number of points in the index post insertion " << end + << std::endl; } - const double elapsedSeconds = timer.elapsed() / 1000000.0; - std::cout << "Initial non-incremental index build time for " << beginning_index_size << " points took " - << elapsedSeconds << " seconds (" << beginning_index_size / elapsedSeconds << " points/second)\n "; - - current_point_offset = beginning_index_size; - - if (points_to_delete_from_beginning > max_points_to_insert) - { - points_to_delete_from_beginning = static_cast(max_points_to_insert); - std::cerr << "WARNING: Reducing points to delete from beginning to " << points_to_delete_from_beginning - << " points since the data file has only that many" << std::endl; + if (checkpoints_per_snapshot > 0 && + last_snapshot_points_threshold != last_point_threshold) { + const auto save_path_inc = get_save_filename( + save_path + ".inc-", points_to_skip, points_to_delete_from_beginning, + last_point_threshold); + // index.save(save_path_inc.c_str(), false); } - if (concurrent) - { - int sub_threads = (thread_count + 1) / 2; - bool delete_launched = false; - std::future delete_task; - - diskann::Timer timer; - - for (size_t start = current_point_offset; start < last_point_threshold; - start += points_per_checkpoint, current_point_offset += points_per_checkpoint) - { - const size_t end = std::min(start + points_per_checkpoint, last_point_threshold); - std::cout << std::endl << "Inserting from " << start << " to " << end << std::endl; - - auto insert_task = std::async(std::launch::async, [&]() { - load_aligned_bin_part(data_path, data, start, end - start); - insert_till_next_checkpoint(index, start, end, sub_threads, data, aligned_dim); - }); - insert_task.wait(); - - if (!delete_launched && end >= start_deletes_after && - end >= points_to_skip + points_to_delete_from_beginning) - { - delete_launched = true; - diskann::IndexWriteParameters delete_params = - diskann::IndexWriteParametersBuilder(params).with_num_threads(sub_threads).build(); - - delete_task = std::async(std::launch::async, [&]() { - delete_from_beginning(index, delete_params, points_to_skip, points_to_delete_from_beginning); - }); - } - } - delete_task.wait(); - - std::cout << "Time Elapsed " << timer.elapsed() / 1000 << "ms\n"; - const auto save_path_inc = get_save_filename(save_path + ".after-concurrent-delete-", points_to_skip, - points_to_delete_from_beginning, last_point_threshold); - index.save(save_path_inc.c_str(), true); - } - else - { - size_t last_snapshot_points_threshold = 0; - size_t num_checkpoints_till_snapshot = checkpoints_per_snapshot; - - for (size_t start = current_point_offset; start < last_point_threshold; - start += points_per_checkpoint, current_point_offset += points_per_checkpoint) - { - const size_t end = std::min(start + points_per_checkpoint, last_point_threshold); - std::cout << std::endl << "Inserting from " << start << " to " << end << std::endl; - - load_aligned_bin_part(data_path, data, start, end - start); - insert_till_next_checkpoint(index, start, end, thread_count, data, aligned_dim); - - if (checkpoints_per_snapshot > 0 && --num_checkpoints_till_snapshot == 0) - { - diskann::Timer save_timer; - - const auto save_path_inc = - get_save_filename(save_path + ".inc-", points_to_skip, points_to_delete_from_beginning, end); - index.save(save_path_inc.c_str(), false); - const double elapsedSeconds = save_timer.elapsed() / 1000000.0; - const size_t points_saved = end - points_to_skip; - - std::cout << "Saved " << points_saved << " points in " << elapsedSeconds << " seconds (" - << points_saved / elapsedSeconds << " points/second)\n"; - - num_checkpoints_till_snapshot = checkpoints_per_snapshot; - last_snapshot_points_threshold = end; - } - - std::cout << "Number of points in the index post insertion " << end << std::endl; - } - - if (checkpoints_per_snapshot > 0 && last_snapshot_points_threshold != last_point_threshold) - { - const auto save_path_inc = get_save_filename(save_path + ".inc-", points_to_skip, - points_to_delete_from_beginning, last_point_threshold); - // index.save(save_path_inc.c_str(), false); - } - - if (points_to_delete_from_beginning > 0) - { - delete_from_beginning(index, params, points_to_skip, points_to_delete_from_beginning); - } - const auto save_path_inc = get_save_filename(save_path + ".after-delete-", points_to_skip, - points_to_delete_from_beginning, last_point_threshold); - index.save(save_path_inc.c_str(), true); + if (points_to_delete_from_beginning > 0) { + delete_from_beginning(index, params, points_to_skip, + points_to_delete_from_beginning); } + const auto save_path_inc = get_save_filename( + save_path + ".after-delete-", points_to_skip, + points_to_delete_from_beginning, last_point_threshold); + index.save(save_path_inc.c_str(), true); + } - diskann::aligned_free(data); + diskann::aligned_free(data); } -int main(int argc, char **argv) -{ - std::string data_type, dist_fn, data_path, index_path_prefix; - 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; - bool concurrent; - - po::options_description desc{"Arguments"}; - try - { - desc.add_options()("help,h", "Print information on arguments"); - desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); - desc.add_options()("dist_fn", po::value(&dist_fn)->required(), "distance function "); - desc.add_options()("data_path", po::value(&data_path)->required(), - "Input data file in bin format"); - desc.add_options()("index_path_prefix", po::value(&index_path_prefix)->required(), - "Path prefix for saving index file components"); - desc.add_options()("max_degree,R", po::value(&R)->default_value(64), "Maximum graph degree"); - desc.add_options()("Lbuild,L", po::value(&L)->default_value(100), - "Build complexity, higher value results in better graphs"); - desc.add_options()("alpha", po::value(&alpha)->default_value(1.2f), - "alpha controls density and diameter of graph, set " - "1 for sparse graph, " - "1.2 or 1.4 for denser graphs with lower diameter"); - desc.add_options()("num_threads,T", po::value(&num_threads)->default_value(omp_get_num_procs()), - "Number of threads used for building index (defaults to " - "omp_get_num_procs())"); - desc.add_options()("points_to_skip", po::value(&points_to_skip)->required(), - "Skip these first set of points from file"); - desc.add_options()("max_points_to_insert", po::value(&max_points_to_insert)->default_value(0), - "These number of points from the file are inserted after " - "points_to_skip"); - desc.add_options()("beginning_index_size", po::value(&beginning_index_size)->required(), - "Batch build will be called on these set of points"); - desc.add_options()("points_per_checkpoint", po::value(&points_per_checkpoint)->required(), - "Insertions are done in batches of points_per_checkpoint"); - desc.add_options()("checkpoints_per_snapshot", po::value(&checkpoints_per_snapshot)->required(), - "Save the index to disk every few checkpoints"); - desc.add_options()("points_to_delete_from_beginning", - po::value(&points_to_delete_from_beginning)->required(), ""); - desc.add_options()("do_concurrent", po::value(&concurrent)->default_value(false), ""); - desc.add_options()("start_deletes_after", po::value(&start_deletes_after)->default_value(0), ""); - desc.add_options()("start_point_norm", po::value(&start_point_norm)->default_value(0), - "Set the start point to a random point on a sphere of this radius"); - desc.add_options()("num_start_points", po::value(&num_start_pts)->default_value(0), - "Set the number of random start (frozen) points to use when " - "inserting and searching"); - - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) - { - std::cout << desc; - return 0; - } - po::notify(vm); - if (beginning_index_size == 0) - if (start_point_norm == 0) - { - std::cout << "When beginning_index_size is 0, use a start " - "point with " - "appropriate norm" - << std::endl; - return -1; - } +int main(int argc, char **argv) { + std::string data_type, dist_fn, data_path, index_path_prefix; + 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; + bool concurrent; + + po::options_description desc{"Arguments"}; + try { + desc.add_options()("help,h", "Print information on arguments"); + desc.add_options()("data_type", + po::value(&data_type)->required(), + "data type "); + desc.add_options()("dist_fn", po::value(&dist_fn)->required(), + "distance function "); + desc.add_options()("data_path", + po::value(&data_path)->required(), + "Input data file in bin format"); + desc.add_options()("index_path_prefix", + po::value(&index_path_prefix)->required(), + "Path prefix for saving index file components"); + desc.add_options()("max_degree,R", + po::value(&R)->default_value(64), + "Maximum graph degree"); + desc.add_options()( + "Lbuild,L", po::value(&L)->default_value(100), + "Build complexity, higher value results in better graphs"); + desc.add_options()("alpha", po::value(&alpha)->default_value(1.2f), + "alpha controls density and diameter of graph, set " + "1 for sparse graph, " + "1.2 or 1.4 for denser graphs with lower diameter"); + desc.add_options()( + "num_threads,T", + po::value(&num_threads)->default_value(omp_get_num_procs()), + "Number of threads used for building index (defaults to " + "omp_get_num_procs())"); + desc.add_options()("points_to_skip", + po::value(&points_to_skip)->required(), + "Skip these first set of points from file"); + desc.add_options()( + "max_points_to_insert", + po::value(&max_points_to_insert)->default_value(0), + "These number of points from the file are inserted after " + "points_to_skip"); + desc.add_options()("beginning_index_size", + po::value(&beginning_index_size)->required(), + "Batch build will be called on these set of points"); + desc.add_options()( + "points_per_checkpoint", + po::value(&points_per_checkpoint)->required(), + "Insertions are done in batches of points_per_checkpoint"); + desc.add_options()( + "checkpoints_per_snapshot", + po::value(&checkpoints_per_snapshot)->required(), + "Save the index to disk every few checkpoints"); + desc.add_options()( + "points_to_delete_from_beginning", + po::value(&points_to_delete_from_beginning)->required(), ""); + desc.add_options()("do_concurrent", + po::value(&concurrent)->default_value(false), ""); + desc.add_options()( + "start_deletes_after", + po::value(&start_deletes_after)->default_value(0), ""); + desc.add_options()( + "start_point_norm", + po::value(&start_point_norm)->default_value(0), + "Set the start point to a random point on a sphere of this radius"); + desc.add_options()( + "num_start_points", + po::value(&num_start_pts)->default_value(0), + "Set the number of random start (frozen) points to use when " + "inserting and searching"); + + po::variables_map vm; + po::store(po::parse_command_line(argc, argv, desc), vm); + if (vm.count("help")) { + std::cout << desc; + return 0; } - catch (const std::exception &ex) - { - std::cerr << ex.what() << '\n'; + po::notify(vm); + if (beginning_index_size == 0) + if (start_point_norm == 0) { + std::cout << "When beginning_index_size is 0, use a start " + "point with " + "appropriate norm" + << std::endl; return -1; - } - - try - { - if (data_type == std::string("int8")) - build_incremental_index(data_path, L, R, alpha, num_threads, 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); - else if (data_type == std::string("uint8")) - build_incremental_index(data_path, L, R, alpha, num_threads, 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); - else if (data_type == std::string("float")) - build_incremental_index(data_path, L, R, alpha, num_threads, 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); - else - std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; - } - catch (const std::exception &e) - { - std::cerr << "Caught exception: " << e.what() << std::endl; - exit(-1); - } - catch (...) - { - std::cerr << "Caught unknown exception" << std::endl; - exit(-1); - } - - return 0; + } + } catch (const std::exception &ex) { + std::cerr << ex.what() << '\n'; + return -1; + } + + try { + if (data_type == std::string("int8")) + build_incremental_index( + data_path, L, R, alpha, num_threads, 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); + else if (data_type == std::string("uint8")) + build_incremental_index( + data_path, L, R, alpha, num_threads, 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); + else if (data_type == std::string("float")) + build_incremental_index( + data_path, L, R, alpha, num_threads, 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); + else + std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; + } catch (const std::exception &e) { + std::cerr << "Caught exception: " << e.what() << std::endl; + exit(-1); + } catch (...) { + std::cerr << "Caught unknown exception" << std::endl; + exit(-1); + } + + return 0; } diff --git a/tests/test_streaming_scenario.cpp b/tests/test_streaming_scenario.cpp index e1fe80c83..ab61bb140 100644 --- a/tests/test_streaming_scenario.cpp +++ b/tests/test_streaming_scenario.cpp @@ -25,354 +25,378 @@ namespace po = boost::program_options; // load_aligned_bin modified to read pieces of the file, but using ifstream // instead of cached_ifstream. template -inline void load_aligned_bin_part(const std::string &bin_file, T *data, size_t offset_points, size_t points_to_read) -{ - std::ifstream reader; - reader.exceptions(std::ios::failbit | std::ios::badbit); - reader.open(bin_file, std::ios::binary | std::ios::ate); - size_t actual_file_size = reader.tellg(); - reader.seekg(0, std::ios::beg); - - int npts_i32, dim_i32; - reader.read((char *)&npts_i32, sizeof(int)); - reader.read((char *)&dim_i32, sizeof(int)); - size_t npts = (uint32_t)npts_i32; - size_t dim = (uint32_t)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; - std::cout << stream.str(); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - - if (offset_points + points_to_read > npts) - { - std::stringstream stream; - stream << "Error. Not enough points in file. Requested " << offset_points << " offset and " << points_to_read - << " points, but have only " << npts << " points" << std::endl; - std::cout << stream.str(); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - - reader.seekg(2 * sizeof(uint32_t) + offset_points * dim * sizeof(T)); - - const size_t rounded_dim = ROUND_UP(dim, 8); - - for (size_t i = 0; i < points_to_read; i++) - { - reader.read((char *)(data + i * rounded_dim), dim * sizeof(T)); - memset(data + i * rounded_dim + dim, 0, (rounded_dim - dim) * sizeof(T)); - } - reader.close(); +inline void load_aligned_bin_part(const std::string &bin_file, T *data, + size_t offset_points, size_t points_to_read) { + std::ifstream reader; + reader.exceptions(std::ios::failbit | std::ios::badbit); + reader.open(bin_file, std::ios::binary | std::ios::ate); + size_t actual_file_size = reader.tellg(); + reader.seekg(0, std::ios::beg); + + int npts_i32, dim_i32; + reader.read((char *)&npts_i32, sizeof(int)); + reader.read((char *)&dim_i32, sizeof(int)); + size_t npts = (uint32_t)npts_i32; + size_t dim = (uint32_t)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; + std::cout << stream.str(); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } + + if (offset_points + points_to_read > npts) { + std::stringstream stream; + stream << "Error. Not enough points in file. Requested " << offset_points + << " offset and " << points_to_read << " points, but have only " + << npts << " points" << std::endl; + std::cout << stream.str(); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } + + reader.seekg(2 * sizeof(uint32_t) + offset_points * dim * sizeof(T)); + + const size_t rounded_dim = ROUND_UP(dim, 8); + + for (size_t i = 0; i < points_to_read; i++) { + reader.read((char *)(data + i * rounded_dim), dim * sizeof(T)); + memset(data + i * rounded_dim + dim, 0, (rounded_dim - dim) * sizeof(T)); + } + reader.close(); } -std::string get_save_filename(const std::string &save_path, size_t active_window, size_t consolidate_interval, - size_t max_points_to_insert) -{ - std::string final_path = save_path; - final_path += "act" + std::to_string(active_window) + "-"; - final_path += "cons" + std::to_string(consolidate_interval) + "-"; - final_path += "max" + std::to_string(max_points_to_insert); - return final_path; +std::string get_save_filename(const std::string &save_path, + size_t active_window, size_t consolidate_interval, + size_t max_points_to_insert) { + std::string final_path = save_path; + final_path += "act" + std::to_string(active_window) + "-"; + final_path += "cons" + std::to_string(consolidate_interval) + "-"; + final_path += "max" + std::to_string(max_points_to_insert); + return final_path; } template -void insert_next_batch(diskann::Index &index, size_t start, size_t end, size_t insert_threads, T *data, - size_t aligned_dim) -{ - try - { - diskann::Timer insert_timer; - std::cout << std::endl << "Inserting from " << start << " to " << end << std::endl; - - size_t num_failed = 0; +void insert_next_batch(diskann::Index &index, size_t start, + size_t end, size_t insert_threads, T *data, + size_t aligned_dim) { + try { + diskann::Timer insert_timer; + std::cout << std::endl + << "Inserting from " << start << " to " << end << std::endl; + + size_t num_failed = 0; #pragma omp parallel for num_threads(insert_threads) schedule(dynamic) reduction(+ : num_failed) - for (int64_t j = start; j < (int64_t)end; j++) - { - if (index.insert_point(&data[(j - start) * aligned_dim], 1 + static_cast(j)) != 0) - { - std::cerr << "Insert failed " << j << std::endl; - num_failed++; - } - } - const double elapsedSeconds = insert_timer.elapsed() / 1000000.0; - std::cout << "Insertion time " << elapsedSeconds << " seconds (" << (end - start) / elapsedSeconds - << " points/second overall, " << (end - start) / elapsedSeconds / insert_threads << " per thread)" - << std::endl; - if (num_failed > 0) - std::cout << num_failed << " of " << end - start << "inserts failed" << std::endl; - } - catch (std::system_error &e) - { - std::cout << "Exiting after catching exception in insertion task: " << e.what() << std::endl; + for (int64_t j = start; j < (int64_t)end; j++) { + if (index.insert_point(&data[(j - start) * aligned_dim], + 1 + static_cast(j)) != 0) { + std::cerr << "Insert failed " << j << std::endl; + num_failed++; + } } + const double elapsedSeconds = insert_timer.elapsed() / 1000000.0; + std::cout << "Insertion time " << elapsedSeconds << " seconds (" + << (end - start) / elapsedSeconds << " points/second overall, " + << (end - start) / elapsedSeconds / insert_threads + << " per thread)" << std::endl; + if (num_failed > 0) + std::cout << num_failed << " of " << end - start << "inserts failed" + << std::endl; + } catch (std::system_error &e) { + std::cout << "Exiting after catching exception in insertion task: " + << e.what() << std::endl; + } } template -void delete_and_consolidate(diskann::Index &index, diskann::IndexWriteParameters &delete_params, - size_t start, size_t end) -{ - try - { - std::cout << std::endl << "Lazy deleting points " << start << " to " << end << "... "; - for (size_t i = start; i < end; ++i) - index.lazy_delete(1 + i); - std::cout << "lazy delete done." << std::endl; - - auto report = index.consolidate_deletes(delete_params); - while (report._status != diskann::consolidation_report::status_code::SUCCESS) - { - int wait_time = 5; - if (report._status == diskann::consolidation_report::status_code::LOCK_FAIL) - { - diskann::cerr << "Unable to acquire consolidate delete lock after " - << "deleting points " << start << " to " << end << ". Will retry in " << wait_time - << "seconds." << std::endl; - } - else if (report._status == diskann::consolidation_report::status_code::INCONSISTENT_COUNT_ERROR) - { - diskann::cerr << "Inconsistent counts in data structure. " - << "Will retry in " << wait_time << "seconds." << std::endl; - } - else - { - std::cerr << "Exiting after unknown error in consolidate delete" << std::endl; - exit(-1); - } - std::this_thread::sleep_for(std::chrono::seconds(wait_time)); - report = index.consolidate_deletes(delete_params); - } - auto points_processed = report._active_points + report._slots_released; - auto deletion_rate = points_processed / report._time; - std::cout << "#active points: " << report._active_points << std::endl - << "max points: " << report._max_points << std::endl - << "empty slots: " << report._empty_slots << std::endl - << "deletes processed: " << report._slots_released << std::endl - << "latest delete size: " << report._delete_set_size << std::endl - << "Deletion rate: " << deletion_rate << "/sec " - << "Deletion rate: " << deletion_rate / delete_params.num_threads << "/thread/sec " << std::endl; - } - catch (std::system_error &e) - { - std::cerr << "Exiting after catching exception in deletion task: " << e.what() << std::endl; +void delete_and_consolidate(diskann::Index &index, + diskann::IndexWriteParameters &delete_params, + size_t start, size_t end) { + try { + std::cout << std::endl + << "Lazy deleting points " << start << " to " << end << "... "; + for (size_t i = start; i < end; ++i) index.lazy_delete(1 + i); + std::cout << "lazy delete done." << std::endl; + + auto report = index.consolidate_deletes(delete_params); + while (report._status != + diskann::consolidation_report::status_code::SUCCESS) { + int wait_time = 5; + if (report._status == + diskann::consolidation_report::status_code::LOCK_FAIL) { + diskann::cerr << "Unable to acquire consolidate delete lock after " + << "deleting points " << start << " to " << end + << ". Will retry in " << wait_time << "seconds." + << std::endl; + } else if (report._status == diskann::consolidation_report::status_code:: + INCONSISTENT_COUNT_ERROR) { + diskann::cerr << "Inconsistent counts in data structure. " + << "Will retry in " << wait_time << "seconds." + << std::endl; + } else { + std::cerr << "Exiting after unknown error in consolidate delete" + << std::endl; exit(-1); + } + std::this_thread::sleep_for(std::chrono::seconds(wait_time)); + report = index.consolidate_deletes(delete_params); } + auto points_processed = report._active_points + report._slots_released; + auto deletion_rate = points_processed / report._time; + std::cout << "#active points: " << report._active_points << std::endl + << "max points: " << report._max_points << std::endl + << "empty slots: " << report._empty_slots << std::endl + << "deletes processed: " << report._slots_released << std::endl + << "latest delete size: " << report._delete_set_size << std::endl + << "Deletion rate: " << deletion_rate << "/sec " + << "Deletion rate: " << deletion_rate / delete_params.num_threads + << "/thread/sec " << std::endl; + } catch (std::system_error &e) { + std::cerr << "Exiting after catching exception in deletion task: " + << e.what() << std::endl; + exit(-1); + } } template -void build_incremental_index(const std::string &data_path, const uint32_t L, const uint32_t R, const float alpha, - const uint32_t insert_threads, const uint32_t consolidate_threads, - size_t max_points_to_insert, size_t active_window, size_t consolidate_interval, - const float start_point_norm, uint32_t num_start_pts, const std::string &save_path) -{ - const uint32_t C = 500; - const bool saturate_graph = false; - - diskann::IndexWriteParameters params = diskann::IndexWriteParametersBuilder(L, R) - .with_max_occlusion_size(C) - .with_alpha(alpha) - .with_saturate_graph(saturate_graph) - .with_num_rounds(1) - .with_num_threads(insert_threads) - .with_num_frozen_points(num_start_pts) - .build(); - - diskann::IndexWriteParameters delete_params = diskann::IndexWriteParametersBuilder(L, R) - .with_max_occlusion_size(C) - .with_alpha(alpha) - .with_saturate_graph(saturate_graph) - .with_num_rounds(1) - .with_num_threads(consolidate_threads) - .build(); - - size_t dim, aligned_dim; - size_t num_points; - - diskann::get_bin_metadata(data_path, num_points, dim); - diskann::cout << "metadata: file " << data_path << " has " << num_points << " points in " << dim << " dims" - << std::endl; - aligned_dim = ROUND_UP(dim, 8); - - if (max_points_to_insert == 0) - { - max_points_to_insert = num_points; - } - - if (num_points < max_points_to_insert) - throw diskann::ANNException(std::string("num_points(") + std::to_string(num_points) + - ") < max_points_to_insert(" + std::to_string(max_points_to_insert) + ")", - -1, __FUNCSIG__, __FILE__, __LINE__); - - if (max_points_to_insert < active_window + consolidate_interval) - throw diskann::ANNException("ERROR: max_points_to_insert < " - "active_window + consolidate_interval", - -1, __FUNCSIG__, __FILE__, __LINE__); - - if (consolidate_interval < max_points_to_insert / 1000) - throw diskann::ANNException("ERROR: consolidate_interval is too small", -1, __FUNCSIG__, __FILE__, __LINE__); - - using TagT = uint32_t; - using LabelT = uint32_t; - const bool enable_tags = true; - - diskann::Index index(diskann::L2, dim, active_window + 4 * consolidate_interval, true, params, L, - insert_threads, enable_tags, true); - index.set_start_points_at_random(static_cast(start_point_norm)); - index.enable_delete(); - - T *data = nullptr; - diskann::alloc_aligned((void **)&data, std::max(consolidate_interval, active_window) * aligned_dim * sizeof(T), - 8 * sizeof(T)); - - std::vector tags(max_points_to_insert); - std::iota(tags.begin(), tags.end(), static_cast(0)); - - diskann::Timer timer; - - std::vector> delete_tasks; - +void build_incremental_index(const std::string &data_path, const uint32_t L, + const uint32_t R, const float alpha, + const uint32_t insert_threads, + const uint32_t consolidate_threads, + size_t max_points_to_insert, size_t active_window, + size_t consolidate_interval, + const float start_point_norm, + uint32_t num_start_pts, + const std::string &save_path) { + const uint32_t C = 500; + const bool saturate_graph = false; + + diskann::IndexWriteParameters params = + diskann::IndexWriteParametersBuilder(L, R) + .with_max_occlusion_size(C) + .with_alpha(alpha) + .with_saturate_graph(saturate_graph) + .with_num_rounds(1) + .with_num_threads(insert_threads) + .with_num_frozen_points(num_start_pts) + .build(); + + diskann::IndexWriteParameters delete_params = + diskann::IndexWriteParametersBuilder(L, R) + .with_max_occlusion_size(C) + .with_alpha(alpha) + .with_saturate_graph(saturate_graph) + .with_num_rounds(1) + .with_num_threads(consolidate_threads) + .build(); + + size_t dim, aligned_dim; + size_t num_points; + + diskann::get_bin_metadata(data_path, num_points, dim); + diskann::cout << "metadata: file " << data_path << " has " << num_points + << " points in " << dim << " dims" << std::endl; + aligned_dim = ROUND_UP(dim, 8); + + if (max_points_to_insert == 0) { + max_points_to_insert = num_points; + } + + if (num_points < max_points_to_insert) + throw diskann::ANNException(std::string("num_points(") + + std::to_string(num_points) + + ") < max_points_to_insert(" + + std::to_string(max_points_to_insert) + ")", + -1, __FUNCSIG__, __FILE__, __LINE__); + + if (max_points_to_insert < active_window + consolidate_interval) + throw diskann::ANNException( + "ERROR: max_points_to_insert < " + "active_window + consolidate_interval", + -1, __FUNCSIG__, __FILE__, __LINE__); + + if (consolidate_interval < max_points_to_insert / 1000) + throw diskann::ANNException("ERROR: consolidate_interval is too small", -1, + __FUNCSIG__, __FILE__, __LINE__); + + using TagT = uint32_t; + using LabelT = uint32_t; + const bool enable_tags = true; + + diskann::Index index( + diskann::L2, dim, active_window + 4 * consolidate_interval, true, params, + L, insert_threads, enable_tags, true); + index.set_start_points_at_random(static_cast(start_point_norm)); + index.enable_delete(); + + T *data = nullptr; + diskann::alloc_aligned( + (void **)&data, + std::max(consolidate_interval, active_window) * aligned_dim * sizeof(T), + 8 * sizeof(T)); + + std::vector tags(max_points_to_insert); + std::iota(tags.begin(), tags.end(), static_cast(0)); + + diskann::Timer timer; + + std::vector> delete_tasks; + + auto insert_task = std::async(std::launch::async, [&]() { + load_aligned_bin_part(data_path, data, 0, active_window); + insert_next_batch(index, 0, active_window, insert_threads, data, + aligned_dim); + }); + insert_task.wait(); + + for (size_t start = active_window; + start + consolidate_interval <= max_points_to_insert; + start += consolidate_interval) { + auto end = std::min(start + consolidate_interval, max_points_to_insert); auto insert_task = std::async(std::launch::async, [&]() { - load_aligned_bin_part(data_path, data, 0, active_window); - insert_next_batch(index, 0, active_window, insert_threads, data, aligned_dim); + load_aligned_bin_part(data_path, data, start, end - start); + insert_next_batch(index, start, end, insert_threads, data, aligned_dim); }); insert_task.wait(); - for (size_t start = active_window; start + consolidate_interval <= max_points_to_insert; - start += consolidate_interval) - { - auto end = std::min(start + consolidate_interval, max_points_to_insert); - auto insert_task = std::async(std::launch::async, [&]() { - load_aligned_bin_part(data_path, data, start, end - start); - insert_next_batch(index, start, end, insert_threads, data, aligned_dim); - }); - insert_task.wait(); - - if (delete_tasks.size() > 0) - delete_tasks[delete_tasks.size() - 1].wait(); - if (start >= active_window + consolidate_interval) - { - auto start_del = start - active_window - consolidate_interval; - auto end_del = start - active_window; - - delete_tasks.emplace_back(std::async( - std::launch::async, [&]() { delete_and_consolidate(index, delete_params, start_del, end_del); })); - } + if (delete_tasks.size() > 0) delete_tasks[delete_tasks.size() - 1].wait(); + if (start >= active_window + consolidate_interval) { + auto start_del = start - active_window - consolidate_interval; + auto end_del = start - active_window; + + delete_tasks.emplace_back(std::async(std::launch::async, [&]() { + delete_and_consolidate(index, delete_params, start_del, end_del); + })); } - if (delete_tasks.size() > 0) - delete_tasks[delete_tasks.size() - 1].wait(); + } + if (delete_tasks.size() > 0) delete_tasks[delete_tasks.size() - 1].wait(); - std::cout << "Time Elapsed " << timer.elapsed() / 1000 << "ms\n"; - const auto save_path_inc = - get_save_filename(save_path + ".after-streaming-", active_window, consolidate_interval, max_points_to_insert); - index.save(save_path_inc.c_str(), true); + std::cout << "Time Elapsed " << timer.elapsed() / 1000 << "ms\n"; + const auto save_path_inc = + get_save_filename(save_path + ".after-streaming-", active_window, + consolidate_interval, max_points_to_insert); + index.save(save_path_inc.c_str(), true); - diskann::aligned_free(data); + diskann::aligned_free(data); } -int main(int argc, char **argv) -{ - std::string data_type, dist_fn, data_path, index_path_prefix; - uint32_t insert_threads, consolidate_threads; - uint32_t R, L, num_start_pts; - float alpha, start_point_norm; - size_t max_points_to_insert, active_window, consolidate_interval; - - po::options_description desc{"Arguments"}; - try - { - desc.add_options()("help,h", "Print information on arguments"); - desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); - desc.add_options()("dist_fn", po::value(&dist_fn)->required(), "distance function "); - desc.add_options()("data_path", po::value(&data_path)->required(), - "Input data file in bin format"); - desc.add_options()("index_path_prefix", po::value(&index_path_prefix)->required(), - "Path prefix for saving index file components"); - desc.add_options()("max_degree,R", po::value(&R)->default_value(64), "Maximum graph degree"); - desc.add_options()("Lbuild,L", po::value(&L)->default_value(100), - "Build complexity, higher value results in better graphs"); - desc.add_options()("alpha", po::value(&alpha)->default_value(1.2f), - "alpha controls density and diameter of graph, set " - "1 for sparse graph, " - "1.2 or 1.4 for denser graphs with lower diameter"); - desc.add_options()("insert_threads", - po::value(&insert_threads)->default_value(omp_get_num_procs() / 2), - "Number of threads used for inserting into the index (defaults to " - "omp_get_num_procs()/2)"); - desc.add_options()("consolidate_threads", - po::value(&consolidate_threads)->default_value(omp_get_num_procs() / 2), - "Number of threads used for consolidating deletes to " - "the index (defaults to omp_get_num_procs()/2)"); - - desc.add_options()("max_points_to_insert", po::value(&max_points_to_insert)->default_value(0), - "The number of points from the file that the program streams " - "over "); - desc.add_options()("active_window", po::value(&active_window)->required(), - "Program maintains an index over an active window of " - "this size that slides through the data"); - desc.add_options()("consolidate_interval", po::value(&consolidate_interval)->required(), - "The program simultaneously adds this number of points to the " - "right of " - "the window while deleting the same number from the left"); - desc.add_options()("start_point_norm", po::value(&start_point_norm)->required(), - "Set the start point to a random point on a sphere of this radius"); - desc.add_options()("num_start_points", po::value(&num_start_pts)->default_value(0), - "Set the number of random start (frozen) points to use when " - "inserting and searching"); - - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) - { - std::cout << desc; - return 0; - } - po::notify(vm); - if (start_point_norm == 0) - { - std::cout << "When beginning_index_size is 0, use a start point with " - "appropriate norm" - << std::endl; - return -1; - } - } - catch (const std::exception &ex) - { - std::cerr << ex.what() << '\n'; - return -1; - } - - try - { - if (data_type == std::string("int8")) - build_incremental_index(data_path, L, R, alpha, insert_threads, consolidate_threads, - max_points_to_insert, active_window, consolidate_interval, start_point_norm, - num_start_pts, index_path_prefix); - else if (data_type == std::string("uint8")) - build_incremental_index(data_path, L, R, alpha, insert_threads, consolidate_threads, - max_points_to_insert, active_window, consolidate_interval, - start_point_norm, num_start_pts, index_path_prefix); - else if (data_type == std::string("float")) - build_incremental_index(data_path, L, R, alpha, insert_threads, consolidate_threads, - max_points_to_insert, active_window, consolidate_interval, start_point_norm, - num_start_pts, index_path_prefix); - else - std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; +int main(int argc, char **argv) { + std::string data_type, dist_fn, data_path, index_path_prefix; + uint32_t insert_threads, consolidate_threads; + uint32_t R, L, num_start_pts; + float alpha, start_point_norm; + size_t max_points_to_insert, active_window, consolidate_interval; + + po::options_description desc{"Arguments"}; + try { + desc.add_options()("help,h", "Print information on arguments"); + desc.add_options()("data_type", + po::value(&data_type)->required(), + "data type "); + desc.add_options()("dist_fn", po::value(&dist_fn)->required(), + "distance function "); + desc.add_options()("data_path", + po::value(&data_path)->required(), + "Input data file in bin format"); + desc.add_options()("index_path_prefix", + po::value(&index_path_prefix)->required(), + "Path prefix for saving index file components"); + desc.add_options()("max_degree,R", + po::value(&R)->default_value(64), + "Maximum graph degree"); + desc.add_options()( + "Lbuild,L", po::value(&L)->default_value(100), + "Build complexity, higher value results in better graphs"); + desc.add_options()("alpha", po::value(&alpha)->default_value(1.2f), + "alpha controls density and diameter of graph, set " + "1 for sparse graph, " + "1.2 or 1.4 for denser graphs with lower diameter"); + desc.add_options()( + "insert_threads", + po::value(&insert_threads) + ->default_value(omp_get_num_procs() / 2), + "Number of threads used for inserting into the index (defaults to " + "omp_get_num_procs()/2)"); + desc.add_options()("consolidate_threads", + po::value(&consolidate_threads) + ->default_value(omp_get_num_procs() / 2), + "Number of threads used for consolidating deletes to " + "the index (defaults to omp_get_num_procs()/2)"); + + desc.add_options()( + "max_points_to_insert", + po::value(&max_points_to_insert)->default_value(0), + "The number of points from the file that the program streams " + "over "); + desc.add_options()("active_window", + po::value(&active_window)->required(), + "Program maintains an index over an active window of " + "this size that slides through the data"); + desc.add_options()( + "consolidate_interval", + po::value(&consolidate_interval)->required(), + "The program simultaneously adds this number of points to the " + "right of " + "the window while deleting the same number from the left"); + desc.add_options()( + "start_point_norm", po::value(&start_point_norm)->required(), + "Set the start point to a random point on a sphere of this radius"); + desc.add_options()( + "num_start_points", + po::value(&num_start_pts)->default_value(0), + "Set the number of random start (frozen) points to use when " + "inserting and searching"); + + po::variables_map vm; + po::store(po::parse_command_line(argc, argv, desc), vm); + if (vm.count("help")) { + std::cout << desc; + return 0; } - catch (const std::exception &e) - { - std::cerr << "Caught exception: " << e.what() << std::endl; - exit(-1); - } - catch (...) - { - std::cerr << "Caught unknown exception" << std::endl; - exit(-1); + po::notify(vm); + if (start_point_norm == 0) { + std::cout << "When beginning_index_size is 0, use a start point with " + "appropriate norm" + << std::endl; + return -1; } - - return 0; + } catch (const std::exception &ex) { + std::cerr << ex.what() << '\n'; + return -1; + } + + try { + if (data_type == std::string("int8")) + build_incremental_index( + data_path, L, R, alpha, insert_threads, consolidate_threads, + max_points_to_insert, active_window, consolidate_interval, + start_point_norm, num_start_pts, index_path_prefix); + else if (data_type == std::string("uint8")) + build_incremental_index( + data_path, L, R, alpha, insert_threads, consolidate_threads, + max_points_to_insert, active_window, consolidate_interval, + start_point_norm, num_start_pts, index_path_prefix); + else if (data_type == std::string("float")) + build_incremental_index( + data_path, L, R, alpha, insert_threads, consolidate_threads, + max_points_to_insert, active_window, consolidate_interval, + start_point_norm, num_start_pts, index_path_prefix); + else + std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; + } catch (const std::exception &e) { + std::cerr << "Caught exception: " << e.what() << std::endl; + exit(-1); + } catch (...) { + std::cerr << "Caught unknown exception" << std::endl; + exit(-1); + } + + return 0; } diff --git a/tests/utils/bin_to_fvecs.cpp b/tests/utils/bin_to_fvecs.cpp index e9a6a8ecc..f35038409 100644 --- a/tests/utils/bin_to_fvecs.cpp +++ b/tests/utils/bin_to_fvecs.cpp @@ -4,60 +4,58 @@ #include #include "util.h" -void block_convert(std::ifstream &writr, std::ofstream &readr, float *read_buf, float *write_buf, uint64_t npts, - uint64_t ndims) -{ - writr.write((char *)read_buf, npts * (ndims * sizeof(float) + sizeof(unsigned))); +void block_convert(std::ifstream &writr, std::ofstream &readr, float *read_buf, + float *write_buf, uint64_t npts, uint64_t ndims) { + writr.write((char *)read_buf, + npts * (ndims * sizeof(float) + sizeof(unsigned))); #pragma omp parallel for - for (uint64_t i = 0; i < npts; i++) - { - memcpy(write_buf + i * ndims, (read_buf + i * (ndims + 1)) + 1, ndims * sizeof(float)); - } - readr.read((char *)write_buf, npts * ndims * sizeof(float)); + for (uint64_t i = 0; i < npts; i++) { + memcpy(write_buf + i * ndims, (read_buf + i * (ndims + 1)) + 1, + ndims * sizeof(float)); + } + readr.read((char *)write_buf, npts * ndims * sizeof(float)); } -int main(int argc, char **argv) -{ - if (argc != 3) - { - std::cout << argv[0] << " input_bin output_fvecs" << std::endl; - exit(-1); - } - std::ifstream readr(argv[1], std::ios::binary); - int npts_s32; - int ndims_s32; - readr.read((char *)&npts_s32, sizeof(int32_t)); - readr.read((char *)&ndims_s32, sizeof(int32_t)); - size_t npts = npts_s32; - size_t ndims = ndims_s32; - uint32_t ndims_u32 = (uint32_t)ndims_s32; - // uint64_t fsize = writr.tellg(); - readr.seekg(0, std::ios::beg); - - unsigned ndims_u32; - writr.write((char *)&ndims_u32, sizeof(unsigned)); - writr.seekg(0, std::ios::beg); - uint64_t ndims = (uint64_t)ndims_u32; - uint64_t npts = fsize / ((ndims + 1) * sizeof(float)); - std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims << std::endl; - - uint64_t blk_size = 131072; - uint64_t nblks = ROUND_UP(npts, blk_size) / blk_size; - std::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]; - for (uint64_t i = 0; i < nblks; i++) - { - uint64_t 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; - } - - delete[] read_buf; - delete[] write_buf; - - writr.close(); - readr.close(); +int main(int argc, char **argv) { + if (argc != 3) { + std::cout << argv[0] << " input_bin output_fvecs" << std::endl; + exit(-1); + } + std::ifstream readr(argv[1], std::ios::binary); + int npts_s32; + int ndims_s32; + readr.read((char *)&npts_s32, sizeof(int32_t)); + readr.read((char *)&ndims_s32, sizeof(int32_t)); + size_t npts = npts_s32; + size_t ndims = ndims_s32; + uint32_t ndims_u32 = (uint32_t)ndims_s32; + // uint64_t fsize = writr.tellg(); + readr.seekg(0, std::ios::beg); + + unsigned ndims_u32; + writr.write((char *)&ndims_u32, sizeof(unsigned)); + writr.seekg(0, std::ios::beg); + uint64_t ndims = (uint64_t)ndims_u32; + uint64_t npts = fsize / ((ndims + 1) * sizeof(float)); + std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims + << std::endl; + + uint64_t blk_size = 131072; + uint64_t nblks = ROUND_UP(npts, blk_size) / blk_size; + std::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]; + for (uint64_t i = 0; i < nblks; i++) { + uint64_t 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; + } + + 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 7851bef6d..0fdb1a902 100644 --- a/tests/utils/bin_to_tsv.cpp +++ b/tests/utils/bin_to_tsv.cpp @@ -5,65 +5,64 @@ #include "utils.h" template -void block_convert(std::ofstream &writer, std::ifstream &reader, T *read_buf, size_t npts, size_t ndims) -{ - reader.read((char *)read_buf, npts * ndims * sizeof(float)); +void block_convert(std::ofstream &writer, std::ifstream &reader, T *read_buf, + size_t npts, size_t ndims) { + reader.read((char *)read_buf, npts * ndims * sizeof(float)); - for (size_t i = 0; i < npts; i++) - { - for (size_t d = 0; d < ndims; d++) - { - writer << read_buf[d + i * ndims]; - if (d < ndims - 1) - writer << "\t"; - else - writer << "\n"; - } + for (size_t i = 0; i < npts; i++) { + for (size_t d = 0; d < ndims; d++) { + writer << read_buf[d + i * ndims]; + if (d < ndims - 1) + writer << "\t"; + else + writer << "\n"; } + } } -int main(int argc, char **argv) -{ - if (argc != 4) - { - std::cout << argv[0] << " input_bin output_tsv" << std::endl; - exit(-1); - } - std::string type_string(argv[1]); - if ((type_string != std::string("float")) && (type_string != std::string("int8")) && - (type_string != std::string("uin8"))) - { - std::cerr << "Error: type not supported. Use float/int8/uint8" << std::endl; - } +int main(int argc, char **argv) { + if (argc != 4) { + std::cout << argv[0] << " input_bin output_tsv" + << std::endl; + exit(-1); + } + std::string type_string(argv[1]); + if ((type_string != std::string("float")) && + (type_string != std::string("int8")) && + (type_string != std::string("uin8"))) { + std::cerr << "Error: type not supported. Use float/int8/uint8" << std::endl; + } - std::ifstream reader(argv[2], std::ios::binary); - uint32_t npts_u32; - uint32_t ndims_u32; - reader.read((char *)&npts_u32, sizeof(uint32_t)); - reader.read((char *)&ndims_u32, sizeof(uint32_t)); - size_t npts = npts_u32; - size_t ndims = ndims_u32; - std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims << std::endl; + std::ifstream reader(argv[2], std::ios::binary); + uint32_t npts_u32; + uint32_t ndims_u32; + reader.read((char *)&npts_u32, sizeof(uint32_t)); + reader.read((char *)&ndims_u32, sizeof(uint32_t)); + size_t npts = npts_u32; + size_t ndims = ndims_u32; + std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims + << std::endl; - size_t blk_size = 131072; - size_t nblks = ROUND_UP(npts, blk_size) / blk_size; + size_t blk_size = 131072; + size_t nblks = ROUND_UP(npts, blk_size) / blk_size; - std::ofstream writer(argv[3]); - char *read_buf = new char[blk_size * ndims * 4]; - for (size_t i = 0; i < nblks; i++) - { - size_t 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); - else if (type_string == std::string("int8")) - 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); - std::cout << "Block #" << i << " written" << std::endl; - } + std::ofstream writer(argv[3]); + char *read_buf = new char[blk_size * ndims * 4]; + for (size_t i = 0; i < nblks; i++) { + size_t 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); + else if (type_string == std::string("int8")) + 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); + std::cout << "Block #" << i << " written" << std::endl; + } - delete[] read_buf; + delete[] read_buf; - writer.close(); - reader.close(); + writer.close(); + reader.close(); } diff --git a/tests/utils/calculate_recall.cpp b/tests/utils/calculate_recall.cpp index a45eb19d7..04d16c7e4 100644 --- a/tests/utils/calculate_recall.cpp +++ b/tests/utils/calculate_recall.cpp @@ -12,43 +12,42 @@ #include "utils.h" #include "disk_utils.h" -int main(int argc, char **argv) -{ - if (argc != 4) - { - std::cout << argv[0] << " " << std::endl; - return -1; - } - uint32_t *gold_std = NULL; - float *gs_dist = nullptr; - uint32_t *our_results = NULL; - float *or_dist = nullptr; - size_t points_num, points_num_gs, points_num_or; - size_t dim_gs; - size_t dim_or; - diskann::load_truthset(argv[1], gold_std, gs_dist, points_num_gs, dim_gs); - diskann::load_truthset(argv[2], our_results, or_dist, points_num_or, dim_or); +int main(int argc, char **argv) { + if (argc != 4) { + std::cout << argv[0] << " " + << std::endl; + return -1; + } + uint32_t *gold_std = NULL; + float *gs_dist = nullptr; + uint32_t *our_results = NULL; + float *or_dist = nullptr; + size_t points_num, points_num_gs, points_num_or; + size_t dim_gs; + size_t dim_or; + diskann::load_truthset(argv[1], gold_std, gs_dist, points_num_gs, dim_gs); + diskann::load_truthset(argv[2], our_results, or_dist, points_num_or, dim_or); - if (points_num_gs != points_num_or) - { - std::cout << "Error. Number of queries mismatch in ground truth and " - "our results" - << std::endl; - return -1; - } - points_num = points_num_gs; + if (points_num_gs != points_num_or) { + std::cout << "Error. Number of queries mismatch in ground truth and " + "our results" + << std::endl; + return -1; + } + points_num = points_num_gs; - uint32_t recall_at = std::atoi(argv[3]); + 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; - 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); + 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; + 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); - // double avg_recall = (recall*1.0)/(points_num*1.0); - std::cout << "Avg. recall@" << recall_at << " is " << recall_val << "\n"; + // double avg_recall = (recall*1.0)/(points_num*1.0); + std::cout << "Avg. recall@" << recall_at << " is " << recall_val << "\n"; } diff --git a/tests/utils/compute_groundtruth.cpp b/tests/utils/compute_groundtruth.cpp index ac448c09b..587cdd58a 100644 --- a/tests/utils/compute_groundtruth.cpp +++ b/tests/utils/compute_groundtruth.cpp @@ -40,532 +40,535 @@ typedef std::string path; namespace po = boost::program_options; -template T div_round_up(const T numerator, const T denominator) -{ - return (numerator % denominator == 0) ? (numerator / denominator) : 1 + (numerator / denominator); +template +T div_round_up(const T numerator, const T denominator) { + return (numerator % denominator == 0) ? (numerator / denominator) + : 1 + (numerator / denominator); } using pairIF = std::pair; -struct cmpmaxstruct -{ - bool operator()(const pairIF &l, const pairIF &r) - { - return l.second < r.second; - }; +struct cmpmaxstruct { + bool operator()(const pairIF &l, const pairIF &r) { + return l.second < r.second; + }; }; -using maxPQIFCS = std::priority_queue, cmpmaxstruct>; +using maxPQIFCS = + std::priority_queue, cmpmaxstruct>; -template T *aligned_malloc(const size_t n, const size_t alignment) -{ +template +T *aligned_malloc(const size_t n, const size_t alignment) { #ifdef _WINDOWS - return (T *)_aligned_malloc(sizeof(T) * n, alignment); + return (T *)_aligned_malloc(sizeof(T) * n, alignment); #else - return static_cast(aligned_alloc(alignment, sizeof(T) * n)); + return static_cast(aligned_alloc(alignment, sizeof(T) * n)); #endif } -inline bool custom_dist(const std::pair &a, const std::pair &b) -{ - return a.second < b.second; +inline bool custom_dist(const std::pair &a, + const std::pair &b) { + return a.second < b.second; } -void compute_l2sq(float *const points_l2sq, const float *const matrix, const int64_t num_points, const int dim) -{ - assert(points_l2sq != NULL); +void compute_l2sq(float *const points_l2sq, const float *const matrix, + const int64_t num_points, const int dim) { + assert(points_l2sq != NULL); #pragma omp parallel for schedule(static, 65536) - for (int64_t d = 0; d < num_points; ++d) - points_l2sq[d] = - cblas_sdot(dim, matrix + (ptrdiff_t)d * (ptrdiff_t)dim, 1, matrix + (ptrdiff_t)d * (ptrdiff_t)dim, 1); + for (int64_t d = 0; d < num_points; ++d) + points_l2sq[d] = cblas_sdot(dim, matrix + (ptrdiff_t)d * (ptrdiff_t)dim, 1, + matrix + (ptrdiff_t)d * (ptrdiff_t)dim, 1); } -void distsq_to_points(const size_t dim, - float *dist_matrix, // Col Major, cols are queries, rows are points - size_t npoints, const float *const points, - const float *const points_l2sq, // points in Col major - size_t nqueries, const float *const queries, - const float *const queries_l2sq, // queries in Col major - float *ones_vec = NULL) // Scratchspace of num_data size and init to 1.0 +void distsq_to_points( + const size_t dim, + float *dist_matrix, // Col Major, cols are queries, rows are points + size_t npoints, const float *const points, + const float *const points_l2sq, // points in Col major + size_t nqueries, const float *const queries, + const float *const queries_l2sq, // queries in Col major + float *ones_vec = NULL) // Scratchspace of num_data size and init to 1.0 { - bool ones_vec_alloc = false; - if (ones_vec == NULL) - { - ones_vec = new float[nqueries > npoints ? nqueries : npoints]; - std::fill_n(ones_vec, nqueries > npoints ? nqueries : npoints, (float)1.0); - ones_vec_alloc = true; - } - cblas_sgemm(CblasColMajor, CblasTrans, CblasNoTrans, npoints, nqueries, dim, (float)-2.0, points, dim, queries, dim, - (float)0.0, dist_matrix, npoints); - cblas_sgemm(CblasColMajor, CblasNoTrans, CblasTrans, npoints, nqueries, 1, (float)1.0, points_l2sq, npoints, - ones_vec, nqueries, (float)1.0, dist_matrix, npoints); - cblas_sgemm(CblasColMajor, CblasNoTrans, CblasTrans, npoints, nqueries, 1, (float)1.0, ones_vec, npoints, - queries_l2sq, nqueries, (float)1.0, dist_matrix, npoints); - if (ones_vec_alloc) - delete[] ones_vec; + bool ones_vec_alloc = false; + if (ones_vec == NULL) { + ones_vec = new float[nqueries > npoints ? nqueries : npoints]; + std::fill_n(ones_vec, nqueries > npoints ? nqueries : npoints, (float)1.0); + ones_vec_alloc = true; + } + cblas_sgemm(CblasColMajor, CblasTrans, CblasNoTrans, npoints, nqueries, dim, + (float)-2.0, points, dim, queries, dim, (float)0.0, dist_matrix, + npoints); + cblas_sgemm(CblasColMajor, CblasNoTrans, CblasTrans, npoints, nqueries, 1, + (float)1.0, points_l2sq, npoints, ones_vec, nqueries, (float)1.0, + dist_matrix, npoints); + cblas_sgemm(CblasColMajor, CblasNoTrans, CblasTrans, npoints, nqueries, 1, + (float)1.0, ones_vec, npoints, queries_l2sq, nqueries, (float)1.0, + dist_matrix, npoints); + if (ones_vec_alloc) delete[] ones_vec; } -void inner_prod_to_points(const size_t dim, - float *dist_matrix, // Col Major, cols are queries, rows are points - size_t npoints, const float *const points, size_t nqueries, const float *const queries, - float *ones_vec = NULL) // Scratchspace of num_data size and init to 1.0 +void inner_prod_to_points( + const size_t dim, + float *dist_matrix, // Col Major, cols are queries, rows are points + size_t npoints, const float *const points, size_t nqueries, + const float *const queries, + float *ones_vec = NULL) // Scratchspace of num_data size and init to 1.0 { - bool ones_vec_alloc = false; - if (ones_vec == NULL) - { - ones_vec = new float[nqueries > npoints ? nqueries : npoints]; - std::fill_n(ones_vec, nqueries > npoints ? nqueries : npoints, (float)1.0); - ones_vec_alloc = true; - } - cblas_sgemm(CblasColMajor, CblasTrans, CblasNoTrans, npoints, nqueries, dim, (float)-1.0, points, dim, queries, dim, - (float)0.0, dist_matrix, npoints); - - if (ones_vec_alloc) - delete[] ones_vec; + bool ones_vec_alloc = false; + if (ones_vec == NULL) { + ones_vec = new float[nqueries > npoints ? nqueries : npoints]; + std::fill_n(ones_vec, nqueries > npoints ? nqueries : npoints, (float)1.0); + ones_vec_alloc = true; + } + cblas_sgemm(CblasColMajor, CblasTrans, CblasNoTrans, npoints, nqueries, dim, + (float)-1.0, points, dim, queries, dim, (float)0.0, dist_matrix, + npoints); + + if (ones_vec_alloc) delete[] ones_vec; } -void exact_knn(const size_t dim, const size_t k, - int *const closest_points, // k * num_queries preallocated, col - // major, queries columns - float *const dist_closest_points, // k * num_queries - // preallocated, Dist to - // corresponding closes_points - size_t npoints, - float *points_in, // points in Col major - size_t nqueries, float *queries_in, - diskann::Metric metric = diskann::Metric::L2) // queries in Col major +void exact_knn( + const size_t dim, const size_t k, + int *const closest_points, // k * num_queries preallocated, col + // major, queries columns + float *const dist_closest_points, // k * num_queries + // preallocated, Dist to + // corresponding closes_points + size_t npoints, + float *points_in, // points in Col major + size_t nqueries, float *queries_in, + diskann::Metric metric = diskann::Metric::L2) // queries in Col major { - float *points_l2sq = new float[npoints]; - float *queries_l2sq = new float[nqueries]; - compute_l2sq(points_l2sq, points_in, npoints, dim); - compute_l2sq(queries_l2sq, queries_in, nqueries, dim); - - float *points = points_in; - float *queries = queries_in; - - if (metric == diskann::Metric::COSINE) - { // we convert cosine distance as - // normalized L2 distnace - points = new float[npoints * dim]; - queries = new float[nqueries * dim]; + float *points_l2sq = new float[npoints]; + float *queries_l2sq = new float[nqueries]; + compute_l2sq(points_l2sq, points_in, npoints, dim); + compute_l2sq(queries_l2sq, queries_in, nqueries, dim); + + float *points = points_in; + float *queries = queries_in; + + if (metric == diskann::Metric::COSINE) { // we convert cosine distance as + // normalized L2 distnace + points = new float[npoints * dim]; + queries = new float[nqueries * dim]; #pragma omp parallel for schedule(static, 4096) - for (int64_t i = 0; i < (int64_t)npoints; i++) - { - float norm = std::sqrt(points_l2sq[i]); - if (norm == 0) - { - norm = std::numeric_limits::epsilon(); - } - for (uint32_t j = 0; j < dim; j++) - { - points[i * dim + j] = points_in[i * dim + j] / norm; - } - } + for (int64_t i = 0; i < (int64_t)npoints; i++) { + float norm = std::sqrt(points_l2sq[i]); + if (norm == 0) { + norm = std::numeric_limits::epsilon(); + } + for (uint32_t j = 0; j < dim; j++) { + points[i * dim + j] = points_in[i * dim + j] / norm; + } + } #pragma omp parallel for schedule(static, 4096) - for (int64_t i = 0; i < (int64_t)nqueries; i++) - { - float norm = std::sqrt(queries_l2sq[i]); - if (norm == 0) - { - norm = std::numeric_limits::epsilon(); - } - for (uint32_t j = 0; j < dim; j++) - { - queries[i * dim + j] = queries_in[i * dim + j] / norm; - } - } - // recalculate norms after normalizing, they should all be one. - compute_l2sq(points_l2sq, points, npoints, dim); - compute_l2sq(queries_l2sq, queries, nqueries, dim); + for (int64_t i = 0; i < (int64_t)nqueries; i++) { + float norm = std::sqrt(queries_l2sq[i]); + if (norm == 0) { + norm = std::numeric_limits::epsilon(); + } + for (uint32_t j = 0; j < dim; j++) { + queries[i * dim + j] = queries_in[i * dim + j] / norm; + } } - - std::cout << "Going to compute " << k << " NNs for " << nqueries << " queries over " << npoints << " points in " - << dim << " dimensions using"; - if (metric == diskann::Metric::INNER_PRODUCT) - std::cout << " MIPS "; - else if (metric == diskann::Metric::COSINE) - std::cout << " Cosine "; - else - std::cout << " L2 "; - std::cout << "distance fn. " << std::endl; - - size_t q_batch_size = (1 << 9); - float *dist_matrix = new float[(size_t)q_batch_size * (size_t)npoints]; - - for (size_t b = 0; b < div_round_up(nqueries, q_batch_size); ++b) - { - int64_t q_b = b * q_batch_size; - int64_t q_e = ((b + 1) * q_batch_size > nqueries) ? nqueries : (b + 1) * q_batch_size; - - if (metric == diskann::Metric::L2 || metric == diskann::Metric::COSINE) - { - 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); - } - else - { - inner_prod_to_points(dim, dist_matrix, npoints, points, q_e - q_b, - queries + (ptrdiff_t)q_b * (ptrdiff_t)dim); - } - std::cout << "Computed distances for queries: [" << q_b << "," << q_e << ")" << std::endl; + // recalculate norms after normalizing, they should all be one. + compute_l2sq(points_l2sq, points, npoints, dim); + compute_l2sq(queries_l2sq, queries, nqueries, dim); + } + + std::cout << "Going to compute " << k << " NNs for " << nqueries + << " queries over " << npoints << " points in " << dim + << " dimensions using"; + if (metric == diskann::Metric::INNER_PRODUCT) + std::cout << " MIPS "; + else if (metric == diskann::Metric::COSINE) + std::cout << " Cosine "; + else + std::cout << " L2 "; + std::cout << "distance fn. " << std::endl; + + size_t q_batch_size = (1 << 9); + float *dist_matrix = new float[(size_t)q_batch_size * (size_t)npoints]; + + for (size_t b = 0; b < div_round_up(nqueries, q_batch_size); ++b) { + int64_t q_b = b * q_batch_size; + int64_t q_e = + ((b + 1) * q_batch_size > nqueries) ? nqueries : (b + 1) * q_batch_size; + + if (metric == diskann::Metric::L2 || metric == diskann::Metric::COSINE) { + 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); + } else { + inner_prod_to_points(dim, dist_matrix, npoints, points, q_e - q_b, + queries + (ptrdiff_t)q_b * (ptrdiff_t)dim); + } + std::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++) - { - maxPQIFCS point_dist; - for (size_t p = 0; p < k; p++) - point_dist.emplace(p, dist_matrix[(ptrdiff_t)p + (ptrdiff_t)(q - q_b) * (ptrdiff_t)npoints]); - for (size_t p = k; p < npoints; p++) - { - if (point_dist.top().second > dist_matrix[(ptrdiff_t)p + (ptrdiff_t)(q - q_b) * (ptrdiff_t)npoints]) - point_dist.emplace(p, dist_matrix[(ptrdiff_t)p + (ptrdiff_t)(q - q_b) * (ptrdiff_t)npoints]); - if (point_dist.size() > k) - point_dist.pop(); - } - for (ptrdiff_t l = 0; l < (ptrdiff_t)k; ++l) - { - closest_points[(ptrdiff_t)(k - 1 - l) + (ptrdiff_t)q * (ptrdiff_t)k] = point_dist.top().first; - dist_closest_points[(ptrdiff_t)(k - 1 - l) + (ptrdiff_t)q * (ptrdiff_t)k] = point_dist.top().second; - point_dist.pop(); - } - assert(std::is_sorted(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; + for (long long q = q_b; q < q_e; q++) { + maxPQIFCS point_dist; + for (size_t p = 0; p < k; p++) + point_dist.emplace( + p, dist_matrix[(ptrdiff_t)p + + (ptrdiff_t)(q - q_b) * (ptrdiff_t)npoints]); + for (size_t p = k; p < npoints; p++) { + if (point_dist.top().second > + dist_matrix[(ptrdiff_t)p + + (ptrdiff_t)(q - q_b) * (ptrdiff_t)npoints]) + point_dist.emplace( + p, dist_matrix[(ptrdiff_t)p + + (ptrdiff_t)(q - q_b) * (ptrdiff_t)npoints]); + if (point_dist.size() > k) point_dist.pop(); + } + for (ptrdiff_t l = 0; l < (ptrdiff_t)k; ++l) { + closest_points[(ptrdiff_t)(k - 1 - l) + (ptrdiff_t)q * (ptrdiff_t)k] = + point_dist.top().first; + dist_closest_points[(ptrdiff_t)(k - 1 - l) + + (ptrdiff_t)q * (ptrdiff_t)k] = + point_dist.top().second; + point_dist.pop(); + } + assert(std::is_sorted( + 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; + } - delete[] dist_matrix; + delete[] dist_matrix; - delete[] points_l2sq; - delete[] queries_l2sq; + delete[] points_l2sq; + delete[] queries_l2sq; - if (metric == diskann::Metric::COSINE) - { - delete[] points; - delete[] queries; - } + if (metric == diskann::Metric::COSINE) { + delete[] points; + delete[] queries; + } } -template inline int get_num_parts(const char *filename) -{ - std::ifstream reader; - reader.exceptions(std::ios::failbit | std::ios::badbit); - reader.open(filename, std::ios::binary); - std::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; - 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; - return num_parts; +template +inline int get_num_parts(const char *filename) { + std::ifstream reader; + reader.exceptions(std::ios::failbit | std::ios::badbit); + reader.open(filename, std::ios::binary); + std::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; + 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; + return num_parts; } template -inline void load_bin_as_float(const char *filename, float *&data, size_t &npts, size_t &ndims, int part_num) -{ - std::ifstream reader; - reader.exceptions(std::ios::failbit | std::ios::badbit); - reader.open(filename, std::ios::binary); - std::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)); - uint64_t start_id = part_num * PARTSIZE; - uint64_t end_id = (std::min)(start_id + PARTSIZE, (uint64_t)npts_i32); - npts = end_id - start_id; - ndims = (uint64_t)ndims_i32; - std::cout << "#pts in part = " << npts << ", #dims = " << ndims << ", size = " << npts * ndims * sizeof(T) << "B" - << std::endl; - - reader.seekg(start_id * ndims * sizeof(T) + 2 * sizeof(uint32_t), std::ios::beg); - T *data_T = new T[npts * ndims]; - reader.read((char *)data_T, sizeof(T) * npts * ndims); - std::cout << "Finished reading part of the bin file." << std::endl; - reader.close(); - data = aligned_malloc(npts * ndims, ALIGNMENT); +inline void load_bin_as_float(const char *filename, float *&data, size_t &npts, + size_t &ndims, int part_num) { + std::ifstream reader; + reader.exceptions(std::ios::failbit | std::ios::badbit); + reader.open(filename, std::ios::binary); + std::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)); + uint64_t start_id = part_num * PARTSIZE; + uint64_t end_id = (std::min)(start_id + PARTSIZE, (uint64_t)npts_i32); + npts = end_id - start_id; + ndims = (uint64_t)ndims_i32; + std::cout << "#pts in part = " << npts << ", #dims = " << ndims + << ", size = " << npts * ndims * sizeof(T) << "B" << std::endl; + + reader.seekg(start_id * ndims * sizeof(T) + 2 * sizeof(uint32_t), + std::ios::beg); + T *data_T = new T[npts * ndims]; + reader.read((char *)data_T, sizeof(T) * npts * ndims); + std::cout << "Finished reading part of the bin file." << std::endl; + reader.close(); + data = aligned_malloc(npts * ndims, ALIGNMENT); #pragma omp parallel for schedule(dynamic, 32768) - for (int64_t i = 0; i < (int64_t)npts; i++) - { - for (int64_t j = 0; j < (int64_t)ndims; j++) - { - float cur_val_float = (float)data_T[i * ndims + j]; - std::memcpy((char *)(data + i * ndims + j), (char *)&cur_val_float, sizeof(float)); - } + for (int64_t i = 0; i < (int64_t)npts; i++) { + for (int64_t j = 0; j < (int64_t)ndims; j++) { + float cur_val_float = (float)data_T[i * ndims + j]; + std::memcpy((char *)(data + i * ndims + j), (char *)&cur_val_float, + sizeof(float)); } - delete[] data_T; - std::cout << "Finished converting part data to float." << std::endl; + } + delete[] data_T; + std::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; - writer.exceptions(std::ios::failbit | std::ios::badbit); - writer.open(filename, std::ios::binary | std::ios::out); - std::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; - - writer.write((char *)data, npts * ndims * sizeof(T)); - writer.close(); - std::cout << "Finished writing bin" << std::endl; +template +inline void save_bin(const std::string filename, T *data, size_t npts, + size_t ndims) { + std::ofstream writer; + writer.exceptions(std::ios::failbit | std::ios::badbit); + writer.open(filename, std::ios::binary | std::ios::out); + std::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; + + writer.write((char *)data, npts * ndims * sizeof(T)); + writer.close(); + std::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) -{ - 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(uint32_t) + 2 * sizeof(int) - << "B" << std::endl; - - writer.write((char *)data, npts * ndims * sizeof(uint32_t)); - writer.write((char *)distances, npts * ndims * sizeof(float)); - writer.close(); - std::cout << "Finished writing truthset" << std::endl; +inline void save_groundtruth_as_one_file(const std::string filename, + int32_t *data, float *distances, + size_t npts, size_t ndims) { + 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(uint32_t) + 2 * sizeof(int) << "B" + << std::endl; + + writer.write((char *)data, npts * ndims * sizeof(uint32_t)); + writer.write((char *)distances, npts * ndims * sizeof(float)); + writer.close(); + std::cout << "Finished writing truthset" << std::endl; } template -std::vector>> processUnfilteredParts(const std::string &base_file, - size_t &nqueries, size_t &npoints, - size_t &dim, size_t &k, float *query_data, - const diskann::Metric &metric, - std::vector &location_to_tag) -{ - float *base_data; - int num_parts = get_num_parts(base_file.c_str()); - std::vector>> res(nqueries); - for (int p = 0; p < num_parts; p++) - { - size_t start_id = p * PARTSIZE; - load_bin_as_float(base_file.c_str(), base_data, npoints, dim, p); - - int *closest_points_part = new int[nqueries * k]; - float *dist_closest_points_part = new float[nqueries * k]; - - uint32_t part_k; - part_k = k < npoints ? k : npoints; - exact_knn(dim, part_k, closest_points_part, dist_closest_points_part, npoints, base_data, nqueries, query_data, - metric); - - for (size_t i = 0; i < nqueries; i++) - { - for (size_t j = 0; j < part_k; j++) - { - if (!location_to_tag.empty()) - if (location_to_tag[closest_points_part[i * k + j] + start_id] == 0) - continue; - - res[i].push_back(std::make_pair((uint32_t)(closest_points_part[i * part_k + j] + start_id), - dist_closest_points_part[i * part_k + j])); - } - } - - delete[] closest_points_part; - delete[] dist_closest_points_part; - - diskann::aligned_free(base_data); - } - return res; -}; - -template -int aux_main(const std::string &base_file, const std::string &query_file, const std::string >_file, size_t k, - const diskann::Metric &metric, const std::string &tags_file = std::string("")) -{ - size_t npoints, nqueries, dim; - - float *query_data; - - load_bin_as_float(query_file.c_str(), query_data, nqueries, dim, 0); - if (nqueries > PARTSIZE) - std::cerr << "WARNING: #Queries provided (" << nqueries << ") is greater than " << PARTSIZE - << ". Computing GT only for the first " << PARTSIZE << " queries." << std::endl; - - // load tags - const bool tags_enabled = tags_file.empty() ? false : true; - std::vector location_to_tag = diskann::loadTags(tags_file, base_file); - - int *closest_points = new int[nqueries * k]; - float *dist_closest_points = new float[nqueries * k]; - - std::vector>> results = - processUnfilteredParts(base_file, nqueries, npoints, dim, k, query_data, metric, location_to_tag); - - for (size_t i = 0; i < nqueries; i++) - { - std::vector> &cur_res = results[i]; - std::sort(cur_res.begin(), cur_res.end(), custom_dist); - size_t j = 0; - for (auto iter : cur_res) - { - if (j == k) - break; - if (tags_enabled) - { - std::uint32_t index_with_tag = location_to_tag[iter.first]; - closest_points[i * k + j] = (int32_t)index_with_tag; - } - else - { - closest_points[i * k + j] = (int32_t)iter.first; - } - - if (metric == diskann::Metric::INNER_PRODUCT) - dist_closest_points[i * k + j] = -iter.second; - else - dist_closest_points[i * k + j] = iter.second; - - ++j; - } - if (j < k) - std::cout << "WARNING: found less than k GT entries for query " << i << std::endl; +std::vector>> processUnfilteredParts( + const std::string &base_file, size_t &nqueries, size_t &npoints, + size_t &dim, size_t &k, float *query_data, const diskann::Metric &metric, + std::vector &location_to_tag) { + float *base_data; + int num_parts = get_num_parts(base_file.c_str()); + std::vector>> res(nqueries); + for (int p = 0; p < num_parts; p++) { + size_t start_id = p * PARTSIZE; + load_bin_as_float(base_file.c_str(), base_data, npoints, dim, p); + + int *closest_points_part = new int[nqueries * k]; + float *dist_closest_points_part = new float[nqueries * k]; + + uint32_t part_k; + part_k = k < npoints ? k : npoints; + exact_knn(dim, part_k, closest_points_part, dist_closest_points_part, + npoints, base_data, nqueries, query_data, metric); + + for (size_t i = 0; i < nqueries; i++) { + for (size_t j = 0; j < part_k; j++) { + if (!location_to_tag.empty()) + if (location_to_tag[closest_points_part[i * k + j] + start_id] == 0) + continue; + + res[i].push_back(std::make_pair( + (uint32_t)(closest_points_part[i * part_k + j] + start_id), + dist_closest_points_part[i * part_k + j])); + } } - save_groundtruth_as_one_file(gt_file, closest_points, dist_closest_points, nqueries, k); - delete[] closest_points; - delete[] dist_closest_points; - diskann::aligned_free(query_data); + delete[] closest_points_part; + delete[] dist_closest_points_part; - return 0; -} - -void load_truthset(const std::string &bin_file, uint32_t *&ids, float *&dists, size_t &npts, size_t &dim) -{ - size_t read_blk_size = 64 * 1024 * 1024; - cached_ifstream reader(bin_file, read_blk_size); - diskann::cout << "Reading truthset file " << bin_file.c_str() << " ..." << std::endl; - size_t actual_file_size = reader.get_file_size(); - - int npts_i32, dim_i32; - reader.read((char *)&npts_i32, sizeof(int)); - reader.read((char *)&dim_i32, sizeof(int)); - npts = (uint32_t)npts_i32; - dim = (uint32_t)dim_i32; - - diskann::cout << "Metadata: #pts = " << npts << ", #dims = " << dim << "... " << std::endl; - - int truthset_type = -1; // 1 means truthset has ids and distances, 2 means - // only ids, -1 is error - size_t expected_file_size_with_dists = 2 * npts * dim * sizeof(uint32_t) + 2 * sizeof(uint32_t); - - if (actual_file_size == expected_file_size_with_dists) - truthset_type = 1; - - size_t expected_file_size_just_ids = npts * dim * sizeof(uint32_t) + 2 * sizeof(uint32_t); - - if (actual_file_size == expected_file_size_just_ids) - truthset_type = 2; - - if (truthset_type == -1) - { - std::stringstream stream; - stream << "Error. File size mismatch. File should have bin format, with " - "npts followed by ngt followed by npts*ngt ids and optionally " - "followed by npts*ngt distance values; actual size: " - << actual_file_size << ", expected: " << expected_file_size_with_dists << " or " - << expected_file_size_just_ids; - diskann::cout << stream.str(); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } + diskann::aligned_free(base_data); + } + return res; +}; - ids = new uint32_t[npts * dim]; - reader.read((char *)ids, npts * dim * sizeof(uint32_t)); +template +int aux_main(const std::string &base_file, const std::string &query_file, + const std::string >_file, size_t k, + const diskann::Metric &metric, + const std::string &tags_file = std::string("")) { + size_t npoints, nqueries, dim; + + float *query_data; + + load_bin_as_float(query_file.c_str(), query_data, nqueries, dim, 0); + if (nqueries > PARTSIZE) + std::cerr << "WARNING: #Queries provided (" << nqueries + << ") is greater than " << PARTSIZE + << ". Computing GT only for the first " << PARTSIZE << " queries." + << std::endl; - if (truthset_type == 1) - { - dists = new float[npts * dim]; - reader.read((char *)dists, npts * dim * sizeof(float)); + // load tags + const bool tags_enabled = tags_file.empty() ? false : true; + std::vector location_to_tag = + diskann::loadTags(tags_file, base_file); + + int *closest_points = new int[nqueries * k]; + float *dist_closest_points = new float[nqueries * k]; + + std::vector>> results = + processUnfilteredParts(base_file, nqueries, npoints, dim, k, + query_data, metric, location_to_tag); + + for (size_t i = 0; i < nqueries; i++) { + std::vector> &cur_res = results[i]; + std::sort(cur_res.begin(), cur_res.end(), custom_dist); + size_t j = 0; + for (auto iter : cur_res) { + if (j == k) break; + if (tags_enabled) { + std::uint32_t index_with_tag = location_to_tag[iter.first]; + closest_points[i * k + j] = (int32_t)index_with_tag; + } else { + closest_points[i * k + j] = (int32_t)iter.first; + } + + if (metric == diskann::Metric::INNER_PRODUCT) + dist_closest_points[i * k + j] = -iter.second; + else + dist_closest_points[i * k + j] = iter.second; + + ++j; } + if (j < k) + std::cout << "WARNING: found less than k GT entries for query " << i + << std::endl; + } + + save_groundtruth_as_one_file(gt_file, closest_points, dist_closest_points, + nqueries, k); + delete[] closest_points; + delete[] dist_closest_points; + diskann::aligned_free(query_data); + + return 0; } -int main(int argc, char **argv) -{ - std::string data_type, dist_fn, base_file, query_file, gt_file, tags_file; - uint64_t K; - - try - { - po::options_description desc{"Arguments"}; - - desc.add_options()("help,h", "Print information on arguments"); - - desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); - desc.add_options()("dist_fn", po::value(&dist_fn)->required(), "distance function "); - desc.add_options()("base_file", po::value(&base_file)->required(), - "File containing the base vectors in binary format"); - desc.add_options()("query_file", po::value(&query_file)->required(), - "File containing the query vectors in binary format"); - desc.add_options()("gt_file", po::value(>_file)->required(), - "File name for the writing ground truth in binary format, please don' append .bin at end if " - "no filter_label or filter_label_file is provided it will save the file with '.bin' at end." - "else it will save the file as filename_label.bin"); - desc.add_options()("K", po::value(&K)->required(), - "Number of ground truth nearest neighbors to compute"); - desc.add_options()("tags_file", po::value(&tags_file)->default_value(std::string()), - "File containing the tags in binary format"); - - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) - { - std::cout << desc; - return 0; - } - po::notify(vm); - } - catch (const std::exception &ex) - { - std::cerr << ex.what() << '\n'; - return -1; - } - - if (data_type != std::string("float") && data_type != std::string("int8") && data_type != std::string("uint8")) - { - std::cout << "Unsupported type. float, int8 and uint8 types are supported." << std::endl; - return -1; - } - - diskann::Metric metric; - if (dist_fn == std::string("l2")) - { - metric = diskann::Metric::L2; - } - else if (dist_fn == std::string("mips")) - { - metric = diskann::Metric::INNER_PRODUCT; - } - else if (dist_fn == std::string("cosine")) - { - metric = diskann::Metric::COSINE; - } - else - { - std::cerr << "Unsupported distance function. Use l2/mips/cosine." << std::endl; - return -1; - } +void load_truthset(const std::string &bin_file, uint32_t *&ids, float *&dists, + size_t &npts, size_t &dim) { + size_t read_blk_size = 64 * 1024 * 1024; + cached_ifstream reader(bin_file, read_blk_size); + diskann::cout << "Reading truthset file " << bin_file.c_str() << " ..." + << std::endl; + size_t actual_file_size = reader.get_file_size(); + + int npts_i32, dim_i32; + reader.read((char *)&npts_i32, sizeof(int)); + reader.read((char *)&dim_i32, sizeof(int)); + npts = (uint32_t)npts_i32; + dim = (uint32_t)dim_i32; + + diskann::cout << "Metadata: #pts = " << npts << ", #dims = " << dim << "... " + << std::endl; + + int truthset_type = -1; // 1 means truthset has ids and distances, 2 means + // only ids, -1 is error + size_t expected_file_size_with_dists = + 2 * npts * dim * sizeof(uint32_t) + 2 * sizeof(uint32_t); + + if (actual_file_size == expected_file_size_with_dists) truthset_type = 1; + + size_t expected_file_size_just_ids = + npts * dim * sizeof(uint32_t) + 2 * sizeof(uint32_t); + + if (actual_file_size == expected_file_size_just_ids) truthset_type = 2; + + if (truthset_type == -1) { + std::stringstream stream; + stream << "Error. File size mismatch. File should have bin format, with " + "npts followed by ngt followed by npts*ngt ids and optionally " + "followed by npts*ngt distance values; actual size: " + << actual_file_size + << ", expected: " << expected_file_size_with_dists << " or " + << expected_file_size_just_ids; + diskann::cout << stream.str(); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } + + ids = new uint32_t[npts * dim]; + reader.read((char *)ids, npts * dim * sizeof(uint32_t)); + + if (truthset_type == 1) { + dists = new float[npts * dim]; + reader.read((char *)dists, npts * dim * sizeof(float)); + } +} - try - { - if (data_type == std::string("float")) - aux_main(base_file, query_file, gt_file, K, metric, tags_file); - if (data_type == std::string("int8")) - aux_main(base_file, query_file, gt_file, K, metric, tags_file); - if (data_type == std::string("uint8")) - aux_main(base_file, query_file, gt_file, K, metric, tags_file); - } - catch (const std::exception &e) - { - std::cout << std::string(e.what()) << std::endl; - diskann::cerr << "Compute GT failed." << std::endl; - return -1; +int main(int argc, char **argv) { + std::string data_type, dist_fn, base_file, query_file, gt_file, tags_file; + uint64_t K; + + try { + po::options_description desc{"Arguments"}; + + desc.add_options()("help,h", "Print information on arguments"); + + desc.add_options()("data_type", + po::value(&data_type)->required(), + "data type "); + desc.add_options()("dist_fn", po::value(&dist_fn)->required(), + "distance function "); + desc.add_options()("base_file", + po::value(&base_file)->required(), + "File containing the base vectors in binary format"); + desc.add_options()("query_file", + po::value(&query_file)->required(), + "File containing the query vectors in binary format"); + desc.add_options()("gt_file", po::value(>_file)->required(), + "File name for the writing ground truth in binary " + "format, please don' append .bin at end if " + "no filter_label or filter_label_file is provided it " + "will save the file with '.bin' at end." + "else it will save the file as filename_label.bin"); + desc.add_options()("K", po::value(&K)->required(), + "Number of ground truth nearest neighbors to compute"); + desc.add_options()( + "tags_file", + po::value(&tags_file)->default_value(std::string()), + "File containing the tags in binary format"); + + po::variables_map vm; + po::store(po::parse_command_line(argc, argv, desc), vm); + if (vm.count("help")) { + std::cout << desc; + return 0; } + po::notify(vm); + } catch (const std::exception &ex) { + std::cerr << ex.what() << '\n'; + return -1; + } + + if (data_type != std::string("float") && data_type != std::string("int8") && + data_type != std::string("uint8")) { + std::cout << "Unsupported type. float, int8 and uint8 types are supported." + << std::endl; + return -1; + } + + diskann::Metric metric; + if (dist_fn == std::string("l2")) { + metric = diskann::Metric::L2; + } else if (dist_fn == std::string("mips")) { + metric = diskann::Metric::INNER_PRODUCT; + } else if (dist_fn == std::string("cosine")) { + metric = diskann::Metric::COSINE; + } else { + std::cerr << "Unsupported distance function. Use l2/mips/cosine." + << std::endl; + return -1; + } + + try { + if (data_type == std::string("float")) + aux_main(base_file, query_file, gt_file, K, metric, tags_file); + if (data_type == std::string("int8")) + aux_main(base_file, query_file, gt_file, K, metric, tags_file); + if (data_type == std::string("uint8")) + aux_main(base_file, query_file, gt_file, K, metric, tags_file); + } catch (const std::exception &e) { + std::cout << std::string(e.what()) << std::endl; + diskann::cerr << "Compute GT failed." << std::endl; + return -1; + } } diff --git a/tests/utils/compute_groundtruth_for_filters.cpp b/tests/utils/compute_groundtruth_for_filters.cpp index d101a07ff..4b8745f4a 100644 --- a/tests/utils/compute_groundtruth_for_filters.cpp +++ b/tests/utils/compute_groundtruth_for_filters.cpp @@ -41,885 +41,878 @@ typedef std::string path; namespace po = boost::program_options; -template T div_round_up(const T numerator, const T denominator) -{ - return (numerator % denominator == 0) ? (numerator / denominator) : 1 + (numerator / denominator); +template +T div_round_up(const T numerator, const T denominator) { + return (numerator % denominator == 0) ? (numerator / denominator) + : 1 + (numerator / denominator); } using pairIF = std::pair; -struct cmpmaxstruct -{ - bool operator()(const pairIF &l, const pairIF &r) - { - return l.second < r.second; - }; +struct cmpmaxstruct { + bool operator()(const pairIF &l, const pairIF &r) { + return l.second < r.second; + }; }; -using maxPQIFCS = std::priority_queue, cmpmaxstruct>; +using maxPQIFCS = + std::priority_queue, cmpmaxstruct>; -template T *aligned_malloc(const size_t n, const size_t alignment) -{ +template +T *aligned_malloc(const size_t n, const size_t alignment) { #ifdef _WINDOWS - return (T *)_aligned_malloc(sizeof(T) * n, alignment); + return (T *)_aligned_malloc(sizeof(T) * n, alignment); #else - return static_cast(aligned_alloc(alignment, sizeof(T) * n)); + return static_cast(aligned_alloc(alignment, sizeof(T) * n)); #endif } -inline bool custom_dist(const std::pair &a, const std::pair &b) -{ - return a.second < b.second; +inline bool custom_dist(const std::pair &a, + const std::pair &b) { + return a.second < b.second; } -void compute_l2sq(float *const points_l2sq, const float *const matrix, const int64_t num_points, const int dim) -{ - assert(points_l2sq != NULL); +void compute_l2sq(float *const points_l2sq, const float *const matrix, + const int64_t num_points, const int dim) { + assert(points_l2sq != NULL); #pragma omp parallel for schedule(static, 65536) - for (int64_t d = 0; d < num_points; ++d) - points_l2sq[d] = - cblas_sdot(dim, matrix + (ptrdiff_t)d * (ptrdiff_t)dim, 1, matrix + (ptrdiff_t)d * (ptrdiff_t)dim, 1); + for (int64_t d = 0; d < num_points; ++d) + points_l2sq[d] = cblas_sdot(dim, matrix + (ptrdiff_t)d * (ptrdiff_t)dim, 1, + matrix + (ptrdiff_t)d * (ptrdiff_t)dim, 1); } -void distsq_to_points(const size_t dim, - float *dist_matrix, // Col Major, cols are queries, rows are points - size_t npoints, const float *const points, - const float *const points_l2sq, // points in Col major - size_t nqueries, const float *const queries, - const float *const queries_l2sq, // queries in Col major - float *ones_vec = NULL) // Scratchspace of num_data size and init to 1.0 +void distsq_to_points( + const size_t dim, + float *dist_matrix, // Col Major, cols are queries, rows are points + size_t npoints, const float *const points, + const float *const points_l2sq, // points in Col major + size_t nqueries, const float *const queries, + const float *const queries_l2sq, // queries in Col major + float *ones_vec = NULL) // Scratchspace of num_data size and init to 1.0 { - bool ones_vec_alloc = false; - if (ones_vec == NULL) - { - ones_vec = new float[nqueries > npoints ? nqueries : npoints]; - std::fill_n(ones_vec, nqueries > npoints ? nqueries : npoints, (float)1.0); - ones_vec_alloc = true; - } - cblas_sgemm(CblasColMajor, CblasTrans, CblasNoTrans, npoints, nqueries, dim, (float)-2.0, points, dim, queries, dim, - (float)0.0, dist_matrix, npoints); - cblas_sgemm(CblasColMajor, CblasNoTrans, CblasTrans, npoints, nqueries, 1, (float)1.0, points_l2sq, npoints, - ones_vec, nqueries, (float)1.0, dist_matrix, npoints); - cblas_sgemm(CblasColMajor, CblasNoTrans, CblasTrans, npoints, nqueries, 1, (float)1.0, ones_vec, npoints, - queries_l2sq, nqueries, (float)1.0, dist_matrix, npoints); - if (ones_vec_alloc) - delete[] ones_vec; + bool ones_vec_alloc = false; + if (ones_vec == NULL) { + ones_vec = new float[nqueries > npoints ? nqueries : npoints]; + std::fill_n(ones_vec, nqueries > npoints ? nqueries : npoints, (float)1.0); + ones_vec_alloc = true; + } + cblas_sgemm(CblasColMajor, CblasTrans, CblasNoTrans, npoints, nqueries, dim, + (float)-2.0, points, dim, queries, dim, (float)0.0, dist_matrix, + npoints); + cblas_sgemm(CblasColMajor, CblasNoTrans, CblasTrans, npoints, nqueries, 1, + (float)1.0, points_l2sq, npoints, ones_vec, nqueries, (float)1.0, + dist_matrix, npoints); + cblas_sgemm(CblasColMajor, CblasNoTrans, CblasTrans, npoints, nqueries, 1, + (float)1.0, ones_vec, npoints, queries_l2sq, nqueries, (float)1.0, + dist_matrix, npoints); + if (ones_vec_alloc) delete[] ones_vec; } -void inner_prod_to_points(const size_t dim, - float *dist_matrix, // Col Major, cols are queries, rows are points - size_t npoints, const float *const points, size_t nqueries, const float *const queries, - float *ones_vec = NULL) // Scratchspace of num_data size and init to 1.0 +void inner_prod_to_points( + const size_t dim, + float *dist_matrix, // Col Major, cols are queries, rows are points + size_t npoints, const float *const points, size_t nqueries, + const float *const queries, + float *ones_vec = NULL) // Scratchspace of num_data size and init to 1.0 { - bool ones_vec_alloc = false; - if (ones_vec == NULL) - { - ones_vec = new float[nqueries > npoints ? nqueries : npoints]; - std::fill_n(ones_vec, nqueries > npoints ? nqueries : npoints, (float)1.0); - ones_vec_alloc = true; - } - cblas_sgemm(CblasColMajor, CblasTrans, CblasNoTrans, npoints, nqueries, dim, (float)-1.0, points, dim, queries, dim, - (float)0.0, dist_matrix, npoints); - - if (ones_vec_alloc) - delete[] ones_vec; + bool ones_vec_alloc = false; + if (ones_vec == NULL) { + ones_vec = new float[nqueries > npoints ? nqueries : npoints]; + std::fill_n(ones_vec, nqueries > npoints ? nqueries : npoints, (float)1.0); + ones_vec_alloc = true; + } + cblas_sgemm(CblasColMajor, CblasTrans, CblasNoTrans, npoints, nqueries, dim, + (float)-1.0, points, dim, queries, dim, (float)0.0, dist_matrix, + npoints); + + if (ones_vec_alloc) delete[] ones_vec; } -void exact_knn(const size_t dim, const size_t k, - int *const closest_points, // k * num_queries preallocated, col - // major, queries columns - float *const dist_closest_points, // k * num_queries - // preallocated, Dist to - // corresponding closes_points - size_t npoints, - float *points_in, // points in Col major - size_t nqueries, float *queries_in, - diskann::Metric metric = diskann::Metric::L2) // queries in Col major +void exact_knn( + const size_t dim, const size_t k, + int *const closest_points, // k * num_queries preallocated, col + // major, queries columns + float *const dist_closest_points, // k * num_queries + // preallocated, Dist to + // corresponding closes_points + size_t npoints, + float *points_in, // points in Col major + size_t nqueries, float *queries_in, + diskann::Metric metric = diskann::Metric::L2) // queries in Col major { - float *points_l2sq = new float[npoints]; - float *queries_l2sq = new float[nqueries]; - compute_l2sq(points_l2sq, points_in, npoints, dim); - compute_l2sq(queries_l2sq, queries_in, nqueries, dim); - - float *points = points_in; - float *queries = queries_in; - - if (metric == diskann::Metric::COSINE) - { // we convert cosine distance as - // normalized L2 distnace - points = new float[npoints * dim]; - queries = new float[nqueries * dim]; + float *points_l2sq = new float[npoints]; + float *queries_l2sq = new float[nqueries]; + compute_l2sq(points_l2sq, points_in, npoints, dim); + compute_l2sq(queries_l2sq, queries_in, nqueries, dim); + + float *points = points_in; + float *queries = queries_in; + + if (metric == diskann::Metric::COSINE) { // we convert cosine distance as + // normalized L2 distnace + points = new float[npoints * dim]; + queries = new float[nqueries * dim]; #pragma omp parallel for schedule(static, 4096) - for (int64_t i = 0; i < (int64_t)npoints; i++) - { - float norm = std::sqrt(points_l2sq[i]); - if (norm == 0) - { - norm = std::numeric_limits::epsilon(); - } - for (uint32_t j = 0; j < dim; j++) - { - points[i * dim + j] = points_in[i * dim + j] / norm; - } - } + for (int64_t i = 0; i < (int64_t)npoints; i++) { + float norm = std::sqrt(points_l2sq[i]); + if (norm == 0) { + norm = std::numeric_limits::epsilon(); + } + for (uint32_t j = 0; j < dim; j++) { + points[i * dim + j] = points_in[i * dim + j] / norm; + } + } #pragma omp parallel for schedule(static, 4096) - for (int64_t i = 0; i < (int64_t)nqueries; i++) - { - float norm = std::sqrt(queries_l2sq[i]); - if (norm == 0) - { - norm = std::numeric_limits::epsilon(); - } - for (uint32_t j = 0; j < dim; j++) - { - queries[i * dim + j] = queries_in[i * dim + j] / norm; - } - } - // recalculate norms after normalizing, they should all be one. - compute_l2sq(points_l2sq, points, npoints, dim); - compute_l2sq(queries_l2sq, queries, nqueries, dim); + for (int64_t i = 0; i < (int64_t)nqueries; i++) { + float norm = std::sqrt(queries_l2sq[i]); + if (norm == 0) { + norm = std::numeric_limits::epsilon(); + } + for (uint32_t j = 0; j < dim; j++) { + queries[i * dim + j] = queries_in[i * dim + j] / norm; + } } - - std::cout << "Going to compute " << k << " NNs for " << nqueries << " queries over " << npoints << " points in " - << dim << " dimensions using"; - if (metric == diskann::Metric::INNER_PRODUCT) - std::cout << " MIPS "; - else if (metric == diskann::Metric::COSINE) - std::cout << " Cosine "; - else - std::cout << " L2 "; - std::cout << "distance fn. " << std::endl; - - size_t q_batch_size = (1 << 9); - float *dist_matrix = new float[(size_t)q_batch_size * (size_t)npoints]; - - for (uint64_t b = 0; b < div_round_up(nqueries, q_batch_size); ++b) - { - int64_t q_b = b * q_batch_size; - int64_t q_e = ((b + 1) * q_batch_size > nqueries) ? nqueries : (b + 1) * q_batch_size; - - if (metric == diskann::Metric::L2 || metric == diskann::Metric::COSINE) - { - 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); - } - else - { - inner_prod_to_points(dim, dist_matrix, npoints, points, q_e - q_b, - queries + (ptrdiff_t)q_b * (ptrdiff_t)dim); - } - std::cout << "Computed distances for queries: [" << q_b << "," << q_e << ")" << std::endl; + // recalculate norms after normalizing, they should all be one. + compute_l2sq(points_l2sq, points, npoints, dim); + compute_l2sq(queries_l2sq, queries, nqueries, dim); + } + + std::cout << "Going to compute " << k << " NNs for " << nqueries + << " queries over " << npoints << " points in " << dim + << " dimensions using"; + if (metric == diskann::Metric::INNER_PRODUCT) + std::cout << " MIPS "; + else if (metric == diskann::Metric::COSINE) + std::cout << " Cosine "; + else + std::cout << " L2 "; + std::cout << "distance fn. " << std::endl; + + size_t q_batch_size = (1 << 9); + float *dist_matrix = new float[(size_t)q_batch_size * (size_t)npoints]; + + for (uint64_t b = 0; b < div_round_up(nqueries, q_batch_size); ++b) { + int64_t q_b = b * q_batch_size; + int64_t q_e = + ((b + 1) * q_batch_size > nqueries) ? nqueries : (b + 1) * q_batch_size; + + if (metric == diskann::Metric::L2 || metric == diskann::Metric::COSINE) { + 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); + } else { + inner_prod_to_points(dim, dist_matrix, npoints, points, q_e - q_b, + queries + (ptrdiff_t)q_b * (ptrdiff_t)dim); + } + std::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++) - { - maxPQIFCS point_dist; - for (uint64_t p = 0; p < k; p++) - point_dist.emplace(p, dist_matrix[(ptrdiff_t)p + (ptrdiff_t)(q - q_b) * (ptrdiff_t)npoints]); - for (size_t p = k; p < npoints; p++) - { - if (point_dist.top().second > dist_matrix[(ptrdiff_t)p + (ptrdiff_t)(q - q_b) * (ptrdiff_t)npoints]) - point_dist.emplace(p, dist_matrix[(ptrdiff_t)p + (ptrdiff_t)(q - q_b) * (ptrdiff_t)npoints]); - if (point_dist.size() > k) - point_dist.pop(); - } - for (ptrdiff_t l = 0; l < (ptrdiff_t)k; ++l) - { - closest_points[(ptrdiff_t)(k - 1 - l) + (ptrdiff_t)q * (ptrdiff_t)k] = point_dist.top().first; - dist_closest_points[(ptrdiff_t)(k - 1 - l) + (ptrdiff_t)q * (ptrdiff_t)k] = point_dist.top().second; - point_dist.pop(); - } - assert(std::is_sorted(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; + for (long long q = q_b; q < q_e; q++) { + maxPQIFCS point_dist; + for (uint64_t p = 0; p < k; p++) + point_dist.emplace( + p, dist_matrix[(ptrdiff_t)p + + (ptrdiff_t)(q - q_b) * (ptrdiff_t)npoints]); + for (size_t p = k; p < npoints; p++) { + if (point_dist.top().second > + dist_matrix[(ptrdiff_t)p + + (ptrdiff_t)(q - q_b) * (ptrdiff_t)npoints]) + point_dist.emplace( + p, dist_matrix[(ptrdiff_t)p + + (ptrdiff_t)(q - q_b) * (ptrdiff_t)npoints]); + if (point_dist.size() > k) point_dist.pop(); + } + for (ptrdiff_t l = 0; l < (ptrdiff_t)k; ++l) { + closest_points[(ptrdiff_t)(k - 1 - l) + (ptrdiff_t)q * (ptrdiff_t)k] = + point_dist.top().first; + dist_closest_points[(ptrdiff_t)(k - 1 - l) + + (ptrdiff_t)q * (ptrdiff_t)k] = + point_dist.top().second; + point_dist.pop(); + } + assert(std::is_sorted( + 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; + } - delete[] dist_matrix; + delete[] dist_matrix; - delete[] points_l2sq; - delete[] queries_l2sq; + delete[] points_l2sq; + delete[] queries_l2sq; - if (metric == diskann::Metric::COSINE) - { - delete[] points; - delete[] queries; - } + if (metric == diskann::Metric::COSINE) { + delete[] points; + delete[] queries; + } } -template inline int get_num_parts(const char *filename) -{ - std::ifstream reader; - reader.exceptions(std::ios::failbit | std::ios::badbit); - reader.open(filename, std::ios::binary); - std::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; - 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; - return num_parts; +template +inline int get_num_parts(const char *filename) { + std::ifstream reader; + reader.exceptions(std::ios::failbit | std::ios::badbit); + reader.open(filename, std::ios::binary); + std::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; + 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; + return num_parts; } template -inline void load_bin_as_float(const char *filename, float *&data, size_t &npts_u64, size_t &ndims_u64, int part_num) -{ - std::ifstream reader; - reader.exceptions(std::ios::failbit | std::ios::badbit); - reader.open(filename, std::ios::binary); - std::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)); - uint64_t start_id = part_num * PARTSIZE; - uint64_t end_id = (std::min)(start_id + PARTSIZE, (uint64_t)npts_i32); - npts_u64 = end_id - start_id; - ndims_u64 = (uint64_t)ndims_i32; - std::cout << "#pts in part = " << npts_u64 << ", #dims = " << ndims_u64 - << ", size = " << npts_u64 * ndims_u64 * sizeof(T) << "B" << std::endl; - - reader.seekg(start_id * ndims_u64 * sizeof(T) + 2 * sizeof(uint32_t), std::ios::beg); - T *data_T = new T[npts_u64 * ndims_u64]; - reader.read((char *)data_T, sizeof(T) * npts_u64 * ndims_u64); - std::cout << "Finished reading part of the bin file." << std::endl; - reader.close(); - data = aligned_malloc(npts_u64 * ndims_u64, ALIGNMENT); +inline void load_bin_as_float(const char *filename, float *&data, + size_t &npts_u64, size_t &ndims_u64, + int part_num) { + std::ifstream reader; + reader.exceptions(std::ios::failbit | std::ios::badbit); + reader.open(filename, std::ios::binary); + std::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)); + uint64_t start_id = part_num * PARTSIZE; + uint64_t end_id = (std::min)(start_id + PARTSIZE, (uint64_t)npts_i32); + npts_u64 = end_id - start_id; + ndims_u64 = (uint64_t)ndims_i32; + std::cout << "#pts in part = " << npts_u64 << ", #dims = " << ndims_u64 + << ", size = " << npts_u64 * ndims_u64 * sizeof(T) << "B" + << std::endl; + + reader.seekg(start_id * ndims_u64 * sizeof(T) + 2 * sizeof(uint32_t), + std::ios::beg); + T *data_T = new T[npts_u64 * ndims_u64]; + reader.read((char *)data_T, sizeof(T) * npts_u64 * ndims_u64); + std::cout << "Finished reading part of the bin file." << std::endl; + reader.close(); + data = aligned_malloc(npts_u64 * ndims_u64, ALIGNMENT); #pragma omp parallel for schedule(dynamic, 32768) - for (int64_t i = 0; i < (int64_t)npts_u64; i++) - { - for (int64_t j = 0; j < (int64_t)ndims_u64; j++) - { - float cur_val_float = (float)data_T[i * ndims_u64 + j]; - std::memcpy((char *)(data + i * ndims_u64 + j), (char *)&cur_val_float, sizeof(float)); - } + for (int64_t i = 0; i < (int64_t)npts_u64; i++) { + for (int64_t j = 0; j < (int64_t)ndims_u64; j++) { + float cur_val_float = (float)data_T[i * ndims_u64 + j]; + std::memcpy((char *)(data + i * ndims_u64 + j), (char *)&cur_val_float, + sizeof(float)); } - delete[] data_T; - std::cout << "Finished converting part data to float." << std::endl; + } + delete[] data_T; + std::cout << "Finished converting part data to float." << std::endl; } template -inline std::vector load_filtered_bin_as_float(const char *filename, float *&data, size_t &npts, size_t &ndims, - int part_num, const char *label_file, - const std::string &filter_label, - const std::string &universal_label, size_t &npoints_filt, - std::vector> &pts_to_labels) -{ - std::ifstream reader(filename, std::ios::binary); - if (reader.fail()) - { - throw diskann::ANNException(std::string("Failed to open file ") + filename, -1); +inline std::vector load_filtered_bin_as_float( + const char *filename, float *&data, size_t &npts, size_t &ndims, + int part_num, const char *label_file, const std::string &filter_label, + const std::string &universal_label, size_t &npoints_filt, + std::vector> &pts_to_labels) { + std::ifstream reader(filename, std::ios::binary); + if (reader.fail()) { + throw diskann::ANNException(std::string("Failed to open file ") + filename, + -1); + } + + std::cout << "Reading bin file " << filename << " ...\n"; + int npts_i32, ndims_i32; + std::vector rev_map; + reader.read((char *)&npts_i32, sizeof(int)); + reader.read((char *)&ndims_i32, sizeof(int)); + uint64_t start_id = part_num * PARTSIZE; + uint64_t end_id = (std::min)(start_id + PARTSIZE, (uint64_t)npts_i32); + npts = end_id - start_id; + ndims = (uint32_t)ndims_i32; + uint64_t nptsuint64_t = (uint64_t)npts; + uint64_t ndimsuint64_t = (uint64_t)ndims; + npoints_filt = 0; + std::cout << "#pts in part = " << npts << ", #dims = " << ndims + << ", size = " << nptsuint64_t * ndimsuint64_t * sizeof(T) << "B" + << std::endl; + std::cout << "start and end ids: " << start_id << ", " << end_id << std::endl; + reader.seekg(start_id * ndims * sizeof(T) + 2 * sizeof(uint32_t), + std::ios::beg); + + 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; + reader.close(); + + data = aligned_malloc(nptsuint64_t * ndimsuint64_t, ALIGNMENT); + + for (int64_t i = 0; i < (int64_t)nptsuint64_t; i++) { + if (std::find(pts_to_labels[start_id + i].begin(), + pts_to_labels[start_id + i].end(), + filter_label) != pts_to_labels[start_id + i].end() || + std::find(pts_to_labels[start_id + i].begin(), + pts_to_labels[start_id + i].end(), + universal_label) != pts_to_labels[start_id + i].end()) { + rev_map.push_back(start_id + i); + for (int64_t j = 0; j < (int64_t)ndimsuint64_t; j++) { + float cur_val_float = (float)data_T[i * ndimsuint64_t + j]; + std::memcpy((char *)(data + npoints_filt * ndimsuint64_t + j), + (char *)&cur_val_float, sizeof(float)); + } + npoints_filt++; } - - std::cout << "Reading bin file " << filename << " ...\n"; - int npts_i32, ndims_i32; - std::vector rev_map; - reader.read((char *)&npts_i32, sizeof(int)); - reader.read((char *)&ndims_i32, sizeof(int)); - uint64_t start_id = part_num * PARTSIZE; - uint64_t end_id = (std::min)(start_id + PARTSIZE, (uint64_t)npts_i32); - npts = end_id - start_id; - ndims = (uint32_t)ndims_i32; - uint64_t nptsuint64_t = (uint64_t)npts; - uint64_t ndimsuint64_t = (uint64_t)ndims; - npoints_filt = 0; - std::cout << "#pts in part = " << npts << ", #dims = " << ndims - << ", size = " << nptsuint64_t * ndimsuint64_t * sizeof(T) << "B" << std::endl; - std::cout << "start and end ids: " << start_id << ", " << end_id << std::endl; - reader.seekg(start_id * ndims * sizeof(T) + 2 * sizeof(uint32_t), std::ios::beg); - - 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; - reader.close(); - - data = aligned_malloc(nptsuint64_t * ndimsuint64_t, ALIGNMENT); - - for (int64_t i = 0; i < (int64_t)nptsuint64_t; i++) - { - if (std::find(pts_to_labels[start_id + i].begin(), pts_to_labels[start_id + i].end(), filter_label) != - pts_to_labels[start_id + i].end() || - std::find(pts_to_labels[start_id + i].begin(), pts_to_labels[start_id + i].end(), universal_label) != - pts_to_labels[start_id + i].end()) - { - rev_map.push_back(start_id + i); - for (int64_t j = 0; j < (int64_t)ndimsuint64_t; j++) - { - float cur_val_float = (float)data_T[i * ndimsuint64_t + j]; - std::memcpy((char *)(data + npoints_filt * ndimsuint64_t + j), (char *)&cur_val_float, sizeof(float)); - } - npoints_filt++; - } - } - delete[] data_T; - std::cout << "Finished converting part data to float.. identified " << npoints_filt - << " points matching the filter." << std::endl; - return rev_map; + } + delete[] data_T; + std::cout << "Finished converting part data to float.. identified " + << npoints_filt << " points matching the filter." << std::endl; + return rev_map; } -template inline void save_bin(const std::string filename, T *data, size_t npts, size_t ndims) -{ - std::ofstream writer; - writer.exceptions(std::ios::failbit | std::ios::badbit); - writer.open(filename, std::ios::binary | std::ios::out); - std::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; - - writer.write((char *)data, npts * ndims * sizeof(T)); - writer.close(); - std::cout << "Finished writing bin" << std::endl; +template +inline void save_bin(const std::string filename, T *data, size_t npts, + size_t ndims) { + std::ofstream writer; + writer.exceptions(std::ios::failbit | std::ios::badbit); + writer.open(filename, std::ios::binary | std::ios::out); + std::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; + + writer.write((char *)data, npts * ndims * sizeof(T)); + writer.close(); + std::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) -{ - 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(uint32_t) + 2 * sizeof(int) - << "B" << std::endl; - - writer.write((char *)data, npts * ndims * sizeof(uint32_t)); - writer.write((char *)distances, npts * ndims * sizeof(float)); - writer.close(); - std::cout << "Finished writing truthset" << std::endl; +inline void save_groundtruth_as_one_file(const std::string filename, + int32_t *data, float *distances, + size_t npts, size_t ndims) { + 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(uint32_t) + 2 * sizeof(int) << "B" + << std::endl; + + writer.write((char *)data, npts * ndims * sizeof(uint32_t)); + writer.write((char *)distances, npts * ndims * sizeof(float)); + writer.close(); + std::cout << "Finished writing truthset" << std::endl; } -inline void parse_label_file_into_vec(size_t &line_cnt, const std::string &map_file, - std::vector> &pts_to_labels) -{ - std::ifstream infile(map_file); - std::string line, token; - std::set labels; - infile.clear(); - infile.seekg(0, std::ios::beg); - while (std::getline(infile, line)) - { - std::istringstream iss(line); - std::vector lbls(0); - - getline(iss, token, '\t'); - std::istringstream new_iss(token); - while (getline(new_iss, token, ',')) - { - token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); - token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); - lbls.push_back(token); - labels.insert(token); - } - if (lbls.size() <= 0) - { - std::cout << "No label found"; - exit(-1); - } - std::sort(lbls.begin(), lbls.end()); - pts_to_labels.push_back(lbls); +inline void parse_label_file_into_vec( + size_t &line_cnt, const std::string &map_file, + std::vector> &pts_to_labels) { + std::ifstream infile(map_file); + std::string line, token; + std::set labels; + infile.clear(); + infile.seekg(0, std::ios::beg); + while (std::getline(infile, line)) { + std::istringstream iss(line); + std::vector lbls(0); + + getline(iss, token, '\t'); + std::istringstream new_iss(token); + while (getline(new_iss, token, ',')) { + token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); + token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); + lbls.push_back(token); + labels.insert(token); + } + if (lbls.size() <= 0) { + std::cout << "No label found"; + exit(-1); } - std::cout << "Identified " << labels.size() << " distinct label(s), and populated labels for " - << pts_to_labels.size() << " points" << std::endl; + std::sort(lbls.begin(), lbls.end()); + pts_to_labels.push_back(lbls); + } + std::cout << "Identified " << labels.size() + << " distinct label(s), and populated labels for " + << pts_to_labels.size() << " points" << std::endl; } template -std::vector>> processUnfilteredParts(const std::string &base_file, - size_t &nqueries, size_t &npoints, - size_t &dim, size_t &k, float *query_data, - const diskann::Metric &metric, - std::vector &location_to_tag) -{ - float *base_data; - int num_parts = get_num_parts(base_file.c_str()); - std::vector>> res(nqueries); - for (int p = 0; p < num_parts; p++) - { - size_t start_id = p * PARTSIZE; - load_bin_as_float(base_file.c_str(), base_data, npoints, dim, p); - - int *closest_points_part = new int[nqueries * k]; - float *dist_closest_points_part = new float[nqueries * k]; - - uint32_t part_k; - part_k = k < npoints ? k : npoints; - exact_knn(dim, part_k, closest_points_part, dist_closest_points_part, npoints, base_data, nqueries, query_data, - metric); - - for (size_t i = 0; i < nqueries; i++) - { - for (uint64_t j = 0; j < part_k; j++) - { - if (!location_to_tag.empty()) - if (location_to_tag[closest_points_part[i * k + j] + start_id] == 0) - continue; - - res[i].push_back(std::make_pair((uint32_t)(closest_points_part[i * part_k + j] + start_id), - dist_closest_points_part[i * part_k + j])); - } - } +std::vector>> processUnfilteredParts( + const std::string &base_file, size_t &nqueries, size_t &npoints, + size_t &dim, size_t &k, float *query_data, const diskann::Metric &metric, + std::vector &location_to_tag) { + float *base_data; + int num_parts = get_num_parts(base_file.c_str()); + std::vector>> res(nqueries); + for (int p = 0; p < num_parts; p++) { + size_t start_id = p * PARTSIZE; + load_bin_as_float(base_file.c_str(), base_data, npoints, dim, p); + + int *closest_points_part = new int[nqueries * k]; + float *dist_closest_points_part = new float[nqueries * k]; + + uint32_t part_k; + part_k = k < npoints ? k : npoints; + exact_knn(dim, part_k, closest_points_part, dist_closest_points_part, + npoints, base_data, nqueries, query_data, metric); + + for (size_t i = 0; i < nqueries; i++) { + for (uint64_t j = 0; j < part_k; j++) { + if (!location_to_tag.empty()) + if (location_to_tag[closest_points_part[i * k + j] + start_id] == 0) + continue; + + res[i].push_back(std::make_pair( + (uint32_t)(closest_points_part[i * part_k + j] + start_id), + dist_closest_points_part[i * part_k + j])); + } + } - delete[] closest_points_part; - delete[] dist_closest_points_part; + delete[] closest_points_part; + delete[] dist_closest_points_part; - diskann::aligned_free(base_data); - } - return res; + diskann::aligned_free(base_data); + } + return res; }; template std::vector>> processFilteredParts( - const std::string &base_file, const std::string &label_file, const std::string &filter_label, - const std::string &universal_label, size_t &nqueries, size_t &npoints, size_t &dim, size_t &k, float *query_data, - const diskann::Metric &metric, std::vector &location_to_tag) -{ - size_t npoints_filt; - float *base_data; - std::vector>> res(nqueries); - int num_parts = get_num_parts(base_file.c_str()); - - std::vector> pts_to_labels; + const std::string &base_file, const std::string &label_file, + const std::string &filter_label, const std::string &universal_label, + size_t &nqueries, size_t &npoints, size_t &dim, size_t &k, + float *query_data, const diskann::Metric &metric, + std::vector &location_to_tag) { + size_t npoints_filt; + float *base_data; + std::vector>> res(nqueries); + int num_parts = get_num_parts(base_file.c_str()); + + std::vector> pts_to_labels; + if (filter_label != "") + parse_label_file_into_vec(npoints, label_file, pts_to_labels); + + for (int p = 0; p < num_parts; p++) { + size_t start_id = p * PARTSIZE; + std::vector rev_map; if (filter_label != "") - parse_label_file_into_vec(npoints, label_file, pts_to_labels); - - for (int p = 0; p < num_parts; p++) - { - size_t start_id = p * PARTSIZE; - std::vector rev_map; - if (filter_label != "") - rev_map = load_filtered_bin_as_float(base_file.c_str(), base_data, npoints, dim, p, label_file.c_str(), - filter_label, universal_label, npoints_filt, pts_to_labels); - int *closest_points_part = new int[nqueries * k]; - float *dist_closest_points_part = new float[nqueries * k]; - - uint32_t part_k; - part_k = k < npoints_filt ? k : npoints_filt; - if (npoints_filt > 0) - { - exact_knn(dim, part_k, closest_points_part, dist_closest_points_part, npoints_filt, base_data, nqueries, - query_data, metric); - } - - for (size_t i = 0; i < nqueries; i++) - { - for (uint64_t j = 0; j < part_k; j++) - { - if (!location_to_tag.empty()) - if (location_to_tag[closest_points_part[i * k + j] + start_id] == 0) - continue; - - res[i].push_back(std::make_pair((uint32_t)(rev_map[closest_points_part[i * part_k + j]]), - dist_closest_points_part[i * part_k + j])); - } - } - - delete[] closest_points_part; - delete[] dist_closest_points_part; - - diskann::aligned_free(base_data); + rev_map = load_filtered_bin_as_float( + base_file.c_str(), base_data, npoints, dim, p, label_file.c_str(), + filter_label, universal_label, npoints_filt, pts_to_labels); + int *closest_points_part = new int[nqueries * k]; + float *dist_closest_points_part = new float[nqueries * k]; + + uint32_t part_k; + part_k = k < npoints_filt ? k : npoints_filt; + if (npoints_filt > 0) { + exact_knn(dim, part_k, closest_points_part, dist_closest_points_part, + npoints_filt, base_data, nqueries, query_data, metric); } - return res; -}; - -template -int aux_main(const std::string &base_file, const std::string &label_file, const std::string &query_file, - const std::string >_file, size_t k, const std::string &universal_label, const diskann::Metric &metric, - const std::string &filter_label, const std::string &tags_file = std::string("")) -{ - size_t npoints, nqueries, dim, npoints_filt; - float *base_data; - float *query_data; + for (size_t i = 0; i < nqueries; i++) { + for (uint64_t j = 0; j < part_k; j++) { + if (!location_to_tag.empty()) + if (location_to_tag[closest_points_part[i * k + j] + start_id] == 0) + continue; - load_bin_as_float(query_file.c_str(), query_data, nqueries, dim, 0); - if (nqueries > PARTSIZE) - std::cerr << "WARNING: #Queries provided (" << nqueries << ") is greater than " << PARTSIZE - << ". Computing GT only for the first " << PARTSIZE << " queries." << std::endl; - - // load tags - const bool tags_enabled = tags_file.empty() ? false : true; - std::vector location_to_tag = diskann::loadTags(tags_file, base_file); - - int *closest_points = new int[nqueries * k]; - float *dist_closest_points = new float[nqueries * k]; - - std::vector>> results; - if (filter_label == "") - { - results = processUnfilteredParts(base_file, nqueries, npoints, dim, k, query_data, metric, location_to_tag); - } - else - { - results = processFilteredParts(base_file, label_file, filter_label, universal_label, nqueries, npoints, dim, - k, query_data, metric, location_to_tag); - } - - for (size_t i = 0; i < nqueries; i++) - { - std::vector> &cur_res = results[i]; - std::sort(cur_res.begin(), cur_res.end(), custom_dist); - size_t j = 0; - for (auto iter : cur_res) - { - if (j == k) - break; - if (tags_enabled) - { - std::uint32_t index_with_tag = location_to_tag[iter.first]; - closest_points[i * k + j] = (int32_t)index_with_tag; - } - else - { - closest_points[i * k + j] = (int32_t)iter.first; - } - - if (metric == diskann::Metric::INNER_PRODUCT) - dist_closest_points[i * k + j] = -iter.second; - else - dist_closest_points[i * k + j] = iter.second; - - ++j; - } - if (j < k) - std::cout << "WARNING: found less than k GT entries for query " << i << std::endl; + res[i].push_back(std::make_pair( + (uint32_t)(rev_map[closest_points_part[i * part_k + j]]), + dist_closest_points_part[i * part_k + j])); + } } - save_groundtruth_as_one_file(gt_file, closest_points, dist_closest_points, nqueries, k); - delete[] closest_points; - delete[] dist_closest_points; - diskann::aligned_free(query_data); + delete[] closest_points_part; + delete[] dist_closest_points_part; - return 0; -} + diskann::aligned_free(base_data); + } + return res; +}; -void load_truthset(const std::string &bin_file, uint32_t *&ids, float *&dists, size_t &npts, size_t &dim) -{ - size_t read_blk_size = 64 * 1024 * 1024; - cached_ifstream reader(bin_file, read_blk_size); - diskann::cout << "Reading truthset file " << bin_file.c_str() << " ..." << std::endl; - size_t actual_file_size = reader.get_file_size(); - - int npts_i32, dim_i32; - reader.read((char *)&npts_i32, sizeof(int)); - reader.read((char *)&dim_i32, sizeof(int)); - npts = (uint32_t)npts_i32; - dim = (uint32_t)dim_i32; - - diskann::cout << "Metadata: #pts = " << npts << ", #dims = " << dim << "... " << std::endl; - - int truthset_type = -1; // 1 means truthset has ids and distances, 2 means - // only ids, -1 is error - size_t expected_file_size_with_dists = 2 * npts * dim * sizeof(uint32_t) + 2 * sizeof(uint32_t); - - if (actual_file_size == expected_file_size_with_dists) - truthset_type = 1; - - size_t expected_file_size_just_ids = npts * dim * sizeof(uint32_t) + 2 * sizeof(uint32_t); - - if (actual_file_size == expected_file_size_just_ids) - truthset_type = 2; - - if (truthset_type == -1) - { - std::stringstream stream; - stream << "Error. File size mismatch. File should have bin format, with " - "npts followed by ngt followed by npts*ngt ids and optionally " - "followed by npts*ngt distance values; actual size: " - << actual_file_size << ", expected: " << expected_file_size_with_dists << " or " - << expected_file_size_just_ids; - diskann::cout << stream.str(); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); +template +int aux_main(const std::string &base_file, const std::string &label_file, + const std::string &query_file, const std::string >_file, + size_t k, const std::string &universal_label, + const diskann::Metric &metric, const std::string &filter_label, + const std::string &tags_file = std::string("")) { + size_t npoints, nqueries, dim, npoints_filt; + + float *base_data; + float *query_data; + + load_bin_as_float(query_file.c_str(), query_data, nqueries, dim, 0); + if (nqueries > PARTSIZE) + std::cerr << "WARNING: #Queries provided (" << nqueries + << ") is greater than " << PARTSIZE + << ". Computing GT only for the first " << PARTSIZE << " queries." + << std::endl; + + // load tags + const bool tags_enabled = tags_file.empty() ? false : true; + std::vector location_to_tag = + diskann::loadTags(tags_file, base_file); + + int *closest_points = new int[nqueries * k]; + float *dist_closest_points = new float[nqueries * k]; + + std::vector>> results; + if (filter_label == "") { + results = processUnfilteredParts(base_file, nqueries, npoints, dim, k, + query_data, metric, location_to_tag); + } else { + results = processFilteredParts(base_file, label_file, filter_label, + universal_label, nqueries, npoints, dim, + k, query_data, metric, location_to_tag); + } + + for (size_t i = 0; i < nqueries; i++) { + std::vector> &cur_res = results[i]; + std::sort(cur_res.begin(), cur_res.end(), custom_dist); + size_t j = 0; + for (auto iter : cur_res) { + if (j == k) break; + if (tags_enabled) { + std::uint32_t index_with_tag = location_to_tag[iter.first]; + closest_points[i * k + j] = (int32_t)index_with_tag; + } else { + closest_points[i * k + j] = (int32_t)iter.first; + } + + if (metric == diskann::Metric::INNER_PRODUCT) + dist_closest_points[i * k + j] = -iter.second; + else + dist_closest_points[i * k + j] = iter.second; + + ++j; } + if (j < k) + std::cout << "WARNING: found less than k GT entries for query " << i + << std::endl; + } + + save_groundtruth_as_one_file(gt_file, closest_points, dist_closest_points, + nqueries, k); + delete[] closest_points; + delete[] dist_closest_points; + diskann::aligned_free(query_data); + + return 0; +} - ids = new uint32_t[npts * dim]; - reader.read((char *)ids, npts * dim * sizeof(uint32_t)); - - if (truthset_type == 1) - { - dists = new float[npts * dim]; - reader.read((char *)dists, npts * dim * sizeof(float)); - } +void load_truthset(const std::string &bin_file, uint32_t *&ids, float *&dists, + size_t &npts, size_t &dim) { + size_t read_blk_size = 64 * 1024 * 1024; + cached_ifstream reader(bin_file, read_blk_size); + diskann::cout << "Reading truthset file " << bin_file.c_str() << " ..." + << std::endl; + size_t actual_file_size = reader.get_file_size(); + + int npts_i32, dim_i32; + reader.read((char *)&npts_i32, sizeof(int)); + reader.read((char *)&dim_i32, sizeof(int)); + npts = (uint32_t)npts_i32; + dim = (uint32_t)dim_i32; + + diskann::cout << "Metadata: #pts = " << npts << ", #dims = " << dim << "... " + << std::endl; + + int truthset_type = -1; // 1 means truthset has ids and distances, 2 means + // only ids, -1 is error + size_t expected_file_size_with_dists = + 2 * npts * dim * sizeof(uint32_t) + 2 * sizeof(uint32_t); + + if (actual_file_size == expected_file_size_with_dists) truthset_type = 1; + + size_t expected_file_size_just_ids = + npts * dim * sizeof(uint32_t) + 2 * sizeof(uint32_t); + + if (actual_file_size == expected_file_size_just_ids) truthset_type = 2; + + if (truthset_type == -1) { + std::stringstream stream; + stream << "Error. File size mismatch. File should have bin format, with " + "npts followed by ngt followed by npts*ngt ids and optionally " + "followed by npts*ngt distance values; actual size: " + << actual_file_size + << ", expected: " << expected_file_size_with_dists << " or " + << expected_file_size_just_ids; + diskann::cout << stream.str(); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, + __LINE__); + } + + ids = new uint32_t[npts * dim]; + reader.read((char *)ids, npts * dim * sizeof(uint32_t)); + + if (truthset_type == 1) { + dists = new float[npts * dim]; + reader.read((char *)dists, npts * dim * sizeof(float)); + } } -int main(int argc, char **argv) -{ - std::string data_type, dist_fn, base_file, query_file, gt_file, tags_file, label_file, filter_label, - universal_label, filter_label_file; - uint64_t K; - - try - { - po::options_description desc{"Arguments"}; - - desc.add_options()("help,h", "Print information on arguments"); - - desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); - desc.add_options()("dist_fn", po::value(&dist_fn)->required(), "distance function "); - desc.add_options()("base_file", po::value(&base_file)->required(), - "File containing the base vectors in binary format"); - desc.add_options()("query_file", po::value(&query_file)->required(), - "File containing the query vectors in binary format"); - desc.add_options()("label_file", po::value(&label_file)->default_value(""), - "Input labels file in txt format if present"); - desc.add_options()("filter_label", po::value(&filter_label)->default_value(""), - "Input filter label if doing filtered groundtruth"); - desc.add_options()("universal_label", po::value(&universal_label)->default_value(""), - "Universal label, if using it, only in conjunction with label_file"); - desc.add_options()("gt_file", po::value(>_file)->required(), - "File name for the writing ground truth in binary format, please don' append .bin at end if " - "no filter_label or filter_label_file is provided it will save the file with '.bin' at end." - "else it will save the file as filename_label.bin"); - desc.add_options()("K", po::value(&K)->required(), - "Number of ground truth nearest neighbors to compute"); - desc.add_options()("tags_file", po::value(&tags_file)->default_value(std::string()), - "File containing the tags in binary format"); - desc.add_options()("filter_label_file", - po::value(&filter_label_file)->default_value(std::string("")), - "Filter file for Queries for Filtered Search "); - - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) - { - std::cout << desc; - return 0; - } - po::notify(vm); - } - catch (const std::exception &ex) - { - std::cerr << ex.what() << '\n'; - return -1; +int main(int argc, char **argv) { + std::string data_type, dist_fn, base_file, query_file, gt_file, tags_file, + label_file, filter_label, universal_label, filter_label_file; + uint64_t K; + + try { + po::options_description desc{"Arguments"}; + + desc.add_options()("help,h", "Print information on arguments"); + + desc.add_options()("data_type", + po::value(&data_type)->required(), + "data type "); + desc.add_options()("dist_fn", po::value(&dist_fn)->required(), + "distance function "); + desc.add_options()("base_file", + po::value(&base_file)->required(), + "File containing the base vectors in binary format"); + desc.add_options()("query_file", + po::value(&query_file)->required(), + "File containing the query vectors in binary format"); + desc.add_options()("label_file", + po::value(&label_file)->default_value(""), + "Input labels file in txt format if present"); + desc.add_options()("filter_label", + po::value(&filter_label)->default_value(""), + "Input filter label if doing filtered groundtruth"); + desc.add_options()( + "universal_label", + po::value(&universal_label)->default_value(""), + "Universal label, if using it, only in conjunction with label_file"); + desc.add_options()("gt_file", po::value(>_file)->required(), + "File name for the writing ground truth in binary " + "format, please don' append .bin at end if " + "no filter_label or filter_label_file is provided it " + "will save the file with '.bin' at end." + "else it will save the file as filename_label.bin"); + desc.add_options()("K", po::value(&K)->required(), + "Number of ground truth nearest neighbors to compute"); + desc.add_options()( + "tags_file", + po::value(&tags_file)->default_value(std::string()), + "File containing the tags in binary format"); + desc.add_options()("filter_label_file", + po::value(&filter_label_file) + ->default_value(std::string("")), + "Filter file for Queries for Filtered Search "); + + po::variables_map vm; + po::store(po::parse_command_line(argc, argv, desc), vm); + if (vm.count("help")) { + std::cout << desc; + return 0; } - - if (data_type != std::string("float") && data_type != std::string("int8") && data_type != std::string("uint8")) - { - std::cout << "Unsupported type. float, int8 and uint8 types are supported." << std::endl; - return -1; + po::notify(vm); + } catch (const std::exception &ex) { + std::cerr << ex.what() << '\n'; + return -1; + } + + if (data_type != std::string("float") && data_type != std::string("int8") && + data_type != std::string("uint8")) { + std::cout << "Unsupported type. float, int8 and uint8 types are supported." + << std::endl; + return -1; + } + + if (filter_label != "" && filter_label_file != "") { + std::cerr + << "Only one of filter_label and query_filters_file should be provided" + << std::endl; + return -1; + } + + diskann::Metric metric; + if (dist_fn == std::string("l2")) { + metric = diskann::Metric::L2; + } else if (dist_fn == std::string("mips")) { + metric = diskann::Metric::INNER_PRODUCT; + } else if (dist_fn == std::string("cosine")) { + metric = diskann::Metric::COSINE; + } else { + std::cerr << "Unsupported distance function. Use l2/mips/cosine." + << std::endl; + return -1; + } + + std::vector filter_labels; + if (filter_label != "") { + filter_labels.push_back(filter_label); + } else if (filter_label_file != "") { + filter_labels = read_file_to_vector_of_strings(filter_label_file, false); + } + + // only when there is no filter label or 1 filter label for all queries + if (filter_labels.size() == 1) { + try { + if (data_type == std::string("float")) + aux_main(base_file, label_file, query_file, gt_file, K, + universal_label, metric, filter_labels[0], tags_file); + if (data_type == std::string("int8")) + aux_main(base_file, label_file, query_file, gt_file, K, + universal_label, metric, filter_labels[0], tags_file); + if (data_type == std::string("uint8")) + aux_main(base_file, label_file, query_file, gt_file, K, + universal_label, metric, filter_labels[0], tags_file); + } catch (const std::exception &e) { + std::cout << std::string(e.what()) << std::endl; + diskann::cerr << "Compute GT failed." << std::endl; + return -1; } - - if (filter_label != "" && filter_label_file != "") - { - std::cerr << "Only one of filter_label and query_filters_file should be provided" << std::endl; - return -1; + } else { // Each query has its own filter label + // Split up data and query bins into label specific ones + tsl::robin_map labels_to_number_of_points; + tsl::robin_map labels_to_number_of_queries; + + label_set all_labels; + for (size_t i = 0; i < filter_labels.size(); i++) { + std::string label = filter_labels[i]; + all_labels.insert(label); + + if (labels_to_number_of_queries.find(label) == + labels_to_number_of_queries.end()) { + labels_to_number_of_queries[label] = 0; + } + labels_to_number_of_queries[label] += 1; } - diskann::Metric metric; - if (dist_fn == std::string("l2")) - { - metric = diskann::Metric::L2; - } - else if (dist_fn == std::string("mips")) - { - metric = diskann::Metric::INNER_PRODUCT; - } - else if (dist_fn == std::string("cosine")) - { - metric = diskann::Metric::COSINE; - } - else - { - std::cerr << "Unsupported distance function. Use l2/mips/cosine." << std::endl; - return -1; + size_t npoints; + std::vector> point_to_labels; + parse_label_file_into_vec(npoints, label_file, point_to_labels); + std::vector point_ids_to_labels(point_to_labels.size()); + std::vector query_ids_to_labels(filter_labels.size()); + + for (size_t i = 0; i < point_to_labels.size(); i++) { + for (size_t j = 0; j < point_to_labels[i].size(); j++) { + std::string label = point_to_labels[i][j]; + if (all_labels.find(label) != all_labels.end()) { + point_ids_to_labels[i].insert(point_to_labels[i][j]); + if (labels_to_number_of_points.find(label) == + labels_to_number_of_points.end()) { + labels_to_number_of_points[label] = 0; + } + labels_to_number_of_points[label] += 1; + } + } } - std::vector filter_labels; - if (filter_label != "") - { - filter_labels.push_back(filter_label); - } - else if (filter_label_file != "") - { - filter_labels = read_file_to_vector_of_strings(filter_label_file, false); + for (size_t i = 0; i < filter_labels.size(); i++) { + query_ids_to_labels[i].insert(filter_labels[i]); } - // only when there is no filter label or 1 filter label for all queries - if (filter_labels.size() == 1) - { - try - { - if (data_type == std::string("float")) - aux_main(base_file, label_file, query_file, gt_file, K, universal_label, metric, - filter_labels[0], tags_file); - if (data_type == std::string("int8")) - aux_main(base_file, label_file, query_file, gt_file, K, universal_label, metric, - filter_labels[0], tags_file); - if (data_type == std::string("uint8")) - aux_main(base_file, label_file, query_file, gt_file, K, universal_label, metric, - filter_labels[0], tags_file); - } - catch (const std::exception &e) - { - std::cout << std::string(e.what()) << std::endl; - diskann::cerr << "Compute GT failed." << std::endl; - return -1; - } + tsl::robin_map> label_id_to_orig_id; + tsl::robin_map> + label_query_id_to_orig_id; + + if (data_type == std::string("float")) { + label_id_to_orig_id = + diskann::generate_label_specific_vector_files_compat( + base_file, labels_to_number_of_points, point_ids_to_labels, + all_labels); + + label_query_id_to_orig_id = + diskann::generate_label_specific_vector_files_compat( + query_file, labels_to_number_of_queries, query_ids_to_labels, + all_labels); // query_filters acts like query_ids_to_labels + } else if (data_type == std::string("int8")) { + label_id_to_orig_id = + diskann::generate_label_specific_vector_files_compat( + base_file, labels_to_number_of_points, point_ids_to_labels, + all_labels); + + label_query_id_to_orig_id = + diskann::generate_label_specific_vector_files_compat( + query_file, labels_to_number_of_queries, query_ids_to_labels, + all_labels); // query_filters acts like query_ids_to_labels + } else if (data_type == std::string("uint8")) { + label_id_to_orig_id = + diskann::generate_label_specific_vector_files_compat( + base_file, labels_to_number_of_points, point_ids_to_labels, + all_labels); + + label_query_id_to_orig_id = + diskann::generate_label_specific_vector_files_compat( + query_file, labels_to_number_of_queries, query_ids_to_labels, + all_labels); // query_filters acts like query_ids_to_labels + } else { + diskann::cerr << "Invalid data type" << std::endl; + return -1; } - else - { // Each query has its own filter label - // Split up data and query bins into label specific ones - tsl::robin_map labels_to_number_of_points; - tsl::robin_map labels_to_number_of_queries; - - label_set all_labels; - for (size_t i = 0; i < filter_labels.size(); i++) - { - std::string label = filter_labels[i]; - all_labels.insert(label); - - if (labels_to_number_of_queries.find(label) == labels_to_number_of_queries.end()) - { - labels_to_number_of_queries[label] = 0; - } - labels_to_number_of_queries[label] += 1; - } - - size_t npoints; - std::vector> point_to_labels; - parse_label_file_into_vec(npoints, label_file, point_to_labels); - std::vector point_ids_to_labels(point_to_labels.size()); - std::vector query_ids_to_labels(filter_labels.size()); - - for (size_t i = 0; i < point_to_labels.size(); i++) - { - for (size_t j = 0; j < point_to_labels[i].size(); j++) - { - std::string label = point_to_labels[i][j]; - if (all_labels.find(label) != all_labels.end()) - { - point_ids_to_labels[i].insert(point_to_labels[i][j]); - if (labels_to_number_of_points.find(label) == labels_to_number_of_points.end()) - { - labels_to_number_of_points[label] = 0; - } - labels_to_number_of_points[label] += 1; - } - } - } - for (size_t i = 0; i < filter_labels.size(); i++) - { - query_ids_to_labels[i].insert(filter_labels[i]); - } - - tsl::robin_map> label_id_to_orig_id; - tsl::robin_map> label_query_id_to_orig_id; + // Generate label specific ground truths + try { + for (const auto &label : all_labels) { + std::string filtered_base_file = base_file + "_" + label; + std::string filtered_query_file = query_file + "_" + label; + std::string filtered_gt_file = gt_file + "_" + label; if (data_type == std::string("float")) - { - label_id_to_orig_id = diskann::generate_label_specific_vector_files_compat( - base_file, labels_to_number_of_points, point_ids_to_labels, all_labels); - - label_query_id_to_orig_id = diskann::generate_label_specific_vector_files_compat( - query_file, labels_to_number_of_queries, query_ids_to_labels, - all_labels); // query_filters acts like query_ids_to_labels - } - else if (data_type == std::string("int8")) - { - label_id_to_orig_id = diskann::generate_label_specific_vector_files_compat( - base_file, labels_to_number_of_points, point_ids_to_labels, all_labels); - - label_query_id_to_orig_id = diskann::generate_label_specific_vector_files_compat( - query_file, labels_to_number_of_queries, query_ids_to_labels, - all_labels); // query_filters acts like query_ids_to_labels - } - else if (data_type == std::string("uint8")) - { - label_id_to_orig_id = diskann::generate_label_specific_vector_files_compat( - base_file, labels_to_number_of_points, point_ids_to_labels, all_labels); - - label_query_id_to_orig_id = diskann::generate_label_specific_vector_files_compat( - query_file, labels_to_number_of_queries, query_ids_to_labels, - all_labels); // query_filters acts like query_ids_to_labels - } - else - { - diskann::cerr << "Invalid data type" << std::endl; - return -1; - } + aux_main(filtered_base_file, "", filtered_query_file, + filtered_gt_file, K, "", metric, ""); + if (data_type == std::string("int8")) + aux_main(filtered_base_file, "", filtered_query_file, + filtered_gt_file, K, "", metric, ""); + if (data_type == std::string("uint8")) + aux_main(filtered_base_file, "", filtered_query_file, + filtered_gt_file, K, "", metric, ""); + } + } catch (const std::exception &e) { + std::cout << std::string(e.what()) << std::endl; + diskann::cerr << "Compute GT failed." << std::endl; + return -1; + } - // Generate label specific ground truths - - try - { - for (const auto &label : all_labels) - { - std::string filtered_base_file = base_file + "_" + label; - std::string filtered_query_file = query_file + "_" + label; - std::string filtered_gt_file = gt_file + "_" + label; - if (data_type == std::string("float")) - aux_main(filtered_base_file, "", filtered_query_file, filtered_gt_file, K, "", metric, ""); - if (data_type == std::string("int8")) - aux_main(filtered_base_file, "", filtered_query_file, filtered_gt_file, K, "", metric, ""); - if (data_type == std::string("uint8")) - aux_main(filtered_base_file, "", filtered_query_file, filtered_gt_file, K, "", metric, ""); - } - } - catch (const std::exception &e) - { - std::cout << std::string(e.what()) << std::endl; - diskann::cerr << "Compute GT failed." << std::endl; - return -1; - } + // Combine the label specific ground truths to produce a single GT file - // Combine the label specific ground truths to produce a single GT file + uint32_t *gt_ids = nullptr; + float *gt_dists = nullptr; + size_t gt_num, gt_dim; - uint32_t *gt_ids = nullptr; - float *gt_dists = nullptr; - size_t gt_num, gt_dim; + std::vector> final_gt_ids; + std::vector> final_gt_dists; - std::vector> final_gt_ids; - std::vector> final_gt_dists; + int query_num = 0; + for (const auto &lbl : all_labels) { + query_num += labels_to_number_of_queries[lbl]; + } - int query_num = 0; - for (const auto &lbl : all_labels) - { - query_num += labels_to_number_of_queries[lbl]; - } + for (int i = 0; i < query_num; i++) { + final_gt_ids.push_back(std::vector(K)); + final_gt_dists.push_back(std::vector(K)); + } - for (int i = 0; i < query_num; i++) - { - final_gt_ids.push_back(std::vector(K)); - final_gt_dists.push_back(std::vector(K)); - } + for (const auto &lbl : all_labels) { + std::string filtered_gt_file = gt_file + "_" + lbl; + load_truthset(filtered_gt_file, gt_ids, gt_dists, gt_num, gt_dim); - for (const auto &lbl : all_labels) - { - std::string filtered_gt_file = gt_file + "_" + lbl; - load_truthset(filtered_gt_file, gt_ids, gt_dists, gt_num, gt_dim); - - for (int i = 0; i < labels_to_number_of_queries[lbl]; i++) - { - int orig_query_id = label_query_id_to_orig_id[lbl][i]; - for (uint64_t j = 0; j < K; j++) - { - final_gt_ids[orig_query_id][j] = label_id_to_orig_id[lbl][gt_ids[i * K + j]]; - final_gt_dists[orig_query_id][j] = gt_dists[i * K + j]; - } - } + for (int i = 0; i < labels_to_number_of_queries[lbl]; i++) { + int orig_query_id = label_query_id_to_orig_id[lbl][i]; + for (uint64_t j = 0; j < K; j++) { + final_gt_ids[orig_query_id][j] = + label_id_to_orig_id[lbl][gt_ids[i * K + j]]; + final_gt_dists[orig_query_id][j] = gt_dists[i * K + j]; } + } + } - int32_t *closest_points = new int32_t[query_num * K]; - float *dist_closest_points = new float[query_num * K]; + int32_t *closest_points = new int32_t[query_num * K]; + float *dist_closest_points = new float[query_num * K]; - for (int i = 0; i < query_num; i++) - { - for (int j = 0; j < K; j++) - { - closest_points[i * K + j] = final_gt_ids[i][j]; - dist_closest_points[i * K + j] = final_gt_dists[i][j]; - } - } + for (int i = 0; i < query_num; i++) { + for (int j = 0; j < K; j++) { + closest_points[i * K + j] = final_gt_ids[i][j]; + dist_closest_points[i * K + j] = final_gt_dists[i][j]; + } + } - save_groundtruth_as_one_file(gt_file, closest_points, dist_closest_points, query_num, K); + save_groundtruth_as_one_file(gt_file, closest_points, dist_closest_points, + query_num, K); - // cleanup artifacts - std::cout << "Cleaning up artifacts..." << std::endl; - tsl::robin_set paths_to_clean{gt_file, base_file, query_file}; - clean_up_artifacts(paths_to_clean, all_labels); - } + // cleanup artifacts + std::cout << "Cleaning up artifacts..." << std::endl; + tsl::robin_set paths_to_clean{gt_file, base_file, query_file}; + clean_up_artifacts(paths_to_clean, all_labels); + } } diff --git a/tests/utils/count_bfs_levels.cpp b/tests/utils/count_bfs_levels.cpp index ddc4eaf0b..86156de3b 100644 --- a/tests/utils/count_bfs_levels.cpp +++ b/tests/utils/count_bfs_levels.cpp @@ -23,59 +23,56 @@ namespace po = boost::program_options; -template void bfs_count(const std::string &index_path, uint32_t data_dims) -{ - using TagT = uint32_t; - using LabelT = uint32_t; - diskann::Index index(diskann::Metric::L2, data_dims, 0, false, false); - std::cout << "Index class instantiated" << std::endl; - index.load(index_path.c_str(), 1, 100); - std::cout << "Index loaded" << std::endl; - index.count_nodes_at_bfs_levels(); +template +void bfs_count(const std::string &index_path, uint32_t data_dims) { + using TagT = uint32_t; + using LabelT = uint32_t; + diskann::Index index(diskann::Metric::L2, data_dims, 0, + false, false); + std::cout << "Index class instantiated" << std::endl; + index.load(index_path.c_str(), 1, 100); + std::cout << "Index loaded" << std::endl; + index.count_nodes_at_bfs_levels(); } -int main(int argc, char **argv) -{ - std::string data_type, index_path_prefix; - uint32_t data_dims; +int main(int argc, char **argv) { + std::string data_type, index_path_prefix; + uint32_t data_dims; - po::options_description desc{"Arguments"}; - try - { - desc.add_options()("help,h", "Print information on arguments"); - desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); - desc.add_options()("index_path_prefix", po::value(&index_path_prefix)->required(), - "Path prefix to the index"); - desc.add_options()("data_dims", po::value(&data_dims)->required(), "Dimensionality of the data"); + po::options_description desc{"Arguments"}; + try { + desc.add_options()("help,h", "Print information on arguments"); + desc.add_options()("data_type", + po::value(&data_type)->required(), + "data type "); + desc.add_options()("index_path_prefix", + po::value(&index_path_prefix)->required(), + "Path prefix to the index"); + desc.add_options()("data_dims", po::value(&data_dims)->required(), + "Dimensionality of the data"); - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) - { - std::cout << desc; - return 0; - } - po::notify(vm); - } - catch (const std::exception &ex) - { - std::cerr << ex.what() << '\n'; - return -1; + po::variables_map vm; + po::store(po::parse_command_line(argc, argv, desc), vm); + if (vm.count("help")) { + std::cout << desc; + return 0; } + po::notify(vm); + } catch (const std::exception &ex) { + std::cerr << ex.what() << '\n'; + return -1; + } - try - { - if (data_type == std::string("int8")) - bfs_count(index_path_prefix, data_dims); - else if (data_type == std::string("uint8")) - bfs_count(index_path_prefix, data_dims); - if (data_type == std::string("float")) - bfs_count(index_path_prefix, data_dims); - } - catch (std::exception &e) - { - std::cout << std::string(e.what()) << std::endl; - diskann::cerr << "Index BFS failed." << std::endl; - return -1; - } + try { + if (data_type == std::string("int8")) + bfs_count(index_path_prefix, data_dims); + else if (data_type == std::string("uint8")) + bfs_count(index_path_prefix, data_dims); + if (data_type == std::string("float")) + bfs_count(index_path_prefix, data_dims); + } catch (std::exception &e) { + std::cout << std::string(e.what()) << std::endl; + diskann::cerr << "Index BFS failed." << std::endl; + return -1; + } } diff --git a/tests/utils/create_disk_layout.cpp b/tests/utils/create_disk_layout.cpp index f494c1227..a6f981db5 100644 --- a/tests/utils/create_disk_layout.cpp +++ b/tests/utils/create_disk_layout.cpp @@ -12,37 +12,34 @@ #include "disk_utils.h" #include "cached_io.h" -template int create_disk_layout(char **argv) -{ - 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); - return 0; +template +int create_disk_layout(char **argv) { + 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); + return 0; } -int main(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); - } +int main(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); + } - int ret_val = -1; - if (std::string(argv[1]) == std::string("float")) - ret_val = create_disk_layout(argv); - else if (std::string(argv[1]) == std::string("int8")) - ret_val = create_disk_layout(argv); - else if (std::string(argv[1]) == std::string("uint8")) - ret_val = create_disk_layout(argv); - else - { - std::cout << "unsupported type. use int8/uint8/float " << std::endl; - ret_val = -2; - } - return ret_val; + int ret_val = -1; + if (std::string(argv[1]) == std::string("float")) + ret_val = create_disk_layout(argv); + else if (std::string(argv[1]) == std::string("int8")) + ret_val = create_disk_layout(argv); + else if (std::string(argv[1]) == std::string("uint8")) + ret_val = create_disk_layout(argv); + else { + std::cout << "unsupported type. use int8/uint8/float " << std::endl; + ret_val = -2; + } + return ret_val; } diff --git a/tests/utils/float_bin_to_int8.cpp b/tests/utils/float_bin_to_int8.cpp index a7632a6cf..15d524850 100644 --- a/tests/utils/float_bin_to_int8.cpp +++ b/tests/utils/float_bin_to_int8.cpp @@ -4,60 +4,59 @@ #include #include "utils.h" -void block_convert(std::ofstream &writer, int8_t *write_buf, std::ifstream &reader, float *read_buf, size_t npts, - size_t ndims, float bias, float scale) -{ - reader.read((char *)read_buf, npts * ndims * sizeof(float)); - - for (size_t i = 0; i < npts; i++) - { - for (size_t d = 0; d < ndims; d++) - { - write_buf[d + i * ndims] = (int8_t)((read_buf[d + i * ndims] - bias) * (254.0 / scale)); - } +void block_convert(std::ofstream &writer, int8_t *write_buf, + std::ifstream &reader, float *read_buf, size_t npts, + size_t ndims, float bias, float scale) { + reader.read((char *)read_buf, npts * ndims * sizeof(float)); + + for (size_t i = 0; i < npts; i++) { + for (size_t d = 0; d < ndims; d++) { + write_buf[d + i * ndims] = + (int8_t)((read_buf[d + i * ndims] - bias) * (254.0 / scale)); } - writer.write((char *)write_buf, npts * ndims); + } + writer.write((char *)write_buf, npts * ndims); } -int main(int argc, char **argv) -{ - if (argc != 5) - { - std::cout << "Usage: " << argv[0] << " input_bin output_tsv bias scale" << std::endl; - exit(-1); - } - - std::ifstream reader(argv[1], std::ios::binary); - uint32_t npts_u32; - uint32_t ndims_u32; - reader.read((char *)&npts_u32, sizeof(uint32_t)); - reader.read((char *)&ndims_u32, sizeof(uint32_t)); - size_t npts = npts_u32; - size_t ndims = ndims_u32; - std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims << std::endl; - - size_t blk_size = 131072; - size_t nblks = ROUND_UP(npts, blk_size) / blk_size; - - std::ofstream writer(argv[2], std::ios::binary); - auto read_buf = new float[blk_size * ndims]; - auto write_buf = new int8_t[blk_size * ndims]; - float bias = atof(argv[3]); - float scale = atof(argv[4]); - - writer.write((char *)(&npts_u32), sizeof(uint32_t)); - writer.write((char *)(&ndims_u32), sizeof(uint32_t)); - - for (size_t i = 0; i < nblks; i++) - { - size_t cblk_size = std::min(npts - i * blk_size, blk_size); - block_convert(writer, write_buf, reader, read_buf, cblk_size, ndims, bias, scale); - std::cout << "Block #" << i << " written" << std::endl; - } - - delete[] read_buf; - delete[] write_buf; - - writer.close(); - reader.close(); +int main(int argc, char **argv) { + if (argc != 5) { + std::cout << "Usage: " << argv[0] << " input_bin output_tsv bias scale" + << std::endl; + exit(-1); + } + + std::ifstream reader(argv[1], std::ios::binary); + uint32_t npts_u32; + uint32_t ndims_u32; + reader.read((char *)&npts_u32, sizeof(uint32_t)); + reader.read((char *)&ndims_u32, sizeof(uint32_t)); + size_t npts = npts_u32; + size_t ndims = ndims_u32; + std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims + << std::endl; + + size_t blk_size = 131072; + size_t nblks = ROUND_UP(npts, blk_size) / blk_size; + + std::ofstream writer(argv[2], std::ios::binary); + auto read_buf = new float[blk_size * ndims]; + auto write_buf = new int8_t[blk_size * ndims]; + float bias = atof(argv[3]); + float scale = atof(argv[4]); + + writer.write((char *)(&npts_u32), sizeof(uint32_t)); + writer.write((char *)(&ndims_u32), sizeof(uint32_t)); + + for (size_t i = 0; i < nblks; i++) { + size_t cblk_size = std::min(npts - i * blk_size, blk_size); + block_convert(writer, write_buf, reader, read_buf, cblk_size, ndims, bias, + scale); + std::cout << "Block #" << i << " written" << std::endl; + } + + delete[] read_buf; + delete[] write_buf; + + writer.close(); + reader.close(); } diff --git a/tests/utils/fvecs_to_bin.cpp b/tests/utils/fvecs_to_bin.cpp index 873ad3b0c..f645b67c9 100644 --- a/tests/utils/fvecs_to_bin.cpp +++ b/tests/utils/fvecs_to_bin.cpp @@ -5,91 +5,88 @@ #include "utils.h" // Convert float types -void block_convert_float(std::ifstream &reader, std::ofstream &writer, float *read_buf, float *write_buf, size_t npts, - size_t ndims) -{ - reader.read((char *)read_buf, npts * (ndims * sizeof(float) + sizeof(uint32_t))); - for (size_t 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)); +void block_convert_float(std::ifstream &reader, std::ofstream &writer, + float *read_buf, float *write_buf, size_t npts, + size_t ndims) { + reader.read((char *)read_buf, + npts * (ndims * sizeof(float) + sizeof(uint32_t))); + for (size_t 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)); } // Convert byte types -void block_convert_byte(std::ifstream &reader, std::ofstream &writer, uint8_t *read_buf, uint8_t *write_buf, - size_t npts, size_t ndims) -{ - reader.read((char *)read_buf, npts * (ndims * sizeof(uint8_t) + sizeof(uint32_t))); - for (size_t i = 0; i < npts; i++) - { - memcpy(write_buf + i * ndims, (read_buf + i * (ndims + sizeof(uint32_t))) + sizeof(uint32_t), - ndims * sizeof(uint8_t)); - } - writer.write((char *)write_buf, npts * ndims * sizeof(uint8_t)); +void block_convert_byte(std::ifstream &reader, std::ofstream &writer, + uint8_t *read_buf, uint8_t *write_buf, size_t npts, + size_t ndims) { + reader.read((char *)read_buf, + npts * (ndims * sizeof(uint8_t) + sizeof(uint32_t))); + for (size_t i = 0; i < npts; i++) { + memcpy(write_buf + i * ndims, + (read_buf + i * (ndims + sizeof(uint32_t))) + sizeof(uint32_t), + ndims * sizeof(uint8_t)); + } + writer.write((char *)write_buf, npts * ndims * sizeof(uint8_t)); } -int main(int argc, char **argv) -{ - if (argc != 4) - { - std::cout << argv[0] << " input_vecs output_bin" << std::endl; - exit(-1); - } +int main(int argc, char **argv) { + if (argc != 4) { + std::cout << argv[0] << " input_vecs output_bin" + << std::endl; + exit(-1); + } - int datasize = sizeof(float); + int datasize = sizeof(float); - if (strcmp(argv[1], "uint8") == 0 || strcmp(argv[1], "int8") == 0) - { - datasize = sizeof(uint8_t); - } - else if (strcmp(argv[1], "float") != 0) - { - std::cout << "Error: type not supported. Use float/int8/uint8" << std::endl; - exit(-1); - } + if (strcmp(argv[1], "uint8") == 0 || strcmp(argv[1], "int8") == 0) { + datasize = sizeof(uint8_t); + } else if (strcmp(argv[1], "float") != 0) { + std::cout << "Error: type not supported. Use float/int8/uint8" << std::endl; + exit(-1); + } - std::ifstream reader(argv[2], std::ios::binary | std::ios::ate); - size_t fsize = reader.tellg(); - reader.seekg(0, std::ios::beg); + std::ifstream reader(argv[2], std::ios::binary | std::ios::ate); + size_t fsize = reader.tellg(); + reader.seekg(0, std::ios::beg); - uint32_t ndims_u32; - reader.read((char *)&ndims_u32, sizeof(uint32_t)); - reader.seekg(0, std::ios::beg); - size_t ndims = (size_t)ndims_u32; - size_t npts = fsize / ((ndims * datasize) + sizeof(uint32_t)); - std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims << std::endl; + uint32_t ndims_u32; + reader.read((char *)&ndims_u32, sizeof(uint32_t)); + reader.seekg(0, std::ios::beg); + size_t ndims = (size_t)ndims_u32; + size_t npts = fsize / ((ndims * datasize) + sizeof(uint32_t)); + std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims + << std::endl; - size_t blk_size = 131072; - size_t nblks = ROUND_UP(npts, blk_size) / blk_size; - std::cout << "# blks: " << nblks << std::endl; - std::ofstream writer(argv[3], std::ios::binary); - int32_t npts_s32 = (int32_t)npts; - int32_t ndims_s32 = (int32_t)ndims; - writer.write((char *)&npts_s32, sizeof(int32_t)); - writer.write((char *)&ndims_s32, sizeof(int32_t)); + size_t blk_size = 131072; + size_t nblks = ROUND_UP(npts, blk_size) / blk_size; + std::cout << "# blks: " << nblks << std::endl; + std::ofstream writer(argv[3], std::ios::binary); + int32_t npts_s32 = (int32_t)npts; + int32_t ndims_s32 = (int32_t)ndims; + writer.write((char *)&npts_s32, sizeof(int32_t)); + writer.write((char *)&ndims_s32, sizeof(int32_t)); - size_t chunknpts = std::min(npts, blk_size); - uint8_t *read_buf = new uint8_t[chunknpts * ((ndims * datasize) + sizeof(uint32_t))]; - uint8_t *write_buf = new uint8_t[chunknpts * ndims * datasize]; + size_t chunknpts = std::min(npts, blk_size); + uint8_t *read_buf = + new uint8_t[chunknpts * ((ndims * datasize) + sizeof(uint32_t))]; + uint8_t *write_buf = new uint8_t[chunknpts * ndims * datasize]; - for (size_t i = 0; i < nblks; i++) - { - size_t cblk_size = std::min(npts - i * blk_size, blk_size); - if (datasize == sizeof(float)) - { - block_convert_float(reader, writer, (float *)read_buf, (float *)write_buf, cblk_size, ndims); - } - else - { - block_convert_byte(reader, writer, read_buf, write_buf, cblk_size, ndims); - } - std::cout << "Block #" << i << " written" << std::endl; + for (size_t i = 0; i < nblks; i++) { + size_t cblk_size = std::min(npts - i * blk_size, blk_size); + if (datasize == sizeof(float)) { + block_convert_float(reader, writer, (float *)read_buf, (float *)write_buf, + cblk_size, ndims); + } else { + block_convert_byte(reader, writer, read_buf, write_buf, cblk_size, ndims); } + std::cout << "Block #" << i << " written" << std::endl; + } - delete[] read_buf; - delete[] write_buf; + delete[] read_buf; + delete[] write_buf; - reader.close(); - writer.close(); + reader.close(); + writer.close(); } diff --git a/tests/utils/fvecs_to_bvecs.cpp b/tests/utils/fvecs_to_bvecs.cpp index f9c2aa71b..1707bfbfe 100644 --- a/tests/utils/fvecs_to_bvecs.cpp +++ b/tests/utils/fvecs_to_bvecs.cpp @@ -4,53 +4,53 @@ #include #include "utils.h" -void block_convert(std::ifstream &reader, std::ofstream &writer, float *read_buf, uint8_t *write_buf, size_t npts, - size_t ndims) -{ - reader.read((char *)read_buf, npts * (ndims * sizeof(float) + sizeof(uint32_t))); - for (size_t i = 0; i < npts; i++) - { - memcpy(write_buf + i * (ndims + 4), read_buf + i * (ndims + 1), sizeof(uint32_t)); - for (size_t d = 0; d < ndims; d++) - write_buf[i * (ndims + 4) + 4 + d] = (uint8_t)read_buf[i * (ndims + 1) + 1 + d]; - } - writer.write((char *)write_buf, npts * (ndims * 1 + 4)); +void block_convert(std::ifstream &reader, std::ofstream &writer, + float *read_buf, uint8_t *write_buf, size_t npts, + size_t ndims) { + reader.read((char *)read_buf, + npts * (ndims * sizeof(float) + sizeof(uint32_t))); + for (size_t i = 0; i < npts; i++) { + memcpy(write_buf + i * (ndims + 4), read_buf + i * (ndims + 1), + sizeof(uint32_t)); + for (size_t d = 0; d < ndims; d++) + write_buf[i * (ndims + 4) + 4 + d] = + (uint8_t)read_buf[i * (ndims + 1) + 1 + d]; + } + writer.write((char *)write_buf, npts * (ndims * 1 + 4)); } -int main(int argc, char **argv) -{ - if (argc != 3) - { - std::cout << argv[0] << " input_fvecs output_bvecs(uint8)" << std::endl; - exit(-1); - } - std::ifstream reader(argv[1], std::ios::binary | std::ios::ate); - size_t fsize = reader.tellg(); - reader.seekg(0, std::ios::beg); +int main(int argc, char **argv) { + if (argc != 3) { + std::cout << argv[0] << " input_fvecs output_bvecs(uint8)" << std::endl; + exit(-1); + } + std::ifstream reader(argv[1], std::ios::binary | std::ios::ate); + size_t fsize = reader.tellg(); + reader.seekg(0, std::ios::beg); - uint32_t ndims_u32; - reader.read((char *)&ndims_u32, sizeof(uint32_t)); - reader.seekg(0, std::ios::beg); - size_t ndims = (size_t)ndims_u32; - size_t npts = fsize / ((ndims + 1) * sizeof(float)); - std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims << std::endl; + uint32_t ndims_u32; + reader.read((char *)&ndims_u32, sizeof(uint32_t)); + reader.seekg(0, std::ios::beg); + size_t ndims = (size_t)ndims_u32; + size_t npts = fsize / ((ndims + 1) * sizeof(float)); + std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims + << std::endl; - size_t blk_size = 131072; - size_t nblks = ROUND_UP(npts, blk_size) / blk_size; - std::cout << "# blks: " << nblks << std::endl; - std::ofstream writer(argv[2], std::ios::binary); - auto read_buf = new float[npts * (ndims + 1)]; - auto write_buf = new uint8_t[npts * (ndims + 4)]; - for (size_t i = 0; i < nblks; i++) - { - size_t 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; - } + size_t blk_size = 131072; + size_t nblks = ROUND_UP(npts, blk_size) / blk_size; + std::cout << "# blks: " << nblks << std::endl; + std::ofstream writer(argv[2], std::ios::binary); + auto read_buf = new float[npts * (ndims + 1)]; + auto write_buf = new uint8_t[npts * (ndims + 4)]; + for (size_t i = 0; i < nblks; i++) { + size_t 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; + } - delete[] read_buf; - delete[] write_buf; + delete[] read_buf; + delete[] write_buf; - reader.close(); - writer.close(); + reader.close(); + writer.close(); } diff --git a/tests/utils/gen_random_slice.cpp b/tests/utils/gen_random_slice.cpp index a4cd96e0a..3914473ab 100644 --- a/tests/utils/gen_random_slice.cpp +++ b/tests/utils/gen_random_slice.cpp @@ -20,39 +20,31 @@ #include #include -template 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])); - gen_random_slice(base_file, output_prefix, sampling_rate); - return 0; +template +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])); + gen_random_slice(base_file, output_prefix, sampling_rate); + return 0; } -int main(int argc, char **argv) -{ - if (argc != 5) - { - std::cout << argv[0] - << " data_type [float/int8/uint8] base_bin_file " - "sample_output_prefix sampling_probability" - << std::endl; - exit(-1); - } +int main(int argc, char **argv) { + if (argc != 5) { + std::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(argv); - } - else if (std::string(argv[1]) == std::string("int8")) - { - aux_main(argv); - } - else if (std::string(argv[1]) == std::string("uint8")) - { - aux_main(argv); - } - else - std::cout << "Unsupported type. Use float/int8/uint8." << std::endl; - return 0; + if (std::string(argv[1]) == std::string("float")) { + aux_main(argv); + } else if (std::string(argv[1]) == std::string("int8")) { + aux_main(argv); + } else if (std::string(argv[1]) == std::string("uint8")) { + aux_main(argv); + } else + std::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 761983129..ff4126c68 100644 --- a/tests/utils/generate_pq.cpp +++ b/tests/utils/generate_pq.cpp @@ -8,63 +8,66 @@ #define KMEANS_ITERS_FOR_PQ 15 template -bool generate_pq(const std::string &data_path, const std::string &index_prefix_path, const size_t num_pq_centers, - const size_t num_pq_chunks, const float sampling_rate, const bool opq) -{ - std::string pq_pivots_path = index_prefix_path + "_pq_pivots.bin"; - std::string pq_compressed_vectors_path = index_prefix_path + "_pq_compressed.bin"; +bool generate_pq(const std::string &data_path, + const std::string &index_prefix_path, + const size_t num_pq_centers, const size_t num_pq_chunks, + const float sampling_rate, const bool opq) { + std::string pq_pivots_path = index_prefix_path + "_pq_pivots.bin"; + std::string pq_compressed_vectors_path = + 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; + // 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; - if (opq) - { - diskann::generate_opq_pivots(train_data, train_size, train_dim, num_pq_centers, num_pq_chunks, pq_pivots_path, - true); - } - else - { - diskann::generate_pq_pivots(train_data, train_size, train_dim, num_pq_centers, num_pq_chunks, - KMEANS_ITERS_FOR_PQ, pq_pivots_path); - } - diskann::generate_pq_data_from_pivots(data_path, num_pq_centers, num_pq_chunks, pq_pivots_path, - pq_compressed_vectors_path, true); + if (opq) { + diskann::generate_opq_pivots(train_data, train_size, train_dim, + num_pq_centers, num_pq_chunks, pq_pivots_path, + true); + } else { + diskann::generate_pq_pivots(train_data, train_size, train_dim, + num_pq_centers, num_pq_chunks, + KMEANS_ITERS_FOR_PQ, pq_pivots_path); + } + diskann::generate_pq_data_from_pivots(data_path, num_pq_centers, + num_pq_chunks, pq_pivots_path, + pq_compressed_vectors_path, true); - delete[] train_data; + delete[] train_data; - return 0; + return 0; } -int main(int argc, char **argv) -{ - if (argc != 7) - { - std::cout << "Usage: \n" - << argv[0] - << " " - " " - " " - << std::endl; - } - else - { - const std::string data_path(argv[2]); - 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 bool opq = atoi(argv[6]) == 0 ? false : true; +int main(int argc, char **argv) { + if (argc != 7) { + std::cout << "Usage: \n" + << argv[0] + << " " + " " + " " + << std::endl; + } else { + const std::string data_path(argv[2]); + 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 bool opq = atoi(argv[6]) == 0 ? false : true; - if (std::string(argv[1]) == std::string("float")) - generate_pq(data_path, index_prefix_path, num_pq_centers, num_pq_chunks, sampling_rate, opq); - else if (std::string(argv[1]) == std::string("int8")) - generate_pq(data_path, index_prefix_path, num_pq_centers, num_pq_chunks, sampling_rate, opq); - else if (std::string(argv[1]) == std::string("uint8")) - generate_pq(data_path, index_prefix_path, num_pq_centers, num_pq_chunks, sampling_rate, opq); - else - std::cout << "Error. wrong file type" << std::endl; - } + if (std::string(argv[1]) == std::string("float")) + generate_pq(data_path, index_prefix_path, num_pq_centers, + num_pq_chunks, sampling_rate, opq); + else if (std::string(argv[1]) == std::string("int8")) + generate_pq(data_path, index_prefix_path, num_pq_centers, + num_pq_chunks, sampling_rate, opq); + else if (std::string(argv[1]) == std::string("uint8")) + generate_pq(data_path, index_prefix_path, num_pq_centers, + num_pq_chunks, sampling_rate, opq); + else + std::cout << "Error. wrong file type" << std::endl; + } } diff --git a/tests/utils/generate_synthetic_labels.cpp b/tests/utils/generate_synthetic_labels.cpp index 9c659c4ca..bbbdd55f5 100644 --- a/tests/utils/generate_synthetic_labels.cpp +++ b/tests/utils/generate_synthetic_labels.cpp @@ -9,199 +9,169 @@ #include "utils.h" namespace po = boost::program_options; -class ZipfDistribution -{ - public: - ZipfDistribution(int num_points, int num_labels) - : num_labels(num_labels), num_points(num_points), - uniform_zero_to_one(std::uniform_real_distribution<>(0.0, 1.0)) - { +class ZipfDistribution { + public: + ZipfDistribution(int num_points, int num_labels) + : num_labels(num_labels), + num_points(num_points), + uniform_zero_to_one(std::uniform_real_distribution<>(0.0, 1.0)) {} + + std::unordered_map createDistributionMap() { + std::unordered_map map; + int primary_label_freq = ceil(num_points * distribution_factor); + for (int i{1}; i < num_labels + 1; i++) { + map[i] = ceil(primary_label_freq / i); } - - std::unordered_map createDistributionMap() - { - std::unordered_map map; - int primary_label_freq = ceil(num_points * distribution_factor); - for (int i{1}; i < num_labels + 1; i++) - { - map[i] = ceil(primary_label_freq / i); - } - return map; - } - - int writeDistribution(std::ofstream &outfile) - { - auto distribution_map = createDistributionMap(); - for (int i{0}; i < num_points; i++) - { - bool label_written = false; - for (auto it = distribution_map.cbegin(), next_it = it; it != distribution_map.cend(); it = next_it) - { - next_it++; - auto label_selection_probability = std::bernoulli_distribution(distribution_factor / (double)it->first); - if (label_selection_probability(rand_engine)) - { - if (label_written) - { - outfile << ','; - } - outfile << it->first; - label_written = true; - // remove label from map if we have used all labels - distribution_map[it->first] -= 1; - if (distribution_map[it->first] == 0) - { - distribution_map.erase(it); - } - } - } - if (!label_written) - { - outfile << 0; - } - if (i < num_points - 1) - { - outfile << '\n'; - } + return map; + } + + int writeDistribution(std::ofstream &outfile) { + auto distribution_map = createDistributionMap(); + for (int i{0}; i < num_points; i++) { + bool label_written = false; + for (auto it = distribution_map.cbegin(), next_it = it; + it != distribution_map.cend(); it = next_it) { + next_it++; + auto label_selection_probability = std::bernoulli_distribution( + distribution_factor / (double)it->first); + if (label_selection_probability(rand_engine)) { + if (label_written) { + outfile << ','; + } + outfile << it->first; + label_written = true; + // remove label from map if we have used all labels + distribution_map[it->first] -= 1; + if (distribution_map[it->first] == 0) { + distribution_map.erase(it); + } } - return 0; + } + if (!label_written) { + outfile << 0; + } + if (i < num_points - 1) { + outfile << '\n'; + } } + return 0; + } - int writeDistribution(std::string filename) - { - std::ofstream outfile(filename); - if (!outfile.is_open()) - { - std::cerr << "Error: could not open output file " << filename << '\n'; - return -1; - } - writeDistribution(outfile); - outfile.close(); + int writeDistribution(std::string filename) { + std::ofstream outfile(filename); + if (!outfile.is_open()) { + std::cerr << "Error: could not open output file " << filename << '\n'; + return -1; } - - private: - int num_labels; - const int num_points; - const double distribution_factor = 0.7; - std::knuth_b rand_engine; - const std::uniform_real_distribution uniform_zero_to_one; + writeDistribution(outfile); + outfile.close(); + } + + private: + int num_labels; + const int num_points; + const double distribution_factor = 0.7; + std::knuth_b rand_engine; + const std::uniform_real_distribution uniform_zero_to_one; }; -int main(int argc, char **argv) -{ - std::string output_file, distribution_type; - size_t num_labels, num_points; - - try - { - po::options_description desc{"Arguments"}; - - desc.add_options()("help,h", "Print information on arguments"); - desc.add_options()("output_file,O", po::value(&output_file)->required(), - "Filename for saving the label file"); - desc.add_options()("num_points,N", po::value(&num_points)->required(), "Number of points in dataset"); - desc.add_options()("num_labels,L", po::value(&num_labels)->required(), - "Number of unique labels, up to 5000"); - desc.add_options()("distribution_type,DT", po::value(&distribution_type)->default_value("random"), - "Distribution function for labels defaults to random"); - - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) - { - std::cout << desc; - return 0; - } - po::notify(vm); +int main(int argc, char **argv) { + std::string output_file, distribution_type; + size_t num_labels, num_points; + + try { + po::options_description desc{"Arguments"}; + + desc.add_options()("help,h", "Print information on arguments"); + desc.add_options()("output_file,O", + po::value(&output_file)->required(), + "Filename for saving the label file"); + desc.add_options()("num_points,N", + po::value(&num_points)->required(), + "Number of points in dataset"); + desc.add_options()("num_labels,L", + po::value(&num_labels)->required(), + "Number of unique labels, up to 5000"); + desc.add_options()( + "distribution_type,DT", + po::value(&distribution_type)->default_value("random"), + "Distribution function for labels defaults " + "to random"); + + po::variables_map vm; + po::store(po::parse_command_line(argc, argv, desc), vm); + if (vm.count("help")) { + std::cout << desc; + return 0; } - catch (const std::exception &ex) - { - std::cerr << ex.what() << '\n'; - return -1; + po::notify(vm); + } catch (const std::exception &ex) { + std::cerr << ex.what() << '\n'; + return -1; + } + + if (num_labels > 5000) { + std::cerr << "Error: num_labels must be 5000 or less" << '\n'; + return -1; + } + + if (num_points <= 0) { + std::cerr << "Error: num_points must be greater than 0" << '\n'; + return -1; + } + + std::cout << "Generating synthetic labels for " << num_points + << " points with " << num_labels << " unique labels" << '\n'; + + try { + std::ofstream outfile(output_file); + if (!outfile.is_open()) { + std::cerr << "Error: could not open output file " << output_file << '\n'; + return -1; } - if (num_labels > 5000) - { - std::cerr << "Error: num_labels must be 5000 or less" << '\n'; - return -1; - } - - if (num_points <= 0) - { - std::cerr << "Error: num_points must be greater than 0" << '\n'; - return -1; - } - - std::cout << "Generating synthetic labels for " << num_points << " points with " << num_labels << " unique labels" - << '\n'; - - try - { - std::ofstream outfile(output_file); - if (!outfile.is_open()) - { - std::cerr << "Error: could not open output file " << output_file << '\n'; - return -1; - } - - if (distribution_type == "zipf") - { - ZipfDistribution zipf(num_points, num_labels); - zipf.writeDistribution(outfile); - } - else if (distribution_type == "random") - { - for (size_t i = 0; i < num_points; i++) - { - bool label_written = false; - for (size_t j = 1; j <= num_labels; j++) - { - // 50% chance to assign each label - if (rand() > (RAND_MAX / 2)) - { - if (label_written) - { - outfile << ','; - } - outfile << j; - label_written = true; - } - } - if (!label_written) - { - outfile << 0; - } - if (i < num_points - 1) - { - outfile << '\n'; - } + if (distribution_type == "zipf") { + ZipfDistribution zipf(num_points, num_labels); + zipf.writeDistribution(outfile); + } else if (distribution_type == "random") { + for (size_t i = 0; i < num_points; i++) { + bool label_written = false; + for (size_t j = 1; j <= num_labels; j++) { + // 50% chance to assign each label + if (rand() > (RAND_MAX / 2)) { + if (label_written) { + outfile << ','; } + outfile << j; + label_written = true; + } } - else if (distribution_type == "one_per_point") - { - std::random_device rd; // obtain a random number from hardware - std::mt19937 gen(rd()); // seed the generator - std::uniform_int_distribution<> distr(0, num_labels); // define the range - - for (size_t i = 0; i < num_points; i++) - { - outfile << distr(gen); - if (i != num_points - 1) - outfile << '\n'; - } + if (!label_written) { + outfile << 0; } - if (outfile.is_open()) - { - outfile.close(); + if (i < num_points - 1) { + outfile << '\n'; } - - std::cout << "Labels written to " << output_file << '\n'; + } + } else if (distribution_type == "one_per_point") { + std::random_device rd; // obtain a random number from hardware + std::mt19937 gen(rd()); // seed the generator + std::uniform_int_distribution<> distr(0, num_labels); // define the range + + for (size_t i = 0; i < num_points; i++) { + outfile << distr(gen); + if (i != num_points - 1) outfile << '\n'; + } } - catch (const std::exception &ex) - { - std::cerr << "Label generation failed: " << ex.what() << '\n'; - return -1; + if (outfile.is_open()) { + outfile.close(); } - return 0; + std::cout << "Labels written to " << output_file << '\n'; + } catch (const std::exception &ex) { + std::cerr << "Label generation failed: " << ex.what() << '\n'; + return -1; + } + + return 0; } \ No newline at end of file diff --git a/tests/utils/int8_to_float.cpp b/tests/utils/int8_to_float.cpp index dcdfddc0d..05662b2a0 100644 --- a/tests/utils/int8_to_float.cpp +++ b/tests/utils/int8_to_float.cpp @@ -4,20 +4,18 @@ #include #include "utils.h" -int main(int argc, char **argv) -{ - if (argc != 3) - { - std::cout << argv[0] << " input_int8_bin output_float_bin" << std::endl; - exit(-1); - } +int main(int argc, char **argv) { + if (argc != 3) { + std::cout << argv[0] << " input_int8_bin output_float_bin" << std::endl; + exit(-1); + } - int8_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; + int8_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/int8_to_float_scale.cpp b/tests/utils/int8_to_float_scale.cpp index 2de1a3a56..e1217c571 100644 --- a/tests/utils/int8_to_float_scale.cpp +++ b/tests/utils/int8_to_float_scale.cpp @@ -4,60 +4,59 @@ #include #include "utils.h" -void block_convert(std::ofstream &writer, float *write_buf, std::ifstream &reader, int8_t *read_buf, size_t npts, - size_t ndims, float bias, float scale) -{ - reader.read((char *)read_buf, npts * ndims * sizeof(int8_t)); - - for (size_t i = 0; i < npts; i++) - { - for (size_t d = 0; d < ndims; d++) - { - write_buf[d + i * ndims] = (((float)read_buf[d + i * ndims] - bias) * scale); - } +void block_convert(std::ofstream &writer, float *write_buf, + std::ifstream &reader, int8_t *read_buf, size_t npts, + size_t ndims, float bias, float scale) { + reader.read((char *)read_buf, npts * ndims * sizeof(int8_t)); + + for (size_t i = 0; i < npts; i++) { + for (size_t d = 0; d < ndims; d++) { + write_buf[d + i * ndims] = + (((float)read_buf[d + i * ndims] - bias) * scale); } - writer.write((char *)write_buf, npts * ndims * sizeof(float)); + } + writer.write((char *)write_buf, npts * ndims * sizeof(float)); } -int main(int argc, char **argv) -{ - if (argc != 5) - { - std::cout << "Usage: " << argv[0] << " input-int8.bin output-float.bin bias scale" << std::endl; - exit(-1); - } - - std::ifstream reader(argv[1], std::ios::binary); - uint32_t npts_u32; - uint32_t ndims_u32; - reader.read((char *)&npts_u32, sizeof(uint32_t)); - reader.read((char *)&ndims_u32, sizeof(uint32_t)); - size_t npts = npts_u32; - size_t ndims = ndims_u32; - std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims << std::endl; - - size_t blk_size = 131072; - size_t nblks = ROUND_UP(npts, blk_size) / blk_size; - - std::ofstream writer(argv[2], std::ios::binary); - auto read_buf = new int8_t[blk_size * ndims]; - auto write_buf = new float[blk_size * ndims]; - float bias = atof(argv[3]); - float scale = atof(argv[4]); - - writer.write((char *)(&npts_u32), sizeof(uint32_t)); - writer.write((char *)(&ndims_u32), sizeof(uint32_t)); - - for (size_t i = 0; i < nblks; i++) - { - size_t cblk_size = std::min(npts - i * blk_size, blk_size); - block_convert(writer, write_buf, reader, read_buf, cblk_size, ndims, bias, scale); - std::cout << "Block #" << i << " written" << std::endl; - } - - delete[] read_buf; - delete[] write_buf; - - writer.close(); - reader.close(); +int main(int argc, char **argv) { + if (argc != 5) { + std::cout << "Usage: " << argv[0] + << " input-int8.bin output-float.bin bias scale" << std::endl; + exit(-1); + } + + std::ifstream reader(argv[1], std::ios::binary); + uint32_t npts_u32; + uint32_t ndims_u32; + reader.read((char *)&npts_u32, sizeof(uint32_t)); + reader.read((char *)&ndims_u32, sizeof(uint32_t)); + size_t npts = npts_u32; + size_t ndims = ndims_u32; + std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims + << std::endl; + + size_t blk_size = 131072; + size_t nblks = ROUND_UP(npts, blk_size) / blk_size; + + std::ofstream writer(argv[2], std::ios::binary); + auto read_buf = new int8_t[blk_size * ndims]; + auto write_buf = new float[blk_size * ndims]; + float bias = atof(argv[3]); + float scale = atof(argv[4]); + + writer.write((char *)(&npts_u32), sizeof(uint32_t)); + writer.write((char *)(&ndims_u32), sizeof(uint32_t)); + + for (size_t i = 0; i < nblks; i++) { + size_t cblk_size = std::min(npts - i * blk_size, blk_size); + block_convert(writer, write_buf, reader, read_buf, cblk_size, ndims, bias, + scale); + std::cout << "Block #" << i << " written" << std::endl; + } + + delete[] read_buf; + delete[] write_buf; + + writer.close(); + reader.close(); } diff --git a/tests/utils/ivecs_to_bin.cpp b/tests/utils/ivecs_to_bin.cpp index ea8a4a3d2..0146330f8 100644 --- a/tests/utils/ivecs_to_bin.cpp +++ b/tests/utils/ivecs_to_bin.cpp @@ -4,55 +4,54 @@ #include #include "utils.h" -void block_convert(std::ifstream &reader, std::ofstream &writer, uint32_t *read_buf, uint32_t *write_buf, size_t npts, - size_t ndims) -{ - reader.read((char *)read_buf, npts * (ndims * sizeof(uint32_t) + sizeof(uint32_t))); - for (size_t i = 0; i < npts; i++) - { - memcpy(write_buf + i * ndims, (read_buf + i * (ndims + 1)) + 1, ndims * sizeof(uint32_t)); - } - writer.write((char *)write_buf, npts * ndims * sizeof(uint32_t)); +void block_convert(std::ifstream &reader, std::ofstream &writer, + uint32_t *read_buf, uint32_t *write_buf, size_t npts, + size_t ndims) { + reader.read((char *)read_buf, + npts * (ndims * sizeof(uint32_t) + sizeof(uint32_t))); + for (size_t i = 0; i < npts; i++) { + memcpy(write_buf + i * ndims, (read_buf + i * (ndims + 1)) + 1, + ndims * sizeof(uint32_t)); + } + writer.write((char *)write_buf, npts * ndims * sizeof(uint32_t)); } -int main(int argc, char **argv) -{ - if (argc != 3) - { - std::cout << argv[0] << " input_ivecs output_bin" << std::endl; - exit(-1); - } - std::ifstream reader(argv[1], std::ios::binary | std::ios::ate); - size_t fsize = reader.tellg(); - reader.seekg(0, std::ios::beg); +int main(int argc, char **argv) { + if (argc != 3) { + std::cout << argv[0] << " input_ivecs output_bin" << std::endl; + exit(-1); + } + std::ifstream reader(argv[1], std::ios::binary | std::ios::ate); + size_t fsize = reader.tellg(); + reader.seekg(0, std::ios::beg); - uint32_t ndims_u32; - reader.read((char *)&ndims_u32, sizeof(uint32_t)); - reader.seekg(0, std::ios::beg); - size_t ndims = (size_t)ndims_u32; - size_t npts = fsize / ((ndims + 1) * sizeof(uint32_t)); - std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims << std::endl; + uint32_t ndims_u32; + reader.read((char *)&ndims_u32, sizeof(uint32_t)); + reader.seekg(0, std::ios::beg); + size_t ndims = (size_t)ndims_u32; + size_t npts = fsize / ((ndims + 1) * sizeof(uint32_t)); + std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims + << std::endl; - size_t blk_size = 131072; - size_t nblks = ROUND_UP(npts, blk_size) / blk_size; - std::cout << "# blks: " << nblks << std::endl; - std::ofstream writer(argv[2], std::ios::binary); - int npts_s32 = (int)npts; - int ndims_s32 = (int)ndims; - writer.write((char *)&npts_s32, sizeof(int)); - writer.write((char *)&ndims_s32, sizeof(int)); - uint32_t *read_buf = new uint32_t[npts * (ndims + 1)]; - uint32_t *write_buf = new uint32_t[npts * ndims]; - for (size_t i = 0; i < nblks; i++) - { - size_t 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; - } + size_t blk_size = 131072; + size_t nblks = ROUND_UP(npts, blk_size) / blk_size; + std::cout << "# blks: " << nblks << std::endl; + std::ofstream writer(argv[2], std::ios::binary); + int npts_s32 = (int)npts; + int ndims_s32 = (int)ndims; + writer.write((char *)&npts_s32, sizeof(int)); + writer.write((char *)&ndims_s32, sizeof(int)); + uint32_t *read_buf = new uint32_t[npts * (ndims + 1)]; + uint32_t *write_buf = new uint32_t[npts * ndims]; + for (size_t i = 0; i < nblks; i++) { + size_t 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; + } - delete[] read_buf; - delete[] write_buf; + delete[] read_buf; + delete[] write_buf; - reader.close(); - writer.close(); + reader.close(); + writer.close(); } diff --git a/tests/utils/merge_shards.cpp b/tests/utils/merge_shards.cpp index 106c15eef..4ae4fa83d 100644 --- a/tests/utils/merge_shards.cpp +++ b/tests/utils/merge_shards.cpp @@ -14,29 +14,28 @@ #include "cached_io.h" #include "utils.h" -int main(int argc, char **argv) -{ - if (argc != 9) - { - std::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] " - "output_medoids_path[8]" - << std::endl; - exit(-1); - } +int main(int argc, char **argv) { + if (argc != 9) { + std::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] " + "output_medoids_path[8]" + << std::endl; + exit(-1); + } - std::string vamana_prefix(argv[1]); - std::string vamana_suffix(argv[2]); - std::string idmaps_prefix(argv[3]); - std::string idmaps_suffix(argv[4]); - uint64_t nshards = (uint64_t)std::atoi(argv[5]); - uint32_t max_degree = (uint64_t)std::atoi(argv[6]); - std::string output_index(argv[7]); - std::string output_medoids(argv[8]); + std::string vamana_prefix(argv[1]); + std::string vamana_suffix(argv[2]); + std::string idmaps_prefix(argv[3]); + std::string idmaps_suffix(argv[4]); + uint64_t nshards = (uint64_t)std::atoi(argv[5]); + uint32_t max_degree = (uint64_t)std::atoi(argv[6]); + std::string output_index(argv[7]); + std::string output_medoids(argv[8]); - return diskann::merge_shards(vamana_prefix, vamana_suffix, idmaps_prefix, idmaps_suffix, nshards, max_degree, - output_index, output_medoids); + return diskann::merge_shards(vamana_prefix, vamana_suffix, idmaps_prefix, + idmaps_suffix, nshards, max_degree, output_index, + output_medoids); } diff --git a/tests/utils/partition_data.cpp b/tests/utils/partition_data.cpp index 2c505315c..0156e7acd 100644 --- a/tests/utils/partition_data.cpp +++ b/tests/utils/partition_data.cpp @@ -8,32 +8,33 @@ // DEPRECATED: NEED TO REPROGRAM -int main(int argc, char **argv) -{ - if (argc != 7) - { - std::cout << "Usage:\n" - << argv[0] - << " datatype " - " " - " " - << std::endl; - exit(-1); - } +int main(int argc, char **argv) { + if (argc != 7) { + std::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 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]); + const std::string data_path(argv[2]); + const std::string prefix_path(argv[3]); + const float sampling_rate = 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]); - if (std::string(argv[1]) == std::string("float")) - partition(data_path, sampling_rate, num_partitions, max_reps, prefix_path, k_index); - else if (std::string(argv[1]) == std::string("int8")) - partition(data_path, sampling_rate, num_partitions, max_reps, prefix_path, k_index); - else if (std::string(argv[1]) == std::string("uint8")) - 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; + if (std::string(argv[1]) == std::string("float")) + partition(data_path, sampling_rate, num_partitions, max_reps, + prefix_path, k_index); + else if (std::string(argv[1]) == std::string("int8")) + partition(data_path, sampling_rate, num_partitions, max_reps, + prefix_path, k_index); + else if (std::string(argv[1]) == std::string("uint8")) + 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; } diff --git a/tests/utils/partition_with_ram_budget.cpp b/tests/utils/partition_with_ram_budget.cpp index 3c546801a..6ef2c9958 100644 --- a/tests/utils/partition_with_ram_budget.cpp +++ b/tests/utils/partition_with_ram_budget.cpp @@ -8,32 +8,33 @@ // DEPRECATED: NEED TO REPROGRAM -int main(int argc, char **argv) -{ - if (argc != 8) - { - std::cout << "Usage:\n" - << argv[0] - << " datatype " - " " - " " - << std::endl; - exit(-1); - } +int main(int argc, char **argv) { + if (argc != 8) { + std::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 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]); + const std::string data_path(argv[2]); + const std::string prefix_path(argv[3]); + const float sampling_rate = 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]); - if (std::string(argv[1]) == std::string("float")) - partition_with_ram_budget(data_path, sampling_rate, ram_budget, graph_degree, prefix_path, k_index); - else if (std::string(argv[1]) == std::string("int8")) - partition_with_ram_budget(data_path, sampling_rate, ram_budget, graph_degree, prefix_path, k_index); - else if (std::string(argv[1]) == std::string("uint8")) - 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; + if (std::string(argv[1]) == std::string("float")) + partition_with_ram_budget(data_path, sampling_rate, ram_budget, + graph_degree, prefix_path, k_index); + else if (std::string(argv[1]) == std::string("int8")) + partition_with_ram_budget(data_path, sampling_rate, ram_budget, + graph_degree, prefix_path, k_index); + else if (std::string(argv[1]) == std::string("uint8")) + 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; } diff --git a/tests/utils/rand_data_gen.cpp b/tests/utils/rand_data_gen.cpp index ea2e67478..d4e81f82e 100644 --- a/tests/utils/rand_data_gen.cpp +++ b/tests/utils/rand_data_gen.cpp @@ -11,199 +11,174 @@ namespace po = boost::program_options; -int block_write_float(std::ofstream &writer, size_t ndims, size_t npts, float norm) -{ - auto vec = new float[ndims]; - - std::random_device rd{}; - std::mt19937 gen{rd()}; - std::normal_distribution<> normal_rand{0, 1}; - - for (size_t i = 0; i < npts; i++) - { - float sum = 0; - for (size_t d = 0; d < ndims; ++d) - vec[d] = normal_rand(gen); - for (size_t d = 0; d < ndims; ++d) - sum += vec[d] * vec[d]; - for (size_t d = 0; d < ndims; ++d) - vec[d] = vec[d] * norm / std::sqrt(sum); - - writer.write((char *)vec, ndims * sizeof(float)); - } +int block_write_float(std::ofstream &writer, size_t ndims, size_t npts, + float norm) { + auto vec = new float[ndims]; - delete[] vec; - return 0; -} + std::random_device rd{}; + std::mt19937 gen{rd()}; + std::normal_distribution<> normal_rand{0, 1}; -int block_write_int8(std::ofstream &writer, size_t ndims, size_t npts, float norm) -{ - auto vec = new float[ndims]; - auto vec_T = new int8_t[ndims]; - - std::random_device rd{}; - std::mt19937 gen{rd()}; - std::normal_distribution<> normal_rand{0, 1}; - - for (size_t i = 0; i < npts; i++) - { - float sum = 0; - for (size_t d = 0; d < ndims; ++d) - vec[d] = normal_rand(gen); - for (size_t d = 0; d < ndims; ++d) - sum += vec[d] * vec[d]; - for (size_t d = 0; d < ndims; ++d) - vec[d] = vec[d] * norm / std::sqrt(sum); - - for (size_t d = 0; d < ndims; ++d) - { - vec_T[d] = std::round(vec[d]); - } - - writer.write((char *)vec_T, ndims * sizeof(int8_t)); - } + for (size_t i = 0; i < npts; i++) { + float sum = 0; + for (size_t d = 0; d < ndims; ++d) vec[d] = normal_rand(gen); + for (size_t d = 0; d < ndims; ++d) sum += vec[d] * vec[d]; + for (size_t d = 0; d < ndims; ++d) vec[d] = vec[d] * norm / std::sqrt(sum); - delete[] vec; - delete[] vec_T; - return 0; + writer.write((char *)vec, ndims * sizeof(float)); + } + + delete[] vec; + return 0; } -int block_write_uint8(std::ofstream &writer, size_t ndims, size_t npts, float norm) -{ - auto vec = new float[ndims]; - auto vec_T = new int8_t[ndims]; - - std::random_device rd{}; - std::mt19937 gen{rd()}; - std::normal_distribution<> normal_rand{0, 1}; - - for (size_t i = 0; i < npts; i++) - { - float sum = 0; - for (size_t d = 0; d < ndims; ++d) - vec[d] = normal_rand(gen); - for (size_t d = 0; d < ndims; ++d) - sum += vec[d] * vec[d]; - for (size_t d = 0; d < ndims; ++d) - vec[d] = vec[d] * norm / std::sqrt(sum); - - for (size_t d = 0; d < ndims; ++d) - { - vec_T[d] = 128 + std::round(vec[d]); - } - - writer.write((char *)vec_T, ndims * sizeof(uint8_t)); +int block_write_int8(std::ofstream &writer, size_t ndims, size_t npts, + float norm) { + auto vec = new float[ndims]; + auto vec_T = new int8_t[ndims]; + + std::random_device rd{}; + std::mt19937 gen{rd()}; + std::normal_distribution<> normal_rand{0, 1}; + + for (size_t i = 0; i < npts; i++) { + float sum = 0; + for (size_t d = 0; d < ndims; ++d) vec[d] = normal_rand(gen); + for (size_t d = 0; d < ndims; ++d) sum += vec[d] * vec[d]; + for (size_t d = 0; d < ndims; ++d) vec[d] = vec[d] * norm / std::sqrt(sum); + + for (size_t d = 0; d < ndims; ++d) { + vec_T[d] = std::round(vec[d]); } - delete[] vec; - delete[] vec_T; - return 0; + writer.write((char *)vec_T, ndims * sizeof(int8_t)); + } + + delete[] vec; + delete[] vec_T; + return 0; } -int main(int argc, char **argv) -{ - std::string data_type, output_file; - size_t ndims, npts; - float norm; - - try - { - po::options_description desc{"Arguments"}; - - desc.add_options()("help,h", "Print information on arguments"); - - desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); - desc.add_options()("output_file", po::value(&output_file)->required(), - "File name for saving the random vectors"); - desc.add_options()("ndims,D", po::value(&ndims)->required(), "Dimensoinality of the vector"); - desc.add_options()("npts,N", po::value(&npts)->required(), "Number of vectors"); - desc.add_options()("norm", po::value(&norm)->required(), "Norm of the vectors"); - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) - { - std::cout << desc; - return 0; - } - po::notify(vm); - } - catch (const std::exception &ex) - { - std::cerr << ex.what() << '\n'; - return -1; - } +int block_write_uint8(std::ofstream &writer, size_t ndims, size_t npts, + float norm) { + auto vec = new float[ndims]; + auto vec_T = new int8_t[ndims]; - if (data_type != std::string("float") && data_type != std::string("int8") && data_type != std::string("uint8")) - { - std::cout << "Unsupported type. float, int8 and uint8 types are supported." << std::endl; - return -1; - } + std::random_device rd{}; + std::mt19937 gen{rd()}; + std::normal_distribution<> normal_rand{0, 1}; - if (norm <= 0.0) - { - std::cerr << "Error: Norm must be a positive number" << std::endl; - return -1; - } + for (size_t i = 0; i < npts; i++) { + float sum = 0; + for (size_t d = 0; d < ndims; ++d) vec[d] = normal_rand(gen); + for (size_t d = 0; d < ndims; ++d) sum += vec[d] * vec[d]; + for (size_t d = 0; d < ndims; ++d) vec[d] = vec[d] * norm / std::sqrt(sum); - if (data_type == std::string("int8") || data_type == std::string("uint8")) - { - if (norm > 127) - { - std::cerr << "Error: for int8/uint8 datatypes, L2 norm can not be " - "greater " - "than 127" - << std::endl; - return -1; - } + for (size_t d = 0; d < ndims; ++d) { + vec_T[d] = 128 + std::round(vec[d]); } - try - { - std::ofstream writer; - writer.exceptions(std::ofstream::failbit | std::ofstream::badbit); - writer.open(output_file, std::ios::binary); - auto npts_u32 = (uint32_t)npts; - auto ndims_u32 = (uint32_t)ndims; - writer.write((char *)&npts_u32, sizeof(uint32_t)); - writer.write((char *)&ndims_u32, sizeof(uint32_t)); - - size_t blk_size = 131072; - size_t nblks = ROUND_UP(npts, blk_size) / blk_size; - std::cout << "# blks: " << nblks << std::endl; - - int ret = 0; - for (size_t i = 0; i < nblks; i++) - { - size_t cblk_size = std::min(npts - i * blk_size, blk_size); - if (data_type == std::string("float")) - { - ret = block_write_float(writer, ndims, cblk_size, norm); - } - else if (data_type == std::string("int8")) - { - ret = block_write_int8(writer, ndims, cblk_size, norm); - } - else if (data_type == std::string("uint8")) - { - ret = block_write_uint8(writer, ndims, cblk_size, norm); - } - if (ret == 0) - std::cout << "Block #" << i << " written" << std::endl; - else - { - writer.close(); - std::cout << "failed to write" << std::endl; - return -1; - } - } - writer.close(); + writer.write((char *)vec_T, ndims * sizeof(uint8_t)); + } + + delete[] vec; + delete[] vec_T; + return 0; +} + +int main(int argc, char **argv) { + std::string data_type, output_file; + size_t ndims, npts; + float norm; + + try { + po::options_description desc{"Arguments"}; + + desc.add_options()("help,h", "Print information on arguments"); + + desc.add_options()("data_type", + po::value(&data_type)->required(), + "data type "); + desc.add_options()("output_file", + po::value(&output_file)->required(), + "File name for saving the random vectors"); + desc.add_options()("ndims,D", po::value(&ndims)->required(), + "Dimensoinality of the vector"); + desc.add_options()("npts,N", po::value(&npts)->required(), + "Number of vectors"); + desc.add_options()("norm", po::value(&norm)->required(), + "Norm of the vectors"); + po::variables_map vm; + po::store(po::parse_command_line(argc, argv, desc), vm); + if (vm.count("help")) { + std::cout << desc; + return 0; + } + po::notify(vm); + } catch (const std::exception &ex) { + std::cerr << ex.what() << '\n'; + return -1; + } + + if (data_type != std::string("float") && data_type != std::string("int8") && + data_type != std::string("uint8")) { + std::cout << "Unsupported type. float, int8 and uint8 types are supported." + << std::endl; + return -1; + } + + if (norm <= 0.0) { + std::cerr << "Error: Norm must be a positive number" << std::endl; + return -1; + } + + if (data_type == std::string("int8") || data_type == std::string("uint8")) { + if (norm > 127) { + std::cerr << "Error: for int8/uint8 datatypes, L2 norm can not be " + "greater " + "than 127" + << std::endl; + return -1; } - catch (const std::exception &e) - { - std::cout << std::string(e.what()) << std::endl; - diskann::cerr << "Index build failed." << std::endl; + } + + try { + std::ofstream writer; + writer.exceptions(std::ofstream::failbit | std::ofstream::badbit); + writer.open(output_file, std::ios::binary); + auto npts_u32 = (uint32_t)npts; + auto ndims_u32 = (uint32_t)ndims; + writer.write((char *)&npts_u32, sizeof(uint32_t)); + writer.write((char *)&ndims_u32, sizeof(uint32_t)); + + size_t blk_size = 131072; + size_t nblks = ROUND_UP(npts, blk_size) / blk_size; + std::cout << "# blks: " << nblks << std::endl; + + int ret = 0; + for (size_t i = 0; i < nblks; i++) { + size_t cblk_size = std::min(npts - i * blk_size, blk_size); + if (data_type == std::string("float")) { + ret = block_write_float(writer, ndims, cblk_size, norm); + } else if (data_type == std::string("int8")) { + ret = block_write_int8(writer, ndims, cblk_size, norm); + } else if (data_type == std::string("uint8")) { + ret = block_write_uint8(writer, ndims, cblk_size, norm); + } + if (ret == 0) + std::cout << "Block #" << i << " written" << std::endl; + else { + writer.close(); + std::cout << "failed to write" << std::endl; return -1; + } } - - return 0; + writer.close(); + } catch (const std::exception &e) { + std::cout << std::string(e.what()) << std::endl; + diskann::cerr << "Index build failed." << std::endl; + return -1; + } + + return 0; } diff --git a/tests/utils/simulate_aggregate_recall.cpp b/tests/utils/simulate_aggregate_recall.cpp index 2e7645ae3..7965a4854 100644 --- a/tests/utils/simulate_aggregate_recall.cpp +++ b/tests/utils/simulate_aggregate_recall.cpp @@ -6,80 +6,73 @@ #include #include -inline float aggregate_recall(const uint32_t k_aggr, const uint32_t k, const uint32_t npart, uint32_t *count, - const std::vector &recalls) -{ - float found = 0; - for (uint32_t i = 0; i < npart; ++i) - { - size_t max_found = std::min(count[i], k); - found += recalls[max_found - 1] * max_found; - } - return found / (float)k_aggr; +inline float aggregate_recall(const uint32_t k_aggr, const uint32_t k, + const uint32_t npart, uint32_t *count, + const std::vector &recalls) { + float found = 0; + for (uint32_t i = 0; i < npart; ++i) { + size_t max_found = std::min(count[i], k); + found += recalls[max_found - 1] * max_found; + } + return found / (float)k_aggr; } -void simulate(const uint32_t k_aggr, const uint32_t k, const uint32_t npart, const uint32_t nsim, - const std::vector &recalls) -{ - std::random_device r; - std::default_random_engine randeng(r()); - std::uniform_int_distribution uniform_dist(0, npart - 1); +void simulate(const uint32_t k_aggr, const uint32_t k, const uint32_t npart, + const uint32_t nsim, const std::vector &recalls) { + std::random_device r; + std::default_random_engine randeng(r()); + std::uniform_int_distribution uniform_dist(0, npart - 1); - uint32_t *count = new uint32_t[npart]; - double aggr_recall = 0; + uint32_t *count = new uint32_t[npart]; + double aggr_recall = 0; - for (uint32_t i = 0; i < nsim; ++i) - { - for (uint32_t p = 0; p < npart; ++p) - { - count[p] = 0; - } - for (uint32_t t = 0; t < k_aggr; ++t) - { - count[uniform_dist(randeng)]++; - } - aggr_recall += aggregate_recall(k_aggr, k, npart, count, recalls); + for (uint32_t i = 0; i < nsim; ++i) { + for (uint32_t p = 0; p < npart; ++p) { + count[p] = 0; + } + for (uint32_t t = 0; t < k_aggr; ++t) { + count[uniform_dist(randeng)]++; } + aggr_recall += aggregate_recall(k_aggr, k, npart, count, recalls); + } - std::cout << "Aggregate recall is " << aggr_recall / (double)nsim << std::endl; - delete[] count; + std::cout << "Aggregate recall is " << aggr_recall / (double)nsim + << std::endl; + delete[] count; } -int main(int argc, char **argv) -{ - if (argc < 6) - { - std::cout << argv[0] << " k_aggregate k_out npart nsim recall@1 recall@2 ... recall@k" << std::endl; - exit(-1); - } +int main(int argc, char **argv) { + if (argc < 6) { + std::cout << argv[0] + << " k_aggregate k_out npart nsim recall@1 recall@2 ... recall@k" + << std::endl; + exit(-1); + } - const uint32_t k_aggr = atoi(argv[1]); - const uint32_t k = atoi(argv[2]); - const uint32_t npart = atoi(argv[3]); - const uint32_t nsim = atoi(argv[4]); + const uint32_t k_aggr = atoi(argv[1]); + const uint32_t k = atoi(argv[2]); + const uint32_t npart = atoi(argv[3]); + const uint32_t nsim = atoi(argv[4]); - std::vector recalls; - for (int ctr = 5; ctr < argc; ctr++) - { - recalls.push_back(atof(argv[ctr])); - } + std::vector recalls; + for (int ctr = 5; ctr < argc; ctr++) { + recalls.push_back(atof(argv[ctr])); + } - if (recalls.size() != k) - { - std::cerr << "Please input k numbers for recall@1, recall@2 .. recall@k" << std::endl; - } - if (k_aggr > npart * k) - { - std::cerr << "k_aggr must be <= k * npart" << std::endl; - exit(-1); - } - if (nsim <= npart * k_aggr) - { - std::cerr << "Choose nsim > npart*k_aggr" << std::endl; - exit(-1); - } + if (recalls.size() != k) { + std::cerr << "Please input k numbers for recall@1, recall@2 .. recall@k" + << std::endl; + } + if (k_aggr > npart * k) { + std::cerr << "k_aggr must be <= k * npart" << std::endl; + exit(-1); + } + if (nsim <= npart * k_aggr) { + std::cerr << "Choose nsim > npart*k_aggr" << std::endl; + exit(-1); + } - simulate(k_aggr, k, npart, nsim, recalls); + simulate(k_aggr, k, npart, nsim, recalls); - return 0; + return 0; } diff --git a/tests/utils/stats_label_data.cpp b/tests/utils/stats_label_data.cpp index 129d5bcb2..1a11238a1 100644 --- a/tests/utils/stats_label_data.cpp +++ b/tests/utils/stats_label_data.cpp @@ -28,120 +28,123 @@ #endif namespace po = boost::program_options; -void stats_analysis(const std::string labels_file, std::string univeral_label, uint32_t density = 10) -{ - std::string token, line; - std::ifstream labels_stream(labels_file); - std::unordered_map label_counts; - std::string label_with_max_points; - uint32_t max_points = 0; - long long sum = 0; - long long point_cnt = 0; - float avg_labels_per_pt, mean_label_size; +void stats_analysis(const std::string labels_file, std::string univeral_label, + uint32_t density = 10) { + std::string token, line; + std::ifstream labels_stream(labels_file); + std::unordered_map label_counts; + std::string label_with_max_points; + uint32_t max_points = 0; + long long sum = 0; + long long point_cnt = 0; + float avg_labels_per_pt, mean_label_size; - std::vector labels_per_point; - uint32_t dense_pts = 0; - if (labels_stream.is_open()) - { - while (getline(labels_stream, line)) - { - point_cnt++; - std::stringstream iss(line); - uint32_t lbl_cnt = 0; - while (getline(iss, token, ',')) - { - lbl_cnt++; - token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); - token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); - if (label_counts.find(token) == label_counts.end()) - label_counts[token] = 0; - label_counts[token]++; - } - if (lbl_cnt >= density) - { - dense_pts++; - } - labels_per_point.emplace_back(lbl_cnt); - } + std::vector labels_per_point; + uint32_t dense_pts = 0; + if (labels_stream.is_open()) { + while (getline(labels_stream, line)) { + point_cnt++; + std::stringstream iss(line); + uint32_t lbl_cnt = 0; + while (getline(iss, token, ',')) { + lbl_cnt++; + token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); + token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); + if (label_counts.find(token) == label_counts.end()) + label_counts[token] = 0; + label_counts[token]++; + } + if (lbl_cnt >= density) { + dense_pts++; + } + labels_per_point.emplace_back(lbl_cnt); } + } - std::cout << "fraction of dense points with >= " << density - << " labels = " << (float)dense_pts / (float)labels_per_point.size() << std::endl; - std::sort(labels_per_point.begin(), labels_per_point.end()); + std::cout << "fraction of dense points with >= " << density + << " labels = " << (float)dense_pts / (float)labels_per_point.size() + << std::endl; + std::sort(labels_per_point.begin(), labels_per_point.end()); - std::vector> label_count_vec; + std::vector> label_count_vec; - for (auto it = label_counts.begin(); it != label_counts.end(); it++) - { - auto &lbl = *it; - label_count_vec.emplace_back(std::make_pair(lbl.first, lbl.second)); - if (lbl.second > max_points) - { - max_points = lbl.second; - label_with_max_points = lbl.first; - } - sum += lbl.second; + for (auto it = label_counts.begin(); it != label_counts.end(); it++) { + auto &lbl = *it; + label_count_vec.emplace_back(std::make_pair(lbl.first, lbl.second)); + if (lbl.second > max_points) { + max_points = lbl.second; + label_with_max_points = lbl.first; } + sum += lbl.second; + } - sort(label_count_vec.begin(), label_count_vec.end(), - [](const std::pair &lhs, const std::pair &rhs) { - return lhs.second < rhs.second; - }); + sort(label_count_vec.begin(), label_count_vec.end(), + [](const std::pair &lhs, + const std::pair &rhs) { + return lhs.second < rhs.second; + }); - for (float p = 0; p < 1; p += 0.05) - { - std::cout << "Percentile " << (100 * p) << "\t" << label_count_vec[(size_t)(p * label_count_vec.size())].first - << " with count=" << label_count_vec[(size_t)(p * label_count_vec.size())].second << std::endl; - } + for (float p = 0; p < 1; p += 0.05) { + std::cout << "Percentile " << (100 * p) << "\t" + << label_count_vec[(size_t)(p * label_count_vec.size())].first + << " with count=" + << label_count_vec[(size_t)(p * label_count_vec.size())].second + << std::endl; + } - std::cout << "Most common label " - << "\t" << label_count_vec[label_count_vec.size() - 1].first - << " with count=" << label_count_vec[label_count_vec.size() - 1].second << std::endl; - if (label_count_vec.size() > 1) - std::cout << "Second common label " - << "\t" << label_count_vec[label_count_vec.size() - 2].first - << " with count=" << label_count_vec[label_count_vec.size() - 2].second << std::endl; - if (label_count_vec.size() > 2) - std::cout << "Third common label " - << "\t" << label_count_vec[label_count_vec.size() - 3].first - << " with count=" << label_count_vec[label_count_vec.size() - 3].second << std::endl; - avg_labels_per_pt = (sum) / (float)point_cnt; - mean_label_size = (sum) / label_counts.size(); - std::cout << "Total number of points = " << point_cnt << ", number of labels = " << label_counts.size() + std::cout << "Most common label " + << "\t" << label_count_vec[label_count_vec.size() - 1].first + << " with count=" + << label_count_vec[label_count_vec.size() - 1].second << std::endl; + if (label_count_vec.size() > 1) + std::cout << "Second common label " + << "\t" << label_count_vec[label_count_vec.size() - 2].first + << " with count=" + << label_count_vec[label_count_vec.size() - 2].second << std::endl; - std::cout << "Average number of labels per point = " << avg_labels_per_pt << std::endl; - std::cout << "Mean label size excluding 0 = " << mean_label_size << std::endl; - std::cout << "Most popular label is " << label_with_max_points << " with " << max_points << " pts" << std::endl; + if (label_count_vec.size() > 2) + std::cout << "Third common label " + << "\t" << label_count_vec[label_count_vec.size() - 3].first + << " with count=" + << label_count_vec[label_count_vec.size() - 3].second + << std::endl; + avg_labels_per_pt = (sum) / (float)point_cnt; + mean_label_size = (sum) / label_counts.size(); + std::cout << "Total number of points = " << point_cnt + << ", number of labels = " << label_counts.size() << std::endl; + std::cout << "Average number of labels per point = " << avg_labels_per_pt + << std::endl; + std::cout << "Mean label size excluding 0 = " << mean_label_size << std::endl; + std::cout << "Most popular label is " << label_with_max_points << " with " + << max_points << " pts" << std::endl; } -int main(int argc, char **argv) -{ - std::string labels_file, universal_label; - uint32_t density; +int main(int argc, char **argv) { + std::string labels_file, universal_label; + uint32_t density; - po::options_description desc{"Arguments"}; - try - { - desc.add_options()("help,h", "Print information on arguments"); - desc.add_options()("labels_file", po::value(&labels_file)->required(), - "path to labels data file."); - desc.add_options()("universal_label", po::value(&universal_label)->required(), - "Universal label used in labels file."); - desc.add_options()("density", po::value(&density)->default_value(1), - "Number of labels each point in labels file, defaults to 1"); - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) - { - std::cout << desc; - return 0; - } - po::notify(vm); - } - catch (const std::exception &e) - { - std::cerr << e.what() << '\n'; - return -1; + po::options_description desc{"Arguments"}; + try { + desc.add_options()("help,h", "Print information on arguments"); + desc.add_options()("labels_file", + po::value(&labels_file)->required(), + "path to labels data file."); + desc.add_options()("universal_label", + po::value(&universal_label)->required(), + "Universal label used in labels file."); + desc.add_options()( + "density", po::value(&density)->default_value(1), + "Number of labels each point in labels file, defaults to 1"); + po::variables_map vm; + po::store(po::parse_command_line(argc, argv, desc), vm); + if (vm.count("help")) { + std::cout << desc; + return 0; } - stats_analysis(labels_file, universal_label, density); + po::notify(vm); + } catch (const std::exception &e) { + std::cerr << e.what() << '\n'; + return -1; + } + stats_analysis(labels_file, universal_label, density); } diff --git a/tests/utils/tsv_to_bin.cpp b/tests/utils/tsv_to_bin.cpp index c590a8f73..43d304f9e 100644 --- a/tests/utils/tsv_to_bin.cpp +++ b/tests/utils/tsv_to_bin.cpp @@ -4,118 +4,105 @@ #include #include "utils.h" -void block_convert_float(std::ifstream &reader, std::ofstream &writer, size_t npts, size_t ndims) -{ - auto read_buf = new float[npts * (ndims + 1)]; - - auto cursor = read_buf; - float val; - - for (size_t i = 0; i < npts; i++) - { - for (size_t d = 0; d < ndims; ++d) - { - reader >> val; - *cursor = val; - cursor++; - } +void block_convert_float(std::ifstream &reader, std::ofstream &writer, + size_t npts, size_t ndims) { + auto read_buf = new float[npts * (ndims + 1)]; + + auto cursor = read_buf; + float val; + + for (size_t i = 0; i < npts; i++) { + for (size_t d = 0; d < ndims; ++d) { + reader >> val; + *cursor = val; + cursor++; } - writer.write((char *)read_buf, npts * ndims * sizeof(float)); - delete[] read_buf; + } + writer.write((char *)read_buf, npts * ndims * sizeof(float)); + delete[] read_buf; } -void block_convert_int8(std::ifstream &reader, std::ofstream &writer, size_t npts, size_t ndims) -{ - auto read_buf = new int8_t[npts * (ndims + 1)]; - - auto cursor = read_buf; - int val; - - for (size_t i = 0; i < npts; i++) - { - for (size_t d = 0; d < ndims; ++d) - { - reader >> val; - *cursor = (int8_t)val; - cursor++; - } - } - writer.write((char *)read_buf, npts * ndims * sizeof(uint8_t)); - delete[] read_buf; -} +void block_convert_int8(std::ifstream &reader, std::ofstream &writer, + size_t npts, size_t ndims) { + auto read_buf = new int8_t[npts * (ndims + 1)]; + + auto cursor = read_buf; + int val; -void block_convert_uint8(std::ifstream &reader, std::ofstream &writer, size_t npts, size_t ndims) -{ - auto read_buf = new uint8_t[npts * (ndims + 1)]; - - auto cursor = read_buf; - int val; - - for (size_t i = 0; i < npts; i++) - { - for (size_t d = 0; d < ndims; ++d) - { - reader >> val; - *cursor = (uint8_t)val; - cursor++; - } + for (size_t i = 0; i < npts; i++) { + for (size_t d = 0; d < ndims; ++d) { + reader >> val; + *cursor = (int8_t)val; + cursor++; } - writer.write((char *)read_buf, npts * ndims * sizeof(uint8_t)); - delete[] read_buf; + } + writer.write((char *)read_buf, npts * ndims * sizeof(uint8_t)); + delete[] read_buf; } -int main(int argc, char **argv) -{ - if (argc != 6) - { - std::cout << argv[0] - << " input_filename.tsv output_filename.bin " - "dim num_pts>" - << std::endl; - exit(-1); - } +void block_convert_uint8(std::ifstream &reader, std::ofstream &writer, + size_t npts, size_t ndims) { + auto read_buf = new uint8_t[npts * (ndims + 1)]; + + auto cursor = read_buf; + int val; - 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; + for (size_t i = 0; i < npts; i++) { + for (size_t d = 0; d < ndims; ++d) { + reader >> val; + *cursor = (uint8_t)val; + cursor++; } + } + writer.write((char *)read_buf, npts * ndims * sizeof(uint8_t)); + delete[] read_buf; +} - size_t ndims = atoi(argv[4]); - size_t npts = atoi(argv[5]); - - std::ifstream reader(argv[2], std::ios::binary | std::ios::ate); - // size_t fsize = reader.tellg(); - reader.seekg(0, std::ios::beg); - reader.seekg(0, std::ios::beg); - - size_t blk_size = 131072; - size_t nblks = ROUND_UP(npts, blk_size) / blk_size; - std::cout << "# blks: " << nblks << std::endl; - std::ofstream writer(argv[3], std::ios::binary); - auto npts_u32 = (uint32_t)npts; - auto ndims_u32 = (uint32_t)ndims; - writer.write((char *)&npts_u32, sizeof(uint32_t)); - writer.write((char *)&ndims_u32, sizeof(uint32_t)); - - for (size_t i = 0; i < nblks; i++) - { - size_t cblk_size = std::min(npts - i * blk_size, blk_size); - if (std::string(argv[1]) == std::string("float")) - { - block_convert_float(reader, writer, cblk_size, ndims); - } - else if (std::string(argv[1]) == std::string("int8")) - { - block_convert_int8(reader, writer, cblk_size, ndims); - } - else if (std::string(argv[1]) == std::string("uint8")) - { - block_convert_uint8(reader, writer, cblk_size, ndims); - } - std::cout << "Block #" << i << " written" << std::endl; +int main(int argc, char **argv) { + if (argc != 6) { + std::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; + } + + size_t ndims = atoi(argv[4]); + size_t npts = atoi(argv[5]); + + std::ifstream reader(argv[2], std::ios::binary | std::ios::ate); + // size_t fsize = reader.tellg(); + reader.seekg(0, std::ios::beg); + reader.seekg(0, std::ios::beg); + + size_t blk_size = 131072; + size_t nblks = ROUND_UP(npts, blk_size) / blk_size; + std::cout << "# blks: " << nblks << std::endl; + std::ofstream writer(argv[3], std::ios::binary); + auto npts_u32 = (uint32_t)npts; + auto ndims_u32 = (uint32_t)ndims; + writer.write((char *)&npts_u32, sizeof(uint32_t)); + writer.write((char *)&ndims_u32, sizeof(uint32_t)); + + for (size_t i = 0; i < nblks; i++) { + size_t cblk_size = std::min(npts - i * blk_size, blk_size); + if (std::string(argv[1]) == std::string("float")) { + block_convert_float(reader, writer, cblk_size, ndims); + } else if (std::string(argv[1]) == std::string("int8")) { + block_convert_int8(reader, writer, cblk_size, ndims); + } else if (std::string(argv[1]) == std::string("uint8")) { + block_convert_uint8(reader, writer, cblk_size, ndims); } + std::cout << "Block #" << i << " written" << std::endl; + } - reader.close(); - writer.close(); + reader.close(); + writer.close(); } diff --git a/tests/utils/uint32_to_uint8.cpp b/tests/utils/uint32_to_uint8.cpp index 87b6fb8ed..3b42e77fd 100644 --- a/tests/utils/uint32_to_uint8.cpp +++ b/tests/utils/uint32_to_uint8.cpp @@ -4,20 +4,18 @@ #include #include "utils.h" -int main(int argc, char **argv) -{ - if (argc != 3) - { - std::cout << argv[0] << " input_uint32_bin output_int8_bin" << std::endl; - exit(-1); - } +int main(int argc, char **argv) { + if (argc != 3) { + std::cout << argv[0] << " input_uint32_bin output_int8_bin" << std::endl; + exit(-1); + } - uint32_t *input; - size_t npts, nd; - diskann::load_bin(argv[1], input, npts, nd); - uint8_t *output = new uint8_t[npts * nd]; - diskann::convert_types(input, output, npts, nd); - diskann::save_bin(argv[2], output, npts, nd); - delete[] output; - delete[] input; + uint32_t *input; + size_t npts, nd; + diskann::load_bin(argv[1], input, npts, nd); + uint8_t *output = new uint8_t[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/uint8_to_float.cpp b/tests/utils/uint8_to_float.cpp index 6415b7c92..52b4ca083 100644 --- a/tests/utils/uint8_to_float.cpp +++ b/tests/utils/uint8_to_float.cpp @@ -4,20 +4,18 @@ #include #include "utils.h" -int main(int argc, char **argv) -{ - if (argc != 3) - { - std::cout << argv[0] << " input_uint8_bin output_float_bin" << std::endl; - exit(-1); - } +int main(int argc, char **argv) { + if (argc != 3) { + std::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; + 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/vector_analysis.cpp b/tests/utils/vector_analysis.cpp index 5e4cb9bf4..9427a4689 100644 --- a/tests/utils/vector_analysis.cpp +++ b/tests/utils/vector_analysis.cpp @@ -20,144 +20,129 @@ #include "partition.h" #include "utils.h" -template int analyze_norm(std::string base_file) -{ - std::cout << "Analyzing data norms" << std::endl; - T *data; - size_t npts, ndims; - diskann::load_bin(base_file, data, npts, ndims); - std::vector norms(npts, 0); +template +int analyze_norm(std::string base_file) { + std::cout << "Analyzing data norms" << std::endl; + T *data; + size_t npts, ndims; + diskann::load_bin(base_file, data, npts, ndims); + std::vector norms(npts, 0); #pragma omp parallel for schedule(dynamic) - for (int64_t i = 0; i < (int64_t)npts; i++) - { - for (size_t d = 0; d < ndims; d++) - norms[i] += data[i * ndims + d] * data[i * ndims + d]; - norms[i] = std::sqrt(norms[i]); - } - std::sort(norms.begin(), norms.end()); - for (int p = 0; p < 100; p += 5) - std::cout << "percentile " << p << ": " << norms[std::floor((p / 100.0) * npts)] << std::endl; - std::cout << "percentile 100" - << ": " << norms[npts - 1] << std::endl; - delete[] data; - return 0; + for (int64_t i = 0; i < (int64_t)npts; i++) { + for (size_t d = 0; d < ndims; d++) + norms[i] += data[i * ndims + d] * data[i * ndims + d]; + norms[i] = std::sqrt(norms[i]); + } + std::sort(norms.begin(), norms.end()); + for (int p = 0; p < 100; p += 5) + std::cout << "percentile " << p << ": " + << norms[std::floor((p / 100.0) * npts)] << std::endl; + std::cout << "percentile 100" + << ": " << norms[npts - 1] << std::endl; + delete[] data; + return 0; } -template int normalize_base(std::string base_file, std::string out_file) -{ - std::cout << "Normalizing base" << std::endl; - T *data; - size_t npts, ndims; - diskann::load_bin(base_file, data, npts, ndims); - // std::vector norms(npts, 0); +template +int normalize_base(std::string base_file, std::string out_file) { + std::cout << "Normalizing base" << std::endl; + T *data; + size_t npts, ndims; + diskann::load_bin(base_file, data, npts, ndims); + // std::vector norms(npts, 0); #pragma omp parallel for schedule(dynamic) - for (int64_t i = 0; i < (int64_t)npts; i++) - { - float pt_norm = 0; - for (size_t d = 0; d < ndims; d++) - pt_norm += data[i * ndims + d] * data[i * ndims + d]; - pt_norm = std::sqrt(pt_norm); - for (size_t d = 0; d < ndims; d++) - data[i * ndims + d] = data[i * ndims + d] / pt_norm; - } - diskann::save_bin(out_file, data, npts, ndims); - delete[] data; - return 0; + for (int64_t i = 0; i < (int64_t)npts; i++) { + float pt_norm = 0; + for (size_t d = 0; d < ndims; d++) + pt_norm += data[i * ndims + d] * data[i * ndims + d]; + pt_norm = std::sqrt(pt_norm); + for (size_t d = 0; d < ndims; d++) + data[i * ndims + d] = data[i * ndims + d] / pt_norm; + } + diskann::save_bin(out_file, data, npts, ndims); + delete[] data; + return 0; } -template int augment_base(std::string base_file, std::string out_file, bool prep_base = true) -{ - std::cout << "Analyzing data norms" << std::endl; - T *data; - size_t npts, ndims; - diskann::load_bin(base_file, data, npts, ndims); - std::vector norms(npts, 0); - float max_norm = 0; +template +int augment_base(std::string base_file, std::string out_file, + bool prep_base = true) { + std::cout << "Analyzing data norms" << std::endl; + T *data; + size_t npts, ndims; + diskann::load_bin(base_file, data, npts, ndims); + std::vector norms(npts, 0); + float max_norm = 0; #pragma omp parallel for schedule(dynamic) - for (int64_t i = 0; i < (int64_t)npts; i++) - { - for (size_t d = 0; d < ndims; d++) - norms[i] += data[i * ndims + d] * data[i * ndims + d]; - max_norm = norms[i] > max_norm ? norms[i] : max_norm; + for (int64_t i = 0; i < (int64_t)npts; i++) { + for (size_t d = 0; d < ndims; d++) + norms[i] += data[i * ndims + d] * data[i * ndims + d]; + max_norm = norms[i] > max_norm ? norms[i] : max_norm; + } + // std::sort(norms.begin(), norms.end()); + max_norm = std::sqrt(max_norm); + std::cout << "Max norm: " << max_norm << std::endl; + T *new_data; + size_t newdims = ndims + 1; + new_data = new T[npts * newdims]; + for (size_t i = 0; i < npts; i++) { + if (prep_base) { + for (size_t j = 0; j < ndims; j++) { + new_data[i * newdims + j] = data[i * ndims + j] / max_norm; + } + float diff = 1 - (norms[i] / (max_norm * max_norm)); + diff = diff <= 0 ? 0 : std::sqrt(diff); + new_data[i * newdims + ndims] = diff; + if (diff <= 0) { + std::cout << i << " has large max norm, investigate if needed. diff = " + << diff << std::endl; + } + } else { + for (size_t j = 0; j < ndims; j++) { + new_data[i * newdims + j] = data[i * ndims + j] / std::sqrt(norms[i]); + } + new_data[i * newdims + ndims] = 0; } - // std::sort(norms.begin(), norms.end()); - max_norm = std::sqrt(max_norm); - std::cout << "Max norm: " << max_norm << std::endl; - T *new_data; - size_t newdims = ndims + 1; - new_data = new T[npts * newdims]; - for (size_t i = 0; i < npts; i++) - { - if (prep_base) - { - for (size_t j = 0; j < ndims; j++) - { - new_data[i * newdims + j] = data[i * ndims + j] / max_norm; - } - float diff = 1 - (norms[i] / (max_norm * max_norm)); - diff = diff <= 0 ? 0 : std::sqrt(diff); - new_data[i * newdims + ndims] = diff; - if (diff <= 0) - { - std::cout << i << " has large max norm, investigate if needed. diff = " << diff << std::endl; - } - } - else - { - for (size_t j = 0; j < ndims; j++) - { - new_data[i * newdims + j] = data[i * ndims + j] / std::sqrt(norms[i]); - } - new_data[i * newdims + ndims] = 0; - } - } - diskann::save_bin(out_file, new_data, npts, newdims); - delete[] new_data; - delete[] data; - return 0; + } + diskann::save_bin(out_file, new_data, npts, newdims); + delete[] new_data; + delete[] data; + return 0; } -template int aux_main(char **argv) -{ - std::string base_file(argv[2]); - uint32_t option = atoi(argv[3]); - if (option == 1) - analyze_norm(base_file); - else if (option == 2) - augment_base(base_file, std::string(argv[4]), true); - else if (option == 3) - augment_base(base_file, std::string(argv[4]), false); - else if (option == 4) - normalize_base(base_file, std::string(argv[4])); - return 0; +template +int aux_main(char **argv) { + std::string base_file(argv[2]); + uint32_t option = atoi(argv[3]); + if (option == 1) + analyze_norm(base_file); + else if (option == 2) + augment_base(base_file, std::string(argv[4]), true); + else if (option == 3) + augment_base(base_file, std::string(argv[4]), false); + else if (option == 4) + normalize_base(base_file, std::string(argv[4])); + return 0; } -int main(int argc, char **argv) -{ - if (argc < 4) - { - std::cout << argv[0] - << " data_type [float/int8/uint8] base_bin_file " - "[option: 1-norm analysis, 2-prep_base_for_mip, " - "3-prep_query_for_mip, 4-normalize-vecs] [out_file for " - "options 2/3/4]" - << std::endl; - exit(-1); - } +int main(int argc, char **argv) { + if (argc < 4) { + std::cout << argv[0] + << " data_type [float/int8/uint8] base_bin_file " + "[option: 1-norm analysis, 2-prep_base_for_mip, " + "3-prep_query_for_mip, 4-normalize-vecs] [out_file for " + "options 2/3/4]" + << std::endl; + exit(-1); + } - if (std::string(argv[1]) == std::string("float")) - { - aux_main(argv); - } - else if (std::string(argv[1]) == std::string("int8")) - { - aux_main(argv); - } - else if (std::string(argv[1]) == std::string("uint8")) - { - aux_main(argv); - } - else - std::cout << "Unsupported type. Use float/int8/uint8." << std::endl; - return 0; + if (std::string(argv[1]) == std::string("float")) { + aux_main(argv); + } else if (std::string(argv[1]) == std::string("int8")) { + aux_main(argv); + } else if (std::string(argv[1]) == std::string("uint8")) { + aux_main(argv); + } else + std::cout << "Unsupported type. Use float/int8/uint8." << std::endl; + return 0; } From 1c127900b4a55318c85cd51f9eb80fb70212ce36 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Wed, 19 Apr 2023 23:53:02 +0530 Subject: [PATCH 58/69] more clang-format --- include/restapi/common.h | 14 +- include/restapi/search_wrapper.h | 211 +++--- include/restapi/server.h | 76 +- python/src/diskann_bindings.cpp | 777 +++++++++++---------- src/restapi/search_wrapper.cpp | 285 ++++---- src/restapi/server.cpp | 434 ++++++------ tests/restapi/client.cpp | 168 +++-- tests/restapi/inmem_server.cpp | 205 +++--- tests/restapi/main.cpp | 106 ++- tests/restapi/multiple_ssdindex_server.cpp | 284 ++++---- tests/restapi/ssd_server.cpp | 210 +++--- 11 files changed, 1361 insertions(+), 1409 deletions(-) diff --git a/include/restapi/common.h b/include/restapi/common.h index b8339635a..640afd323 100644 --- a/include/restapi/common.h +++ b/include/restapi/common.h @@ -6,13 +6,15 @@ #include #include -namespace diskann -{ +namespace diskann { // Constants -static const std::string VECTOR_KEY = "query", K_KEY = "k", INDICES_KEY = "indices", DISTANCES_KEY = "distances", - TAGS_KEY = "tags", QUERY_ID_KEY = "query_id", ERROR_MESSAGE_KEY = "error", L_KEY = "Ls", - TIME_TAKEN_KEY = "time_taken_in_us", PARTITION_KEY = "partition", +static const std::string VECTOR_KEY = "query", K_KEY = "k", + INDICES_KEY = "indices", DISTANCES_KEY = "distances", + TAGS_KEY = "tags", QUERY_ID_KEY = "query_id", + ERROR_MESSAGE_KEY = "error", L_KEY = "Ls", + TIME_TAKEN_KEY = "time_taken_in_us", + PARTITION_KEY = "partition", UNKNOWN_ERROR = "unknown_error"; const unsigned int DEFAULT_L = 100; -} // namespace diskann \ No newline at end of file +} // namespace diskann \ No newline at end of file diff --git a/include/restapi/search_wrapper.h b/include/restapi/search_wrapper.h index ebd067d8a..9a8c71f14 100644 --- a/include/restapi/search_wrapper.h +++ b/include/restapi/search_wrapper.h @@ -10,131 +10,108 @@ #include #include -namespace diskann -{ -class SearchResult -{ - public: - SearchResult(unsigned int K, unsigned int elapsed_time_in_ms, const unsigned *const indices, - const float *const distances, const std::string *const tags = nullptr, - const unsigned *const partitions = nullptr); - - const std::vector &get_indices() const - { - return _indices; - } - const std::vector &get_distances() const - { - return _distances; - } - bool tags_enabled() const - { - return _tags_enabled; - } - const std::vector &get_tags() const - { - return _tags; - } - bool partitions_enabled() const - { - return _partitions_enabled; - } - const std::vector &get_partitions() const - { - return _partitions; - } - unsigned get_time() const - { - return _search_time_in_ms; - } - - private: - unsigned int _K; - unsigned int _search_time_in_ms; - std::vector _indices; - std::vector _distances; - - bool _tags_enabled; - std::vector _tags; - - bool _partitions_enabled; - std::vector _partitions; +namespace diskann { +class SearchResult { + public: + SearchResult(unsigned int K, unsigned int elapsed_time_in_ms, + const unsigned *const indices, const float *const distances, + const std::string *const tags = nullptr, + const unsigned *const partitions = nullptr); + + const std::vector &get_indices() const { return _indices; } + const std::vector &get_distances() const { return _distances; } + bool tags_enabled() const { return _tags_enabled; } + const std::vector &get_tags() const { return _tags; } + bool partitions_enabled() const { return _partitions_enabled; } + const std::vector &get_partitions() const { return _partitions; } + unsigned get_time() const { return _search_time_in_ms; } + + private: + unsigned int _K; + unsigned int _search_time_in_ms; + std::vector _indices; + std::vector _distances; + + bool _tags_enabled; + std::vector _tags; + + bool _partitions_enabled; + std::vector _partitions; }; -class SearchNotImplementedException : public std::logic_error -{ - private: - std::string _errormsg; - - public: - SearchNotImplementedException(const char *type) : std::logic_error("Not Implemented") - { - _errormsg = "Search with data type "; - _errormsg += std::string(type); - _errormsg += " not implemented : "; - _errormsg += __FUNCTION__; - } - - virtual const char *what() const throw() - { - return _errormsg.c_str(); - } +class SearchNotImplementedException : public std::logic_error { + private: + std::string _errormsg; + + public: + SearchNotImplementedException(const char *type) + : std::logic_error("Not Implemented") { + _errormsg = "Search with data type "; + _errormsg += std::string(type); + _errormsg += " not implemented : "; + _errormsg += __FUNCTION__; + } + + virtual const char *what() const throw() { return _errormsg.c_str(); } }; -class BaseSearch -{ - public: - BaseSearch(const std::string &tagsFile = nullptr); - virtual SearchResult search(const float *query, const unsigned int dimensions, const unsigned int K, - const unsigned int Ls) - { - throw SearchNotImplementedException("float"); - } - virtual SearchResult search(const int8_t *query, const unsigned int dimensions, const unsigned int K, - const unsigned int Ls) - { - throw SearchNotImplementedException("int8_t"); - } - - virtual SearchResult search(const uint8_t *query, const unsigned int dimensions, const unsigned int K, - const unsigned int Ls) - { - throw SearchNotImplementedException("uint8_t"); - } - - void lookup_tags(const unsigned K, const unsigned *indices, std::string *ret_tags); - - protected: - bool _tags_enabled; - std::vector _tags_str; +class BaseSearch { + public: + BaseSearch(const std::string &tagsFile = nullptr); + virtual SearchResult search(const float *query, const unsigned int dimensions, + const unsigned int K, const unsigned int Ls) { + throw SearchNotImplementedException("float"); + } + virtual SearchResult search(const int8_t *query, + const unsigned int dimensions, + const unsigned int K, const unsigned int Ls) { + throw SearchNotImplementedException("int8_t"); + } + + virtual SearchResult search(const uint8_t *query, + const unsigned int dimensions, + const unsigned int K, const unsigned int Ls) { + throw SearchNotImplementedException("uint8_t"); + } + + void lookup_tags(const unsigned K, const unsigned *indices, + std::string *ret_tags); + + protected: + bool _tags_enabled; + std::vector _tags_str; }; -template class InMemorySearch : public BaseSearch -{ - public: - InMemorySearch(const std::string &baseFile, const std::string &indexFile, const std::string &tagsFile, Metric m, - uint32_t num_threads, uint32_t search_l); - virtual ~InMemorySearch(); +template +class InMemorySearch : public BaseSearch { + public: + InMemorySearch(const std::string &baseFile, const std::string &indexFile, + const std::string &tagsFile, Metric m, uint32_t num_threads, + uint32_t search_l); + virtual ~InMemorySearch(); - SearchResult search(const T *query, const unsigned int dimensions, const unsigned int K, const unsigned int Ls); + SearchResult search(const T *query, const unsigned int dimensions, + const unsigned int K, const unsigned int Ls); - private: - unsigned int _dimensions, _numPoints; - std::unique_ptr> _index; + private: + unsigned int _dimensions, _numPoints; + std::unique_ptr> _index; }; -template class PQFlashSearch : public BaseSearch -{ - public: - PQFlashSearch(const std::string &indexPrefix, const unsigned num_nodes_to_cache, const unsigned num_threads, - const std::string &tagsFile, Metric m); - virtual ~PQFlashSearch(); - - SearchResult search(const T *query, const unsigned int dimensions, const unsigned int K, const unsigned int Ls); - - private: - unsigned int _dimensions, _numPoints; - std::unique_ptr> _index; - std::shared_ptr reader; +template +class PQFlashSearch : public BaseSearch { + public: + PQFlashSearch(const std::string &indexPrefix, + const unsigned num_nodes_to_cache, const unsigned num_threads, + const std::string &tagsFile, Metric m); + virtual ~PQFlashSearch(); + + SearchResult search(const T *query, const unsigned int dimensions, + const unsigned int K, const unsigned int Ls); + + private: + unsigned int _dimensions, _numPoints; + std::unique_ptr> _index; + std::shared_ptr reader; }; -} // namespace diskann +} // namespace diskann diff --git a/include/restapi/server.h b/include/restapi/server.h index 1d75847a2..2dcc0ebc1 100644 --- a/include/restapi/server.h +++ b/include/restapi/server.h @@ -6,40 +6,44 @@ #include #include -namespace diskann -{ -class Server -{ - public: - Server(web::uri &url, std::vector> &multi_searcher, - const std::string &typestring); - virtual ~Server(); - - pplx::task open(); - pplx::task close(); - - protected: - template void handle_post(web::http::http_request message); - - template - web::json::value toJsonArray(const std::vector &v, std::function valConverter); - web::json::value prepareResponse(const int64_t &queryId, const int k); - - template - void parseJson(const utility::string_t &body, unsigned int &k, int64_t &queryId, T *&queryVector, - unsigned int &dimensions, unsigned &Ls); - - web::json::value idsToJsonArray(const diskann::SearchResult &result); - web::json::value distancesToJsonArray(const diskann::SearchResult &result); - web::json::value tagsToJsonArray(const diskann::SearchResult &result); - web::json::value partitionsToJsonArray(const diskann::SearchResult &result); - - SearchResult aggregate_results(const unsigned K, const std::vector &results); - - private: - bool _isDebug; - std::unique_ptr _listener; - const bool _multi_search; - std::vector> _multi_searcher; +namespace diskann { +class Server { + public: + Server(web::uri &url, + std::vector> &multi_searcher, + const std::string &typestring); + virtual ~Server(); + + pplx::task open(); + pplx::task close(); + + protected: + template + void handle_post(web::http::http_request message); + + template + web::json::value toJsonArray( + const std::vector &v, + std::function valConverter); + web::json::value prepareResponse(const int64_t &queryId, const int k); + + template + void parseJson(const utility::string_t &body, unsigned int &k, + int64_t &queryId, T *&queryVector, unsigned int &dimensions, + unsigned &Ls); + + web::json::value idsToJsonArray(const diskann::SearchResult &result); + web::json::value distancesToJsonArray(const diskann::SearchResult &result); + web::json::value tagsToJsonArray(const diskann::SearchResult &result); + web::json::value partitionsToJsonArray(const diskann::SearchResult &result); + + SearchResult aggregate_results( + const unsigned K, const std::vector &results); + + private: + bool _isDebug; + std::unique_ptr _listener; + const bool _multi_search; + std::vector> _multi_searcher; }; -} // namespace diskann +} // namespace diskann diff --git a/python/src/diskann_bindings.cpp b/python/src/diskann_bindings.cpp index 6a977cac7..96e4709f9 100644 --- a/python/src/diskann_bindings.cpp +++ b/python/src/diskann_bindings.cpp @@ -30,417 +30,458 @@ PYBIND11_MAKE_OPAQUE(std::vector); namespace py = pybind11; using namespace diskann; -template struct DiskANNIndex -{ - PQFlashIndex *pq_flash_index; - std::shared_ptr reader; +template +struct DiskANNIndex { + PQFlashIndex *pq_flash_index; + std::shared_ptr reader; - DiskANNIndex(diskann::Metric metric) - { + DiskANNIndex(diskann::Metric metric) { #ifdef _WINDOWS - reader = std::make_shared(); + reader = std::make_shared(); #else - reader = std::make_shared(); + reader = std::make_shared(); #endif - pq_flash_index = new PQFlashIndex(reader, metric); - } + pq_flash_index = new PQFlashIndex(reader, metric); + } - ~DiskANNIndex() - { - delete pq_flash_index; - } + ~DiskANNIndex() { delete pq_flash_index; } - auto get_metric() - { - return pq_flash_index->get_metric(); - } + auto get_metric() { return pq_flash_index->get_metric(); } - void cache_bfs_levels(size_t num_nodes_to_cache) - { - std::vector node_list; - pq_flash_index->cache_bfs_levels(num_nodes_to_cache, node_list); - pq_flash_index->load_cache_list(node_list); - } + void cache_bfs_levels(size_t num_nodes_to_cache) { + std::vector node_list; + pq_flash_index->cache_bfs_levels(num_nodes_to_cache, node_list); + pq_flash_index->load_cache_list(node_list); + } - void cache_sample_paths(size_t num_nodes_to_cache, const std::string &warmup_query_file, uint32_t num_threads) - { - if (!file_exists(warmup_query_file)) - { - return; - } - - std::vector node_list; - pq_flash_index->generate_cache_list_from_sample_queries(warmup_query_file, 15, 4, num_nodes_to_cache, - num_threads, node_list); - pq_flash_index->load_cache_list(node_list); + void cache_sample_paths(size_t num_nodes_to_cache, + const std::string &warmup_query_file, + uint32_t num_threads) { + if (!file_exists(warmup_query_file)) { + return; } - int load_index(const std::string &index_path_prefix, const int num_threads, const size_t num_nodes_to_cache, - int cache_mechanism) - { - int load_success = pq_flash_index->load(num_threads, index_path_prefix.c_str()); - if (load_success != 0) - { - throw std::runtime_error("load_index failed."); - } - if (cache_mechanism == 0) - { - // Nothing to do - } - else if (cache_mechanism == 1) - { - std::string sample_file = index_path_prefix + std::string("_sample_data.bin"); - cache_sample_paths(num_nodes_to_cache, sample_file, num_threads); - } - else if (cache_mechanism == 2) - { - cache_bfs_levels(num_nodes_to_cache); - } - return 0; + std::vector node_list; + pq_flash_index->generate_cache_list_from_sample_queries( + warmup_query_file, 15, 4, num_nodes_to_cache, num_threads, node_list); + pq_flash_index->load_cache_list(node_list); + } + + int load_index(const std::string &index_path_prefix, const int num_threads, + const size_t num_nodes_to_cache, int cache_mechanism) { + int load_success = + pq_flash_index->load(num_threads, index_path_prefix.c_str()); + if (load_success != 0) { + throw std::runtime_error("load_index failed."); } + if (cache_mechanism == 0) { + // Nothing to do + } else if (cache_mechanism == 1) { + std::string sample_file = + index_path_prefix + std::string("_sample_data.bin"); + cache_sample_paths(num_nodes_to_cache, sample_file, num_threads); + } else if (cache_mechanism == 2) { + cache_bfs_levels(num_nodes_to_cache); + } + return 0; + } - auto search(py::array_t &query, const uint64_t knn, - const uint64_t l_search, const uint64_t beam_width) - { - py::array_t ids(knn); - py::array_t dists(knn); + auto search(py::array_t &query, + const uint64_t knn, const uint64_t l_search, + const uint64_t beam_width) { + py::array_t ids(knn); + py::array_t dists(knn); - std::vector u32_ids(knn); - std::vector u64_ids(knn); - QueryStats stats; + std::vector u32_ids(knn); + std::vector u64_ids(knn); + QueryStats stats; - pq_flash_index->cached_beam_search(query.data(), knn, l_search, u64_ids.data(), dists.mutable_data(), - beam_width, false, &stats); + pq_flash_index->cached_beam_search(query.data(), knn, l_search, + u64_ids.data(), dists.mutable_data(), + beam_width, false, &stats); - auto r = ids.mutable_unchecked<1>(); - for (uint64_t i = 0; i < knn; ++i) - r(i) = (unsigned)u64_ids[i]; + auto r = ids.mutable_unchecked<1>(); + for (uint64_t i = 0; i < knn; ++i) r(i) = (unsigned)u64_ids[i]; - return std::make_pair(ids, dists); - } + return std::make_pair(ids, dists); + } - auto batch_search(py::array_t &queries, const uint64_t num_queries, - const uint64_t knn, const uint64_t l_search, const uint64_t beam_width, const int num_threads) - { - py::array_t ids({num_queries, knn}); - py::array_t dists({num_queries, knn}); + auto batch_search( + py::array_t &queries, + const uint64_t num_queries, const uint64_t knn, const uint64_t l_search, + const uint64_t beam_width, const int num_threads) { + py::array_t ids({num_queries, knn}); + py::array_t dists({num_queries, knn}); - omp_set_num_threads(num_threads); + omp_set_num_threads(num_threads); - std::vector u64_ids(knn * num_queries); + std::vector u64_ids(knn * num_queries); #pragma omp parallel for schedule(dynamic, 1) - for (int64_t i = 0; i < (int64_t)num_queries; i++) - { - pq_flash_index->cached_beam_search(queries.data(i), knn, l_search, u64_ids.data() + i * knn, - dists.mutable_data(i), beam_width); - } - - auto r = ids.mutable_unchecked(); - for (uint64_t i = 0; i < num_queries; ++i) - for (uint64_t j = 0; j < knn; ++j) - r(i, j) = (unsigned)u64_ids[i * knn + j]; - - return std::make_pair(ids, dists); + for (int64_t i = 0; i < (int64_t)num_queries; i++) { + pq_flash_index->cached_beam_search(queries.data(i), knn, l_search, + u64_ids.data() + i * knn, + dists.mutable_data(i), beam_width); } + + auto r = ids.mutable_unchecked(); + for (uint64_t i = 0; i < num_queries; ++i) + for (uint64_t j = 0; j < knn; ++j) + r(i, j) = (unsigned)u64_ids[i * knn + j]; + + return std::make_pair(ids, dists); + } }; typedef uint32_t IdT; typedef uint32_t filterT; -template struct DynamicInMemIndex -{ - Index *_index; - const IndexWriteParameters write_params; - - DynamicInMemIndex(Metric m, const size_t dim, const size_t max_points, const IndexWriteParameters &index_parameters, - const uint32_t initial_search_list_size, const uint32_t search_threads, - const bool concurrent_consolidate) - : write_params(index_parameters) - { - _index = new Index(m, dim, max_points, - true, // dynamic_index - index_parameters, // used for insert - initial_search_list_size, // used to prepare the scratch space for searching. can / may be - // expanded if the search asks for a larger L. - search_threads, // also used for the scratch space - true, // enable_tags - concurrent_consolidate, - false, // pq_dist_build - 0, // num_pq_chunks - false); // use_opq = false - } - - ~DynamicInMemIndex() - { - delete _index; - } - - int insert(py::array_t &vector, const IdT id) - { - return _index->insert_point(vector.data(), id); - } - - int mark_deleted(const IdT id) - { - return _index->lazy_delete(id); - } - - auto search(py::array_t &query, const uint64_t knn, - const uint64_t l_search) - { - py::array_t ids(knn); - py::array_t dists(knn); - std::vector empty_vector; - _index->search_with_tags(query.data(), knn, l_search, ids.mutable_data(), dists.mutable_data(), empty_vector); - return std::make_pair(ids, dists); - } - - auto batch_search(py::array_t &queries, const uint64_t num_queries, - const uint64_t knn, const uint64_t l_search, const int num_threads) - { - py::array_t ids({num_queries, knn}); - py::array_t dists({num_queries, knn}); - std::vector empty_vector; - - omp_set_num_threads(num_threads); +template +struct DynamicInMemIndex { + Index *_index; + const IndexWriteParameters write_params; + + DynamicInMemIndex(Metric m, const size_t dim, const size_t max_points, + const IndexWriteParameters &index_parameters, + const uint32_t initial_search_list_size, + const uint32_t search_threads, + const bool concurrent_consolidate) + : write_params(index_parameters) { + _index = new Index( + m, dim, max_points, + true, // dynamic_index + index_parameters, // used for insert + initial_search_list_size, // used to prepare the scratch space for + // searching. can / may be expanded if the + // search asks for a larger L. + search_threads, // also used for the scratch space + true, // enable_tags + concurrent_consolidate, + false, // pq_dist_build + 0, // num_pq_chunks + false); // use_opq = false + } + + ~DynamicInMemIndex() { delete _index; } + + int insert(py::array_t &vector, + const IdT id) { + return _index->insert_point(vector.data(), id); + } + + int mark_deleted(const IdT id) { return _index->lazy_delete(id); } + + auto search(py::array_t &query, + const uint64_t knn, const uint64_t l_search) { + py::array_t ids(knn); + py::array_t dists(knn); + std::vector empty_vector; + _index->search_with_tags(query.data(), knn, l_search, ids.mutable_data(), + dists.mutable_data(), empty_vector); + return std::make_pair(ids, dists); + } + + auto batch_search( + py::array_t &queries, + const uint64_t num_queries, const uint64_t knn, const uint64_t l_search, + const int num_threads) { + py::array_t ids({num_queries, knn}); + py::array_t dists({num_queries, knn}); + std::vector empty_vector; + + omp_set_num_threads(num_threads); #pragma omp parallel for schedule(dynamic, 1) - for (int64_t i = 0; i < (int64_t)num_queries; i++) - { - _index->search_with_tags(queries.data(i), knn, l_search, ids.mutable_data(i), dists.mutable_data(i), - empty_vector); - } - - return std::make_pair(ids, dists); + for (int64_t i = 0; i < (int64_t)num_queries; i++) { + _index->search_with_tags(queries.data(i), knn, l_search, + ids.mutable_data(i), dists.mutable_data(i), + empty_vector); } - auto consolidate_delete() - { - return _index->consolidate_deletes(write_params); - } -}; - -template struct StaticInMemIndex -{ - Index *_index; - - StaticInMemIndex(Metric m, const std::string &data_path, IndexWriteParameters &index_parameters) - { - size_t ndims, npoints; - diskann::get_bin_metadata(data_path, npoints, ndims); - _index = new Index(m, ndims, npoints, - false, // not a dynamic_index - false, // no enable_tags/ids - false, // no concurrent_consolidate, - false, // pq_dist_build - 0, // num_pq_chunks - false, // use_opq = false - 0); // num_frozen_pts = 0 - _index->build(data_path.c_str(), npoints, index_parameters); - } + return std::make_pair(ids, dists); + } - ~StaticInMemIndex() - { - delete _index; - } - - auto search(py::array_t &query, const uint64_t knn, - const uint64_t l_search) - { - py::array_t ids(knn); - py::array_t dists(knn); - std::vector empty_vector; - _index->search(query.data(), knn, l_search, ids.mutable_data(), dists.mutable_data()); - return std::make_pair(ids, dists); - } - - auto batch_search(py::array_t &queries, const uint64_t num_queries, - const uint64_t knn, const uint64_t l_search, const int num_threads) - { - py::array_t ids({num_queries, knn}); - py::array_t dists({num_queries, knn}); - std::vector empty_vector; + auto consolidate_delete() { + return _index->consolidate_deletes(write_params); + } +}; - omp_set_num_threads(num_threads); +template +struct StaticInMemIndex { + Index *_index; + + StaticInMemIndex(Metric m, const std::string &data_path, + IndexWriteParameters &index_parameters) { + size_t ndims, npoints; + diskann::get_bin_metadata(data_path, npoints, ndims); + _index = new Index(m, ndims, npoints, + false, // not a dynamic_index + false, // no enable_tags/ids + false, // no concurrent_consolidate, + false, // pq_dist_build + 0, // num_pq_chunks + false, // use_opq = false + 0); // num_frozen_pts = 0 + _index->build(data_path.c_str(), npoints, index_parameters); + } + + ~StaticInMemIndex() { delete _index; } + + auto search(py::array_t &query, + const uint64_t knn, const uint64_t l_search) { + py::array_t ids(knn); + py::array_t dists(knn); + std::vector empty_vector; + _index->search(query.data(), knn, l_search, ids.mutable_data(), + dists.mutable_data()); + return std::make_pair(ids, dists); + } + + auto batch_search( + py::array_t &queries, + const uint64_t num_queries, const uint64_t knn, const uint64_t l_search, + const int num_threads) { + py::array_t ids({num_queries, knn}); + py::array_t dists({num_queries, knn}); + std::vector empty_vector; + + omp_set_num_threads(num_threads); #pragma omp parallel for schedule(dynamic, 1) - for (int64_t i = 0; i < (int64_t)num_queries; i++) - { - _index->search(queries.data(i), knn, l_search, ids.mutable_data(i), dists.mutable_data(i)); - } - - return std::make_pair(ids, dists); + for (int64_t i = 0; i < (int64_t)num_queries; i++) { + _index->search(queries.data(i), knn, l_search, ids.mutable_data(i), + dists.mutable_data(i)); } + + return std::make_pair(ids, dists); + } }; -PYBIND11_MODULE(_diskannpy, m) -{ - m.doc() = "DiskANN Python Bindings"; +PYBIND11_MODULE(_diskannpy, m) { + m.doc() = "DiskANN Python Bindings"; #ifdef VERSION_INFO - m.attr("__version__") = VERSION_INFO; + m.attr("__version__") = VERSION_INFO; #else - m.attr("__version__") = "dev"; + m.attr("__version__") = "dev"; #endif - py::enum_(m, "Metric") - .value("L2", Metric::L2) - .value("INNER_PRODUCT", Metric::INNER_PRODUCT) - .export_values(); - - py::class_>(m, "DiskANNStaticInMemFloatIndex") - .def(py::init([](diskann::Metric metric, const std::string &data_path, IndexWriteParameters &index_parameters) { - return std::unique_ptr>( - new StaticInMemIndex(metric, data_path, index_parameters)); - })) - .def("search", &StaticInMemIndex::search, py::arg("query"), py::arg("knn"), py::arg("l_search")) - .def("batch_search", &StaticInMemIndex::batch_search, py::arg("queries"), py::arg("num_queries"), - py::arg("knn"), py::arg("l_search"), py::arg("num_threads")); - - py::class_>(m, "DiskANNStaticInMemInt8Index") - .def(py::init([](diskann::Metric metric, const std::string &data_path, IndexWriteParameters &index_parameters) { - return std::unique_ptr>( - new StaticInMemIndex(metric, data_path, index_parameters)); - })) - .def("search", &StaticInMemIndex::search, py::arg("query"), py::arg("knn"), py::arg("l_search")) - .def("batch_search", &StaticInMemIndex::batch_search, py::arg("queries"), py::arg("num_queries"), - py::arg("knn"), py::arg("l_search"), py::arg("num_threads")); - - py::class_>(m, "DiskANNStaticInMemUint8Index") - .def(py::init([](diskann::Metric metric, const std::string &data_path, IndexWriteParameters &index_parameters) { - return std::unique_ptr>( - new StaticInMemIndex(metric, data_path, index_parameters)); - })) - .def("search", &StaticInMemIndex::search, py::arg("query"), py::arg("knn"), py::arg("l_search")) - .def("batch_search", &StaticInMemIndex::batch_search, py::arg("queries"), py::arg("num_queries"), - py::arg("knn"), py::arg("l_search"), py::arg("num_threads")); - - py::class_>(m, "DiskANNDynamicInMemFloatIndex") - .def(py::init([](diskann::Metric metric, const size_t dim, const size_t max_points, - const IndexWriteParameters &index_parameters, const uint32_t initial_search_list_size, - const uint32_t search_threads, const bool concurrent_consolidate) { + py::enum_(m, "Metric") + .value("L2", Metric::L2) + .value("INNER_PRODUCT", Metric::INNER_PRODUCT) + .export_values(); + + py::class_>(m, "DiskANNStaticInMemFloatIndex") + .def(py::init([](diskann::Metric metric, const std::string &data_path, + IndexWriteParameters &index_parameters) { + return std::unique_ptr>( + new StaticInMemIndex(metric, data_path, index_parameters)); + })) + .def("search", &StaticInMemIndex::search, py::arg("query"), + py::arg("knn"), py::arg("l_search")) + .def("batch_search", &StaticInMemIndex::batch_search, + py::arg("queries"), py::arg("num_queries"), py::arg("knn"), + py::arg("l_search"), py::arg("num_threads")); + + py::class_>(m, "DiskANNStaticInMemInt8Index") + .def(py::init([](diskann::Metric metric, const std::string &data_path, + IndexWriteParameters &index_parameters) { + return std::unique_ptr>( + new StaticInMemIndex(metric, data_path, index_parameters)); + })) + .def("search", &StaticInMemIndex::search, py::arg("query"), + py::arg("knn"), py::arg("l_search")) + .def("batch_search", &StaticInMemIndex::batch_search, + py::arg("queries"), py::arg("num_queries"), py::arg("knn"), + py::arg("l_search"), py::arg("num_threads")); + + py::class_>(m, "DiskANNStaticInMemUint8Index") + .def(py::init([](diskann::Metric metric, const std::string &data_path, + IndexWriteParameters &index_parameters) { + return std::unique_ptr>( + new StaticInMemIndex(metric, data_path, index_parameters)); + })) + .def("search", &StaticInMemIndex::search, py::arg("query"), + py::arg("knn"), py::arg("l_search")) + .def("batch_search", &StaticInMemIndex::batch_search, + py::arg("queries"), py::arg("num_queries"), py::arg("knn"), + py::arg("l_search"), py::arg("num_threads")); + + py::class_>(m, "DiskANNDynamicInMemFloatIndex") + .def(py::init( + [](diskann::Metric metric, const size_t dim, const size_t max_points, + const IndexWriteParameters &index_parameters, + const uint32_t initial_search_list_size, + const uint32_t search_threads, const bool concurrent_consolidate) { return std::unique_ptr>( - new DynamicInMemIndex(metric, dim, max_points, index_parameters, initial_search_list_size, - search_threads, concurrent_consolidate)); - })) - .def("search", &DynamicInMemIndex::search, py::arg("query"), py::arg("knn"), py::arg("l_search")) - .def("batch_search", &DynamicInMemIndex::batch_search, py::arg("queries"), py::arg("num_queries"), - py::arg("knn"), py::arg("l_search"), py::arg("num_threads")) - .def("insert", &DynamicInMemIndex::insert, py::arg("vector"), py::arg("id")) - .def("mark_deleted", &DynamicInMemIndex::mark_deleted, py::arg("id")) - .def("consolidate_delete", &DynamicInMemIndex::consolidate_delete); - - py::class_>(m, "DiskANNDynamicInMemInt8Index") - .def(py::init([](diskann::Metric metric, const size_t dim, const size_t max_points, - const IndexWriteParameters &index_parameters, const uint32_t initial_search_list_size, - const uint32_t search_threads, const bool concurrent_consolidate) { + new DynamicInMemIndex( + metric, dim, max_points, index_parameters, + initial_search_list_size, search_threads, + concurrent_consolidate)); + })) + .def("search", &DynamicInMemIndex::search, py::arg("query"), + py::arg("knn"), py::arg("l_search")) + .def("batch_search", &DynamicInMemIndex::batch_search, + py::arg("queries"), py::arg("num_queries"), py::arg("knn"), + py::arg("l_search"), py::arg("num_threads")) + .def("insert", &DynamicInMemIndex::insert, py::arg("vector"), + py::arg("id")) + .def("mark_deleted", &DynamicInMemIndex::mark_deleted, + py::arg("id")) + .def("consolidate_delete", &DynamicInMemIndex::consolidate_delete); + + py::class_>(m, "DiskANNDynamicInMemInt8Index") + .def(py::init( + [](diskann::Metric metric, const size_t dim, const size_t max_points, + const IndexWriteParameters &index_parameters, + const uint32_t initial_search_list_size, + const uint32_t search_threads, const bool concurrent_consolidate) { return std::unique_ptr>( - new DynamicInMemIndex(metric, dim, max_points, index_parameters, initial_search_list_size, - search_threads, concurrent_consolidate)); - })) - .def("search", &DynamicInMemIndex::search, py::arg("query"), py::arg("knn"), py::arg("l_search")) - .def("batch_search", &DynamicInMemIndex::batch_search, py::arg("queries"), py::arg("num_queries"), - py::arg("knn"), py::arg("l_search"), py::arg("num_threads")) - .def("insert", &DynamicInMemIndex::insert, py::arg("vector"), py::arg("id")) - .def("mark_deleted", &DynamicInMemIndex::mark_deleted, py::arg("id")) - .def("consolidate_delete", &DynamicInMemIndex::consolidate_delete); - - py::class_>(m, "DiskANNDynamicInMemUint8Index") - .def(py::init([](diskann::Metric metric, const size_t dim, const size_t max_points, - const IndexWriteParameters &index_parameters, const uint32_t initial_search_list_size, - const uint32_t search_threads, const bool concurrent_consolidate) { + new DynamicInMemIndex( + metric, dim, max_points, index_parameters, + initial_search_list_size, search_threads, + concurrent_consolidate)); + })) + .def("search", &DynamicInMemIndex::search, py::arg("query"), + py::arg("knn"), py::arg("l_search")) + .def("batch_search", &DynamicInMemIndex::batch_search, + py::arg("queries"), py::arg("num_queries"), py::arg("knn"), + py::arg("l_search"), py::arg("num_threads")) + .def("insert", &DynamicInMemIndex::insert, py::arg("vector"), + py::arg("id")) + .def("mark_deleted", &DynamicInMemIndex::mark_deleted, + py::arg("id")) + .def("consolidate_delete", + &DynamicInMemIndex::consolidate_delete); + + py::class_>(m, "DiskANNDynamicInMemUint8Index") + .def(py::init( + [](diskann::Metric metric, const size_t dim, const size_t max_points, + const IndexWriteParameters &index_parameters, + const uint32_t initial_search_list_size, + const uint32_t search_threads, const bool concurrent_consolidate) { return std::unique_ptr>( - new DynamicInMemIndex(metric, dim, max_points, index_parameters, initial_search_list_size, - search_threads, concurrent_consolidate)); - })) - .def("search", &DynamicInMemIndex::search, py::arg("query"), py::arg("knn"), py::arg("l_search")) - .def("batch_search", &DynamicInMemIndex::batch_search, py::arg("queries"), py::arg("num_queries"), - py::arg("knn"), py::arg("l_search"), py::arg("num_threads")) - .def("insert", &DynamicInMemIndex::insert, py::arg("vector"), py::arg("id")) - .def("mark_deleted", &DynamicInMemIndex::mark_deleted, py::arg("id")) - .def("consolidate_delete", &DynamicInMemIndex::consolidate_delete); - - py::class_>(m, "DiskANNFloatIndex") - .def(py::init([](diskann::Metric metric) { - return std::unique_ptr>(new DiskANNIndex(metric)); - })) - .def("cache_bfs_levels", &DiskANNIndex::cache_bfs_levels, py::arg("num_nodes_to_cache")) - .def("load_index", &DiskANNIndex::load_index, py::arg("index_path_prefix"), py::arg("num_threads"), - py::arg("num_nodes_to_cache"), py::arg("cache_mechanism") = 1) - .def("search", &DiskANNIndex::search, py::arg("query"), py::arg("knn"), py::arg("l_search"), - py::arg("beam_width")) - .def("batch_search", &DiskANNIndex::batch_search, py::arg("queries"), py::arg("num_queries"), - py::arg("knn"), py::arg("l_search"), py::arg("beam_width"), py::arg("num_threads")) - .def( - "build", - [](DiskANNIndex &self, const char *data_file_path, const char *index_prefix_path, unsigned R, - unsigned L, double final_index_ram_limit, double indexing_ram_budget, unsigned num_threads, - unsigned pq_disk_bytes) { - std::string params = std::to_string(R) + " " + std::to_string(L) + " " + - std::to_string(final_index_ram_limit) + " " + std::to_string(indexing_ram_budget) + - " " + std::to_string(num_threads); - if (pq_disk_bytes > 0) - { - params = params + " " + std::to_string(pq_disk_bytes); - } - diskann::build_disk_index(data_file_path, index_prefix_path, params.c_str(), self.get_metric()); - }, - py::arg("data_file_path"), py::arg("index_prefix_path"), py::arg("R"), py::arg("L"), - py::arg("final_index_ram_limit"), py::arg("indexing_ram_limit"), py::arg("num_threads"), - py::arg("pq_disk_bytes") = 0); - - py::class_>(m, "DiskANNInt8Index") - .def(py::init([](diskann::Metric metric) { - return std::unique_ptr>(new DiskANNIndex(metric)); - })) - .def("cache_bfs_levels", &DiskANNIndex::cache_bfs_levels, py::arg("num_nodes_to_cache")) - .def("load_index", &DiskANNIndex::load_index, py::arg("index_path_prefix"), py::arg("num_threads"), - py::arg("num_nodes_to_cache"), py::arg("cache_mechanism") = 1) - .def("search", &DiskANNIndex::search, py::arg("query"), py::arg("knn"), py::arg("l_search"), - py::arg("beam_width")) - .def("batch_search", &DiskANNIndex::batch_search, py::arg("queries"), py::arg("num_queries"), - py::arg("knn"), py::arg("l_search"), py::arg("beam_width"), py::arg("num_threads")) - .def( - "build", - [](DiskANNIndex &self, const char *data_file_path, const char *index_prefix_path, unsigned R, - unsigned L, double final_index_ram_limit, double indexing_ram_budget, unsigned num_threads, - unsigned pq_disk_bytes) { - std::string params = std::to_string(R) + " " + std::to_string(L) + " " + - std::to_string(final_index_ram_limit) + " " + std::to_string(indexing_ram_budget) + - " " + std::to_string(num_threads); - if (pq_disk_bytes > 0) - params = params + " " + std::to_string(pq_disk_bytes); - diskann::build_disk_index(data_file_path, index_prefix_path, params.c_str(), self.get_metric()); - }, - py::arg("data_file_path"), py::arg("index_prefix_path"), py::arg("R"), py::arg("L"), - py::arg("final_index_ram_limit"), py::arg("indexing_ram_limit"), py::arg("num_threads"), - py::arg("pq_disk_bytes") = 0); - - py::class_>(m, "DiskANNUInt8Index") - .def(py::init([](diskann::Metric metric) { - return std::unique_ptr>(new DiskANNIndex(metric)); - })) - .def("cache_bfs_levels", &DiskANNIndex::cache_bfs_levels, py::arg("num_nodes_to_cache")) - .def("load_index", &DiskANNIndex::load_index, py::arg("index_path_prefix"), py::arg("num_threads"), - py::arg("num_nodes_to_cache"), py::arg("cache_mechanism") = 1) - .def("search", &DiskANNIndex::search, py::arg("query"), py::arg("knn"), py::arg("l_search"), - py::arg("beam_width")) - .def("batch_search", &DiskANNIndex::batch_search, py::arg("queries"), py::arg("num_queries"), - py::arg("knn"), py::arg("l_search"), py::arg("beam_width"), py::arg("num_threads")) - .def( - "build", - [](DiskANNIndex &self, const char *data_file_path, const char *index_prefix_path, unsigned R, - unsigned L, double final_index_ram_limit, double indexing_ram_budget, unsigned num_threads, - unsigned pq_disk_bytes) { - std::string params = std::to_string(R) + " " + std::to_string(L) + " " + - std::to_string(final_index_ram_limit) + " " + std::to_string(indexing_ram_budget) + - " " + std::to_string(num_threads); - if (pq_disk_bytes > 0) - params = params + " " + std::to_string(pq_disk_bytes); - diskann::build_disk_index(data_file_path, index_prefix_path, params.c_str(), - self.get_metric()); - }, - py::arg("data_file_path"), py::arg("index_prefix_path"), py::arg("R"), py::arg("L"), - py::arg("final_index_ram_limit"), py::arg("indexing_ram_limit"), py::arg("num_threads"), - py::arg("pq_disk_bytes") = 0); + new DynamicInMemIndex( + metric, dim, max_points, index_parameters, + initial_search_list_size, search_threads, + concurrent_consolidate)); + })) + .def("search", &DynamicInMemIndex::search, py::arg("query"), + py::arg("knn"), py::arg("l_search")) + .def("batch_search", &DynamicInMemIndex::batch_search, + py::arg("queries"), py::arg("num_queries"), py::arg("knn"), + py::arg("l_search"), py::arg("num_threads")) + .def("insert", &DynamicInMemIndex::insert, py::arg("vector"), + py::arg("id")) + .def("mark_deleted", &DynamicInMemIndex::mark_deleted, + py::arg("id")) + .def("consolidate_delete", + &DynamicInMemIndex::consolidate_delete); + + py::class_>(m, "DiskANNFloatIndex") + .def(py::init([](diskann::Metric metric) { + return std::unique_ptr>( + new DiskANNIndex(metric)); + })) + .def("cache_bfs_levels", &DiskANNIndex::cache_bfs_levels, + py::arg("num_nodes_to_cache")) + .def("load_index", &DiskANNIndex::load_index, + py::arg("index_path_prefix"), py::arg("num_threads"), + py::arg("num_nodes_to_cache"), py::arg("cache_mechanism") = 1) + .def("search", &DiskANNIndex::search, py::arg("query"), + py::arg("knn"), py::arg("l_search"), py::arg("beam_width")) + .def("batch_search", &DiskANNIndex::batch_search, + py::arg("queries"), py::arg("num_queries"), py::arg("knn"), + py::arg("l_search"), py::arg("beam_width"), py::arg("num_threads")) + .def( + "build", + [](DiskANNIndex &self, const char *data_file_path, + const char *index_prefix_path, unsigned R, unsigned L, + double final_index_ram_limit, double indexing_ram_budget, + unsigned num_threads, unsigned pq_disk_bytes) { + std::string params = std::to_string(R) + " " + std::to_string(L) + + " " + std::to_string(final_index_ram_limit) + + " " + std::to_string(indexing_ram_budget) + + " " + std::to_string(num_threads); + if (pq_disk_bytes > 0) { + params = params + " " + std::to_string(pq_disk_bytes); + } + diskann::build_disk_index(data_file_path, index_prefix_path, + params.c_str(), self.get_metric()); + }, + py::arg("data_file_path"), py::arg("index_prefix_path"), py::arg("R"), + py::arg("L"), py::arg("final_index_ram_limit"), + py::arg("indexing_ram_limit"), py::arg("num_threads"), + py::arg("pq_disk_bytes") = 0); + + py::class_>(m, "DiskANNInt8Index") + .def(py::init([](diskann::Metric metric) { + return std::unique_ptr>( + new DiskANNIndex(metric)); + })) + .def("cache_bfs_levels", &DiskANNIndex::cache_bfs_levels, + py::arg("num_nodes_to_cache")) + .def("load_index", &DiskANNIndex::load_index, + py::arg("index_path_prefix"), py::arg("num_threads"), + py::arg("num_nodes_to_cache"), py::arg("cache_mechanism") = 1) + .def("search", &DiskANNIndex::search, py::arg("query"), + py::arg("knn"), py::arg("l_search"), py::arg("beam_width")) + .def("batch_search", &DiskANNIndex::batch_search, + py::arg("queries"), py::arg("num_queries"), py::arg("knn"), + py::arg("l_search"), py::arg("beam_width"), py::arg("num_threads")) + .def( + "build", + [](DiskANNIndex &self, const char *data_file_path, + const char *index_prefix_path, unsigned R, unsigned L, + double final_index_ram_limit, double indexing_ram_budget, + unsigned num_threads, unsigned pq_disk_bytes) { + std::string params = std::to_string(R) + " " + std::to_string(L) + + " " + std::to_string(final_index_ram_limit) + + " " + std::to_string(indexing_ram_budget) + + " " + std::to_string(num_threads); + if (pq_disk_bytes > 0) + params = params + " " + std::to_string(pq_disk_bytes); + diskann::build_disk_index(data_file_path, index_prefix_path, + params.c_str(), + self.get_metric()); + }, + py::arg("data_file_path"), py::arg("index_prefix_path"), py::arg("R"), + py::arg("L"), py::arg("final_index_ram_limit"), + py::arg("indexing_ram_limit"), py::arg("num_threads"), + py::arg("pq_disk_bytes") = 0); + + py::class_>(m, "DiskANNUInt8Index") + .def(py::init([](diskann::Metric metric) { + return std::unique_ptr>( + new DiskANNIndex(metric)); + })) + .def("cache_bfs_levels", &DiskANNIndex::cache_bfs_levels, + py::arg("num_nodes_to_cache")) + .def("load_index", &DiskANNIndex::load_index, + py::arg("index_path_prefix"), py::arg("num_threads"), + py::arg("num_nodes_to_cache"), py::arg("cache_mechanism") = 1) + .def("search", &DiskANNIndex::search, py::arg("query"), + py::arg("knn"), py::arg("l_search"), py::arg("beam_width")) + .def("batch_search", &DiskANNIndex::batch_search, + py::arg("queries"), py::arg("num_queries"), py::arg("knn"), + py::arg("l_search"), py::arg("beam_width"), py::arg("num_threads")) + .def( + "build", + [](DiskANNIndex &self, const char *data_file_path, + const char *index_prefix_path, unsigned R, unsigned L, + double final_index_ram_limit, double indexing_ram_budget, + unsigned num_threads, unsigned pq_disk_bytes) { + std::string params = std::to_string(R) + " " + std::to_string(L) + + " " + std::to_string(final_index_ram_limit) + + " " + std::to_string(indexing_ram_budget) + + " " + std::to_string(num_threads); + if (pq_disk_bytes > 0) + params = params + " " + std::to_string(pq_disk_bytes); + diskann::build_disk_index( + data_file_path, index_prefix_path, params.c_str(), + self.get_metric()); + }, + py::arg("data_file_path"), py::arg("index_prefix_path"), py::arg("R"), + py::arg("L"), py::arg("final_index_ram_limit"), + py::arg("indexing_ram_limit"), py::arg("num_threads"), + py::arg("pq_disk_bytes") = 0); } diff --git a/src/restapi/search_wrapper.cpp b/src/restapi/search_wrapper.cpp index dc9f5734e..5fe049bf6 100644 --- a/src/restapi/search_wrapper.cpp +++ b/src/restapi/search_wrapper.cpp @@ -21,189 +21,182 @@ #endif #endif -namespace diskann -{ +namespace diskann { const unsigned int DEFAULT_W = 1; -SearchResult::SearchResult(unsigned int K, unsigned int elapsed_time_in_ms, const unsigned *const indices, - const float *const distances, const std::string *const tags, +SearchResult::SearchResult(unsigned int K, unsigned int elapsed_time_in_ms, + const unsigned *const indices, + const float *const distances, + const std::string *const tags, const unsigned *const partitions) - : _K(K), _search_time_in_ms(elapsed_time_in_ms) -{ - for (unsigned i = 0; i < K; ++i) - { - this->_indices.push_back(indices[i]); - this->_distances.push_back(distances[i]); - if (tags != NULL) - this->_tags.push_back(tags[i]); - if (partitions != NULL) - this->_partitions.push_back(partitions[i]); - } - if (tags != nullptr) - this->_tags_enabled = true; - else - this->_tags_enabled = false; - - if (partitions != nullptr) - this->_partitions_enabled = true; - else - this->_partitions_enabled = false; + : _K(K), _search_time_in_ms(elapsed_time_in_ms) { + for (unsigned i = 0; i < K; ++i) { + this->_indices.push_back(indices[i]); + this->_distances.push_back(distances[i]); + if (tags != NULL) this->_tags.push_back(tags[i]); + if (partitions != NULL) this->_partitions.push_back(partitions[i]); + } + if (tags != nullptr) + this->_tags_enabled = true; + else + this->_tags_enabled = false; + + if (partitions != nullptr) + this->_partitions_enabled = true; + else + this->_partitions_enabled = false; } -BaseSearch::BaseSearch(const std::string &tagsFile) -{ - if (tagsFile.size() != 0) - { - std::ifstream in(tagsFile); +BaseSearch::BaseSearch(const std::string &tagsFile) { + if (tagsFile.size() != 0) { + std::ifstream in(tagsFile); - if (!in.is_open()) - { - std::cerr << "Could not open " << tagsFile << std::endl; - } + if (!in.is_open()) { + std::cerr << "Could not open " << tagsFile << std::endl; + } - std::string tag; - while (std::getline(in, tag)) - { - _tags_str.push_back(tag); - } + std::string tag; + while (std::getline(in, tag)) { + _tags_str.push_back(tag); + } - _tags_enabled = true; + _tags_enabled = true; - std::cout << "Loaded " << _tags_str.size() << " tags from " << tagsFile << std::endl; - } - else - { - _tags_enabled = false; - } + std::cout << "Loaded " << _tags_str.size() << " tags from " << tagsFile + << std::endl; + } else { + _tags_enabled = false; + } } -void BaseSearch::lookup_tags(const unsigned K, const unsigned *indices, std::string *ret_tags) -{ - if (_tags_enabled == false) - throw std::runtime_error("Can not look up tags as they are not enabled."); - else - { - for (unsigned k = 0; k < K; ++k) - { - if (indices[k] > _tags_str.size()) - throw std::runtime_error("In tag lookup, index exceeded the number of tags"); - else - ret_tags[k] = _tags_str[indices[k]]; - } +void BaseSearch::lookup_tags(const unsigned K, const unsigned *indices, + std::string *ret_tags) { + if (_tags_enabled == false) + throw std::runtime_error("Can not look up tags as they are not enabled."); + else { + for (unsigned k = 0; k < K; ++k) { + if (indices[k] > _tags_str.size()) + throw std::runtime_error( + "In tag lookup, index exceeded the number of tags"); + else + ret_tags[k] = _tags_str[indices[k]]; } + } } template -InMemorySearch::InMemorySearch(const std::string &baseFile, const std::string &indexFile, - const std::string &tagsFile, Metric m, uint32_t num_threads, uint32_t search_l) - : BaseSearch(tagsFile) -{ - size_t dimensions, total_points = 0; - diskann::get_bin_metadata(baseFile, total_points, dimensions); - _index = std::unique_ptr>(new diskann::Index(m, dimensions, total_points, false)); - - _index->load(indexFile.c_str(), num_threads, search_l); +InMemorySearch::InMemorySearch(const std::string &baseFile, + const std::string &indexFile, + const std::string &tagsFile, Metric m, + uint32_t num_threads, uint32_t search_l) + : BaseSearch(tagsFile) { + size_t dimensions, total_points = 0; + diskann::get_bin_metadata(baseFile, total_points, dimensions); + _index = std::unique_ptr>( + new diskann::Index(m, dimensions, total_points, false)); + + _index->load(indexFile.c_str(), num_threads, search_l); } template -SearchResult InMemorySearch::search(const T *query, const unsigned int dimensions, const unsigned int K, - const unsigned int Ls) -{ - unsigned int *indices = new unsigned int[K]; - float *distances = new float[K]; - - auto startTime = std::chrono::high_resolution_clock::now(); - _index->search(query, K, Ls, indices, distances); - auto duration = - std::chrono::duration_cast(std::chrono::high_resolution_clock::now() - startTime) - .count(); - - std::string *tags = nullptr; - if (_tags_enabled) - { - tags = new std::string[K]; - lookup_tags(K, indices, tags); - } - - SearchResult result(K, (unsigned int)duration, indices, distances, tags); - - delete[] indices; - delete[] distances; - return result; +SearchResult InMemorySearch::search(const T *query, + const unsigned int dimensions, + const unsigned int K, + const unsigned int Ls) { + unsigned int *indices = new unsigned int[K]; + float *distances = new float[K]; + + auto startTime = std::chrono::high_resolution_clock::now(); + _index->search(query, K, Ls, indices, distances); + auto duration = std::chrono::duration_cast( + std::chrono::high_resolution_clock::now() - startTime) + .count(); + + std::string *tags = nullptr; + if (_tags_enabled) { + tags = new std::string[K]; + lookup_tags(K, indices, tags); + } + + SearchResult result(K, (unsigned int)duration, indices, distances, tags); + + delete[] indices; + delete[] distances; + return result; } -template InMemorySearch::~InMemorySearch() -{ -} +template +InMemorySearch::~InMemorySearch() {} template -PQFlashSearch::PQFlashSearch(const std::string &indexPrefix, const unsigned num_nodes_to_cache, - const unsigned num_threads, const std::string &tagsFile, Metric m) - : BaseSearch(tagsFile) -{ +PQFlashSearch::PQFlashSearch(const std::string &indexPrefix, + const unsigned num_nodes_to_cache, + const unsigned num_threads, + const std::string &tagsFile, Metric m) + : BaseSearch(tagsFile) { #ifdef _WINDOWS #ifndef USE_BING_INFRA - reader.reset(new WindowsAlignedFileReader()); + reader.reset(new WindowsAlignedFileReader()); #else - reader.reset(new diskann::BingAlignedFileReader()); + reader.reset(new diskann::BingAlignedFileReader()); #endif #else - auto ptr = new LinuxAlignedFileReader(); - reader.reset(ptr); + auto ptr = new LinuxAlignedFileReader(); + reader.reset(ptr); #endif - std::string index_prefix_path(indexPrefix); - std::string disk_index_file = index_prefix_path + "_disk.index"; - std::string warmup_query_file = index_prefix_path + "_sample_data.bin"; + std::string index_prefix_path(indexPrefix); + std::string disk_index_file = index_prefix_path + "_disk.index"; + std::string warmup_query_file = index_prefix_path + "_sample_data.bin"; - _index = std::unique_ptr>(new diskann::PQFlashIndex(reader, m)); + _index = std::unique_ptr>( + new diskann::PQFlashIndex(reader, m)); - int res = _index->load(num_threads, index_prefix_path.c_str()); + int res = _index->load(num_threads, index_prefix_path.c_str()); - if (res != 0) - { - std::cerr << "Unable to load index. Status code: " << res << "." << std::endl; - } + if (res != 0) { + std::cerr << "Unable to load index. Status code: " << res << "." + << std::endl; + } - std::vector node_list; - std::cout << "Caching " << num_nodes_to_cache << " BFS nodes around medoid(s)" << std::endl; - _index->cache_bfs_levels(num_nodes_to_cache, node_list); - _index->load_cache_list(node_list); - omp_set_num_threads(num_threads); + std::vector node_list; + std::cout << "Caching " << num_nodes_to_cache << " BFS nodes around medoid(s)" + << std::endl; + _index->cache_bfs_levels(num_nodes_to_cache, node_list); + _index->load_cache_list(node_list); + omp_set_num_threads(num_threads); } template -SearchResult PQFlashSearch::search(const T *query, const unsigned int dimensions, const unsigned int K, - const unsigned int Ls) -{ - uint64_t *indices_u64 = new uint64_t[K]; - unsigned *indices = new unsigned[K]; - float *distances = new float[K]; - - auto startTime = std::chrono::high_resolution_clock::now(); - _index->cached_beam_search(query, K, Ls, indices_u64, distances, DEFAULT_W); - auto duration = - std::chrono::duration_cast(std::chrono::high_resolution_clock::now() - startTime) - .count(); - for (unsigned k = 0; k < K; ++k) - indices[k] = indices_u64[k]; - - std::string *tags = nullptr; - if (_tags_enabled) - { - tags = new std::string[K]; - lookup_tags(K, indices, tags); - } - SearchResult result(K, (unsigned int)duration, indices, distances, tags); - delete[] indices_u64; - delete[] indices; - delete[] distances; - return result; +SearchResult PQFlashSearch::search(const T *query, + const unsigned int dimensions, + const unsigned int K, + const unsigned int Ls) { + uint64_t *indices_u64 = new uint64_t[K]; + unsigned *indices = new unsigned[K]; + float *distances = new float[K]; + + auto startTime = std::chrono::high_resolution_clock::now(); + _index->cached_beam_search(query, K, Ls, indices_u64, distances, DEFAULT_W); + auto duration = std::chrono::duration_cast( + std::chrono::high_resolution_clock::now() - startTime) + .count(); + for (unsigned k = 0; k < K; ++k) indices[k] = indices_u64[k]; + + std::string *tags = nullptr; + if (_tags_enabled) { + tags = new std::string[K]; + lookup_tags(K, indices, tags); + } + SearchResult result(K, (unsigned int)duration, indices, distances, tags); + delete[] indices_u64; + delete[] indices; + delete[] distances; + return result; } -template PQFlashSearch::~PQFlashSearch() -{ -} +template +PQFlashSearch::~PQFlashSearch() {} template class InMemorySearch; template class InMemorySearch; @@ -212,4 +205,4 @@ template class InMemorySearch; template class PQFlashSearch; template class PQFlashSearch; template class PQFlashSearch; -} // namespace diskann +} // namespace diskann diff --git a/src/restapi/server.cpp b/src/restapi/server.cpp index f79b0affb..6d365913c 100644 --- a/src/restapi/server.cpp +++ b/src/restapi/server.cpp @@ -11,261 +11,239 @@ #include -namespace diskann -{ - -Server::Server(web::uri &uri, std::vector> &multi_searcher, - const std::string &typestring) - : _multi_search(multi_searcher.size() > 1 ? true : false) -{ - for (auto &searcher : multi_searcher) - _multi_searcher.push_back(std::move(searcher)); - - _listener = std::unique_ptr( - new web::http::experimental::listener::http_listener(uri)); - if (typestring == std::string("float")) - { - _listener->support(std::bind(&Server::handle_post, this, std::placeholders::_1)); - } - else if (typestring == std::string("int8_t")) - { - _listener->support(web::http::methods::POST, - std::bind(&Server::handle_post, this, std::placeholders::_1)); - } - else if (typestring == std::string("uint8_t")) - { - _listener->support(web::http::methods::POST, - std::bind(&Server::handle_post, this, std::placeholders::_1)); - } - else - { - throw "Unsupported type in server constuctor"; - } +namespace diskann { + +Server::Server( + web::uri &uri, + std::vector> &multi_searcher, + const std::string &typestring) + : _multi_search(multi_searcher.size() > 1 ? true : false) { + for (auto &searcher : multi_searcher) + _multi_searcher.push_back(std::move(searcher)); + + _listener = std::unique_ptr( + new web::http::experimental::listener::http_listener(uri)); + if (typestring == std::string("float")) { + _listener->support( + std::bind(&Server::handle_post, this, std::placeholders::_1)); + } else if (typestring == std::string("int8_t")) { + _listener->support( + web::http::methods::POST, + std::bind(&Server::handle_post, this, std::placeholders::_1)); + } else if (typestring == std::string("uint8_t")) { + _listener->support( + web::http::methods::POST, + std::bind(&Server::handle_post, this, std::placeholders::_1)); + } else { + throw "Unsupported type in server constuctor"; + } } -Server::~Server() -{ -} +Server::~Server() {} -pplx::task Server::open() -{ - return _listener->open(); -} -pplx::task Server::close() -{ - return _listener->close(); -} +pplx::task Server::open() { return _listener->open(); } +pplx::task Server::close() { return _listener->close(); } -diskann::SearchResult Server::aggregate_results(const unsigned K, const std::vector &results) -{ - if (_multi_search) - { - auto best_indices = new unsigned[K]; - auto best_distances = new float[K]; - auto best_partitions = new unsigned[K]; - auto best_tags = results[0].tags_enabled() ? new std::string[K] : nullptr; +diskann::SearchResult Server::aggregate_results( + const unsigned K, const std::vector &results) { + if (_multi_search) { + auto best_indices = new unsigned[K]; + auto best_distances = new float[K]; + auto best_partitions = new unsigned[K]; + auto best_tags = results[0].tags_enabled() ? new std::string[K] : nullptr; - auto numsearchers = _multi_searcher.size(); - std::vector pos(numsearchers, 0); + auto numsearchers = _multi_searcher.size(); + std::vector pos(numsearchers, 0); - for (size_t k = 0; k < K; ++k) - { - float best_distance = std::numeric_limits::max(); - unsigned best_partition = 0; + for (size_t k = 0; k < K; ++k) { + float best_distance = std::numeric_limits::max(); + unsigned best_partition = 0; - for (size_t i = 0; i < numsearchers; ++i) - { - if (results[i].get_distances()[pos[i]] < best_distance) - { - best_distance = results[i].get_distances()[pos[i]]; - best_partition = i; - } - } - best_distances[k] = best_distance; - best_indices[k] = results[best_partition].get_indices()[pos[best_partition]]; - best_partitions[k] = best_partition; - if (results[best_partition].tags_enabled()) - best_tags[k] = results[best_partition].get_tags()[pos[best_partition]]; - std::cout << best_partition << " " << pos[best_partition] << std::endl; - pos[best_partition]++; + for (size_t i = 0; i < numsearchers; ++i) { + if (results[i].get_distances()[pos[i]] < best_distance) { + best_distance = results[i].get_distances()[pos[i]]; + best_partition = i; } - - unsigned int total_time = 0; - for (size_t i = 0; i < numsearchers; ++i) - total_time += results[i].get_time(); - diskann::SearchResult result = - SearchResult(K, total_time, best_indices, best_distances, best_tags, best_partitions); - - delete[] best_indices; - delete[] best_distances; - delete[] best_partitions; - delete[] best_tags; - - return result; + } + best_distances[k] = best_distance; + best_indices[k] = + results[best_partition].get_indices()[pos[best_partition]]; + best_partitions[k] = best_partition; + if (results[best_partition].tags_enabled()) + best_tags[k] = results[best_partition].get_tags()[pos[best_partition]]; + std::cout << best_partition << " " << pos[best_partition] << std::endl; + pos[best_partition]++; } - else - { - return results[0]; - } -} - -template void Server::handle_post(web::http::http_request message) -{ - message.extract_string(true) - .then([=](utility::string_t body) { - int64_t queryId = -1; - unsigned int K = 0; - try - { - T *queryVector = nullptr; - unsigned int dimensions = 0; - unsigned int Ls; - parseJson(body, K, queryId, queryVector, dimensions, Ls); - auto startTime = std::chrono::high_resolution_clock::now(); - std::vector results; - - for (auto &searcher : _multi_searcher) - results.push_back(searcher->search(queryVector, dimensions, (unsigned int)K, Ls)); - diskann::SearchResult result = aggregate_results(K, results); - diskann::aligned_free(queryVector); - web::json::value response = prepareResponse(queryId, K); - response[INDICES_KEY] = idsToJsonArray(result); - response[DISTANCES_KEY] = distancesToJsonArray(result); - if (result.tags_enabled()) - response[TAGS_KEY] = tagsToJsonArray(result); - if (result.partitions_enabled()) - response[PARTITION_KEY] = partitionsToJsonArray(result); - - response[TIME_TAKEN_KEY] = std::chrono::duration_cast( - std::chrono::high_resolution_clock::now() - startTime) - .count(); + unsigned int total_time = 0; + for (size_t i = 0; i < numsearchers; ++i) + total_time += results[i].get_time(); + diskann::SearchResult result = + SearchResult(K, total_time, best_indices, best_distances, best_tags, + best_partitions); + + delete[] best_indices; + delete[] best_distances; + delete[] best_partitions; + delete[] best_tags; + + return result; + } else { + return results[0]; + } +} - std::cout << "Responding to: " << queryId << std::endl; - return std::make_pair(web::http::status_codes::OK, response); - } - catch (const std::exception &ex) - { - std::cerr << "Exception while processing query: " << queryId << ":" << ex.what() << std::endl; - web::json::value response = prepareResponse(queryId, K); - response[ERROR_MESSAGE_KEY] = web::json::value::string(ex.what()); - return std::make_pair(web::http::status_codes::InternalError, response); - } - catch (...) - { - std::cerr << "Uncaught exception while processing query: " << queryId; - web::json::value response = prepareResponse(queryId, K); - response[ERROR_MESSAGE_KEY] = web::json::value::string(UNKNOWN_ERROR); - return std::make_pair(web::http::status_codes::InternalError, response); - } - }) - .then([=](std::pair response_status) { - try - { - message.reply(response_status.first, response_status.second).wait(); - } - catch (const std::exception &ex) - { - std::cerr << "Exception while processing reply: " << ex.what() << std::endl; +template +void Server::handle_post(web::http::http_request message) { + message.extract_string(true) + .then([=](utility::string_t body) { + int64_t queryId = -1; + unsigned int K = 0; + try { + T *queryVector = nullptr; + unsigned int dimensions = 0; + unsigned int Ls; + parseJson(body, K, queryId, queryVector, dimensions, Ls); + + auto startTime = std::chrono::high_resolution_clock::now(); + std::vector results; + + for (auto &searcher : _multi_searcher) + results.push_back( + searcher->search(queryVector, dimensions, (unsigned int)K, Ls)); + diskann::SearchResult result = aggregate_results(K, results); + diskann::aligned_free(queryVector); + web::json::value response = prepareResponse(queryId, K); + response[INDICES_KEY] = idsToJsonArray(result); + response[DISTANCES_KEY] = distancesToJsonArray(result); + if (result.tags_enabled()) + response[TAGS_KEY] = tagsToJsonArray(result); + if (result.partitions_enabled()) + response[PARTITION_KEY] = partitionsToJsonArray(result); + + response[TIME_TAKEN_KEY] = + std::chrono::duration_cast( + std::chrono::high_resolution_clock::now() - startTime) + .count(); + + std::cout << "Responding to: " << queryId << std::endl; + return std::make_pair(web::http::status_codes::OK, response); + } catch (const std::exception &ex) { + std::cerr << "Exception while processing query: " << queryId << ":" + << ex.what() << std::endl; + web::json::value response = prepareResponse(queryId, K); + response[ERROR_MESSAGE_KEY] = web::json::value::string(ex.what()); + return std::make_pair(web::http::status_codes::InternalError, + response); + } catch (...) { + std::cerr << "Uncaught exception while processing query: " << queryId; + web::json::value response = prepareResponse(queryId, K); + response[ERROR_MESSAGE_KEY] = web::json::value::string(UNKNOWN_ERROR); + return std::make_pair(web::http::status_codes::InternalError, + response); + } + }) + .then( + [=](std::pair response_status) { + try { + message.reply(response_status.first, response_status.second) + .wait(); + } catch (const std::exception &ex) { + std::cerr << "Exception while processing reply: " << ex.what() + << std::endl; }; - }); + }); } -web::json::value Server::prepareResponse(const int64_t &queryId, const int k) -{ - web::json::value response = web::json::value::object(); - response[QUERY_ID_KEY] = queryId; - response[K_KEY] = k; +web::json::value Server::prepareResponse(const int64_t &queryId, const int k) { + web::json::value response = web::json::value::object(); + response[QUERY_ID_KEY] = queryId; + response[K_KEY] = k; - return response; + return response; } template -void Server::parseJson(const utility::string_t &body, unsigned int &k, int64_t &queryId, T *&queryVector, - unsigned int &dimensions, unsigned &Ls) -{ - std::cout << body << std::endl; - web::json::value val = web::json::value::parse(body); - web::json::array queryArr = val.at(VECTOR_KEY).as_array(); - queryId = val.has_field(QUERY_ID_KEY) ? val.at(QUERY_ID_KEY).as_number().to_int64() : -1; - Ls = val.has_field(L_KEY) ? val.at(L_KEY).as_number().to_uint32() : DEFAULT_L; - k = val.at(K_KEY).as_integer(); - - if (k <= 0 || k > Ls) - { - throw new std::invalid_argument("Num of expected NN (k) must be greater than zero and less than or " - "equal to Ls."); - } - if (queryArr.size() == 0) - { - throw new std::invalid_argument("Query vector has zero elements."); - } - - dimensions = static_cast(queryArr.size()); - unsigned new_dim = ROUND_UP(dimensions, 8); - diskann::alloc_aligned((void **)&queryVector, new_dim * sizeof(T), 8 * sizeof(T)); - memset(queryVector, 0, new_dim * sizeof(float)); - for (size_t i = 0; i < queryArr.size(); i++) - { - queryVector[i] = (float)queryArr[i].as_double(); - } +void Server::parseJson(const utility::string_t &body, unsigned int &k, + int64_t &queryId, T *&queryVector, + unsigned int &dimensions, unsigned &Ls) { + std::cout << body << std::endl; + web::json::value val = web::json::value::parse(body); + web::json::array queryArr = val.at(VECTOR_KEY).as_array(); + queryId = val.has_field(QUERY_ID_KEY) + ? val.at(QUERY_ID_KEY).as_number().to_int64() + : -1; + Ls = val.has_field(L_KEY) ? val.at(L_KEY).as_number().to_uint32() : DEFAULT_L; + k = val.at(K_KEY).as_integer(); + + if (k <= 0 || k > Ls) { + throw new std::invalid_argument( + "Num of expected NN (k) must be greater than zero and less than or " + "equal to Ls."); + } + if (queryArr.size() == 0) { + throw new std::invalid_argument("Query vector has zero elements."); + } + + dimensions = static_cast(queryArr.size()); + unsigned new_dim = ROUND_UP(dimensions, 8); + diskann::alloc_aligned((void **)&queryVector, new_dim * sizeof(T), + 8 * sizeof(T)); + memset(queryVector, 0, new_dim * sizeof(float)); + for (size_t i = 0; i < queryArr.size(); i++) { + queryVector[i] = (float)queryArr[i].as_double(); + } } template -web::json::value Server::toJsonArray(const std::vector &v, std::function valConverter) -{ - web::json::value rslts = web::json::value::array(); - for (size_t i = 0; i < v.size(); i++) - { - auto jsonVal = valConverter(v[i]); - rslts[i] = jsonVal; - } - return rslts; +web::json::value Server::toJsonArray( + const std::vector &v, + std::function valConverter) { + web::json::value rslts = web::json::value::array(); + for (size_t i = 0; i < v.size(); i++) { + auto jsonVal = valConverter(v[i]); + rslts[i] = jsonVal; + } + return rslts; } -web::json::value Server::idsToJsonArray(const diskann::SearchResult &result) -{ - web::json::value idArray = web::json::value::array(); - auto ids = result.get_indices(); - for (size_t i = 0; i < ids.size(); i++) - { - auto idVal = web::json::value::number(ids[i]); - idArray[i] = idVal; - } - std::cout << "Vector size: " << ids.size() << std::endl; - return idArray; +web::json::value Server::idsToJsonArray(const diskann::SearchResult &result) { + web::json::value idArray = web::json::value::array(); + auto ids = result.get_indices(); + for (size_t i = 0; i < ids.size(); i++) { + auto idVal = web::json::value::number(ids[i]); + idArray[i] = idVal; + } + std::cout << "Vector size: " << ids.size() << std::endl; + return idArray; } -web::json::value Server::distancesToJsonArray(const diskann::SearchResult &result) -{ - web::json::value distArray = web::json::value::array(); - auto distances = result.get_distances(); - for (size_t i = 0; i < distances.size(); i++) - { - distArray[i] = web::json::value::number(distances[i]); - } - return distArray; +web::json::value Server::distancesToJsonArray( + const diskann::SearchResult &result) { + web::json::value distArray = web::json::value::array(); + auto distances = result.get_distances(); + for (size_t i = 0; i < distances.size(); i++) { + distArray[i] = web::json::value::number(distances[i]); + } + return distArray; } -web::json::value Server::tagsToJsonArray(const diskann::SearchResult &result) -{ - web::json::value tagArray = web::json::value::array(); - auto tags = result.get_tags(); - for (size_t i = 0; i < tags.size(); i++) - { - tagArray[i] = web::json::value::string(tags[i]); - } - return tagArray; +web::json::value Server::tagsToJsonArray(const diskann::SearchResult &result) { + web::json::value tagArray = web::json::value::array(); + auto tags = result.get_tags(); + for (size_t i = 0; i < tags.size(); i++) { + tagArray[i] = web::json::value::string(tags[i]); + } + return tagArray; } -web::json::value Server::partitionsToJsonArray(const diskann::SearchResult &result) -{ - web::json::value partitionArray = web::json::value::array(); - auto partitions = result.get_partitions(); - for (size_t i = 0; i < partitions.size(); i++) - { - partitionArray[i] = web::json::value::number(partitions[i]); - } - return partitionArray; +web::json::value Server::partitionsToJsonArray( + const diskann::SearchResult &result) { + web::json::value partitionArray = web::json::value::array(); + auto partitions = result.get_partitions(); + for (size_t i = 0; i < partitions.size(); i++) { + partitionArray[i] = web::json::value::number(partitions[i]); + } + return partitionArray; } -}; // namespace diskann \ No newline at end of file +}; // namespace diskann \ No newline at end of file diff --git a/tests/restapi/client.cpp b/tests/restapi/client.cpp index fdf4414dd..40d83ac41 100644 --- a/tests/restapi/client.cpp +++ b/tests/restapi/client.cpp @@ -23,102 +23,92 @@ using namespace diskann; namespace po = boost::program_options; template -void query_loop(const std::string &ip_addr_port, const std::string &query_file, const unsigned nq, const unsigned Ls, - const unsigned k_value) -{ - web::http::client::http_client client(U(ip_addr_port)); +void query_loop(const std::string &ip_addr_port, const std::string &query_file, + const unsigned nq, const unsigned Ls, const unsigned k_value) { + web::http::client::http_client client(U(ip_addr_port)); - T *data; - size_t npts = 1, ndims = 128, rounded_dim = 128; - diskann::load_aligned_bin(query_file, data, npts, ndims, rounded_dim); + T *data; + size_t npts = 1, ndims = 128, rounded_dim = 128; + diskann::load_aligned_bin(query_file, data, npts, ndims, rounded_dim); - for (unsigned i = 0; i < nq; ++i) - { - T *vec = data + i * rounded_dim; - web::http::http_request http_query(methods::POST); - web::json::value queryJson = web::json::value::object(); - queryJson[QUERY_ID_KEY] = i; - queryJson[K_KEY] = k_value; - queryJson[L_KEY] = Ls; - for (size_t i = 0; i < ndims; ++i) - { - queryJson[VECTOR_KEY][i] = web::json::value::number(vec[i]); - } - http_query.set_body(queryJson); - - client.request(http_query) - .then([](web::http::http_response response) -> pplx::task { - if (response.status_code() == status_codes::OK) - { - return response.extract_string(); - } - std::cerr << "Query failed" << std::endl; - return pplx::task_from_result(utility::string_t()); - }) - .then([](pplx::task previousTask) { - try - { - std::cout << previousTask.get() << std::endl; - } - catch (http_exception const &e) - { - std::wcout << e.what() << std::endl; - } - }) - .wait(); + for (unsigned i = 0; i < nq; ++i) { + T *vec = data + i * rounded_dim; + web::http::http_request http_query(methods::POST); + web::json::value queryJson = web::json::value::object(); + queryJson[QUERY_ID_KEY] = i; + queryJson[K_KEY] = k_value; + queryJson[L_KEY] = Ls; + for (size_t i = 0; i < ndims; ++i) { + queryJson[VECTOR_KEY][i] = web::json::value::number(vec[i]); } + http_query.set_body(queryJson); + + client.request(http_query) + .then([](web::http::http_response response) + -> pplx::task { + if (response.status_code() == status_codes::OK) { + return response.extract_string(); + } + std::cerr << "Query failed" << std::endl; + return pplx::task_from_result(utility::string_t()); + }) + .then([](pplx::task previousTask) { + try { + std::cout << previousTask.get() << std::endl; + } catch (http_exception const &e) { + std::wcout << e.what() << std::endl; + } + }) + .wait(); + } } -int main(int argc, char *argv[]) -{ - std::string data_type, query_file, address; - uint32_t num_queries; - uint32_t l_search, k_value; +int main(int argc, char *argv[]) { + std::string data_type, query_file, address; + uint32_t num_queries; + uint32_t l_search, k_value; - po::options_description desc{"Arguments"}; - try - { - desc.add_options()("help,h", "Print information on arguments"); - desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); - desc.add_options()("address", po::value(&address)->required(), "Web server address"); - desc.add_options()("query_file", po::value(&query_file)->required(), - "File containing the queries to search"); - desc.add_options()("num_queries,Q", po::value(&num_queries)->required(), - "Number of queries to search"); - desc.add_options()("l_search", po::value(&l_search)->required(), "Value of L"); - desc.add_options()("k_value,K", po::value(&k_value)->default_value(10), "Value of K (default 10)"); - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) - { - std::cout << desc; - return 0; - } - po::notify(vm); - } - catch (const std::exception &ex) - { - std::cerr << ex.what() << std::endl; - return -1; + po::options_description desc{"Arguments"}; + try { + desc.add_options()("help,h", "Print information on arguments"); + desc.add_options()("data_type", + po::value(&data_type)->required(), + "data type "); + desc.add_options()("address", po::value(&address)->required(), + "Web server address"); + desc.add_options()("query_file", + po::value(&query_file)->required(), + "File containing the queries to search"); + desc.add_options()("num_queries,Q", + po::value(&num_queries)->required(), + "Number of queries to search"); + desc.add_options()("l_search", po::value(&l_search)->required(), + "Value of L"); + desc.add_options()("k_value,K", + po::value(&k_value)->default_value(10), + "Value of K (default 10)"); + po::variables_map vm; + po::store(po::parse_command_line(argc, argv, desc), vm); + if (vm.count("help")) { + std::cout << desc; + return 0; } + po::notify(vm); + } catch (const std::exception &ex) { + std::cerr << ex.what() << std::endl; + return -1; + } - if (data_type == std::string("float")) - { - query_loop(address, query_file, num_queries, l_search, k_value); - } - else if (data_type == std::string("int8")) - { - query_loop(address, query_file, num_queries, l_search, k_value); - } - else if (data_type == std::string("uint8")) - { - query_loop(address, query_file, num_queries, l_search, k_value); - } - else - { - std::cerr << "Unsupported type " << argv[2] << std::endl; - return -1; - } + if (data_type == std::string("float")) { + query_loop(address, query_file, num_queries, l_search, k_value); + } else if (data_type == std::string("int8")) { + query_loop(address, query_file, num_queries, l_search, k_value); + } else if (data_type == std::string("uint8")) { + query_loop(address, query_file, num_queries, l_search, k_value); + } else { + std::cerr << "Unsupported type " << argv[2] << std::endl; + return -1; + } - return 0; + return 0; } \ No newline at end of file diff --git a/tests/restapi/inmem_server.cpp b/tests/restapi/inmem_server.cpp index 11da541ff..fb0fb6223 100644 --- a/tests/restapi/inmem_server.cpp +++ b/tests/restapi/inmem_server.cpp @@ -17,122 +17,115 @@ namespace po = boost::program_options; std::unique_ptr g_httpServer(nullptr); std::vector> g_inMemorySearch; -void setup(const utility::string_t &address, const std::string &typestring) -{ - web::http::uri_builder uriBldr(address); - auto uri = uriBldr.to_uri(); +void setup(const utility::string_t &address, const std::string &typestring) { + web::http::uri_builder uriBldr(address); + auto uri = uriBldr.to_uri(); - std::cout << "Attempting to start server on " << uri.to_string() << std::endl; + std::cout << "Attempting to start server on " << uri.to_string() << std::endl; - g_httpServer = std::unique_ptr(new Server(uri, g_inMemorySearch, typestring)); - std::cout << "Created a server object" << std::endl; + g_httpServer = + std::unique_ptr(new Server(uri, g_inMemorySearch, typestring)); + std::cout << "Created a server object" << std::endl; - g_httpServer->open().wait(); - ucout << U"Listening for requests on: " << address << std::endl; + g_httpServer->open().wait(); + ucout << U"Listening for requests on: " << address << std::endl; } -void teardown(const utility::string_t &address) -{ - g_httpServer->close().wait(); +void teardown(const utility::string_t &address) { + g_httpServer->close().wait(); } -int main(int argc, char *argv[]) -{ - std::string data_type, index_file, data_file, address, dist_fn, tags_file; - uint32_t num_threads; - uint32_t l_search; +int main(int argc, char *argv[]) { + std::string data_type, index_file, data_file, address, dist_fn, tags_file; + uint32_t num_threads; + uint32_t l_search; - po::options_description desc{"Arguments"}; - try - { - desc.add_options()("help,h", "Print information on arguments"); - desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); - desc.add_options()("address", po::value(&address)->required(), "Web server address"); - desc.add_options()("data_file", po::value(&data_file)->required(), - "File containing the data found in the index"); - desc.add_options()("index_path_prefix", po::value(&index_file)->required(), - "Path prefix for saving index file components"); - desc.add_options()("num_threads,T", po::value(&num_threads)->required(), - "Number of threads used for building index"); - desc.add_options()("l_search", po::value(&l_search)->required(), "Value of L"); - desc.add_options()("dist_fn", po::value(&dist_fn)->default_value("l2"), - "distance function "); - desc.add_options()("tags_file", po::value(&tags_file)->default_value(std::string()), - "Tags file location"); - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) - { - std::cout << desc; - return 0; - } - po::notify(vm); - } - catch (const std::exception &ex) - { - std::cerr << ex.what() << std::endl; - return -1; - } - diskann::Metric metric; - if (dist_fn == std::string("l2")) - metric = diskann::Metric::L2; - else if (dist_fn == std::string("mips")) - metric = diskann::Metric::INNER_PRODUCT; - else - { - std::cout << "Error. Only l2 and mips distance functions are supported" << std::endl; - return -1; + po::options_description desc{"Arguments"}; + try { + desc.add_options()("help,h", "Print information on arguments"); + desc.add_options()("data_type", + po::value(&data_type)->required(), + "data type "); + desc.add_options()("address", po::value(&address)->required(), + "Web server address"); + desc.add_options()("data_file", + po::value(&data_file)->required(), + "File containing the data found in the index"); + desc.add_options()("index_path_prefix", + po::value(&index_file)->required(), + "Path prefix for saving index file components"); + desc.add_options()("num_threads,T", + po::value(&num_threads)->required(), + "Number of threads used for building index"); + desc.add_options()("l_search", po::value(&l_search)->required(), + "Value of L"); + desc.add_options()("dist_fn", + po::value(&dist_fn)->default_value("l2"), + "distance function "); + desc.add_options()( + "tags_file", + po::value(&tags_file)->default_value(std::string()), + "Tags file location"); + po::variables_map vm; + po::store(po::parse_command_line(argc, argv, desc), vm); + if (vm.count("help")) { + std::cout << desc; + return 0; } + po::notify(vm); + } catch (const std::exception &ex) { + std::cerr << ex.what() << std::endl; + return -1; + } + diskann::Metric metric; + if (dist_fn == std::string("l2")) + metric = diskann::Metric::L2; + else if (dist_fn == std::string("mips")) + metric = diskann::Metric::INNER_PRODUCT; + else { + std::cout << "Error. Only l2 and mips distance functions are supported" + << std::endl; + return -1; + } - if (data_type == std::string("float")) - { - auto searcher = std::unique_ptr( - new diskann::InMemorySearch(data_file, index_file, tags_file, metric, num_threads, l_search)); - g_inMemorySearch.push_back(std::move(searcher)); - } - else if (data_type == std::string("int8")) - { - auto searcher = std::unique_ptr( - new diskann::InMemorySearch(data_file, index_file, tags_file, metric, num_threads, l_search)); - g_inMemorySearch.push_back(std::move(searcher)); - } - else if (data_type == std::string("uint8")) - { - auto searcher = std::unique_ptr( - new diskann::InMemorySearch(data_file, index_file, tags_file, metric, num_threads, l_search)); - g_inMemorySearch.push_back(std::move(searcher)); - } - else - { - std::cerr << "Unsupported data type " << argv[2] << std::endl; - } + if (data_type == std::string("float")) { + auto searcher = + std::unique_ptr(new diskann::InMemorySearch( + data_file, index_file, tags_file, metric, num_threads, l_search)); + g_inMemorySearch.push_back(std::move(searcher)); + } else if (data_type == std::string("int8")) { + auto searcher = std::unique_ptr( + new diskann::InMemorySearch(data_file, index_file, tags_file, + metric, num_threads, l_search)); + g_inMemorySearch.push_back(std::move(searcher)); + } else if (data_type == std::string("uint8")) { + auto searcher = std::unique_ptr( + new diskann::InMemorySearch(data_file, index_file, tags_file, + metric, num_threads, l_search)); + g_inMemorySearch.push_back(std::move(searcher)); + } else { + std::cerr << "Unsupported data type " << argv[2] << std::endl; + } - while (1) - { - try - { - setup(address, data_type); - std::cout << "Type 'exit' (case-sensitive) to exit" << std::endl; - std::string line; - std::getline(std::cin, line); - if (line == "exit") - { - teardown(address); - g_httpServer->close().wait(); - exit(0); - } - } - catch (const std::exception &ex) - { - std::cerr << "Exception occurred: " << ex.what() << std::endl; - std::cerr << "Restarting HTTP server"; - teardown(address); - } - catch (...) - { - std::cerr << "Unknown exception occurreed" << std::endl; - std::cerr << "Restarting HTTP server"; - teardown(address); - } + while (1) { + try { + setup(address, data_type); + std::cout << "Type 'exit' (case-sensitive) to exit" << std::endl; + std::string line; + std::getline(std::cin, line); + if (line == "exit") { + teardown(address); + g_httpServer->close().wait(); + exit(0); + } + } catch (const std::exception &ex) { + std::cerr << "Exception occurred: " << ex.what() << std::endl; + std::cerr << "Restarting HTTP server"; + teardown(address); + } catch (...) { + std::cerr << "Unknown exception occurreed" << std::endl; + std::cerr << "Restarting HTTP server"; + teardown(address); } + } } diff --git a/tests/restapi/main.cpp b/tests/restapi/main.cpp index cb48d6787..f1c5a78c7 100644 --- a/tests/restapi/main.cpp +++ b/tests/restapi/main.cpp @@ -9,75 +9,65 @@ std::unique_ptr g_httpServer(nullptr); std::unique_ptr g_inMemorySearch(nullptr); -void setup(const utility::string_t &address) -{ - web::http::uri_builder uriBldr(address); - auto uri = uriBldr.to_uri(); +void setup(const utility::string_t &address) { + web::http::uri_builder uriBldr(address); + auto uri = uriBldr.to_uri(); - std::wcout << L"Attempting to start server on " << uri.to_string() << std::endl; + std::wcout << L"Attempting to start server on " << uri.to_string() + << std::endl; - g_httpServer = std::unique_ptr(new Server(uri, g_inMemorySearch)); - g_httpServer->open().wait(); + g_httpServer = std::unique_ptr(new Server(uri, g_inMemorySearch)); + g_httpServer->open().wait(); - ucout << U"Listening for requests on: " << address << std::endl; + ucout << U"Listening for requests on: " << address << std::endl; } -void teardown(const utility::string_t &address) -{ - g_httpServer->close().wait(); +void teardown(const utility::string_t &address) { + g_httpServer->close().wait(); } -void loadIndex(const char *indexFile, const char *baseFile, const char *idsFile) -{ - auto nsgSearch = new diskann::InMemorySearch(baseFile, indexFile, idsFile, diskann::L2); - g_inMemorySearch = std::unique_ptr(nsgSearch); +void loadIndex(const char *indexFile, const char *baseFile, + const char *idsFile) { + auto nsgSearch = + new diskann::InMemorySearch(baseFile, indexFile, idsFile, diskann::L2); + g_inMemorySearch = std::unique_ptr(nsgSearch); } -std::wstring getHostingAddress(const char *hostNameAndPort) -{ - wchar_t buffer[4096]; - mbstowcs_s(nullptr, buffer, sizeof(buffer) / sizeof(buffer[0]), hostNameAndPort, - sizeof(buffer) / sizeof(buffer[0])); - return std::wstring(buffer); +std::wstring getHostingAddress(const char *hostNameAndPort) { + wchar_t buffer[4096]; + mbstowcs_s(nullptr, buffer, sizeof(buffer) / sizeof(buffer[0]), + hostNameAndPort, sizeof(buffer) / sizeof(buffer[0])); + return std::wstring(buffer); } -int main(int argc, char *argv[]) -{ - if (argc != 5) - { - std::cout << "Usage: nsg_server " - " " - << std::endl; - exit(1); - } +int main(int argc, char *argv[]) { + if (argc != 5) { + std::cout << "Usage: nsg_server " + " " + << std::endl; + exit(1); + } - auto address = getHostingAddress(argv[1]); - loadIndex(argv[2], argv[3], argv[4]); - while (1) - { - try - { - setup(address); - std::cout << "Type 'exit' (case-sensitive) to exit" << std::endl; - std::string line; - std::getline(std::cin, line); - if (line == "exit") - { - teardown(address); - exit(0); - } - } - catch (const std::exception &ex) - { - std::cerr << "Exception occurred: " << ex.what() << std::endl; - std::cerr << "Restarting HTTP server"; - teardown(address); - } - catch (...) - { - std::cerr << "Unknown exception occurreed" << std::endl; - std::cerr << "Restarting HTTP server"; - teardown(address); - } + auto address = getHostingAddress(argv[1]); + loadIndex(argv[2], argv[3], argv[4]); + while (1) { + try { + setup(address); + std::cout << "Type 'exit' (case-sensitive) to exit" << std::endl; + std::string line; + std::getline(std::cin, line); + if (line == "exit") { + teardown(address); + exit(0); + } + } catch (const std::exception &ex) { + std::cerr << "Exception occurred: " << ex.what() << std::endl; + std::cerr << "Restarting HTTP server"; + teardown(address); + } catch (...) { + std::cerr << "Unknown exception occurreed" << std::endl; + std::cerr << "Restarting HTTP server"; + teardown(address); } + } } diff --git a/tests/restapi/multiple_ssdindex_server.cpp b/tests/restapi/multiple_ssdindex_server.cpp index 89cb06fcb..c550cc0ca 100644 --- a/tests/restapi/multiple_ssdindex_server.cpp +++ b/tests/restapi/multiple_ssdindex_server.cpp @@ -18,165 +18,153 @@ namespace po = boost::program_options; std::unique_ptr g_httpServer(nullptr); std::vector> g_ssdSearch; -void setup(const utility::string_t &address, const std::string &typestring) -{ - web::http::uri_builder uriBldr(address); - auto uri = uriBldr.to_uri(); +void setup(const utility::string_t &address, const std::string &typestring) { + web::http::uri_builder uriBldr(address); + auto uri = uriBldr.to_uri(); - std::cout << "Attempting to start server on " << uri.to_string() << std::endl; + std::cout << "Attempting to start server on " << uri.to_string() << std::endl; - g_httpServer = std::unique_ptr(new Server(uri, g_ssdSearch, typestring)); - std::cout << "Created a server object" << std::endl; + g_httpServer = + std::unique_ptr(new Server(uri, g_ssdSearch, typestring)); + std::cout << "Created a server object" << std::endl; - g_httpServer->open().wait(); - ucout << U"Listening for requests on: " << address << std::endl; + g_httpServer->open().wait(); + ucout << U"Listening for requests on: " << address << std::endl; } -void teardown(const utility::string_t &address) -{ - g_httpServer->close().wait(); +void teardown(const utility::string_t &address) { + g_httpServer->close().wait(); } -int main(int argc, char *argv[]) -{ - std::string data_type, index_prefix_paths, address, dist_fn, tags_file; - uint32_t num_nodes_to_cache; - uint32_t num_threads; - - po::options_description desc{"Arguments"}; - try - { - desc.add_options()("help,h", "Print information on arguments"); - desc.add_options()("address", po::value(&address)->required(), "Web server address"); - desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); - desc.add_options()("index_prefix_paths", po::value(&index_prefix_paths)->required(), - "Path prefix for loading index file components"); - desc.add_options()("num_nodes_to_cache", po::value(&num_nodes_to_cache)->default_value(0), - "Number of nodes to cache during search"); - desc.add_options()("num_threads,T", po::value(&num_threads)->default_value(omp_get_num_procs()), - "Number of threads used for building index (defaults to " - "omp_get_num_procs())"); - desc.add_options()("dist_fn", po::value(&dist_fn)->default_value("l2"), - "distance function "); - desc.add_options()("tags_file", po::value(&tags_file)->default_value(std::string()), - "Tags file location"); - - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) - { - std::cout << desc; - return 0; - } - po::notify(vm); +int main(int argc, char *argv[]) { + std::string data_type, index_prefix_paths, address, dist_fn, tags_file; + uint32_t num_nodes_to_cache; + uint32_t num_threads; + + po::options_description desc{"Arguments"}; + try { + desc.add_options()("help,h", "Print information on arguments"); + desc.add_options()("address", po::value(&address)->required(), + "Web server address"); + desc.add_options()("data_type", + po::value(&data_type)->required(), + "data type "); + desc.add_options()("index_prefix_paths", + po::value(&index_prefix_paths)->required(), + "Path prefix for loading index file components"); + desc.add_options()( + "num_nodes_to_cache", + po::value(&num_nodes_to_cache)->default_value(0), + "Number of nodes to cache during search"); + desc.add_options()( + "num_threads,T", + po::value(&num_threads)->default_value(omp_get_num_procs()), + "Number of threads used for building index (defaults to " + "omp_get_num_procs())"); + desc.add_options()("dist_fn", + po::value(&dist_fn)->default_value("l2"), + "distance function "); + desc.add_options()( + "tags_file", + po::value(&tags_file)->default_value(std::string()), + "Tags file location"); + + po::variables_map vm; + po::store(po::parse_command_line(argc, argv, desc), vm); + if (vm.count("help")) { + std::cout << desc; + return 0; } - catch (const std::exception &ex) - { - std::cerr << ex.what() << std::endl; - return -1; + po::notify(vm); + } catch (const std::exception &ex) { + std::cerr << ex.what() << std::endl; + return -1; + } + + diskann::Metric metric; + if (dist_fn == std::string("l2")) + metric = diskann::Metric::L2; + else if (dist_fn == std::string("mips")) + metric = diskann::Metric::INNER_PRODUCT; + else { + std::cout << "Error. Only l2 and mips distance functions are supported" + << std::endl; + return -1; + } + + std::vector> index_tag_paths; + std::ifstream index_in(index_prefix_paths); + if (!index_in.is_open()) { + std::cerr << "Could not open " << index_prefix_paths << std::endl; + exit(-1); + } + std::ifstream tags_in(tags_file); + if (!tags_in.is_open()) { + std::cerr << "Could not open " << tags_file << std::endl; + exit(-1); + } + std::string prefix, tagfile; + while (std::getline(index_in, prefix)) { + if (std::getline(tags_in, tagfile)) { + index_tag_paths.push_back(std::make_pair(prefix, tagfile)); + } else { + std::cerr << "The number of tags specified does not match the number of " + "indices specified" + << std::endl; + exit(-1); } - - diskann::Metric metric; - if (dist_fn == std::string("l2")) - metric = diskann::Metric::L2; - else if (dist_fn == std::string("mips")) - metric = diskann::Metric::INNER_PRODUCT; - else - { - std::cout << "Error. Only l2 and mips distance functions are supported" << std::endl; - return -1; - } - - std::vector> index_tag_paths; - std::ifstream index_in(index_prefix_paths); - if (!index_in.is_open()) - { - std::cerr << "Could not open " << index_prefix_paths << std::endl; - exit(-1); - } - std::ifstream tags_in(tags_file); - if (!tags_in.is_open()) - { - std::cerr << "Could not open " << tags_file << std::endl; - exit(-1); + } + index_in.close(); + tags_in.close(); + + if (data_type == std::string("float")) { + for (auto &index_tag : index_tag_paths) { + auto searcher = std::unique_ptr( + new diskann::PQFlashSearch(index_tag.first.c_str(), + num_nodes_to_cache, num_threads, + index_tag.second.c_str(), metric)); + g_ssdSearch.push_back(std::move(searcher)); } - std::string prefix, tagfile; - while (std::getline(index_in, prefix)) - { - if (std::getline(tags_in, tagfile)) - { - index_tag_paths.push_back(std::make_pair(prefix, tagfile)); - } - else - { - std::cerr << "The number of tags specified does not match the number of " - "indices specified" - << std::endl; - exit(-1); - } + } else if (data_type == std::string("int8")) { + for (auto &index_tag : index_tag_paths) { + auto searcher = std::unique_ptr( + new diskann::PQFlashSearch(index_tag.first.c_str(), + num_nodes_to_cache, num_threads, + index_tag.second.c_str(), metric)); + g_ssdSearch.push_back(std::move(searcher)); } - index_in.close(); - tags_in.close(); - - if (data_type == std::string("float")) - { - for (auto &index_tag : index_tag_paths) - { - auto searcher = std::unique_ptr(new diskann::PQFlashSearch( - index_tag.first.c_str(), num_nodes_to_cache, num_threads, index_tag.second.c_str(), metric)); - g_ssdSearch.push_back(std::move(searcher)); - } + } else if (data_type == std::string("uint8")) { + for (auto &index_tag : index_tag_paths) { + auto searcher = std::unique_ptr( + new diskann::PQFlashSearch( + index_tag.first.c_str(), num_nodes_to_cache, num_threads, + index_tag.second.c_str(), metric)); + g_ssdSearch.push_back(std::move(searcher)); } - else if (data_type == std::string("int8")) - { - for (auto &index_tag : index_tag_paths) - { - auto searcher = std::unique_ptr(new diskann::PQFlashSearch( - index_tag.first.c_str(), num_nodes_to_cache, num_threads, index_tag.second.c_str(), metric)); - g_ssdSearch.push_back(std::move(searcher)); - } - } - else if (data_type == std::string("uint8")) - { - for (auto &index_tag : index_tag_paths) - { - auto searcher = std::unique_ptr(new diskann::PQFlashSearch( - index_tag.first.c_str(), num_nodes_to_cache, num_threads, index_tag.second.c_str(), metric)); - g_ssdSearch.push_back(std::move(searcher)); - } - } - else - { - std::cerr << "Unsupported data type " << data_type << std::endl; - exit(-1); - } - - while (1) - { - try - { - setup(address, data_type); - std::cout << "Type 'exit' (case-sensitive) to exit" << std::endl; - std::string line; - std::getline(std::cin, line); - if (line == "exit") - { - teardown(address); - g_httpServer->close().wait(); - exit(0); - } - } - catch (const std::exception &ex) - { - std::cerr << "Exception occurred: " << ex.what() << std::endl; - std::cerr << "Restarting HTTP server"; - teardown(address); - } - catch (...) - { - std::cerr << "Unknown exception occurreed" << std::endl; - std::cerr << "Restarting HTTP server"; - teardown(address); - } + } else { + std::cerr << "Unsupported data type " << data_type << std::endl; + exit(-1); + } + + while (1) { + try { + setup(address, data_type); + std::cout << "Type 'exit' (case-sensitive) to exit" << std::endl; + std::string line; + std::getline(std::cin, line); + if (line == "exit") { + teardown(address); + g_httpServer->close().wait(); + exit(0); + } + } catch (const std::exception &ex) { + std::cerr << "Exception occurred: " << ex.what() << std::endl; + std::cerr << "Restarting HTTP server"; + teardown(address); + } catch (...) { + std::cerr << "Unknown exception occurreed" << std::endl; + std::cerr << "Restarting HTTP server"; + teardown(address); } + } } diff --git a/tests/restapi/ssd_server.cpp b/tests/restapi/ssd_server.cpp index d17997374..9364e122e 100644 --- a/tests/restapi/ssd_server.cpp +++ b/tests/restapi/ssd_server.cpp @@ -18,124 +18,120 @@ namespace po = boost::program_options; std::unique_ptr g_httpServer(nullptr); std::vector> g_ssdSearch; -void setup(const utility::string_t &address, const std::string &typestring) -{ - web::http::uri_builder uriBldr(address); - auto uri = uriBldr.to_uri(); +void setup(const utility::string_t &address, const std::string &typestring) { + web::http::uri_builder uriBldr(address); + auto uri = uriBldr.to_uri(); - std::cout << "Attempting to start server on " << uri.to_string() << std::endl; + std::cout << "Attempting to start server on " << uri.to_string() << std::endl; - g_httpServer = std::unique_ptr(new Server(uri, g_ssdSearch, typestring)); - std::cout << "Created a server object" << std::endl; + g_httpServer = + std::unique_ptr(new Server(uri, g_ssdSearch, typestring)); + std::cout << "Created a server object" << std::endl; - g_httpServer->open().wait(); - ucout << U"Listening for requests on: " << address << std::endl; + g_httpServer->open().wait(); + ucout << U"Listening for requests on: " << address << std::endl; } -void teardown(const utility::string_t &address) -{ - g_httpServer->close().wait(); +void teardown(const utility::string_t &address) { + g_httpServer->close().wait(); } -int main(int argc, char *argv[]) -{ - std::string data_type, index_path_prefix, address, dist_fn, tags_file; - uint32_t num_nodes_to_cache; - uint32_t num_threads; +int main(int argc, char *argv[]) { + std::string data_type, index_path_prefix, address, dist_fn, tags_file; + uint32_t num_nodes_to_cache; + uint32_t num_threads; - po::options_description desc{"Arguments"}; - try - { - desc.add_options()("help,h", "Print information on arguments"); - desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); - desc.add_options()("address", po::value(&address)->required(), "Web server address"); - desc.add_options()("index_path_prefix", po::value(&index_path_prefix)->required(), - "Path prefix for loading index file components"); - desc.add_options()("num_nodes_to_cache", po::value(&num_nodes_to_cache)->default_value(0), - "Number of nodes to cache during search"); - desc.add_options()("num_threads,T", po::value(&num_threads)->default_value(omp_get_num_procs()), - "Number of threads used for building index (defaults to " - "omp_get_num_procs())"); - desc.add_options()("dist_fn", po::value(&dist_fn)->default_value("l2"), - "distance function "); - desc.add_options()("tags_file", po::value(&tags_file)->default_value(std::string()), - "Tags file location"); - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) - { - std::cout << desc; - return 0; - } - po::notify(vm); - } - catch (const std::exception &ex) - { - std::cerr << ex.what() << std::endl; - return -1; + po::options_description desc{"Arguments"}; + try { + desc.add_options()("help,h", "Print information on arguments"); + desc.add_options()("data_type", + po::value(&data_type)->required(), + "data type "); + desc.add_options()("address", po::value(&address)->required(), + "Web server address"); + desc.add_options()("index_path_prefix", + po::value(&index_path_prefix)->required(), + "Path prefix for loading index file components"); + desc.add_options()( + "num_nodes_to_cache", + po::value(&num_nodes_to_cache)->default_value(0), + "Number of nodes to cache during search"); + desc.add_options()( + "num_threads,T", + po::value(&num_threads)->default_value(omp_get_num_procs()), + "Number of threads used for building index (defaults to " + "omp_get_num_procs())"); + desc.add_options()("dist_fn", + po::value(&dist_fn)->default_value("l2"), + "distance function "); + desc.add_options()( + "tags_file", + po::value(&tags_file)->default_value(std::string()), + "Tags file location"); + po::variables_map vm; + po::store(po::parse_command_line(argc, argv, desc), vm); + if (vm.count("help")) { + std::cout << desc; + return 0; } + po::notify(vm); + } catch (const std::exception &ex) { + std::cerr << ex.what() << std::endl; + return -1; + } - diskann::Metric metric; - if (dist_fn == std::string("l2")) - metric = diskann::Metric::L2; - else if (dist_fn == std::string("mips")) - metric = diskann::Metric::INNER_PRODUCT; - else - { - std::cout << "Error. Only l2 and mips distance functions are supported" << std::endl; - return -1; - } + diskann::Metric metric; + if (dist_fn == std::string("l2")) + metric = diskann::Metric::L2; + else if (dist_fn == std::string("mips")) + metric = diskann::Metric::INNER_PRODUCT; + else { + std::cout << "Error. Only l2 and mips distance functions are supported" + << std::endl; + return -1; + } - if (data_type == std::string("float")) - { - auto searcher = std::unique_ptr( - new diskann::PQFlashSearch(index_path_prefix, num_nodes_to_cache, num_threads, tags_file, metric)); - g_ssdSearch.push_back(std::move(searcher)); - } - else if (data_type == std::string("int8")) - { - auto searcher = std::unique_ptr( - new diskann::PQFlashSearch(index_path_prefix, num_nodes_to_cache, num_threads, tags_file, metric)); - g_ssdSearch.push_back(std::move(searcher)); - } - else if (data_type == std::string("uint8")) - { - auto searcher = std::unique_ptr( - new diskann::PQFlashSearch(index_path_prefix, num_nodes_to_cache, num_threads, tags_file, metric)); - g_ssdSearch.push_back(std::move(searcher)); - } - else - { - std::cerr << "Unsupported data type " << argv[2] << std::endl; - exit(-1); - } + if (data_type == std::string("float")) { + auto searcher = std::unique_ptr( + new diskann::PQFlashSearch(index_path_prefix, num_nodes_to_cache, + num_threads, tags_file, metric)); + g_ssdSearch.push_back(std::move(searcher)); + } else if (data_type == std::string("int8")) { + auto searcher = + std::unique_ptr(new diskann::PQFlashSearch( + index_path_prefix, num_nodes_to_cache, num_threads, tags_file, + metric)); + g_ssdSearch.push_back(std::move(searcher)); + } else if (data_type == std::string("uint8")) { + auto searcher = std::unique_ptr( + new diskann::PQFlashSearch(index_path_prefix, + num_nodes_to_cache, num_threads, + tags_file, metric)); + g_ssdSearch.push_back(std::move(searcher)); + } else { + std::cerr << "Unsupported data type " << argv[2] << std::endl; + exit(-1); + } - while (1) - { - try - { - setup(address, data_type); - std::cout << "Type 'exit' (case-sensitive) to exit" << std::endl; - std::string line; - std::getline(std::cin, line); - if (line == "exit") - { - teardown(address); - g_httpServer->close().wait(); - exit(0); - } - } - catch (const std::exception &ex) - { - std::cerr << "Exception occurred: " << ex.what() << std::endl; - std::cerr << "Restarting HTTP server"; - teardown(address); - } - catch (...) - { - std::cerr << "Unknown exception occurreed" << std::endl; - std::cerr << "Restarting HTTP server"; - teardown(address); - } + while (1) { + try { + setup(address, data_type); + std::cout << "Type 'exit' (case-sensitive) to exit" << std::endl; + std::string line; + std::getline(std::cin, line); + if (line == "exit") { + teardown(address); + g_httpServer->close().wait(); + exit(0); + } + } catch (const std::exception &ex) { + std::cerr << "Exception occurred: " << ex.what() << std::endl; + std::cerr << "Restarting HTTP server"; + teardown(address); + } catch (...) { + std::cerr << "Unknown exception occurreed" << std::endl; + std::cerr << "Restarting HTTP server"; + teardown(address); } + } } From a4a5cb5ebc2839aae2862441bc51ec6de436d69d Mon Sep 17 00:00:00 2001 From: ravishankar Date: Wed, 19 Apr 2023 23:58:41 +0530 Subject: [PATCH 59/69] clang-format 3 --- python/src/diskann_bindings.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/src/diskann_bindings.cpp b/python/src/diskann_bindings.cpp index 96e4709f9..c7b63233d 100644 --- a/python/src/diskann_bindings.cpp +++ b/python/src/diskann_bindings.cpp @@ -154,8 +154,8 @@ struct DynamicInMemIndex { initial_search_list_size, // used to prepare the scratch space for // searching. can / may be expanded if the // search asks for a larger L. - search_threads, // also used for the scratch space - true, // enable_tags + search_threads, // also used for the scratch space + true, // enable_tags concurrent_consolidate, false, // pq_dist_build 0, // num_pq_chunks From 4ce5abf8182d45ca4d570eb91a74a15b7acef5ac Mon Sep 17 00:00:00 2001 From: ravishankar Date: Thu, 20 Apr 2023 11:13:05 +0530 Subject: [PATCH 60/69] reverted clang to Microsoft --- .clang-format | 2 +- include/abstract_data_store.h | 202 +- include/abstract_graph_store.h | 28 +- include/aligned_file_reader.h | 139 +- include/ann_exception.h | 33 +- include/boost_dynamic_bitset_fwd.h | 9 +- include/cached_io.h | 348 +- include/concurrent_queue.h | 175 +- include/cosine_similarity.h | 392 +- include/defaults.h | 10 +- include/disk_utils.h | 98 +- include/distance.h | 370 +- include/exceptions.h | 15 +- include/filter_utils.h | 275 +- include/in_mem_data_store.h | 122 +- include/in_mem_graph_store.h | 20 +- include/index.h | 794 ++- include/linux_aligned_file_reader.h | 56 +- include/locking.h | 5 +- include/logger.h | 15 +- include/logger_impl.h | 65 +- include/math_utils.h | 48 +- include/memory_mapper.h | 34 +- include/natural_number_map.h | 130 +- include/natural_number_set.h | 57 +- include/neighbor.h | 191 +- include/parameters.h | 216 +- include/partition.h | 39 +- include/percentile_stats.h | 78 +- include/pq.h | 174 +- include/pq_flash_index.h | 356 +- include/restapi/common.h | 14 +- include/restapi/search_wrapper.h | 211 +- include/restapi/server.h | 76 +- include/scratch.h | 320 +- include/simd_utils.h | 137 +- include/timer.h | 46 +- include/types.h | 5 +- include/utils.h | 1716 ++--- include/windows_aligned_file_reader.h | 64 +- include/windows_slim_lock.h | 76 +- python/src/diskann_bindings.cpp | 778 ++- src/abstract_data_store.cpp | 47 +- src/ann_exception.cpp | 43 +- src/disk_utils.cpp | 2572 ++++---- src/distance.cpp | 1092 ++-- src/dll/dllmain.cpp | 13 +- src/filter_utils.cpp | 452 +- src/in_mem_data_store.cpp | 565 +- src/in_mem_graph_store.cpp | 28 +- src/index.cpp | 5550 +++++++++-------- src/linux_aligned_file_reader.cpp | 333 +- src/logger.cpp | 112 +- src/math_utils.cpp | 668 +- src/memory_mapper.cpp | 146 +- src/natural_number_map.cpp | 132 +- src/natural_number_set.cpp | 76 +- src/partition.cpp | 1121 ++-- src/pq.cpp | 1739 +++--- src/pq_flash_index.cpp | 2521 ++++---- src/restapi/search_wrapper.cpp | 285 +- src/restapi/server.cpp | 434 +- src/scratch.cpp | 189 +- src/utils.cpp | 814 ++- src/windows_aligned_file_reader.cpp | 283 +- tests/build_disk_index.cpp | 341 +- tests/build_memory_index.cpp | 345 +- tests/build_stitched_index.cpp | 706 +-- tests/range_search_disk_index.cpp | 598 +- tests/restapi/client.cpp | 168 +- tests/restapi/inmem_server.cpp | 205 +- tests/restapi/main.cpp | 106 +- tests/restapi/multiple_ssdindex_server.cpp | 284 +- tests/restapi/ssd_server.cpp | 210 +- tests/search_disk_index.cpp | 824 ++- tests/search_memory_index.cpp | 756 +-- tests/test_insert_deletes_consolidate.cpp | 796 ++- tests/test_streaming_scenario.cpp | 672 +- tests/utils/bin_to_fvecs.cpp | 104 +- tests/utils/bin_to_tsv.cpp | 103 +- tests/utils/calculate_recall.cpp | 69 +- tests/utils/compute_groundtruth.cpp | 947 ++- .../utils/compute_groundtruth_for_filters.cpp | 1565 ++--- tests/utils/count_bfs_levels.cpp | 95 +- tests/utils/create_disk_layout.cpp | 57 +- tests/utils/float_bin_to_int8.cpp | 105 +- tests/utils/fvecs_to_bin.cpp | 141 +- tests/utils/fvecs_to_bvecs.cpp | 86 +- tests/utils/gen_random_slice.cpp | 56 +- tests/utils/generate_pq.cpp | 107 +- tests/utils/generate_synthetic_labels.cpp | 327 +- tests/utils/int8_to_float.cpp | 28 +- tests/utils/int8_to_float_scale.cpp | 105 +- tests/utils/ivecs_to_bin.cpp | 89 +- tests/utils/merge_shards.cpp | 45 +- tests/utils/partition_data.cpp | 53 +- tests/utils/partition_with_ram_budget.cpp | 53 +- tests/utils/rand_data_gen.cpp | 331 +- tests/utils/simulate_aggregate_recall.cpp | 119 +- tests/utils/stats_label_data.cpp | 207 +- tests/utils/tsv_to_bin.cpp | 191 +- tests/utils/uint32_to_uint8.cpp | 28 +- tests/utils/uint8_to_float.cpp | 28 +- tests/utils/vector_analysis.cpp | 239 +- 104 files changed, 19378 insertions(+), 19065 deletions(-) diff --git a/.clang-format b/.clang-format index 7ff715f79..ad3192fd6 100644 --- a/.clang-format +++ b/.clang-format @@ -1,5 +1,5 @@ --- -BasedOnStyle: Google +BasedOnStyle: Microsoft --- Language: Cpp SortIncludes: false diff --git a/include/abstract_data_store.h b/include/abstract_data_store.h index 4662fc7dd..71ce319fc 100644 --- a/include/abstract_data_store.h +++ b/include/abstract_data_store.h @@ -9,110 +9,102 @@ #include "types.h" #include "windows_customizations.h" -namespace diskann { - -template -class AbstractDataStore { - public: - AbstractDataStore(const location_t capacity, const size_t dim); - - // Return number of points returned - virtual location_t load(const std::string &filename) = 0; - - // Why does store take num_pts? Since store only has capacity, but we allow - // resizing we can end up in a situation where the store has spare capacity. - // To optimize disk utilization, we pass the number of points that are "true" - // points, so that the store can discard the empty locations before saving. - virtual size_t save(const std::string &filename, - const location_t num_pts) = 0; - - DISKANN_DLLEXPORT virtual location_t capacity() const; - - DISKANN_DLLEXPORT virtual size_t get_dims() const; - - // Implementers can choose to return _dim if they are not - // concerned about memory alignment. - // Some distance metrics (like l2) need data vectors to be aligned, so we - // align the dimension by padding zeros. - virtual size_t get_aligned_dim() const = 0; - - // populate the store with vectors (either from a pointer or bin file), - // potentially after pre-processing the vectors if the metric deems so - // e.g., normalizing vectors for cosine distance over floating-point vectors - // useful for bulk or static index building. - virtual void populate_data(const data_t *vectors, - const location_t num_pts) = 0; - virtual void populate_data(const std::string &filename, - const size_t offset) = 0; - - // save the first num_pts many vectors back to bin file - // note: cannot undo the pre-processing done in populate data - virtual void extract_data_to_bin(const std::string &filename, - const location_t num_pts) = 0; - - // Returns the updated capacity of the datastore. Clients should check - // if resize actually changed the capacity to new_num_points before - // proceeding with operations. See the code below: - // auto new_capcity = data_store->resize(new_num_points); - // if ( new_capacity >= new_num_points) { - // //PROCEED - // else - // //ERROR. - virtual location_t resize(const location_t new_num_points); - - // operations on vectors - // like populate_data function, but over one vector at a time useful for - // streaming setting - virtual void get_vector(const location_t i, data_t *dest) const = 0; - virtual void set_vector(const location_t i, const data_t *const vector) = 0; - virtual void prefetch_vector(const location_t loc) = 0; - - // internal shuffle operations to move around vectors - // will bulk-move all the vectors in [old_start_loc, old_start_loc + - // num_points) to [new_start_loc, new_start_loc + num_points) and set the old - // positions to zero vectors. - virtual void move_vectors(const location_t old_start_loc, - const location_t new_start_loc, - const location_t num_points) = 0; - - // same as above, without resetting the vectors in [from_loc, from_loc + - // num_points) to zero - virtual void copy_vectors(const location_t from_loc, const location_t to_loc, - const location_t num_points) = 0; - - // metric specific operations - - virtual float get_distance(const data_t *query, - const location_t loc) const = 0; - virtual void get_distance(const data_t *query, const location_t *locations, - const uint32_t location_count, - float *distances) const = 0; - virtual float get_distance(const location_t loc1, - const location_t loc2) const = 0; - - // stats of the data stored in store - // Returns the point in the dataset that is closest to the mean of all points - // in the dataset - virtual location_t calculate_medoid() const = 0; - - // search helpers - // if the base data is aligned per the request of the metric, this will tell - // how to align the query vector in a consistent manner - virtual size_t get_alignment_factor() const = 0; - - protected: - // Expand the datastore to new_num_points. Returns the new capacity created, - // which should be == new_num_points in the normal case. Implementers can also - // return _capacity to indicate that there are not implementing this method. - virtual location_t expand(const location_t new_num_points) = 0; - - // Shrink the datastore to new_num_points. It is NOT an error if shrink - // doesn't reduce the capacity so callers need to check this correctly. See - // also for "default" implementation - virtual location_t shrink(const location_t new_num_points) = 0; - - location_t _capacity; - size_t _dim; +namespace diskann +{ + +template class AbstractDataStore +{ + public: + AbstractDataStore(const location_t capacity, const size_t dim); + + // Return number of points returned + virtual location_t load(const std::string &filename) = 0; + + // Why does store take num_pts? Since store only has capacity, but we allow + // resizing we can end up in a situation where the store has spare capacity. + // To optimize disk utilization, we pass the number of points that are "true" + // points, so that the store can discard the empty locations before saving. + virtual size_t save(const std::string &filename, const location_t num_pts) = 0; + + DISKANN_DLLEXPORT virtual location_t capacity() const; + + DISKANN_DLLEXPORT virtual size_t get_dims() const; + + // Implementers can choose to return _dim if they are not + // concerned about memory alignment. + // Some distance metrics (like l2) need data vectors to be aligned, so we + // align the dimension by padding zeros. + virtual size_t get_aligned_dim() const = 0; + + // populate the store with vectors (either from a pointer or bin file), + // potentially after pre-processing the vectors if the metric deems so + // e.g., normalizing vectors for cosine distance over floating-point vectors + // useful for bulk or static index building. + virtual void populate_data(const data_t *vectors, const location_t num_pts) = 0; + virtual void populate_data(const std::string &filename, const size_t offset) = 0; + + // save the first num_pts many vectors back to bin file + // note: cannot undo the pre-processing done in populate data + virtual void extract_data_to_bin(const std::string &filename, const location_t num_pts) = 0; + + // Returns the updated capacity of the datastore. Clients should check + // if resize actually changed the capacity to new_num_points before + // proceeding with operations. See the code below: + // auto new_capcity = data_store->resize(new_num_points); + // if ( new_capacity >= new_num_points) { + // //PROCEED + // else + // //ERROR. + virtual location_t resize(const location_t new_num_points); + + // operations on vectors + // like populate_data function, but over one vector at a time useful for + // streaming setting + virtual void get_vector(const location_t i, data_t *dest) const = 0; + virtual void set_vector(const location_t i, const data_t *const vector) = 0; + virtual void prefetch_vector(const location_t loc) = 0; + + // internal shuffle operations to move around vectors + // will bulk-move all the vectors in [old_start_loc, old_start_loc + + // num_points) to [new_start_loc, new_start_loc + num_points) and set the old + // positions to zero vectors. + virtual void move_vectors(const location_t old_start_loc, const location_t new_start_loc, + const location_t num_points) = 0; + + // same as above, without resetting the vectors in [from_loc, from_loc + + // num_points) to zero + virtual void copy_vectors(const location_t from_loc, const location_t to_loc, const location_t num_points) = 0; + + // metric specific operations + + virtual float get_distance(const data_t *query, const location_t loc) const = 0; + virtual void get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, + float *distances) const = 0; + virtual float get_distance(const location_t loc1, const location_t loc2) const = 0; + + // stats of the data stored in store + // Returns the point in the dataset that is closest to the mean of all points + // in the dataset + virtual location_t calculate_medoid() const = 0; + + // search helpers + // if the base data is aligned per the request of the metric, this will tell + // how to align the query vector in a consistent manner + virtual size_t get_alignment_factor() const = 0; + + protected: + // Expand the datastore to new_num_points. Returns the new capacity created, + // which should be == new_num_points in the normal case. Implementers can also + // return _capacity to indicate that there are not implementing this method. + virtual location_t expand(const location_t new_num_points) = 0; + + // Shrink the datastore to new_num_points. It is NOT an error if shrink + // doesn't reduce the capacity so callers need to check this correctly. See + // also for "default" implementation + virtual location_t shrink(const location_t new_num_points) = 0; + + location_t _capacity; + size_t _dim; }; -} // namespace diskann +} // namespace diskann diff --git a/include/abstract_graph_store.h b/include/abstract_graph_store.h index 6f0eecdd7..f7735b79a 100644 --- a/include/abstract_graph_store.h +++ b/include/abstract_graph_store.h @@ -8,22 +8,24 @@ #include "types.h" -namespace diskann { +namespace diskann +{ -class AbstractGraphStore { - public: - AbstractGraphStore(const size_t max_pts) : _capacity(max_pts) {} +class AbstractGraphStore +{ + public: + AbstractGraphStore(const size_t max_pts) : _capacity(max_pts) + { + } - virtual int load(const std::string &index_path_prefix) = 0; - virtual int store(const std::string &index_path_prefix) = 0; + virtual int load(const std::string &index_path_prefix) = 0; + virtual int store(const std::string &index_path_prefix) = 0; - virtual void get_adj_list(const location_t i, - std::vector &neighbors) = 0; - virtual void set_adj_list(const location_t i, - std::vector &neighbors) = 0; + virtual void get_adj_list(const location_t i, std::vector &neighbors) = 0; + virtual void set_adj_list(const location_t i, std::vector &neighbors) = 0; - private: - size_t _capacity; + private: + size_t _capacity; }; -} // namespace diskann +} // namespace diskann diff --git a/include/aligned_file_reader.h b/include/aligned_file_reader.h index cf5d5fd78..f5e2af5c3 100644 --- a/include/aligned_file_reader.h +++ b/include/aligned_file_reader.h @@ -18,10 +18,11 @@ typedef io_context_t IOContext; #include #ifndef USE_BING_INFRA -struct IOContext { - HANDLE fhandle = NULL; - HANDLE iocp = NULL; - std::vector reqs; +struct IOContext +{ + HANDLE fhandle = NULL; + HANDLE iocp = NULL; + std::vector reqs; }; #else #include "IDiskPriorityIO.h" @@ -31,25 +32,32 @@ struct IOContext { // errors. // Because of such callous copying, we have to use ptr->atomic instead // of atomic, as atomic is not copyable. -struct IOContext { - enum Status { READ_WAIT = 0, READ_SUCCESS, READ_FAILED, PROCESS_COMPLETE }; - - std::shared_ptr m_pDiskIO = nullptr; - std::shared_ptr > m_pRequests; - std::shared_ptr > m_pRequestsStatus; - - // waitonaddress on this memory to wait for IO completion signal - // reader should signal this memory after IO completion - // TODO: WindowsAlignedFileReader can be modified to take advantage of this - // and can largely share code with the file reader for Bing. - mutable volatile long m_completeCount = 0; - - IOContext() - : m_pRequestsStatus(new std::vector()), - m_pRequests(new std::vector()) { - (*m_pRequestsStatus).reserve(MAX_IO_DEPTH); - (*m_pRequests).reserve(MAX_IO_DEPTH); - } +struct IOContext +{ + enum Status + { + READ_WAIT = 0, + READ_SUCCESS, + READ_FAILED, + PROCESS_COMPLETE + }; + + std::shared_ptr m_pDiskIO = nullptr; + std::shared_ptr> m_pRequests; + std::shared_ptr> m_pRequestsStatus; + + // waitonaddress on this memory to wait for IO completion signal + // reader should signal this memory after IO completion + // TODO: WindowsAlignedFileReader can be modified to take advantage of this + // and can largely share code with the file reader for Bing. + mutable volatile long m_completeCount = 0; + + IOContext() + : m_pRequestsStatus(new std::vector()), m_pRequests(new std::vector()) + { + (*m_pRequestsStatus).reserve(MAX_IO_DEPTH); + (*m_pRequests).reserve(MAX_IO_DEPTH); + } }; #endif @@ -63,47 +71,50 @@ struct IOContext { #include "utils.h" // NOTE :: all 3 fields must be 512-aligned -struct AlignedRead { - uint64_t offset; // where to read from - uint64_t len; // how much to read - void *buf; // where to read into - - AlignedRead() : offset(0), len(0), buf(nullptr) {} - - AlignedRead(uint64_t offset, uint64_t len, void *buf) - : offset(offset), len(len), buf(buf) { - assert(IS_512_ALIGNED(offset)); - assert(IS_512_ALIGNED(len)); - assert(IS_512_ALIGNED(buf)); - // assert(malloc_usable_size(buf) >= len); - } +struct AlignedRead +{ + uint64_t offset; // where to read from + uint64_t len; // how much to read + void *buf; // where to read into + + AlignedRead() : offset(0), len(0), buf(nullptr) + { + } + + AlignedRead(uint64_t offset, uint64_t len, void *buf) : offset(offset), len(len), buf(buf) + { + assert(IS_512_ALIGNED(offset)); + assert(IS_512_ALIGNED(len)); + assert(IS_512_ALIGNED(buf)); + // assert(malloc_usable_size(buf) >= len); + } }; -class AlignedFileReader { - protected: - tsl::robin_map ctx_map; - std::mutex ctx_mut; - - public: - // returns the thread-specific context - // returns (io_context_t)(-1) if thread is not registered - virtual IOContext &get_ctx() = 0; - - virtual ~AlignedFileReader(){}; - - // register thread-id for a context - virtual void register_thread() = 0; - // 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 close() = 0; - - // process batch of aligned requests in parallel - // NOTE :: blocking call - virtual void read(std::vector &read_reqs, IOContext &ctx, - bool async = false) = 0; +class AlignedFileReader +{ + protected: + tsl::robin_map ctx_map; + std::mutex ctx_mut; + + public: + // returns the thread-specific context + // returns (io_context_t)(-1) if thread is not registered + virtual IOContext &get_ctx() = 0; + + virtual ~AlignedFileReader(){}; + + // register thread-id for a context + virtual void register_thread() = 0; + // 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 close() = 0; + + // process batch of aligned requests in parallel + // NOTE :: blocking call + virtual void read(std::vector &read_reqs, IOContext &ctx, bool async = false) = 0; }; diff --git a/include/ann_exception.h b/include/ann_exception.h index 5d721e31a..6b81373c1 100644 --- a/include/ann_exception.h +++ b/include/ann_exception.h @@ -11,25 +11,24 @@ #define __FUNCSIG__ __PRETTY_FUNCTION__ #endif -namespace diskann { +namespace diskann +{ -class ANNException : public std::runtime_error { - public: - DISKANN_DLLEXPORT ANNException(const std::string &message, int errorCode); - DISKANN_DLLEXPORT ANNException(const std::string &message, int errorCode, - const std::string &funcSig, - const std::string &fileName, uint32_t lineNum); +class ANNException : public std::runtime_error +{ + public: + DISKANN_DLLEXPORT ANNException(const std::string &message, int errorCode); + DISKANN_DLLEXPORT ANNException(const std::string &message, int errorCode, const std::string &funcSig, + const std::string &fileName, uint32_t lineNum); - private: - int _errorCode; + private: + int _errorCode; }; -class FileException : public ANNException { - public: - DISKANN_DLLEXPORT FileException(const std::string &filename, - std::system_error &e, - const std::string &funcSig, - const std::string &fileName, - uint32_t lineNum); +class FileException : public ANNException +{ + public: + DISKANN_DLLEXPORT FileException(const std::string &filename, std::system_error &e, const std::string &funcSig, + const std::string &fileName, uint32_t lineNum); }; -} // namespace diskann +} // namespace diskann diff --git a/include/boost_dynamic_bitset_fwd.h b/include/boost_dynamic_bitset_fwd.h index d0dd72c13..5aebb2bc2 100644 --- a/include/boost_dynamic_bitset_fwd.h +++ b/include/boost_dynamic_bitset_fwd.h @@ -3,10 +3,9 @@ #pragma once -namespace boost { +namespace boost +{ #ifndef BOOST_DYNAMIC_BITSET_FWD_HPP -template > -class dynamic_bitset; +template > class dynamic_bitset; #endif -} // namespace boost +} // namespace boost diff --git a/include/cached_io.h b/include/cached_io.h index ac31283a8..daef2f2f7 100644 --- a/include/cached_io.h +++ b/include/cached_io.h @@ -11,175 +11,207 @@ #include "ann_exception.h" // sequential cached reads -class cached_ifstream { - public: - cached_ifstream() {} - cached_ifstream(const std::string &filename, uint64_t cacheSize) - : cache_size(cacheSize), cur_off(0) { - reader.exceptions(std::ifstream::failbit | std::ifstream::badbit); - this->open(filename, cache_size); - } - ~cached_ifstream() { - delete[] cache_buf; - reader.close(); - } - - void open(const std::string &filename, uint64_t cacheSize) { - this->cur_off = 0; - - try { - reader.open(filename, std::ios::binary | std::ios::ate); - fsize = reader.tellg(); - reader.seekg(0, std::ios::beg); - assert(reader.is_open()); - assert(cacheSize > 0); - cacheSize = (std::min)(cacheSize, fsize); - this->cache_size = cacheSize; - cache_buf = new char[cacheSize]; - reader.read(cache_buf, cacheSize); - diskann::cout << "Opened: " << filename.c_str() << ", size: " << fsize - << ", cache_size: " << cacheSize << std::endl; - } catch (std::system_error &e) { - throw diskann::FileException(filename, e, __FUNCSIG__, __FILE__, - __LINE__); +class cached_ifstream +{ + public: + cached_ifstream() + { } - } - - size_t get_file_size() { return fsize; } - - void read(char *read_buf, uint64_t n_bytes) { - assert(cache_buf != nullptr); - assert(read_buf != nullptr); - - if (n_bytes <= (cache_size - cur_off)) { - // case 1: cache contains all data - memcpy(read_buf, cache_buf + cur_off, n_bytes); - cur_off += n_bytes; - } else { - // case 2: cache contains some data - uint64_t cached_bytes = cache_size - cur_off; - if (n_bytes - cached_bytes > fsize - reader.tellg()) { - std::stringstream stream; - stream << "Reading beyond end of file" << std::endl; - stream << "n_bytes: " << n_bytes << " cached_bytes: " << cached_bytes - << " fsize: " << fsize << " current pos:" << reader.tellg() - << std::endl; - diskann::cout << stream.str() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); - } - memcpy(read_buf, cache_buf + cur_off, cached_bytes); - - // go to disk and fetch more data - reader.read(read_buf + cached_bytes, n_bytes - cached_bytes); - // reset cur off - cur_off = cache_size; - - uint64_t size_left = fsize - reader.tellg(); - - if (size_left >= cache_size) { - reader.read(cache_buf, cache_size); - cur_off = 0; - } - // note that if size_left < cache_size, then cur_off = cache_size, - // so subsequent reads will all be directly from file + cached_ifstream(const std::string &filename, uint64_t cacheSize) : cache_size(cacheSize), cur_off(0) + { + reader.exceptions(std::ifstream::failbit | std::ifstream::badbit); + this->open(filename, cache_size); + } + ~cached_ifstream() + { + delete[] cache_buf; + reader.close(); + } + + void open(const std::string &filename, uint64_t cacheSize) + { + this->cur_off = 0; + + try + { + reader.open(filename, std::ios::binary | std::ios::ate); + fsize = reader.tellg(); + reader.seekg(0, std::ios::beg); + assert(reader.is_open()); + assert(cacheSize > 0); + cacheSize = (std::min)(cacheSize, fsize); + this->cache_size = cacheSize; + cache_buf = new char[cacheSize]; + reader.read(cache_buf, cacheSize); + diskann::cout << "Opened: " << filename.c_str() << ", size: " << fsize << ", cache_size: " << cacheSize + << std::endl; + } + catch (std::system_error &e) + { + throw diskann::FileException(filename, e, __FUNCSIG__, __FILE__, __LINE__); + } + } + + size_t get_file_size() + { + return fsize; + } + + void read(char *read_buf, uint64_t n_bytes) + { + assert(cache_buf != nullptr); + assert(read_buf != nullptr); + + if (n_bytes <= (cache_size - cur_off)) + { + // case 1: cache contains all data + memcpy(read_buf, cache_buf + cur_off, n_bytes); + cur_off += n_bytes; + } + else + { + // case 2: cache contains some data + uint64_t cached_bytes = cache_size - cur_off; + if (n_bytes - cached_bytes > fsize - reader.tellg()) + { + std::stringstream stream; + stream << "Reading beyond end of file" << std::endl; + stream << "n_bytes: " << n_bytes << " cached_bytes: " << cached_bytes << " fsize: " << fsize + << " current pos:" << reader.tellg() << std::endl; + diskann::cout << stream.str() << std::endl; + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } + memcpy(read_buf, cache_buf + cur_off, cached_bytes); + + // go to disk and fetch more data + reader.read(read_buf + cached_bytes, n_bytes - cached_bytes); + // reset cur off + cur_off = cache_size; + + uint64_t size_left = fsize - reader.tellg(); + + if (size_left >= cache_size) + { + reader.read(cache_buf, cache_size); + cur_off = 0; + } + // note that if size_left < cache_size, then cur_off = cache_size, + // so subsequent reads will all be directly from file + } } - } - - private: - // underlying ifstream - std::ifstream reader; - // # bytes to cache in one shot read - uint64_t cache_size = 0; - // underlying buf for cache - char *cache_buf = nullptr; - // offset into cache_buf for cur_pos - uint64_t cur_off = 0; - // file size - uint64_t fsize = 0; + + private: + // underlying ifstream + std::ifstream reader; + // # bytes to cache in one shot read + uint64_t cache_size = 0; + // underlying buf for cache + char *cache_buf = nullptr; + // offset into cache_buf for cur_pos + uint64_t cur_off = 0; + // file size + uint64_t fsize = 0; }; // sequential cached writes -class cached_ofstream { - public: - cached_ofstream(const std::string &filename, uint64_t cache_size) - : cache_size(cache_size), cur_off(0) { - writer.exceptions(std::ifstream::failbit | std::ifstream::badbit); - try { - writer.open(filename, std::ios::binary); - assert(writer.is_open()); - assert(cache_size > 0); - cache_buf = new char[cache_size]; - diskann::cout << "Opened: " << filename.c_str() - << ", cache_size: " << cache_size << std::endl; - } catch (std::system_error &e) { - throw diskann::FileException(filename, e, __FUNCSIG__, __FILE__, - __LINE__); +class cached_ofstream +{ + public: + cached_ofstream(const std::string &filename, uint64_t cache_size) : cache_size(cache_size), cur_off(0) + { + writer.exceptions(std::ifstream::failbit | std::ifstream::badbit); + try + { + writer.open(filename, std::ios::binary); + assert(writer.is_open()); + assert(cache_size > 0); + cache_buf = new char[cache_size]; + diskann::cout << "Opened: " << filename.c_str() << ", cache_size: " << cache_size << std::endl; + } + catch (std::system_error &e) + { + throw diskann::FileException(filename, e, __FUNCSIG__, __FILE__, __LINE__); + } } - } - ~cached_ofstream() { this->close(); } + ~cached_ofstream() + { + this->close(); + } - void close() { - // dump any remaining data in memory - if (cur_off > 0) { - this->flush_cache(); + void close() + { + // dump any remaining data in memory + if (cur_off > 0) + { + this->flush_cache(); + } + + if (cache_buf != nullptr) + { + delete[] cache_buf; + cache_buf = nullptr; + } + + if (writer.is_open()) + writer.close(); + diskann::cout << "Finished writing " << fsize << "B" << std::endl; } - if (cache_buf != nullptr) { - delete[] cache_buf; - cache_buf = nullptr; + size_t get_file_size() + { + return fsize; + } + // writes n_bytes from write_buf to the underlying ofstream/cache + void write(char *write_buf, uint64_t n_bytes) + { + assert(cache_buf != nullptr); + if (n_bytes <= (cache_size - cur_off)) + { + // case 1: cache can take all data + memcpy(cache_buf + cur_off, write_buf, n_bytes); + cur_off += n_bytes; + } + else + { + // case 2: cache cant take all data + // go to disk and write existing cache data + writer.write(cache_buf, cur_off); + fsize += cur_off; + // write the new data to disk + writer.write(write_buf, n_bytes); + fsize += n_bytes; + // memset all cache data and reset cur_off + memset(cache_buf, 0, cache_size); + cur_off = 0; + } } - if (writer.is_open()) writer.close(); - diskann::cout << "Finished writing " << fsize << "B" << std::endl; - } - - size_t get_file_size() { return fsize; } - // writes n_bytes from write_buf to the underlying ofstream/cache - void write(char *write_buf, uint64_t n_bytes) { - assert(cache_buf != nullptr); - if (n_bytes <= (cache_size - cur_off)) { - // case 1: cache can take all data - memcpy(cache_buf + cur_off, write_buf, n_bytes); - cur_off += n_bytes; - } else { - // case 2: cache cant take all data - // go to disk and write existing cache data - writer.write(cache_buf, cur_off); - fsize += cur_off; - // write the new data to disk - writer.write(write_buf, n_bytes); - fsize += n_bytes; - // memset all cache data and reset cur_off - memset(cache_buf, 0, cache_size); - cur_off = 0; + void flush_cache() + { + assert(cache_buf != nullptr); + writer.write(cache_buf, cur_off); + fsize += cur_off; + memset(cache_buf, 0, cache_size); + cur_off = 0; } - } - - void flush_cache() { - assert(cache_buf != nullptr); - writer.write(cache_buf, cur_off); - fsize += cur_off; - memset(cache_buf, 0, cache_size); - cur_off = 0; - } - - void reset() { - flush_cache(); - writer.seekp(0); - } - - private: - // underlying ofstream - std::ofstream writer; - // # bytes to cache for one shot write - uint64_t cache_size = 0; - // underlying buf for cache - char *cache_buf = nullptr; - // offset into cache_buf for cur_pos - uint64_t cur_off = 0; - - // file size - uint64_t fsize = 0; + + void reset() + { + flush_cache(); + writer.seekp(0); + } + + private: + // underlying ofstream + std::ofstream writer; + // # bytes to cache for one shot write + uint64_t cache_size = 0; + // underlying buf for cache + char *cache_buf = nullptr; + // offset into cache_buf for cur_pos + uint64_t cur_off = 0; + + // file size + uint64_t fsize = 0; }; diff --git a/include/concurrent_queue.h b/include/concurrent_queue.h index a375a7466..1e57bbf0f 100644 --- a/include/concurrent_queue.h +++ b/include/concurrent_queue.h @@ -11,91 +11,122 @@ #include #include -namespace diskann { +namespace diskann +{ -template -class ConcurrentQueue { - typedef std::chrono::microseconds chrono_us_t; - typedef std::unique_lock mutex_locker; +template class ConcurrentQueue +{ + typedef std::chrono::microseconds chrono_us_t; + typedef std::unique_lock mutex_locker; - std::queue q; - std::mutex mut; - std::mutex push_mut; - std::mutex pop_mut; - std::condition_variable push_cv; - std::condition_variable pop_cv; - T null_T; + std::queue q; + std::mutex mut; + std::mutex push_mut; + std::mutex pop_mut; + std::condition_variable push_cv; + std::condition_variable pop_cv; + T null_T; - public: - ConcurrentQueue() {} + public: + ConcurrentQueue() + { + } - ConcurrentQueue(T nullT) { this->null_T = nullT; } + ConcurrentQueue(T nullT) + { + this->null_T = nullT; + } - ~ConcurrentQueue() { - this->push_cv.notify_all(); - this->pop_cv.notify_all(); - } + ~ConcurrentQueue() + { + this->push_cv.notify_all(); + this->pop_cv.notify_all(); + } - // queue stats - uint64_t size() { - mutex_locker lk(this->mut); - uint64_t ret = q.size(); - lk.unlock(); - return ret; - } + // queue stats + uint64_t size() + { + mutex_locker lk(this->mut); + uint64_t ret = q.size(); + lk.unlock(); + return ret; + } - bool empty() { return (this->size() == 0); } + bool empty() + { + return (this->size() == 0); + } - // PUSH BACK - void push(T &new_val) { - mutex_locker lk(this->mut); - this->q.push(new_val); - lk.unlock(); - } + // PUSH BACK + void push(T &new_val) + { + mutex_locker lk(this->mut); + this->q.push(new_val); + lk.unlock(); + } - template - void insert(Iterator iter_begin, Iterator iter_end) { - mutex_locker lk(this->mut); - for (Iterator it = iter_begin; it != iter_end; it++) { - this->q.push(*it); + template void insert(Iterator iter_begin, Iterator iter_end) + { + mutex_locker lk(this->mut); + for (Iterator it = iter_begin; it != iter_end; it++) + { + this->q.push(*it); + } + lk.unlock(); } - lk.unlock(); - } - // POP FRONT - T pop() { - mutex_locker lk(this->mut); - if (this->q.empty()) { - lk.unlock(); - return this->null_T; - } else { - T ret = this->q.front(); - this->q.pop(); - // diskann::cout << "thread_id: " << std::this_thread::get_id() << - // ", ctx: " - // << ret.ctx << "\n"; - lk.unlock(); - return ret; + // POP FRONT + T pop() + { + mutex_locker lk(this->mut); + if (this->q.empty()) + { + lk.unlock(); + return this->null_T; + } + else + { + T ret = this->q.front(); + this->q.pop(); + // diskann::cout << "thread_id: " << std::this_thread::get_id() << + // ", ctx: " + // << ret.ctx << "\n"; + lk.unlock(); + return ret; + } } - } - // register for notifications - void wait_for_push_notify(chrono_us_t wait_time = chrono_us_t{10}) { - mutex_locker lk(this->push_mut); - this->push_cv.wait_for(lk, wait_time); - lk.unlock(); - } + // register for notifications + void wait_for_push_notify(chrono_us_t wait_time = chrono_us_t{10}) + { + mutex_locker lk(this->push_mut); + this->push_cv.wait_for(lk, wait_time); + lk.unlock(); + } - void wait_for_pop_notify(chrono_us_t wait_time = chrono_us_t{10}) { - mutex_locker lk(this->pop_mut); - this->pop_cv.wait_for(lk, wait_time); - lk.unlock(); - } + void wait_for_pop_notify(chrono_us_t wait_time = chrono_us_t{10}) + { + mutex_locker lk(this->pop_mut); + this->pop_cv.wait_for(lk, wait_time); + lk.unlock(); + } - // just notify functions - void push_notify_one() { this->push_cv.notify_one(); } - void push_notify_all() { this->push_cv.notify_all(); } - void pop_notify_one() { this->pop_cv.notify_one(); } - void pop_notify_all() { this->pop_cv.notify_all(); } + // just notify functions + void push_notify_one() + { + this->push_cv.notify_one(); + } + void push_notify_all() + { + this->push_cv.notify_all(); + } + void pop_notify_one() + { + this->pop_cv.notify_one(); + } + void pop_notify_all() + { + this->pop_cv.notify_all(); + } }; -} // namespace diskann +} // namespace diskann diff --git a/include/cosine_similarity.h b/include/cosine_similarity.h index 5c2ec3d50..dc51f6c0a 100644 --- a/include/cosine_similarity.h +++ b/include/cosine_similarity.h @@ -38,194 +38,202 @@ extern bool Avx2SupportedCPU; * */ -namespace diskann { +namespace diskann +{ using namespace std; #define PORTABLE_ALIGN16 __declspec(align(16)) -static float NormScalarProductSIMD2(const int8_t *pVect1, const int8_t *pVect2, - uint32_t qty) { - if (Avx2SupportedCPU) { - __m256 cos, p1Len, p2Len; - cos = p1Len = p2Len = _mm256_setzero_ps(); - while (qty >= 32) { - __m256i rx = _mm256_load_si256((__m256i *)pVect1), - ry = _mm256_load_si256((__m256i *)pVect2); - cos = _mm256_add_ps(cos, _mm256_mul_epi8(rx, ry)); - p1Len = _mm256_add_ps(p1Len, _mm256_mul_epi8(rx, rx)); - p2Len = _mm256_add_ps(p2Len, _mm256_mul_epi8(ry, ry)); - pVect1 += 32; - pVect2 += 32; - qty -= 32; +static float NormScalarProductSIMD2(const int8_t *pVect1, const int8_t *pVect2, uint32_t qty) +{ + if (Avx2SupportedCPU) + { + __m256 cos, p1Len, p2Len; + cos = p1Len = p2Len = _mm256_setzero_ps(); + while (qty >= 32) + { + __m256i rx = _mm256_load_si256((__m256i *)pVect1), ry = _mm256_load_si256((__m256i *)pVect2); + cos = _mm256_add_ps(cos, _mm256_mul_epi8(rx, ry)); + p1Len = _mm256_add_ps(p1Len, _mm256_mul_epi8(rx, rx)); + p2Len = _mm256_add_ps(p2Len, _mm256_mul_epi8(ry, ry)); + pVect1 += 32; + pVect2 += 32; + qty -= 32; + } + while (qty > 0) + { + __m128i rx = _mm_load_si128((__m128i *)pVect1), ry = _mm_load_si128((__m128i *)pVect2); + cos = _mm256_add_ps(cos, _mm256_mul32_pi8(rx, ry)); + p1Len = _mm256_add_ps(p1Len, _mm256_mul32_pi8(rx, rx)); + p2Len = _mm256_add_ps(p2Len, _mm256_mul32_pi8(ry, ry)); + pVect1 += 4; + pVect2 += 4; + qty -= 4; + } + cos = _mm256_hadd_ps(_mm256_hadd_ps(cos, cos), cos); + p1Len = _mm256_hadd_ps(_mm256_hadd_ps(p1Len, p1Len), p1Len); + p2Len = _mm256_hadd_ps(_mm256_hadd_ps(p2Len, p2Len), p2Len); + float denominator = max(numeric_limits::min() * 2, sqrt(p1Len.m256_f32[0] + p1Len.m256_f32[4]) * + sqrt(p2Len.m256_f32[0] + p2Len.m256_f32[4])); + float cosine = (cos.m256_f32[0] + cos.m256_f32[4]) / denominator; + + return max(float(-1), min(float(1), cosine)); } - while (qty > 0) { - __m128i rx = _mm_load_si128((__m128i *)pVect1), - ry = _mm_load_si128((__m128i *)pVect2); - cos = _mm256_add_ps(cos, _mm256_mul32_pi8(rx, ry)); - p1Len = _mm256_add_ps(p1Len, _mm256_mul32_pi8(rx, rx)); - p2Len = _mm256_add_ps(p2Len, _mm256_mul32_pi8(ry, ry)); - pVect1 += 4; - pVect2 += 4; - qty -= 4; + + __m128 cos, p1Len, p2Len; + cos = p1Len = p2Len = _mm_setzero_ps(); + __m128i rx, ry; + while (qty >= 16) + { + rx = _mm_load_si128((__m128i *)pVect1); + ry = _mm_load_si128((__m128i *)pVect2); + cos = _mm_add_ps(cos, _mm_mul_epi8(rx, ry)); + p1Len = _mm_add_ps(p1Len, _mm_mul_epi8(rx, rx)); + p2Len = _mm_add_ps(p2Len, _mm_mul_epi8(ry, ry)); + pVect1 += 16; + pVect2 += 16; + qty -= 16; + } + while (qty > 0) + { + rx = _mm_load_si128((__m128i *)pVect1); + ry = _mm_load_si128((__m128i *)pVect2); + cos = _mm_add_ps(cos, _mm_mul32_pi8(rx, ry)); + p1Len = _mm_add_ps(p1Len, _mm_mul32_pi8(rx, rx)); + p2Len = _mm_add_ps(p2Len, _mm_mul32_pi8(ry, ry)); + pVect1 += 4; + pVect2 += 4; + qty -= 4; } - cos = _mm256_hadd_ps(_mm256_hadd_ps(cos, cos), cos); - p1Len = _mm256_hadd_ps(_mm256_hadd_ps(p1Len, p1Len), p1Len); - p2Len = _mm256_hadd_ps(_mm256_hadd_ps(p2Len, p2Len), p2Len); - float denominator = max(numeric_limits::min() * 2, - sqrt(p1Len.m256_f32[0] + p1Len.m256_f32[4]) * - sqrt(p2Len.m256_f32[0] + p2Len.m256_f32[4])); - float cosine = (cos.m256_f32[0] + cos.m256_f32[4]) / denominator; - - return max(float(-1), min(float(1), cosine)); - } - - __m128 cos, p1Len, p2Len; - cos = p1Len = p2Len = _mm_setzero_ps(); - __m128i rx, ry; - while (qty >= 16) { - rx = _mm_load_si128((__m128i *)pVect1); - ry = _mm_load_si128((__m128i *)pVect2); - cos = _mm_add_ps(cos, _mm_mul_epi8(rx, ry)); - p1Len = _mm_add_ps(p1Len, _mm_mul_epi8(rx, rx)); - p2Len = _mm_add_ps(p2Len, _mm_mul_epi8(ry, ry)); - pVect1 += 16; - pVect2 += 16; - qty -= 16; - } - while (qty > 0) { - rx = _mm_load_si128((__m128i *)pVect1); - ry = _mm_load_si128((__m128i *)pVect2); - cos = _mm_add_ps(cos, _mm_mul32_pi8(rx, ry)); - p1Len = _mm_add_ps(p1Len, _mm_mul32_pi8(rx, rx)); - p2Len = _mm_add_ps(p2Len, _mm_mul32_pi8(ry, ry)); - pVect1 += 4; - pVect2 += 4; - qty -= 4; - } - cos = _mm_hadd_ps(_mm_hadd_ps(cos, cos), cos); - p1Len = _mm_hadd_ps(_mm_hadd_ps(p1Len, p1Len), p1Len); - p2Len = _mm_hadd_ps(_mm_hadd_ps(p2Len, p2Len), p2Len); - float norm1 = p1Len.m128_f32[0]; - float norm2 = p2Len.m128_f32[0]; - - static const float eps = numeric_limits::min() * 2; - - if (norm1 < eps) { /* - * This shouldn't normally happen for this space, but - * if it does, we don't want to get NANs - */ - if (norm2 < eps) { - return 1; + cos = _mm_hadd_ps(_mm_hadd_ps(cos, cos), cos); + p1Len = _mm_hadd_ps(_mm_hadd_ps(p1Len, p1Len), p1Len); + p2Len = _mm_hadd_ps(_mm_hadd_ps(p2Len, p2Len), p2Len); + float norm1 = p1Len.m128_f32[0]; + float norm2 = p2Len.m128_f32[0]; + + static const float eps = numeric_limits::min() * 2; + + if (norm1 < eps) + { /* + * This shouldn't normally happen for this space, but + * if it does, we don't want to get NANs + */ + if (norm2 < eps) + { + return 1; + } + return 0; } - return 0; - } - /* - * Sometimes due to rounding errors, we get values > 1 or < -1. - * This throws off other functions that use scalar product, e.g., acos - */ - return max(float(-1), - min(float(1), cos.m128_f32[0] / sqrt(norm1) / sqrt(norm2))); + /* + * Sometimes due to rounding errors, we get values > 1 or < -1. + * This throws off other functions that use scalar product, e.g., acos + */ + return max(float(-1), min(float(1), cos.m128_f32[0] / sqrt(norm1) / sqrt(norm2))); } -static float NormScalarProductSIMD(const float *pVect1, const float *pVect2, - uint32_t qty) { - // Didn't get significant performance gain compared with 128bit version. - static const float eps = numeric_limits::min() * 2; - - if (Avx2SupportedCPU) { - uint32_t qty8 = qty / 8; - - const float *pEnd1 = pVect1 + 8 * qty8; - const float *pEnd2 = pVect1 + qty; - - __m256 v1, v2; - __m256 sum_prod = _mm256_set_ps(0, 0, 0, 0, 0, 0, 0, 0); - __m256 sum_square1 = sum_prod; - __m256 sum_square2 = sum_prod; - - while (pVect1 < pEnd1) { - v1 = _mm256_loadu_ps(pVect1); - pVect1 += 8; - v2 = _mm256_loadu_ps(pVect2); - pVect2 += 8; - sum_prod = _mm256_add_ps(sum_prod, _mm256_mul_ps(v1, v2)); - sum_square1 = _mm256_add_ps(sum_square1, _mm256_mul_ps(v1, v1)); - sum_square2 = _mm256_add_ps(sum_square2, _mm256_mul_ps(v2, v2)); +static float NormScalarProductSIMD(const float *pVect1, const float *pVect2, uint32_t qty) +{ + // Didn't get significant performance gain compared with 128bit version. + static const float eps = numeric_limits::min() * 2; + + if (Avx2SupportedCPU) + { + uint32_t qty8 = qty / 8; + + const float *pEnd1 = pVect1 + 8 * qty8; + const float *pEnd2 = pVect1 + qty; + + __m256 v1, v2; + __m256 sum_prod = _mm256_set_ps(0, 0, 0, 0, 0, 0, 0, 0); + __m256 sum_square1 = sum_prod; + __m256 sum_square2 = sum_prod; + + while (pVect1 < pEnd1) + { + v1 = _mm256_loadu_ps(pVect1); + pVect1 += 8; + v2 = _mm256_loadu_ps(pVect2); + pVect2 += 8; + sum_prod = _mm256_add_ps(sum_prod, _mm256_mul_ps(v1, v2)); + sum_square1 = _mm256_add_ps(sum_square1, _mm256_mul_ps(v1, v1)); + sum_square2 = _mm256_add_ps(sum_square2, _mm256_mul_ps(v2, v2)); + } + + float PORTABLE_ALIGN16 TmpResProd[8]; + float PORTABLE_ALIGN16 TmpResSquare1[8]; + float PORTABLE_ALIGN16 TmpResSquare2[8]; + + _mm256_store_ps(TmpResProd, sum_prod); + _mm256_store_ps(TmpResSquare1, sum_square1); + _mm256_store_ps(TmpResSquare2, sum_square2); + + float sum = 0.0f; + float norm1 = 0.0f; + float norm2 = 0.0f; + for (uint32_t i = 0; i < 8; ++i) + { + sum += TmpResProd[i]; + norm1 += TmpResSquare1[i]; + norm2 += TmpResSquare2[i]; + } + + while (pVect1 < pEnd2) + { + sum += (*pVect1) * (*pVect2); + norm1 += (*pVect1) * (*pVect1); + norm2 += (*pVect2) * (*pVect2); + + ++pVect1; + ++pVect2; + } + + if (norm1 < eps) + { + return norm2 < eps ? 1.0f : 0.0f; + } + + return max(float(-1), min(float(1), sum / sqrt(norm1) / sqrt(norm2))); } - float PORTABLE_ALIGN16 TmpResProd[8]; - float PORTABLE_ALIGN16 TmpResSquare1[8]; - float PORTABLE_ALIGN16 TmpResSquare2[8]; - - _mm256_store_ps(TmpResProd, sum_prod); - _mm256_store_ps(TmpResSquare1, sum_square1); - _mm256_store_ps(TmpResSquare2, sum_square2); - - float sum = 0.0f; - float norm1 = 0.0f; - float norm2 = 0.0f; - for (uint32_t i = 0; i < 8; ++i) { - sum += TmpResProd[i]; - norm1 += TmpResSquare1[i]; - norm2 += TmpResSquare2[i]; + __m128 v1, v2; + __m128 sum_prod = _mm_set1_ps(0); + __m128 sum_square1 = sum_prod; + __m128 sum_square2 = sum_prod; + + while (qty >= 4) + { + v1 = _mm_loadu_ps(pVect1); + pVect1 += 4; + v2 = _mm_loadu_ps(pVect2); + pVect2 += 4; + sum_prod = _mm_add_ps(sum_prod, _mm_mul_ps(v1, v2)); + sum_square1 = _mm_add_ps(sum_square1, _mm_mul_ps(v1, v1)); + sum_square2 = _mm_add_ps(sum_square2, _mm_mul_ps(v2, v2)); + + qty -= 4; } - while (pVect1 < pEnd2) { - sum += (*pVect1) * (*pVect2); - norm1 += (*pVect1) * (*pVect1); - norm2 += (*pVect2) * (*pVect2); - - ++pVect1; - ++pVect2; - } + float sum = sum_prod.m128_f32[0] + sum_prod.m128_f32[1] + sum_prod.m128_f32[2] + sum_prod.m128_f32[3]; + float norm1 = sum_square1.m128_f32[0] + sum_square1.m128_f32[1] + sum_square1.m128_f32[2] + sum_square1.m128_f32[3]; + float norm2 = sum_square2.m128_f32[0] + sum_square2.m128_f32[1] + sum_square2.m128_f32[2] + sum_square2.m128_f32[3]; - if (norm1 < eps) { - return norm2 < eps ? 1.0f : 0.0f; + if (norm1 < eps) + { + return norm2 < eps ? 1.0f : 0.0f; } return max(float(-1), min(float(1), sum / sqrt(norm1) / sqrt(norm2))); - } - - __m128 v1, v2; - __m128 sum_prod = _mm_set1_ps(0); - __m128 sum_square1 = sum_prod; - __m128 sum_square2 = sum_prod; - - while (qty >= 4) { - v1 = _mm_loadu_ps(pVect1); - pVect1 += 4; - v2 = _mm_loadu_ps(pVect2); - pVect2 += 4; - sum_prod = _mm_add_ps(sum_prod, _mm_mul_ps(v1, v2)); - sum_square1 = _mm_add_ps(sum_square1, _mm_mul_ps(v1, v1)); - sum_square2 = _mm_add_ps(sum_square2, _mm_mul_ps(v2, v2)); - - qty -= 4; - } - - float sum = sum_prod.m128_f32[0] + sum_prod.m128_f32[1] + - sum_prod.m128_f32[2] + sum_prod.m128_f32[3]; - float norm1 = sum_square1.m128_f32[0] + sum_square1.m128_f32[1] + - sum_square1.m128_f32[2] + sum_square1.m128_f32[3]; - float norm2 = sum_square2.m128_f32[0] + sum_square2.m128_f32[1] + - sum_square2.m128_f32[2] + sum_square2.m128_f32[3]; - - if (norm1 < eps) { - return norm2 < eps ? 1.0f : 0.0f; - } - - return max(float(-1), min(float(1), sum / sqrt(norm1) / sqrt(norm2))); } -static float NormScalarProductSIMD2(const float *pVect1, const float *pVect2, - uint32_t qty) { - return NormScalarProductSIMD(pVect1, pVect2, qty); +static float NormScalarProductSIMD2(const float *pVect1, const float *pVect2, uint32_t qty) +{ + return NormScalarProductSIMD(pVect1, pVect2, qty); } -template -static float CosineSimilarity2(const T *p1, const T *p2, uint32_t qty) { - return std::max(0.0f, 1.0f - NormScalarProductSIMD2(p1, p2, qty)); +template static float CosineSimilarity2(const T *p1, const T *p2, uint32_t qty) +{ + return std::max(0.0f, 1.0f - NormScalarProductSIMD2(p1, p2, qty)); } // static template float CosineSimilarity2<__int8>(const __int8* pVect1, @@ -234,19 +242,22 @@ static float CosineSimilarity2(const T *p1, const T *p2, uint32_t qty) { // static template float CosineSimilarity2(const float* pVect1, // const float* pVect2, size_t qty); -template -static void CosineSimilarityNormalize(T *pVector, uint32_t qty) { - T sum = 0; - for (uint32_t i = 0; i < qty; ++i) { - sum += pVector[i] * pVector[i]; - } - sum = 1 / sqrt(sum); - if (sum == 0) { - sum = numeric_limits::min(); - } - for (uint32_t i = 0; i < qty; ++i) { - pVector[i] *= sum; - } +template static void CosineSimilarityNormalize(T *pVector, uint32_t qty) +{ + T sum = 0; + for (uint32_t i = 0; i < qty; ++i) + { + sum += pVector[i] * pVector[i]; + } + sum = 1 / sqrt(sum); + if (sum == 0) + { + sum = numeric_limits::min(); + } + for (uint32_t i = 0; i < qty; ++i) + { + pVector[i] *= sum; + } } // template static void CosineSimilarityNormalize(float* pVector, @@ -254,22 +265,19 @@ static void CosineSimilarityNormalize(T *pVector, uint32_t qty) { // template static void CosineSimilarityNormalize(double* pVector, // size_t qty); -template <> -void CosineSimilarityNormalize(__int8 * /*pVector*/, uint32_t /*qty*/) { - throw std::runtime_error( - "For int8 type vector, you can not use cosine distance!"); +template <> void CosineSimilarityNormalize(__int8 * /*pVector*/, uint32_t /*qty*/) +{ + throw std::runtime_error("For int8 type vector, you can not use cosine distance!"); } -template <> -void CosineSimilarityNormalize(__int16 * /*pVector*/, uint32_t /*qty*/) { - throw std::runtime_error( - "For int16 type vector, you can not use cosine distance!"); +template <> void CosineSimilarityNormalize(__int16 * /*pVector*/, uint32_t /*qty*/) +{ + throw std::runtime_error("For int16 type vector, you can not use cosine distance!"); } -template <> -void CosineSimilarityNormalize(int * /*pVector*/, uint32_t /*qty*/) { - throw std::runtime_error( - "For int type vector, you can not use cosine distance!"); +template <> void CosineSimilarityNormalize(int * /*pVector*/, uint32_t /*qty*/) +{ + throw std::runtime_error("For int type vector, you can not use cosine distance!"); } -} // namespace diskann +} // namespace diskann #endif \ No newline at end of file diff --git a/include/defaults.h b/include/defaults.h index 87e9323e8..6b14bc2b3 100644 --- a/include/defaults.h +++ b/include/defaults.h @@ -4,8 +4,10 @@ #pragma once #include -namespace diskann { -namespace defaults { +namespace diskann +{ +namespace defaults +{ const float ALPHA = 1.2f; const uint32_t NUM_THREADS = 0; const uint32_t NUM_ROUNDS = 2; @@ -18,5 +20,5 @@ const uint32_t MAX_DEGREE = 64; const uint32_t BUILD_LIST_SIZE = 100; const uint32_t SATURATE_GRAPH = false; const uint32_t SEARCH_LIST_SIZE = 100; -} // namespace defaults -} // namespace diskann +} // namespace defaults +} // namespace diskann diff --git a/include/disk_utils.h b/include/disk_utils.h index 65d0a51aa..08f046dcd 100644 --- a/include/disk_utils.h +++ b/include/disk_utils.h @@ -31,7 +31,8 @@ typedef int FileHandle; #include "utils.h" #include "windows_customizations.h" -namespace diskann { +namespace diskann +{ const size_t MAX_SAMPLE_POINTS_FOR_WARMUP = 100000; const double PQ_TRAINING_SET_FRACTION = 0.1; const double SPACE_FOR_CACHED_NODES_IN_GB = 0.25; @@ -40,87 +41,68 @@ const uint32_t NUM_NODES_TO_CACHE = 250000; const uint32_t WARMUP_L = 20; const uint32_t NUM_KMEANS_REPS = 12; -template -class PQFlashIndex; +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 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 size_t calculate_num_pq_chunks(double final_index_ram_limit, size_t points_num, uint32_t dim); -DISKANN_DLLEXPORT void read_idmap(const std::string &fname, - std::vector &ivecs); +DISKANN_DLLEXPORT void read_idmap(const std::string &fname, std::vector &ivecs); #ifdef EXEC_ENV_OLS template -DISKANN_DLLEXPORT T *load_warmup(MemoryMappedFiles &files, - const std::string &cache_warmup_file, - uint64_t &warmup_num, uint64_t warmup_dim, - uint64_t warmup_aligned_dim); +DISKANN_DLLEXPORT T *load_warmup(MemoryMappedFiles &files, const std::string &cache_warmup_file, uint64_t &warmup_num, + uint64_t warmup_dim, uint64_t warmup_aligned_dim); #else template -DISKANN_DLLEXPORT T *load_warmup(const std::string &cache_warmup_file, - uint64_t &warmup_num, uint64_t warmup_dim, +DISKANN_DLLEXPORT T *load_warmup(const std::string &cache_warmup_file, uint64_t &warmup_num, uint64_t warmup_dim, uint64_t warmup_aligned_dim); #endif -DISKANN_DLLEXPORT int merge_shards( - const std::string &vamana_prefix, const std::string &vamana_suffix, - const std::string &idmaps_prefix, const std::string &idmaps_suffix, - const uint64_t nshards, uint32_t max_degree, - const std::string &output_vamana, const std::string &medoids_file, - bool use_filters = false, - const std::string &labels_to_medoids_file = std::string("")); +DISKANN_DLLEXPORT int merge_shards(const std::string &vamana_prefix, const std::string &vamana_suffix, + const std::string &idmaps_prefix, const std::string &idmaps_suffix, + const uint64_t nshards, uint32_t max_degree, const std::string &output_vamana, + const std::string &medoids_file, bool use_filters = false, + const std::string &labels_to_medoids_file = std::string("")); -DISKANN_DLLEXPORT void extract_shard_labels( - const std::string &in_label_file, const std::string &shard_ids_bin, - const std::string &shard_label_file); +DISKANN_DLLEXPORT void extract_shard_labels(const std::string &in_label_file, const std::string &shard_ids_bin, + const std::string &shard_label_file); template -DISKANN_DLLEXPORT std::string preprocess_base_file( - const std::string &infile, const std::string &indexPrefix, - diskann::Metric &distMetric); +DISKANN_DLLEXPORT std::string preprocess_base_file(const std::string &infile, const std::string &indexPrefix, + diskann::Metric &distMetric); template -DISKANN_DLLEXPORT int build_merged_vamana_index( - std::string base_file, diskann::Metric _compareMetric, uint32_t L, - uint32_t R, double sampling_rate, double ram_budget, - std::string mem_index_path, std::string medoids_file, - std::string centroids_file, size_t build_pq_bytes, bool use_opq, - uint32_t num_threads, bool use_filters = false, - const std::string &label_file = std::string(""), - const std::string &labels_to_medoids_file = std::string(""), - const std::string &universal_label = "", const uint32_t Lf = 0); +DISKANN_DLLEXPORT int build_merged_vamana_index(std::string base_file, diskann::Metric _compareMetric, uint32_t L, + uint32_t R, double sampling_rate, double ram_budget, + std::string mem_index_path, std::string medoids_file, + std::string centroids_file, size_t build_pq_bytes, bool use_opq, + uint32_t num_threads, bool use_filters = false, + const std::string &label_file = std::string(""), + const std::string &labels_to_medoids_file = std::string(""), + const std::string &universal_label = "", const uint32_t Lf = 0); template -DISKANN_DLLEXPORT uint32_t optimize_beamwidth( - std::unique_ptr> &_pFlashIndex, - T *tuning_sample, uint64_t tuning_sample_num, - uint64_t tuning_sample_aligned_dim, uint32_t L, uint32_t nthreads, - uint32_t start_bw = 2); +DISKANN_DLLEXPORT uint32_t optimize_beamwidth(std::unique_ptr> &_pFlashIndex, + T *tuning_sample, uint64_t tuning_sample_num, + uint64_t tuning_sample_aligned_dim, uint32_t L, uint32_t nthreads, + uint32_t start_bw = 2); template DISKANN_DLLEXPORT int build_disk_index( - const char *dataFilePath, const char *indexFilePath, - const char *indexBuildParameters, diskann::Metric _compareMetric, - bool use_opq = false, - const std::string &codebook_prefix = - "", // default is empty for no codebook pass in + const char *dataFilePath, const char *indexFilePath, const char *indexBuildParameters, + diskann::Metric _compareMetric, bool use_opq = false, + const std::string &codebook_prefix = "", // default is empty for no codebook pass in bool use_filters = false, - const std::string &label_file = - std::string(""), // default is empty string for no label_file - const std::string &universal_label = "", - const uint32_t filter_threshold = 0, - const uint32_t Lf = 0); // default is empty string for no universal label + const std::string &label_file = std::string(""), // default is empty string for no label_file + const std::string &universal_label = "", const uint32_t filter_threshold = 0, + const uint32_t Lf = 0); // default is empty string for no universal label template -DISKANN_DLLEXPORT void create_disk_layout( - const std::string base_file, const std::string mem_index_file, - const std::string output_file, - const std::string reorder_data_file = std::string("")); +DISKANN_DLLEXPORT void create_disk_layout(const std::string base_file, const std::string mem_index_file, + const std::string output_file, + const std::string reorder_data_file = std::string("")); -} // namespace diskann +} // namespace diskann diff --git a/include/distance.h b/include/distance.h index b76a3ec30..8b20e586b 100644 --- a/include/distance.h +++ b/include/distance.h @@ -2,210 +2,234 @@ #include "windows_customizations.h" #include -namespace diskann { -enum Metric { L2 = 0, INNER_PRODUCT = 1, COSINE = 2, FAST_L2 = 3 }; - -template -class Distance { - public: - DISKANN_DLLEXPORT Distance(diskann::Metric dist_metric) - : _distance_metric(dist_metric) {} - - // distance comparison function - DISKANN_DLLEXPORT virtual float compare(const T *a, const T *b, - uint32_t length) const = 0; - - // Needed only for COSINE-BYTE and INNER_PRODUCT-BYTE - DISKANN_DLLEXPORT virtual float compare(const T *a, const T *b, - const float normA, const float normB, - uint32_t length) const; - - // For MIPS, normalization adds an extra dimension to the vectors. - // This function lets callers know if the normalization process - // changes the dimension. - DISKANN_DLLEXPORT virtual uint32_t post_normalization_dimension( - uint32_t orig_dimension) const; - - DISKANN_DLLEXPORT virtual diskann::Metric get_metric() const; - - // This is for efficiency. If no normalization is required, the callers - // can simply ignore the normalize_data_for_build() function. - DISKANN_DLLEXPORT virtual bool preprocessing_required() const; - - // Check the preprocessing_required() function before calling this. - // Clients can call the function like this: - // - // if (metric->preprocessing_required()){ - // T* normalized_data_batch; - // Split data into batches of batch_size and for each, call: - // metric->preprocess_base_points(data_batch, batch_size); - // - // TODO: This does not take into account the case for SSD inner product - // where the dimensions change after normalization. - DISKANN_DLLEXPORT virtual void preprocess_base_points( - T *original_data, const size_t orig_dim, const size_t num_points); - - // Invokes normalization for a single vector during search. The scratch space - // has to be created by the caller keeping track of the fact that - // normalization might change the dimension of the query vector. - DISKANN_DLLEXPORT virtual void preprocess_query(const T *query_vec, - const size_t query_dim, - T *scratch_query); - - // If an algorithm has a requirement that some data be aligned to a certain - // boundary it can use this function to indicate that requirement. Currently, - // we are setting it to 8 because that works well for AVX2. If we have AVX512 - // implementations of distance algos, they might have to set this to 16 - // (depending on how they are implemented) - DISKANN_DLLEXPORT virtual size_t get_required_alignment() const; - - // Providing a default implementation for the virtual destructor because we - // don't expect most metric implementations to need it. - DISKANN_DLLEXPORT virtual ~Distance(); - - protected: - diskann::Metric _distance_metric; - size_t _alignment_factor = 8; +namespace diskann +{ +enum Metric +{ + L2 = 0, + INNER_PRODUCT = 1, + COSINE = 2, + FAST_L2 = 3 }; -class DistanceCosineInt8 : public Distance { - public: - DistanceCosineInt8() : Distance(diskann::Metric::COSINE) {} - DISKANN_DLLEXPORT virtual float compare(const int8_t *a, const int8_t *b, - uint32_t length) const; +template class Distance +{ + public: + DISKANN_DLLEXPORT Distance(diskann::Metric dist_metric) : _distance_metric(dist_metric) + { + } + + // distance comparison function + DISKANN_DLLEXPORT virtual float compare(const T *a, const T *b, uint32_t length) const = 0; + + // Needed only for COSINE-BYTE and INNER_PRODUCT-BYTE + DISKANN_DLLEXPORT virtual float compare(const T *a, const T *b, const float normA, const float normB, + uint32_t length) const; + + // For MIPS, normalization adds an extra dimension to the vectors. + // This function lets callers know if the normalization process + // changes the dimension. + DISKANN_DLLEXPORT virtual uint32_t post_normalization_dimension(uint32_t orig_dimension) const; + + DISKANN_DLLEXPORT virtual diskann::Metric get_metric() const; + + // This is for efficiency. If no normalization is required, the callers + // can simply ignore the normalize_data_for_build() function. + DISKANN_DLLEXPORT virtual bool preprocessing_required() const; + + // Check the preprocessing_required() function before calling this. + // Clients can call the function like this: + // + // if (metric->preprocessing_required()){ + // T* normalized_data_batch; + // Split data into batches of batch_size and for each, call: + // metric->preprocess_base_points(data_batch, batch_size); + // + // TODO: This does not take into account the case for SSD inner product + // where the dimensions change after normalization. + DISKANN_DLLEXPORT virtual void preprocess_base_points(T *original_data, const size_t orig_dim, + const size_t num_points); + + // Invokes normalization for a single vector during search. The scratch space + // has to be created by the caller keeping track of the fact that + // normalization might change the dimension of the query vector. + DISKANN_DLLEXPORT virtual void preprocess_query(const T *query_vec, const size_t query_dim, T *scratch_query); + + // If an algorithm has a requirement that some data be aligned to a certain + // boundary it can use this function to indicate that requirement. Currently, + // we are setting it to 8 because that works well for AVX2. If we have AVX512 + // implementations of distance algos, they might have to set this to 16 + // (depending on how they are implemented) + DISKANN_DLLEXPORT virtual size_t get_required_alignment() const; + + // Providing a default implementation for the virtual destructor because we + // don't expect most metric implementations to need it. + DISKANN_DLLEXPORT virtual ~Distance(); + + protected: + diskann::Metric _distance_metric; + size_t _alignment_factor = 8; }; -class DistanceL2Int8 : public Distance { - public: - DistanceL2Int8() : Distance(diskann::Metric::L2) {} - DISKANN_DLLEXPORT virtual float compare(const int8_t *a, const int8_t *b, - uint32_t size) const; +class DistanceCosineInt8 : public Distance +{ + public: + DistanceCosineInt8() : Distance(diskann::Metric::COSINE) + { + } + DISKANN_DLLEXPORT virtual float compare(const int8_t *a, const int8_t *b, uint32_t length) const; +}; + +class DistanceL2Int8 : public Distance +{ + public: + DistanceL2Int8() : Distance(diskann::Metric::L2) + { + } + DISKANN_DLLEXPORT virtual float compare(const int8_t *a, const int8_t *b, uint32_t size) const; }; // AVX implementations. Borrowed from HNSW code. -class AVXDistanceL2Int8 : public Distance { - public: - AVXDistanceL2Int8() : Distance(diskann::Metric::L2) {} - DISKANN_DLLEXPORT virtual float compare(const int8_t *a, const int8_t *b, - uint32_t length) const; +class AVXDistanceL2Int8 : public Distance +{ + public: + AVXDistanceL2Int8() : Distance(diskann::Metric::L2) + { + } + DISKANN_DLLEXPORT virtual float compare(const int8_t *a, const int8_t *b, uint32_t length) const; }; -class DistanceCosineFloat : public Distance { - public: - DistanceCosineFloat() : Distance(diskann::Metric::COSINE) {} - DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, - uint32_t length) const; +class DistanceCosineFloat : public Distance +{ + public: + DistanceCosineFloat() : Distance(diskann::Metric::COSINE) + { + } + DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t length) const; }; -class DistanceL2Float : public Distance { - public: - DistanceL2Float() : Distance(diskann::Metric::L2) {} +class DistanceL2Float : public Distance +{ + public: + DistanceL2Float() : Distance(diskann::Metric::L2) + { + } #ifdef _WINDOWS - DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, - uint32_t size) const; + DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t size) const; #else - DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, - uint32_t size) const - __attribute__((hot)); + DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t size) const __attribute__((hot)); #endif }; -class AVXDistanceL2Float : public Distance { - public: - AVXDistanceL2Float() : Distance(diskann::Metric::L2) {} - DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, - uint32_t length) const; +class AVXDistanceL2Float : public Distance +{ + public: + AVXDistanceL2Float() : Distance(diskann::Metric::L2) + { + } + DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t length) const; }; -template -class SlowDistanceL2 : public Distance { - public: - SlowDistanceL2() : Distance(diskann::Metric::L2) {} - DISKANN_DLLEXPORT virtual float compare(const T *a, const T *b, - uint32_t length) const; +template class SlowDistanceL2 : public Distance +{ + public: + SlowDistanceL2() : Distance(diskann::Metric::L2) + { + } + DISKANN_DLLEXPORT virtual float compare(const T *a, const T *b, uint32_t length) const; }; -class SlowDistanceCosineUInt8 : public Distance { - public: - SlowDistanceCosineUInt8() : Distance(diskann::Metric::COSINE) {} - DISKANN_DLLEXPORT virtual float compare(const uint8_t *a, const uint8_t *b, - uint32_t length) const; +class SlowDistanceCosineUInt8 : public Distance +{ + public: + SlowDistanceCosineUInt8() : Distance(diskann::Metric::COSINE) + { + } + DISKANN_DLLEXPORT virtual float compare(const uint8_t *a, const uint8_t *b, uint32_t length) const; }; -class DistanceL2UInt8 : public Distance { - public: - DistanceL2UInt8() : Distance(diskann::Metric::L2) {} - DISKANN_DLLEXPORT virtual float compare(const uint8_t *a, const uint8_t *b, - uint32_t size) const; +class DistanceL2UInt8 : public Distance +{ + public: + DistanceL2UInt8() : Distance(diskann::Metric::L2) + { + } + DISKANN_DLLEXPORT virtual float compare(const uint8_t *a, const uint8_t *b, uint32_t size) const; }; -template -class DistanceInnerProduct : public Distance { - public: - DistanceInnerProduct() : Distance(diskann::Metric::INNER_PRODUCT) {} - - DistanceInnerProduct(diskann::Metric metric) : Distance(metric) {} - inline float inner_product(const T *a, const T *b, unsigned size) const; - - inline float compare(const T *a, const T *b, unsigned size) const { - float result = inner_product(a, b, size); - // if (result < 0) - // return std::numeric_limits::max(); - // else - return -result; - } +template class DistanceInnerProduct : public Distance +{ + public: + DistanceInnerProduct() : Distance(diskann::Metric::INNER_PRODUCT) + { + } + + DistanceInnerProduct(diskann::Metric metric) : Distance(metric) + { + } + inline float inner_product(const T *a, const T *b, unsigned size) const; + + inline float compare(const T *a, const T *b, unsigned size) const + { + float result = inner_product(a, b, size); + // if (result < 0) + // return std::numeric_limits::max(); + // else + return -result; + } }; -template -class DistanceFastL2 : public DistanceInnerProduct { - // currently defined only for float. - // templated for future use. - public: - DistanceFastL2() : DistanceInnerProduct(diskann::Metric::FAST_L2) {} - float norm(const T *a, unsigned size) const; - float compare(const T *a, const T *b, float norm, unsigned size) const; +template class DistanceFastL2 : public DistanceInnerProduct +{ + // currently defined only for float. + // templated for future use. + public: + DistanceFastL2() : DistanceInnerProduct(diskann::Metric::FAST_L2) + { + } + float norm(const T *a, unsigned size) const; + float compare(const T *a, const T *b, float norm, unsigned size) const; }; -class AVXDistanceInnerProductFloat : public Distance { - public: - AVXDistanceInnerProductFloat() - : Distance(diskann::Metric::INNER_PRODUCT) {} - DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, - uint32_t length) const; +class AVXDistanceInnerProductFloat : public Distance +{ + public: + AVXDistanceInnerProductFloat() : Distance(diskann::Metric::INNER_PRODUCT) + { + } + DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t length) const; }; -class AVXNormalizedCosineDistanceFloat : public Distance { - private: - AVXDistanceInnerProductFloat _innerProduct; - - protected: - void normalize_and_copy(const float *a, uint32_t length, float *a_norm) const; - - public: - AVXNormalizedCosineDistanceFloat() - : Distance(diskann::Metric::COSINE) {} - DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, - uint32_t length) const { - // Inner product returns negative values to indicate distance. - // This will ensure that cosine is between -1 and 1. - return 1.0f + _innerProduct.compare(a, b, length); - } - DISKANN_DLLEXPORT virtual uint32_t post_normalization_dimension( - uint32_t orig_dimension) const override; - - DISKANN_DLLEXPORT virtual bool preprocessing_required() const; - - DISKANN_DLLEXPORT virtual void preprocess_base_points( - float *original_data, const size_t orig_dim, - const size_t num_points) override; - - DISKANN_DLLEXPORT virtual void preprocess_query( - const float *query_vec, const size_t query_dim, - float *scratch_query_vector) override; +class AVXNormalizedCosineDistanceFloat : public Distance +{ + private: + AVXDistanceInnerProductFloat _innerProduct; + + protected: + void normalize_and_copy(const float *a, uint32_t length, float *a_norm) const; + + public: + AVXNormalizedCosineDistanceFloat() : Distance(diskann::Metric::COSINE) + { + } + DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t length) const + { + // Inner product returns negative values to indicate distance. + // This will ensure that cosine is between -1 and 1. + return 1.0f + _innerProduct.compare(a, b, length); + } + DISKANN_DLLEXPORT virtual uint32_t post_normalization_dimension(uint32_t orig_dimension) const override; + + DISKANN_DLLEXPORT virtual bool preprocessing_required() const; + + DISKANN_DLLEXPORT virtual void preprocess_base_points(float *original_data, const size_t orig_dim, + const size_t num_points) override; + + DISKANN_DLLEXPORT virtual void preprocess_query(const float *query_vec, const size_t query_dim, + float *scratch_query_vector) override; }; -template -Distance *get_distance_function(Metric m); +template Distance *get_distance_function(Metric m); -} // namespace diskann +} // namespace diskann diff --git a/include/exceptions.h b/include/exceptions.h index b23dae72c..99e4e7361 100644 --- a/include/exceptions.h +++ b/include/exceptions.h @@ -4,11 +4,14 @@ #pragma once #include -namespace diskann { +namespace diskann +{ -class NotImplementedException : public std::logic_error { - public: - NotImplementedException() - : std::logic_error("Function not yet implemented.") {} +class NotImplementedException : public std::logic_error +{ + public: + NotImplementedException() : std::logic_error("Function not yet implemented.") + { + } }; -} // namespace diskann +} // namespace diskann diff --git a/include/filter_utils.h b/include/filter_utils.h index 073543f1f..b7114d264 100644 --- a/include/filter_utils.h +++ b/include/filter_utils.h @@ -44,32 +44,24 @@ typedef tsl::robin_set label_set; typedef std::string path; // structs for returning multiple items from a function -typedef std::tuple, - tsl::robin_map, - tsl::robin_set> +typedef std::tuple, tsl::robin_map, tsl::robin_set> parse_label_file_return_values; -typedef std::tuple>, uint64_t> - load_label_index_return_values; +typedef std::tuple>, uint64_t> load_label_index_return_values; -namespace diskann { +namespace diskann +{ template -DISKANN_DLLEXPORT void generate_label_indices(path input_data_path, - path final_index_path_prefix, - label_set all_labels, unsigned R, - unsigned L, float alpha, - unsigned num_threads); +DISKANN_DLLEXPORT void generate_label_indices(path input_data_path, path final_index_path_prefix, label_set all_labels, + unsigned R, unsigned L, float alpha, unsigned num_threads); -DISKANN_DLLEXPORT load_label_index_return_values -load_label_index(path label_index_path, uint32_t label_number_of_points); +DISKANN_DLLEXPORT load_label_index_return_values load_label_index(path label_index_path, + uint32_t label_number_of_points); -DISKANN_DLLEXPORT parse_label_file_return_values -parse_label_file(path label_data_path, std::string universal_label); +DISKANN_DLLEXPORT parse_label_file_return_values parse_label_file(path label_data_path, std::string universal_label); template -DISKANN_DLLEXPORT tsl::robin_map> -generate_label_specific_vector_files_compat( - path input_data_path, - tsl::robin_map labels_to_number_of_points, +DISKANN_DLLEXPORT tsl::robin_map> generate_label_specific_vector_files_compat( + path input_data_path, tsl::robin_map labels_to_number_of_points, std::vector point_ids_to_labels, label_set all_labels); /* @@ -83,138 +75,139 @@ generate_label_specific_vector_files_compat( * input_data_path + "_" + label */ template -inline tsl::robin_map> -generate_label_specific_vector_files( - path input_data_path, - tsl::robin_map labels_to_number_of_points, - std::vector point_ids_to_labels, label_set all_labels) { - auto file_writing_timer = std::chrono::high_resolution_clock::now(); - diskann::MemoryMapper input_data(input_data_path); - char *input_start = input_data.getBuf(); - - uint32_t number_of_points, dimension; - std::memcpy(&number_of_points, input_start, sizeof(uint32_t)); - std::memcpy(&dimension, input_start + sizeof(uint32_t), sizeof(uint32_t)); - const uint32_t VECTOR_SIZE = dimension * sizeof(T); - const size_t METADATA = 2 * sizeof(uint32_t); - if (number_of_points != point_ids_to_labels.size()) { - std::cerr << "Error: number of points in labels file and data file differ." - << std::endl; - throw; - } - - tsl::robin_map label_to_iovec_map; - tsl::robin_map label_to_curr_iovec; - tsl::robin_map> label_id_to_orig_id; - - // setup iovec list for each label - for (const auto &lbl : all_labels) { - iovec *label_iovecs = - (iovec *)malloc(labels_to_number_of_points[lbl] * sizeof(iovec)); - if (label_iovecs == nullptr) { - throw; +inline tsl::robin_map> generate_label_specific_vector_files( + path input_data_path, tsl::robin_map labels_to_number_of_points, + std::vector point_ids_to_labels, label_set all_labels) +{ + auto file_writing_timer = std::chrono::high_resolution_clock::now(); + diskann::MemoryMapper input_data(input_data_path); + char *input_start = input_data.getBuf(); + + uint32_t number_of_points, dimension; + std::memcpy(&number_of_points, input_start, sizeof(uint32_t)); + std::memcpy(&dimension, input_start + sizeof(uint32_t), sizeof(uint32_t)); + const uint32_t VECTOR_SIZE = dimension * sizeof(T); + const size_t METADATA = 2 * sizeof(uint32_t); + if (number_of_points != point_ids_to_labels.size()) + { + std::cerr << "Error: number of points in labels file and data file differ." << std::endl; + throw; } - label_to_iovec_map[lbl] = label_iovecs; - label_to_curr_iovec[lbl] = 0; - label_id_to_orig_id[lbl].reserve(labels_to_number_of_points[lbl]); - } - - // each point added to corresponding per-label iovec list - for (uint32_t point_id = 0; point_id < number_of_points; point_id++) { - char *curr_point = input_start + METADATA + (VECTOR_SIZE * point_id); - iovec curr_iovec; - - curr_iovec.iov_base = curr_point; - curr_iovec.iov_len = VECTOR_SIZE; - for (const auto &lbl : point_ids_to_labels[point_id]) { - *(label_to_iovec_map[lbl] + label_to_curr_iovec[lbl]) = curr_iovec; - label_to_curr_iovec[lbl]++; - label_id_to_orig_id[lbl].push_back(point_id); + + tsl::robin_map label_to_iovec_map; + tsl::robin_map label_to_curr_iovec; + tsl::robin_map> label_id_to_orig_id; + + // setup iovec list for each label + for (const auto &lbl : all_labels) + { + iovec *label_iovecs = (iovec *)malloc(labels_to_number_of_points[lbl] * sizeof(iovec)); + if (label_iovecs == nullptr) + { + throw; + } + label_to_iovec_map[lbl] = label_iovecs; + label_to_curr_iovec[lbl] = 0; + label_id_to_orig_id[lbl].reserve(labels_to_number_of_points[lbl]); } - } - - // write each label iovec to resp. file - for (const auto &lbl : all_labels) { - int label_input_data_fd; - path curr_label_input_data_path(input_data_path + "_" + lbl); - uint32_t curr_num_pts = labels_to_number_of_points[lbl]; - - label_input_data_fd = - open(curr_label_input_data_path.c_str(), - O_CREAT | O_WRONLY | O_TRUNC | O_APPEND, (mode_t)0644); - if (label_input_data_fd == -1) throw; - - // write metadata - uint32_t metadata[2] = {curr_num_pts, dimension}; - int return_value = - write(label_input_data_fd, metadata, sizeof(uint32_t) * 2); - if (return_value == -1) { - throw; + + // each point added to corresponding per-label iovec list + for (uint32_t point_id = 0; point_id < number_of_points; point_id++) + { + char *curr_point = input_start + METADATA + (VECTOR_SIZE * point_id); + iovec curr_iovec; + + curr_iovec.iov_base = curr_point; + curr_iovec.iov_len = VECTOR_SIZE; + for (const auto &lbl : point_ids_to_labels[point_id]) + { + *(label_to_iovec_map[lbl] + label_to_curr_iovec[lbl]) = curr_iovec; + label_to_curr_iovec[lbl]++; + label_id_to_orig_id[lbl].push_back(point_id); + } } - // limits on number of iovec structs per writev means we need to perform - // multiple writevs - size_t i = 0; - while (curr_num_pts > IOV_MAX) { - return_value = writev(label_input_data_fd, - (label_to_iovec_map[lbl] + (IOV_MAX * i)), IOV_MAX); - if (return_value == -1) { + // write each label iovec to resp. file + for (const auto &lbl : all_labels) + { + int label_input_data_fd; + path curr_label_input_data_path(input_data_path + "_" + lbl); + uint32_t curr_num_pts = labels_to_number_of_points[lbl]; + + label_input_data_fd = + open(curr_label_input_data_path.c_str(), O_CREAT | O_WRONLY | O_TRUNC | O_APPEND, (mode_t)0644); + if (label_input_data_fd == -1) + throw; + + // write metadata + uint32_t metadata[2] = {curr_num_pts, dimension}; + int return_value = write(label_input_data_fd, metadata, sizeof(uint32_t) * 2); + if (return_value == -1) + { + throw; + } + + // limits on number of iovec structs per writev means we need to perform + // multiple writevs + size_t i = 0; + while (curr_num_pts > IOV_MAX) + { + return_value = writev(label_input_data_fd, (label_to_iovec_map[lbl] + (IOV_MAX * i)), IOV_MAX); + if (return_value == -1) + { + close(label_input_data_fd); + throw; + } + curr_num_pts -= IOV_MAX; + i += 1; + } + return_value = writev(label_input_data_fd, (label_to_iovec_map[lbl] + (IOV_MAX * i)), curr_num_pts); + if (return_value == -1) + { + close(label_input_data_fd); + throw; + } + + free(label_to_iovec_map[lbl]); close(label_input_data_fd); - throw; - } - curr_num_pts -= IOV_MAX; - i += 1; - } - return_value = - writev(label_input_data_fd, (label_to_iovec_map[lbl] + (IOV_MAX * i)), - curr_num_pts); - if (return_value == -1) { - close(label_input_data_fd); - throw; } - free(label_to_iovec_map[lbl]); - close(label_input_data_fd); - } - - std::chrono::duration file_writing_time = - std::chrono::high_resolution_clock::now() - file_writing_timer; - std::cout << "generated " << all_labels.size() - << " label-specific vector files for index building in time " - << file_writing_time.count() << "\n" - << std::endl; + std::chrono::duration file_writing_time = std::chrono::high_resolution_clock::now() - file_writing_timer; + std::cout << "generated " << all_labels.size() << " label-specific vector files for index building in time " + << file_writing_time.count() << "\n" + << std::endl; - return label_id_to_orig_id; + return label_id_to_orig_id; } -inline std::vector loadTags(const std::string &tags_file, - const std::string &base_file) { - const bool tags_enabled = tags_file.empty() ? false : true; - std::vector location_to_tag; - if (tags_enabled) { - size_t tag_file_ndims, tag_file_npts; - std::uint32_t *tag_data; - diskann::load_bin(tags_file, tag_data, tag_file_npts, - tag_file_ndims); - if (tag_file_ndims != 1) { - diskann::cerr << "tags file error" << std::endl; - throw diskann::ANNException("tag file error", -1, __FUNCSIG__, __FILE__, - __LINE__); +inline std::vector loadTags(const std::string &tags_file, const std::string &base_file) +{ + const bool tags_enabled = tags_file.empty() ? false : true; + std::vector location_to_tag; + if (tags_enabled) + { + size_t tag_file_ndims, tag_file_npts; + std::uint32_t *tag_data; + diskann::load_bin(tags_file, tag_data, tag_file_npts, tag_file_ndims); + if (tag_file_ndims != 1) + { + diskann::cerr << "tags file error" << std::endl; + throw diskann::ANNException("tag file error", -1, __FUNCSIG__, __FILE__, __LINE__); + } + + // check if the point count match + size_t base_file_npts, base_file_ndims; + diskann::get_bin_metadata(base_file, base_file_npts, base_file_ndims); + if (base_file_npts != tag_file_npts) + { + diskann::cerr << "point num in tags file mismatch" << std::endl; + throw diskann::ANNException("point num in tags file mismatch", -1, __FUNCSIG__, __FILE__, __LINE__); + } + + location_to_tag.assign(tag_data, tag_data + tag_file_npts); + delete[] tag_data; } - - // check if the point count match - size_t base_file_npts, base_file_ndims; - diskann::get_bin_metadata(base_file, base_file_npts, base_file_ndims); - if (base_file_npts != tag_file_npts) { - diskann::cerr << "point num in tags file mismatch" << std::endl; - throw diskann::ANNException("point num in tags file mismatch", -1, - __FUNCSIG__, __FILE__, __LINE__); - } - - location_to_tag.assign(tag_data, tag_data + tag_file_npts); - delete[] tag_data; - } - return location_to_tag; + return location_to_tag; } -} // namespace diskann +} // namespace diskann diff --git a/include/in_mem_data_store.h b/include/in_mem_data_store.h index db8b0bf20..f86aa0edc 100644 --- a/include/in_mem_data_store.h +++ b/include/in_mem_data_store.h @@ -15,77 +15,67 @@ #include "natural_number_map.h" #include "natural_number_set.h" -namespace diskann { -template -class InMemDataStore : public AbstractDataStore { - public: - InMemDataStore(const location_t capacity, const size_t dim, - std::shared_ptr> distance_metric); - virtual ~InMemDataStore(); - - virtual location_t load(const std::string &filename) override; - virtual size_t save(const std::string &filename, - const location_t num_points) override; - - virtual size_t get_aligned_dim() const override; - - // Populate internal data from unaligned data while doing alignment and any - // normalization that is required. - virtual void populate_data(const data_t *vectors, - const location_t num_pts) override; - virtual void populate_data(const std::string &filename, - const size_t offset) override; - - virtual void extract_data_to_bin(const std::string &filename, - const location_t num_pts) override; - - virtual void get_vector(const location_t i, data_t *target) const override; - virtual void set_vector(const location_t i, - const data_t *const vector) override; - virtual void prefetch_vector(const location_t loc) override; - - virtual void move_vectors(const location_t old_location_start, - const location_t new_location_start, - const location_t num_points) override; - virtual void copy_vectors(const location_t from_loc, const location_t to_loc, - const location_t num_points) override; - - virtual float get_distance(const data_t *query, - const location_t loc) const override; - virtual float get_distance(const location_t loc1, - const location_t loc2) const override; - virtual void get_distance(const data_t *query, const location_t *locations, - const uint32_t location_count, - float *distances) const override; - - virtual location_t calculate_medoid() const override; - - virtual Distance *get_dist_fn(); - - virtual size_t get_alignment_factor() const override; - - protected: - virtual location_t expand(const location_t new_size) override; - virtual location_t shrink(const location_t new_size) override; - - virtual location_t load_impl(const std::string &filename); +namespace diskann +{ +template class InMemDataStore : public AbstractDataStore +{ + public: + InMemDataStore(const location_t capacity, const size_t dim, std::shared_ptr> distance_metric); + virtual ~InMemDataStore(); + + virtual location_t load(const std::string &filename) override; + virtual size_t save(const std::string &filename, const location_t num_points) override; + + virtual size_t get_aligned_dim() const override; + + // Populate internal data from unaligned data while doing alignment and any + // normalization that is required. + virtual void populate_data(const data_t *vectors, const location_t num_pts) override; + virtual void populate_data(const std::string &filename, const size_t offset) override; + + virtual void extract_data_to_bin(const std::string &filename, const location_t num_pts) override; + + virtual void get_vector(const location_t i, data_t *target) const override; + virtual void set_vector(const location_t i, const data_t *const vector) override; + virtual void prefetch_vector(const location_t loc) override; + + virtual void move_vectors(const location_t old_location_start, const location_t new_location_start, + const location_t num_points) override; + virtual void copy_vectors(const location_t from_loc, const location_t to_loc, const location_t num_points) override; + + virtual float get_distance(const data_t *query, const location_t loc) const override; + virtual float get_distance(const location_t loc1, const location_t loc2) const override; + virtual void get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, + float *distances) const override; + + virtual location_t calculate_medoid() const override; + + virtual Distance *get_dist_fn(); + + virtual size_t get_alignment_factor() const override; + + protected: + virtual location_t expand(const location_t new_size) override; + virtual location_t shrink(const location_t new_size) override; + + virtual location_t load_impl(const std::string &filename); #ifdef EXEC_ENV_OLS - virtual location_t load_impl(AlignedFileReader &reader); + virtual location_t load_impl(AlignedFileReader &reader); #endif - private: - data_t *_data = nullptr; + private: + data_t *_data = nullptr; - size_t _aligned_dim; + size_t _aligned_dim; - // It may seem weird to put distance metric along with the data store class, - // but this gives us perf benefits as the datastore can do distance - // computations during search and compute norms of vectors internally without - // have to copy data back and forth. - std::shared_ptr> _distance_fn; + // It may seem weird to put distance metric along with the data store class, + // but this gives us perf benefits as the datastore can do distance + // computations during search and compute norms of vectors internally without + // have to copy data back and forth. + std::shared_ptr> _distance_fn; - // in case we need to save vector norms for optimization - std::shared_ptr _pre_computed_norms; + // in case we need to save vector norms for optimization + std::shared_ptr _pre_computed_norms; }; -} // namespace diskann \ No newline at end of file +} // namespace diskann \ No newline at end of file diff --git a/include/in_mem_graph_store.h b/include/in_mem_graph_store.h index a5a907b38..98a9e4dc5 100644 --- a/include/in_mem_graph_store.h +++ b/include/in_mem_graph_store.h @@ -5,17 +5,19 @@ #include "abstract_graph_store.h" -namespace diskann { +namespace diskann +{ -class InMemGraphStore : public AbstractGraphStore { - public: - InMemGraphStore(const size_t max_pts); +class InMemGraphStore : public AbstractGraphStore +{ + public: + InMemGraphStore(const size_t max_pts); - int load(const std::string &index_path_prefix); - int store(const std::string &index_path_prefix); + int load(const std::string &index_path_prefix); + int store(const std::string &index_path_prefix); - void get_adj_list(const location_t i, std::vector &neighbors); - void set_adj_list(const location_t i, std::vector &neighbors); + void get_adj_list(const location_t i, std::vector &neighbors); + void set_adj_list(const location_t i, std::vector &neighbors); }; -} // namespace diskann +} // namespace diskann diff --git a/include/index.h b/include/index.h index aa6415e3a..9090580f0 100644 --- a/include/index.h +++ b/include/index.h @@ -24,437 +24,395 @@ #define EXPAND_IF_FULL 0 #define DEFAULT_MAXC 750 -namespace diskann { - -inline double estimate_ram_usage(size_t size, uint32_t dim, uint32_t datasize, - uint32_t degree) { - double size_of_data = ((double)size) * ROUND_UP(dim, 8) * datasize; - double size_of_graph = - ((double)size) * degree * sizeof(uint32_t) * GRAPH_SLACK_FACTOR; - double size_of_locks = ((double)size) * sizeof(non_recursive_mutex); - double size_of_outer_vector = ((double)size) * sizeof(ptrdiff_t); - - return OVERHEAD_FACTOR * - (size_of_data + size_of_graph + size_of_locks + size_of_outer_vector); +namespace diskann +{ + +inline double estimate_ram_usage(size_t size, uint32_t dim, uint32_t datasize, uint32_t degree) +{ + double size_of_data = ((double)size) * ROUND_UP(dim, 8) * datasize; + double size_of_graph = ((double)size) * degree * sizeof(uint32_t) * GRAPH_SLACK_FACTOR; + double size_of_locks = ((double)size) * sizeof(non_recursive_mutex); + double size_of_outer_vector = ((double)size) * sizeof(ptrdiff_t); + + return OVERHEAD_FACTOR * (size_of_data + size_of_graph + size_of_locks + size_of_outer_vector); } -struct consolidation_report { - enum status_code { - SUCCESS = 0, - FAIL = 1, - LOCK_FAIL = 2, - INCONSISTENT_COUNT_ERROR = 3 - }; - status_code _status; - size_t _active_points, _max_points, _empty_slots, _slots_released, - _delete_set_size, _num_calls_to_process_delete; - double _time; - - consolidation_report(status_code status, size_t active_points, - size_t max_points, size_t empty_slots, - size_t slots_released, size_t delete_set_size, - size_t num_calls_to_process_delete, double time_secs) - : _status(status), - _active_points(active_points), - _max_points(max_points), - _empty_slots(empty_slots), - _slots_released(slots_released), - _delete_set_size(delete_set_size), - _num_calls_to_process_delete(num_calls_to_process_delete), - _time(time_secs) {} +struct consolidation_report +{ + enum status_code + { + SUCCESS = 0, + FAIL = 1, + LOCK_FAIL = 2, + INCONSISTENT_COUNT_ERROR = 3 + }; + status_code _status; + size_t _active_points, _max_points, _empty_slots, _slots_released, _delete_set_size, _num_calls_to_process_delete; + double _time; + + consolidation_report(status_code status, size_t active_points, size_t max_points, size_t empty_slots, + size_t slots_released, size_t delete_set_size, size_t num_calls_to_process_delete, + double time_secs) + : _status(status), _active_points(active_points), _max_points(max_points), _empty_slots(empty_slots), + _slots_released(slots_released), _delete_set_size(delete_set_size), + _num_calls_to_process_delete(num_calls_to_process_delete), _time(time_secs) + { + } }; -template -class Index { - /************************************************************************** - * - * Public functions acquire one or more of _update_lock, _consolidate_lock, - * _tag_lock, _delete_lock before calling protected functions which DO NOT - * acquire these locks. They might acquire locks on _locks[i] - * - **************************************************************************/ - - public: - // Constructor for Bulk operations and for creating the index object solely - // for loading a prexisting index. - DISKANN_DLLEXPORT Index( - Metric m, const size_t dim, const size_t max_points = 1, - 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 size_t num_frozen_pts = 0); - - // Constructor for incremental index - DISKANN_DLLEXPORT Index( - Metric m, const size_t dim, const size_t max_points, - const bool dynamic_index, const IndexWriteParameters &indexParameters, - const uint32_t initial_search_list_size, const uint32_t search_threads, - 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); - - DISKANN_DLLEXPORT ~Index(); - - // Saves graph, data, metadata and associated tags. - DISKANN_DLLEXPORT void save(const char *filename, - bool compact_before_save = false); - - // Load functions +template class Index +{ + /************************************************************************** + * + * Public functions acquire one or more of _update_lock, _consolidate_lock, + * _tag_lock, _delete_lock before calling protected functions which DO NOT + * acquire these locks. They might acquire locks on _locks[i] + * + **************************************************************************/ + + public: + // Constructor for Bulk operations and for creating the index object solely + // for loading a prexisting index. + DISKANN_DLLEXPORT Index(Metric m, const size_t dim, const size_t max_points = 1, 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 size_t num_frozen_pts = 0); + + // Constructor for incremental index + DISKANN_DLLEXPORT Index(Metric m, const size_t dim, const size_t max_points, const bool dynamic_index, + const IndexWriteParameters &indexParameters, const uint32_t initial_search_list_size, + const uint32_t search_threads, 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); + + DISKANN_DLLEXPORT ~Index(); + + // Saves graph, data, metadata and associated tags. + DISKANN_DLLEXPORT void save(const char *filename, bool compact_before_save = false); + + // Load functions #ifdef EXEC_ENV_OLS - DISKANN_DLLEXPORT void load(AlignedFileReader &reader, uint32_t num_threads, - uint32_t search_l); + DISKANN_DLLEXPORT void load(AlignedFileReader &reader, uint32_t num_threads, uint32_t search_l); #else - // Reads the number of frozen points from graph's metadata file section. - DISKANN_DLLEXPORT static size_t get_graph_num_frozen_points( - const std::string &graph_file); + // Reads the number of frozen points from graph's metadata file section. + DISKANN_DLLEXPORT static size_t get_graph_num_frozen_points(const std::string &graph_file); - DISKANN_DLLEXPORT void load(const char *index_file, uint32_t num_threads, - uint32_t search_l); + DISKANN_DLLEXPORT void load(const char *index_file, uint32_t num_threads, uint32_t search_l); #endif - // get some private variables - DISKANN_DLLEXPORT size_t get_num_points(); - DISKANN_DLLEXPORT size_t get_max_points(); - - // Batch build from a file. Optionally pass tags vector. - DISKANN_DLLEXPORT void build( - const char *filename, const size_t num_points_to_load, - IndexWriteParameters ¶meters, - const std::vector &tags = std::vector()); - - // Batch build from a file. Optionally pass tags file. - DISKANN_DLLEXPORT void build(const char *filename, - const size_t num_points_to_load, - IndexWriteParameters ¶meters, - const char *tag_filename); - - // Batch build from a data array, which must pad vectors to aligned_dim - DISKANN_DLLEXPORT void build(const T *data, const size_t num_points_to_load, - IndexWriteParameters ¶meters, - const std::vector &tags); - - // Filtered Support - DISKANN_DLLEXPORT void build_filtered_index( - const char *filename, const std::string &label_file, - const size_t num_points_to_load, IndexWriteParameters ¶meters, - const std::vector &tags = std::vector()); - - DISKANN_DLLEXPORT void set_universal_label(const LabelT &label); - - // Get converted integer label from string to int map (_label_map) - DISKANN_DLLEXPORT LabelT get_converted_label(const std::string &raw_label); - - // Set starting point of an index before inserting any points incrementally. - // The data count should be equal to _num_frozen_pts * _aligned_dim. - DISKANN_DLLEXPORT void set_start_points(const T *data, size_t data_count); - // Set starting points to random points on a sphere of certain radius. - // A fixed random seed can be specified for scenarios where it's important - // to have higher consistency between index builds. - DISKANN_DLLEXPORT void set_start_points_at_random(T radius, - uint32_t random_seed = 0); - - // For FastL2 search on a static index, we interleave the data with graph - DISKANN_DLLEXPORT void optimize_index_layout(); - - // For FastL2 search on optimized layout - DISKANN_DLLEXPORT void search_with_optimized_layout(const T *query, size_t K, - size_t L, - uint32_t *indices); - - // Added search overload that takes L as parameter, so that we - // can customize L on a per-query basis without tampering with "Parameters" - template - DISKANN_DLLEXPORT std::pair search( - const T *query, const size_t K, const uint32_t L, IDType *indices, - float *distances = nullptr); - - // Initialize space for res_vectors before calling. - DISKANN_DLLEXPORT size_t search_with_tags(const T *query, const uint64_t K, - const uint32_t L, TagT *tags, - float *distances, - std::vector &res_vectors); - - // Filter support search - template - DISKANN_DLLEXPORT std::pair search_with_filters( - const T *query, const LabelT &filter_label, const size_t K, - const uint32_t L, IndexType *indices, float *distances); - - // Will fail if tag already in the index or if tag=0. - DISKANN_DLLEXPORT int insert_point(const T *point, const TagT tag); - - // call this before issuing deletions to sets relevant flags - DISKANN_DLLEXPORT int enable_delete(); - - // Record deleted point now and restructure graph later. Return -1 if tag - // not found, 0 if OK. - DISKANN_DLLEXPORT int lazy_delete(const TagT &tag); - - // Record deleted points now and restructure graph later. Add to failed_tags - // if tag not found. - DISKANN_DLLEXPORT void lazy_delete(const std::vector &tags, - std::vector &failed_tags); - - // Call after a series of lazy deletions - // Returns number of live points left after consolidation - // If _conc_consolidates is set in the ctor, then this call can be invoked - // alongside inserts and lazy deletes, else it acquires _update_lock - DISKANN_DLLEXPORT consolidation_report - consolidate_deletes(const IndexWriteParameters ¶meters); - - DISKANN_DLLEXPORT void prune_all_neighbors(const uint32_t max_degree, - const uint32_t max_occlusion, - const float alpha); - - DISKANN_DLLEXPORT bool is_index_saved(); - - // 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_points(uint32_t old_location_start, - uint32_t new_location_start, - uint32_t num_locations); - - // DISKANN_DLLEXPORT void save_index_as_one_file(bool flag); - - DISKANN_DLLEXPORT void get_active_tags(tsl::robin_set &active_tags); - - // memory should be allocated for vec before calling this function - DISKANN_DLLEXPORT int get_vector_by_tag(TagT &tag, T *vec); - - DISKANN_DLLEXPORT void print_status(); - - DISKANN_DLLEXPORT void count_nodes_at_bfs_levels(); - - // 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: - // No copy/assign. - Index(const Index &) = delete; - Index &operator=(const Index &) = delete; - - // Use after _data and _nd have been populated - // Acquire exclusive _update_lock before calling - void build_with_data_populated(IndexWriteParameters ¶meters, + // get some private variables + DISKANN_DLLEXPORT size_t get_num_points(); + DISKANN_DLLEXPORT size_t get_max_points(); + + // Batch build from a file. Optionally pass tags vector. + DISKANN_DLLEXPORT void build(const char *filename, const size_t num_points_to_load, + IndexWriteParameters ¶meters, const std::vector &tags = std::vector()); + + // Batch build from a file. Optionally pass tags file. + DISKANN_DLLEXPORT void build(const char *filename, const size_t num_points_to_load, + IndexWriteParameters ¶meters, const char *tag_filename); + + // Batch build from a data array, which must pad vectors to aligned_dim + DISKANN_DLLEXPORT void build(const T *data, const size_t num_points_to_load, IndexWriteParameters ¶meters, const std::vector &tags); - // generates 1 frozen point that will never be deleted from the graph - // This is not visible to the user - void generate_frozen_point(); - - // determines navigating node of the graph by calculating medoid of datafopt - uint32_t calculate_entry_point(); - - void parse_label_file(const std::string &label_file, size_t &num_pts_labels); - - std::unordered_map load_label_map( - const std::string &map_file); - - // Returns the locations of start point and frozen points suitable for use - // with iterate_to_fixed_point. - std::vector get_init_ids(); - - std::pair iterate_to_fixed_point( - const T *node_coords, const uint32_t Lindex, - const std::vector &init_ids, InMemQueryScratch *scratch, - bool use_filter, const std::vector &filters, - bool search_invocation); - - void search_for_point_and_prune(int location, uint32_t Lindex, - std::vector &pruned_list, - InMemQueryScratch *scratch, - bool use_filter = false, - uint32_t filteredLindex = 0); - - void prune_neighbors(const uint32_t location, std::vector &pool, - std::vector &pruned_list, - InMemQueryScratch *scratch); - - void prune_neighbors(const uint32_t location, std::vector &pool, - const uint32_t range, const uint32_t max_candidate_size, - const float alpha, std::vector &pruned_list, - InMemQueryScratch *scratch); - - // Prunes candidates in @pool to a shorter list @result - // @pool must be sorted before calling - void occlude_list( - const uint32_t location, std::vector &pool, const float alpha, - const uint32_t degree, const uint32_t maxc, std::vector &result, - InMemQueryScratch *scratch, - const tsl::robin_set *const delete_set_ptr = nullptr); - - // add reverse links from all the visited nodes to node n. - void inter_insert(uint32_t n, std::vector &pruned_list, - const uint32_t range, InMemQueryScratch *scratch); - - void inter_insert(uint32_t n, std::vector &pruned_list, - InMemQueryScratch *scratch); - - // Acquire exclusive _update_lock before calling - void link(IndexWriteParameters ¶meters); - - // Acquire exclusive _tag_lock and _delete_lock before calling - int reserve_location(); - - // Acquire exclusive _tag_lock before calling - size_t release_location(int location); - size_t release_locations(const tsl::robin_set &locations); - - // Resize the index when no slots are left for insertion. - // Acquire exclusive _update_lock and _tag_lock before calling. - void resize(size_t new_max_points); - - // Acquire unique lock on _update_lock, _consolidate_lock, _tag_lock - // and _delete_lock before calling these functions. - // Renumber nodes, update tag and location maps and compact the - // graph, mode = _consolidated_order in case of lazy deletion and - // _compacted_order in case of eager deletion - DISKANN_DLLEXPORT void compact_data(); - DISKANN_DLLEXPORT void compact_frozen_point(); - - // Remove deleted nodes from adjacency list of node loc - // Replace removed neighbors with second order neighbors. - // Also acquires _locks[i] for i = loc and out-neighbors of loc. - void process_delete(const tsl::robin_set &old_delete_set, - size_t loc, const uint32_t range, const uint32_t maxc, - const float alpha, InMemQueryScratch *scratch); - - void initialize_query_scratch(uint32_t num_threads, uint32_t search_l, - uint32_t indexing_l, uint32_t r, uint32_t maxc, - size_t dim); - - // Do not call without acquiring appropriate locks - // call public member functions save and load to invoke these. - DISKANN_DLLEXPORT size_t save_graph(std::string filename); - DISKANN_DLLEXPORT size_t save_data(std::string filename); - DISKANN_DLLEXPORT size_t save_tags(std::string filename); - DISKANN_DLLEXPORT size_t save_delete_list(const std::string &filename); + // Filtered Support + DISKANN_DLLEXPORT void build_filtered_index(const char *filename, const std::string &label_file, + const size_t num_points_to_load, IndexWriteParameters ¶meters, + const std::vector &tags = std::vector()); + + DISKANN_DLLEXPORT void set_universal_label(const LabelT &label); + + // Get converted integer label from string to int map (_label_map) + DISKANN_DLLEXPORT LabelT get_converted_label(const std::string &raw_label); + + // Set starting point of an index before inserting any points incrementally. + // The data count should be equal to _num_frozen_pts * _aligned_dim. + DISKANN_DLLEXPORT void set_start_points(const T *data, size_t data_count); + // Set starting points to random points on a sphere of certain radius. + // A fixed random seed can be specified for scenarios where it's important + // to have higher consistency between index builds. + DISKANN_DLLEXPORT void set_start_points_at_random(T radius, uint32_t random_seed = 0); + + // For FastL2 search on a static index, we interleave the data with graph + DISKANN_DLLEXPORT void optimize_index_layout(); + + // For FastL2 search on optimized layout + DISKANN_DLLEXPORT void search_with_optimized_layout(const T *query, size_t K, size_t L, uint32_t *indices); + + // Added search overload that takes L as parameter, so that we + // can customize L on a per-query basis without tampering with "Parameters" + template + DISKANN_DLLEXPORT std::pair search(const T *query, const size_t K, const uint32_t L, + IDType *indices, float *distances = nullptr); + + // Initialize space for res_vectors before calling. + DISKANN_DLLEXPORT size_t search_with_tags(const T *query, const uint64_t K, const uint32_t L, TagT *tags, + float *distances, std::vector &res_vectors); + + // Filter support search + template + DISKANN_DLLEXPORT std::pair search_with_filters(const T *query, const LabelT &filter_label, + const size_t K, const uint32_t L, + IndexType *indices, float *distances); + + // Will fail if tag already in the index or if tag=0. + DISKANN_DLLEXPORT int insert_point(const T *point, const TagT tag); + + // call this before issuing deletions to sets relevant flags + DISKANN_DLLEXPORT int enable_delete(); + + // Record deleted point now and restructure graph later. Return -1 if tag + // not found, 0 if OK. + DISKANN_DLLEXPORT int lazy_delete(const TagT &tag); + + // Record deleted points now and restructure graph later. Add to failed_tags + // if tag not found. + DISKANN_DLLEXPORT void lazy_delete(const std::vector &tags, std::vector &failed_tags); + + // Call after a series of lazy deletions + // Returns number of live points left after consolidation + // If _conc_consolidates is set in the ctor, then this call can be invoked + // alongside inserts and lazy deletes, else it acquires _update_lock + DISKANN_DLLEXPORT consolidation_report consolidate_deletes(const IndexWriteParameters ¶meters); + + DISKANN_DLLEXPORT void prune_all_neighbors(const uint32_t max_degree, const uint32_t max_occlusion, + const float alpha); + + DISKANN_DLLEXPORT bool is_index_saved(); + + // 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_points(uint32_t old_location_start, uint32_t new_location_start, + uint32_t num_locations); + + // DISKANN_DLLEXPORT void save_index_as_one_file(bool flag); + + DISKANN_DLLEXPORT void get_active_tags(tsl::robin_set &active_tags); + + // memory should be allocated for vec before calling this function + DISKANN_DLLEXPORT int get_vector_by_tag(TagT &tag, T *vec); + + DISKANN_DLLEXPORT void print_status(); + + DISKANN_DLLEXPORT void count_nodes_at_bfs_levels(); + + // 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: + // No copy/assign. + Index(const Index &) = delete; + Index &operator=(const Index &) = delete; + + // Use after _data and _nd have been populated + // Acquire exclusive _update_lock before calling + void build_with_data_populated(IndexWriteParameters ¶meters, const std::vector &tags); + + // generates 1 frozen point that will never be deleted from the graph + // This is not visible to the user + void generate_frozen_point(); + + // determines navigating node of the graph by calculating medoid of datafopt + uint32_t calculate_entry_point(); + + void parse_label_file(const std::string &label_file, size_t &num_pts_labels); + + std::unordered_map load_label_map(const std::string &map_file); + + // Returns the locations of start point and frozen points suitable for use + // with iterate_to_fixed_point. + std::vector get_init_ids(); + + std::pair iterate_to_fixed_point(const T *node_coords, const uint32_t Lindex, + const std::vector &init_ids, + InMemQueryScratch *scratch, bool use_filter, + const std::vector &filters, bool search_invocation); + + void search_for_point_and_prune(int location, uint32_t Lindex, std::vector &pruned_list, + InMemQueryScratch *scratch, bool use_filter = false, + uint32_t filteredLindex = 0); + + void prune_neighbors(const uint32_t location, std::vector &pool, std::vector &pruned_list, + InMemQueryScratch *scratch); + + void prune_neighbors(const uint32_t location, std::vector &pool, const uint32_t range, + const uint32_t max_candidate_size, const float alpha, std::vector &pruned_list, + InMemQueryScratch *scratch); + + // Prunes candidates in @pool to a shorter list @result + // @pool must be sorted before calling + void occlude_list(const uint32_t location, std::vector &pool, const float alpha, const uint32_t degree, + const uint32_t maxc, std::vector &result, InMemQueryScratch *scratch, + const tsl::robin_set *const delete_set_ptr = nullptr); + + // add reverse links from all the visited nodes to node n. + void inter_insert(uint32_t n, std::vector &pruned_list, const uint32_t range, + InMemQueryScratch *scratch); + + void inter_insert(uint32_t n, std::vector &pruned_list, InMemQueryScratch *scratch); + + // Acquire exclusive _update_lock before calling + void link(IndexWriteParameters ¶meters); + + // Acquire exclusive _tag_lock and _delete_lock before calling + int reserve_location(); + + // Acquire exclusive _tag_lock before calling + size_t release_location(int location); + size_t release_locations(const tsl::robin_set &locations); + + // Resize the index when no slots are left for insertion. + // Acquire exclusive _update_lock and _tag_lock before calling. + void resize(size_t new_max_points); + + // Acquire unique lock on _update_lock, _consolidate_lock, _tag_lock + // and _delete_lock before calling these functions. + // Renumber nodes, update tag and location maps and compact the + // graph, mode = _consolidated_order in case of lazy deletion and + // _compacted_order in case of eager deletion + DISKANN_DLLEXPORT void compact_data(); + DISKANN_DLLEXPORT void compact_frozen_point(); + + // Remove deleted nodes from adjacency list of node loc + // Replace removed neighbors with second order neighbors. + // Also acquires _locks[i] for i = loc and out-neighbors of loc. + void process_delete(const tsl::robin_set &old_delete_set, size_t loc, const uint32_t range, + const uint32_t maxc, const float alpha, InMemQueryScratch *scratch); + + void initialize_query_scratch(uint32_t num_threads, uint32_t search_l, uint32_t indexing_l, uint32_t r, + uint32_t maxc, size_t dim); + + // Do not call without acquiring appropriate locks + // call public member functions save and load to invoke these. + DISKANN_DLLEXPORT size_t save_graph(std::string filename); + DISKANN_DLLEXPORT size_t save_data(std::string filename); + DISKANN_DLLEXPORT size_t save_tags(std::string filename); + DISKANN_DLLEXPORT size_t save_delete_list(const std::string &filename); #ifdef EXEC_ENV_OLS - DISKANN_DLLEXPORT size_t load_graph(AlignedFileReader &reader, - size_t expected_num_points); - DISKANN_DLLEXPORT size_t load_data(AlignedFileReader &reader); - DISKANN_DLLEXPORT size_t load_tags(AlignedFileReader &reader); - DISKANN_DLLEXPORT size_t load_delete_set(AlignedFileReader &reader); + DISKANN_DLLEXPORT size_t load_graph(AlignedFileReader &reader, size_t expected_num_points); + DISKANN_DLLEXPORT size_t load_data(AlignedFileReader &reader); + DISKANN_DLLEXPORT size_t load_tags(AlignedFileReader &reader); + DISKANN_DLLEXPORT size_t load_delete_set(AlignedFileReader &reader); #else - DISKANN_DLLEXPORT size_t load_graph(const std::string filename, - size_t expected_num_points); - DISKANN_DLLEXPORT size_t load_data(std::string filename0); - DISKANN_DLLEXPORT size_t load_tags(const std::string tag_file_name); - DISKANN_DLLEXPORT size_t load_delete_set(const std::string &filename); + DISKANN_DLLEXPORT size_t load_graph(const std::string filename, size_t expected_num_points); + DISKANN_DLLEXPORT size_t load_data(std::string filename0); + DISKANN_DLLEXPORT size_t load_tags(const std::string tag_file_name); + DISKANN_DLLEXPORT size_t load_delete_set(const std::string &filename); #endif - private: - // Distance functions - Metric _dist_metric = diskann::L2; - std::shared_ptr> _distance; - - // Data - std::unique_ptr> _data_store; - char *_opt_graph = nullptr; - - // Graph related data structures - std::vector> _final_graph; - - // Dimensions - size_t _dim = 0; - 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 - - // _num_frozen_pts is the number of points which are used as initial - // candidates when iterating to closest point(s). These are not visible - // externally and won't be returned by search. At least 1 frozen point is - // needed for a dynamic index. The frozen points have consecutive locations. - // See also _start below. - size_t _num_frozen_pts = 0; - size_t _max_range_of_loaded_graph = 0; - size_t _node_size; - size_t _data_len; - size_t _neighbor_len; - - uint32_t _max_observed_degree = 0; - // Start point of the search. When _num_frozen_pts is greater than zero, - // this is the location of the first frozen point. Otherwise, this is a - // location of one of the points in index. - uint32_t _start = 0; - - bool _has_built = false; - bool _saturate_graph = false; - bool _save_as_one_file = false; // plan to support in next version - bool _dynamic_index = false; - bool _enable_tags = false; - bool _normalize_vecs = false; // Using normalied L2 for cosine. - - // Filter Support - - bool _filtered_index = false; - std::vector> _pts_to_labels; - tsl::robin_set _labels; - std::string _labels_file; - std::unordered_map _label_to_medoid_id; - std::unordered_map _medoid_counts; - bool _use_universal_label = false; - LabelT _universal_label = 0; - uint32_t _filterIndexingQueueSize; - std::unordered_map _label_map; - - // Indexing parameters - uint32_t _indexingQueueSize; - uint32_t _indexingRange; - uint32_t _indexingMaxC; - float _indexingAlpha; - - // Query scratch data structures - ConcurrentQueue *> _query_scratch; - - // Flags for PQ based distance calculation - bool _pq_dist = false; - bool _use_opq = false; - size_t _num_pq_chunks = 0; - uint8_t *_pq_data = nullptr; - bool _pq_generated = false; - FixedChunkPQTable _pq_table; - - // - // Data structures, locks and flags for dynamic indexing and tags - // - - // lazy_delete removes entry from _location_to_tag and _tag_to_location. If - // _location_to_tag does not resolve a location, infer that it was deleted. - tsl::sparse_map _tag_to_location; - natural_number_map _location_to_tag; - - // _empty_slots has unallocated slots and those freed by consolidate_delete. - // _delete_set has locations marked deleted by lazy_delete. Will not be - // immediately available for insert. consolidate_delete will release these - // slots to _empty_slots. - natural_number_set _empty_slots; - std::unique_ptr> _delete_set; - - bool _data_compacted = true; // true if data has been compacted - bool _is_saved = false; // Checking if the index is already saved. - bool _conc_consolidate = false; // use _lock while searching - - // Acquire locks in the order below when acquiring multiple locks - std::shared_timed_mutex // RW mutex between save/load (exclusive lock) and - _update_lock; // search/inserts/deletes/consolidate (shared lock) - std::shared_timed_mutex // Ensure only one consolidate or compact_data is - _consolidate_lock; // ever active - std::shared_timed_mutex // RW lock for _tag_to_location, - _tag_lock; // _location_to_tag, _empty_slots, _nd, _max_points - std::shared_timed_mutex // RW Lock on _delete_set and _data_compacted - _delete_lock; // variable - - // Per node lock, cardinality=_max_points - std::vector _locks; - - static const float INDEX_GROWTH_FACTOR; + private: + // Distance functions + Metric _dist_metric = diskann::L2; + std::shared_ptr> _distance; + + // Data + std::unique_ptr> _data_store; + char *_opt_graph = nullptr; + + // Graph related data structures + std::vector> _final_graph; + + // Dimensions + size_t _dim = 0; + 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 + + // _num_frozen_pts is the number of points which are used as initial + // candidates when iterating to closest point(s). These are not visible + // externally and won't be returned by search. At least 1 frozen point is + // needed for a dynamic index. The frozen points have consecutive locations. + // See also _start below. + size_t _num_frozen_pts = 0; + size_t _max_range_of_loaded_graph = 0; + size_t _node_size; + size_t _data_len; + size_t _neighbor_len; + + uint32_t _max_observed_degree = 0; + // Start point of the search. When _num_frozen_pts is greater than zero, + // this is the location of the first frozen point. Otherwise, this is a + // location of one of the points in index. + uint32_t _start = 0; + + bool _has_built = false; + bool _saturate_graph = false; + bool _save_as_one_file = false; // plan to support in next version + bool _dynamic_index = false; + bool _enable_tags = false; + bool _normalize_vecs = false; // Using normalied L2 for cosine. + + // Filter Support + + bool _filtered_index = false; + std::vector> _pts_to_labels; + tsl::robin_set _labels; + std::string _labels_file; + std::unordered_map _label_to_medoid_id; + std::unordered_map _medoid_counts; + bool _use_universal_label = false; + LabelT _universal_label = 0; + uint32_t _filterIndexingQueueSize; + std::unordered_map _label_map; + + // Indexing parameters + uint32_t _indexingQueueSize; + uint32_t _indexingRange; + uint32_t _indexingMaxC; + float _indexingAlpha; + + // Query scratch data structures + ConcurrentQueue *> _query_scratch; + + // Flags for PQ based distance calculation + bool _pq_dist = false; + bool _use_opq = false; + size_t _num_pq_chunks = 0; + uint8_t *_pq_data = nullptr; + bool _pq_generated = false; + FixedChunkPQTable _pq_table; + + // + // Data structures, locks and flags for dynamic indexing and tags + // + + // lazy_delete removes entry from _location_to_tag and _tag_to_location. If + // _location_to_tag does not resolve a location, infer that it was deleted. + tsl::sparse_map _tag_to_location; + natural_number_map _location_to_tag; + + // _empty_slots has unallocated slots and those freed by consolidate_delete. + // _delete_set has locations marked deleted by lazy_delete. Will not be + // immediately available for insert. consolidate_delete will release these + // slots to _empty_slots. + natural_number_set _empty_slots; + std::unique_ptr> _delete_set; + + bool _data_compacted = true; // true if data has been compacted + bool _is_saved = false; // Checking if the index is already saved. + bool _conc_consolidate = false; // use _lock while searching + + // Acquire locks in the order below when acquiring multiple locks + std::shared_timed_mutex // RW mutex between save/load (exclusive lock) and + _update_lock; // search/inserts/deletes/consolidate (shared lock) + std::shared_timed_mutex // Ensure only one consolidate or compact_data is + _consolidate_lock; // ever active + std::shared_timed_mutex // RW lock for _tag_to_location, + _tag_lock; // _location_to_tag, _empty_slots, _nd, _max_points + std::shared_timed_mutex // RW Lock on _delete_set and _data_compacted + _delete_lock; // variable + + // Per node lock, cardinality=_max_points + std::vector _locks; + + static const float INDEX_GROWTH_FACTOR; }; -} // namespace diskann +} // namespace diskann diff --git a/include/linux_aligned_file_reader.h b/include/linux_aligned_file_reader.h index fd78cac79..7620e3194 100644 --- a/include/linux_aligned_file_reader.h +++ b/include/linux_aligned_file_reader.h @@ -6,34 +6,34 @@ #include "aligned_file_reader.h" -class LinuxAlignedFileReader : public AlignedFileReader { - private: - uint64_t file_sz; - FileHandle file_desc; - io_context_t bad_ctx = (io_context_t)-1; - - public: - LinuxAlignedFileReader(); - ~LinuxAlignedFileReader(); - - IOContext &get_ctx(); - - // register thread-id for a context - void register_thread(); - - // 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 close(); - - // process batch of aligned requests in parallel - // NOTE :: blocking call - void read(std::vector &read_reqs, IOContext &ctx, - bool async = false); +class LinuxAlignedFileReader : public AlignedFileReader +{ + private: + uint64_t file_sz; + FileHandle file_desc; + io_context_t bad_ctx = (io_context_t)-1; + + public: + LinuxAlignedFileReader(); + ~LinuxAlignedFileReader(); + + IOContext &get_ctx(); + + // register thread-id for a context + void register_thread(); + + // 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 close(); + + // process batch of aligned requests in parallel + // NOTE :: blocking call + void read(std::vector &read_reqs, IOContext &ctx, bool async = false); }; #endif diff --git a/include/locking.h b/include/locking.h index 6f0c51477..2a70f4ffa 100644 --- a/include/locking.h +++ b/include/locking.h @@ -7,7 +7,8 @@ #include "windows_slim_lock.h" #endif -namespace diskann { +namespace diskann +{ #ifdef _WINDOWS using non_recursive_mutex = windows_exclusive_slim_lock; using LockGuard = windows_exclusive_slim_lock_guard; @@ -15,4 +16,4 @@ using LockGuard = windows_exclusive_slim_lock_guard; using non_recursive_mutex = std::mutex; using LockGuard = std::lock_guard; #endif -} // namespace diskann +} // namespace diskann diff --git a/include/logger.h b/include/logger.h index 5fe04ac0c..28a9f619c 100644 --- a/include/logger.h +++ b/include/logger.h @@ -6,14 +6,19 @@ #include #include "windows_customizations.h" -namespace diskann { +namespace diskann +{ DISKANN_DLLEXPORT extern std::basic_ostream cout; DISKANN_DLLEXPORT extern std::basic_ostream cerr; -enum class DISKANN_DLLEXPORT LogLevel { LL_Info = 0, LL_Error, LL_Count }; +enum class DISKANN_DLLEXPORT LogLevel +{ + LL_Info = 0, + LL_Error, + LL_Count +}; #ifdef EXEC_ENV_OLS -DISKANN_DLLEXPORT void SetCustomLogger( - std::function logger); +DISKANN_DLLEXPORT void SetCustomLogger(std::function logger); #endif -} // namespace diskann +} // namespace diskann diff --git a/include/logger_impl.h b/include/logger_impl.h index 4196f32bf..af6108c06 100644 --- a/include/logger_impl.h +++ b/include/logger_impl.h @@ -9,29 +9,32 @@ #include "ann_exception.h" #include "logger.h" -namespace diskann { -class ANNStreamBuf : public std::basic_streambuf { - public: - DISKANN_DLLEXPORT explicit ANNStreamBuf(FILE *fp); - DISKANN_DLLEXPORT ~ANNStreamBuf(); - - DISKANN_DLLEXPORT bool is_open() const { - return true; // because stdout and stderr are always open. - } - DISKANN_DLLEXPORT void close(); - DISKANN_DLLEXPORT virtual int underflow(); - DISKANN_DLLEXPORT virtual int overflow(int c); - DISKANN_DLLEXPORT virtual int sync(); - - private: - FILE *_fp; - char *_buf; - int _bufIndex; - std::mutex _mutex; - LogLevel _logLevel; - - int flush(); - void logImpl(char *str, int numchars); +namespace diskann +{ +class ANNStreamBuf : public std::basic_streambuf +{ + public: + DISKANN_DLLEXPORT explicit ANNStreamBuf(FILE *fp); + DISKANN_DLLEXPORT ~ANNStreamBuf(); + + DISKANN_DLLEXPORT bool is_open() const + { + return true; // because stdout and stderr are always open. + } + DISKANN_DLLEXPORT void close(); + DISKANN_DLLEXPORT virtual int underflow(); + DISKANN_DLLEXPORT virtual int overflow(int c); + DISKANN_DLLEXPORT virtual int sync(); + + private: + FILE *_fp; + char *_buf; + int _bufIndex; + std::mutex _mutex; + LogLevel _logLevel; + + int flush(); + void logImpl(char *str, int numchars); // Why the two buffer-sizes? If we are running normally, we are basically // interacting with a character output system, so we short-circuit the @@ -48,15 +51,15 @@ class ANNStreamBuf : public std::basic_streambuf { // 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 = 1024; #else - // Allocating an arbitrarily small buffer here because the overflow() and - // other function implementations push the BUFFER_SIZE chars into the - // buffer before flushing to fwrite. - static const int BUFFER_SIZE = 4; + // Allocating an arbitrarily small buffer here because the overflow() and + // other function implementations push the BUFFER_SIZE chars into the + // buffer before flushing to fwrite. + static const int BUFFER_SIZE = 4; #endif - ANNStreamBuf(const ANNStreamBuf &); - ANNStreamBuf &operator=(const ANNStreamBuf &); + ANNStreamBuf(const ANNStreamBuf &); + ANNStreamBuf &operator=(const ANNStreamBuf &); }; -} // namespace diskann +} // namespace diskann diff --git a/include/math_utils.h b/include/math_utils.h index 333c00c2c..83d189f70 100644 --- a/include/math_utils.h +++ b/include/math_utils.h @@ -6,18 +6,17 @@ #include "common_includes.h" #include "utils.h" -namespace math_utils { +namespace math_utils +{ float calc_distance(float *vec_1, float *vec_2, size_t dim); // compute l2-squared norms of data stored in row major num_points * dim, // needs // to be pre-allocated -void compute_vecs_l2sq(float *vecs_l2sq, float *data, const size_t num_points, - const size_t dim); +void compute_vecs_l2sq(float *vecs_l2sq, float *data, const size_t num_points, const size_t dim); -void rotate_data_randomly(float *data, size_t num_points, size_t dim, - float *rot_mat, float *&new_mat, +void rotate_data_randomly(float *data, size_t num_points, size_t dim, float *rot_mat, float *&new_mat, bool transpose_rot = false); // calculate closest center to data of num_points * dim (row major) @@ -29,11 +28,10 @@ void rotate_data_randomly(float *data, size_t num_points, size_t dim, // squared distances // Ideally used only by compute_closest_centers -void compute_closest_centers_in_block( - const float *const data, const size_t num_points, const size_t dim, - const float *const centers, const size_t num_centers, - const float *const docs_l2sq, const float *const centers_l2sq, - uint32_t *center_index, float *const dist_matrix, size_t k = 1); +void compute_closest_centers_in_block(const float *const data, const size_t num_points, const size_t dim, + const float *const centers, const size_t num_centers, + const float *const docs_l2sq, const float *const centers_l2sq, + uint32_t *center_index, float *const dist_matrix, size_t k = 1); // Given data in num_points * new_dim row major // Pivots stored in full_pivot_data as k * new_dim row major @@ -45,23 +43,21 @@ void compute_closest_centers_in_block( // those // values -void compute_closest_centers(float *data, size_t num_points, size_t dim, - float *pivot_data, size_t num_centers, size_t k, - uint32_t *closest_centers_ivf, - std::vector *inverted_index = NULL, +void compute_closest_centers(float *data, size_t num_points, size_t dim, float *pivot_data, size_t num_centers, + size_t k, uint32_t *closest_centers_ivf, std::vector *inverted_index = NULL, float *pts_norms_squared = NULL); // if to_subtract is 1, will subtract nearest center from each row. Else will // add. Output will be in data_load iself. // Nearest centers need to be provided in closst_centers. -void process_residuals(float *data_load, size_t num_points, size_t dim, - float *cur_pivot_data, size_t num_centers, +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); -} // namespace math_utils +} // namespace math_utils -namespace kmeans { +namespace kmeans +{ // run Lloyds one iteration // Given data in row major num_points * dim, and centers in row major @@ -71,8 +67,7 @@ namespace kmeans { // If closest_centers == NULL, will allocate memory and return. // Similarly, if closest_docs == NULL, will allocate memory and return. -float lloyds_iter(float *data, size_t num_points, size_t dim, float *centers, - size_t num_centers, float *docs_l2sq, +float lloyds_iter(float *data, size_t num_points, size_t dim, float *centers, size_t num_centers, float *docs_l2sq, std::vector *closest_docs, uint32_t *&closest_center); // Run Lloyds until max_reps or stopping criterion @@ -81,15 +76,12 @@ float lloyds_iter(float *data, size_t num_points, size_t dim, float *centers, // vector [num_centers], and closest_center = new size_t[num_points] // Final centers are output in centers as row major num_centers * dim // -float run_lloyds(float *data, size_t num_points, size_t dim, float *centers, - const size_t num_centers, const size_t max_reps, - std::vector *closest_docs, uint32_t *closest_center); +float run_lloyds(float *data, size_t num_points, size_t dim, float *centers, const size_t num_centers, + const size_t max_reps, std::vector *closest_docs, uint32_t *closest_center); // assumes already memory allocated for pivot_data as new // float[num_centers*dim] and select randomly num_centers points as pivots -void selecting_pivots(float *data, size_t num_points, size_t dim, - float *pivot_data, size_t num_centers); +void selecting_pivots(float *data, size_t num_points, size_t dim, float *pivot_data, size_t num_centers); -void kmeanspp_selecting_pivots(float *data, size_t num_points, size_t dim, - float *pivot_data, size_t num_centers); -} // namespace kmeans +void kmeanspp_selecting_pivots(float *data, size_t num_points, size_t dim, float *pivot_data, size_t num_centers); +} // namespace kmeans diff --git a/include/memory_mapper.h b/include/memory_mapper.h index bf78d1687..75faca1bb 100644 --- a/include/memory_mapper.h +++ b/include/memory_mapper.h @@ -15,27 +15,29 @@ #endif #include -namespace diskann { -class MemoryMapper { - private: +namespace diskann +{ +class MemoryMapper +{ + private: #ifndef _WINDOWS - int _fd; + int _fd; #else - HANDLE _bareFile; - HANDLE _fd; + HANDLE _bareFile; + HANDLE _fd; #endif - char *_buf; - size_t _fileSize; - const char *_fileName; + char *_buf; + size_t _fileSize; + const char *_fileName; - public: - MemoryMapper(const char *filename); - MemoryMapper(const std::string &filename); + public: + MemoryMapper(const char *filename); + MemoryMapper(const std::string &filename); - char *getBuf(); - size_t getFileSize(); + char *getBuf(); + size_t getFileSize(); - ~MemoryMapper(); + ~MemoryMapper(); }; -} // namespace diskann +} // namespace diskann diff --git a/include/natural_number_map.h b/include/natural_number_map.h index 6868e5551..820ac3fdf 100644 --- a/include/natural_number_map.h +++ b/include/natural_number_map.h @@ -9,7 +9,8 @@ #include -namespace diskann { +namespace diskann +{ // A map whose key is a natural number (from 0 onwards) and maps to a value. // Made as both memory and performance efficient map for scenario such as // DiskANN location-to-tag map. There, the pool of numbers is consecutive from @@ -21,67 +22,68 @@ namespace diskann { // Thread-safety: this class is not thread-safe in general. // Exception: multiple read-only operations are safe on the object only if // there are no writers to it in parallel. -template -class natural_number_map { - public: - static_assert(std::is_trivial::value, "Key must be a trivial type"); - // Some of the class member prototypes are done with this assumption to - // minimize verbosity since it's the only use case. - static_assert(std::is_trivial::value, "Value must be a trivial type"); - - // Represents a reference to a element in the map. Used while iterating - // over map entries. - struct position { - size_t _key; - // The number of keys that were enumerated when iterating through the - // map so far. Used to early-terminate enumeration when ithere are no - // more entries in the map. - size_t _keys_already_enumerated; - - // Returns whether it's valid to access the element at this position in - // the map. - bool is_valid() const; - }; - - natural_number_map(); - - void reserve(size_t count); - size_t size() const; - - void set(Key key, Value value); - void erase(Key key); - - bool contains(Key key) const; - bool try_get(Key key, Value &value) const; - - // Returns the value at the specified position. Prerequisite: position is - // valid. - Value get(const position &pos) const; - - // Finds the first element in the map, if any. Invalidated by changes in the - // map. - position find_first() const; - - // Finds the next element in the map after the specified position. - // Invalidated by changes in the map. - position find_next(const position &after_position) const; - - void clear(); - - private: - // Number of entries in the map. Not the same as size() of the - // _values_vector below. - size_t _size; - - // Array of values. The key is the index of the value. - std::vector _values_vector; - - // Values that are in the set have the corresponding bit index set - // to 1. - // - // Use a pointer here to allow for forward declaration of dynamic_bitset - // in public headers to avoid making boost a dependency for clients - // of DiskANN. - std::unique_ptr> _values_bitset; +template class natural_number_map +{ + public: + static_assert(std::is_trivial::value, "Key must be a trivial type"); + // Some of the class member prototypes are done with this assumption to + // minimize verbosity since it's the only use case. + static_assert(std::is_trivial::value, "Value must be a trivial type"); + + // Represents a reference to a element in the map. Used while iterating + // over map entries. + struct position + { + size_t _key; + // The number of keys that were enumerated when iterating through the + // map so far. Used to early-terminate enumeration when ithere are no + // more entries in the map. + size_t _keys_already_enumerated; + + // Returns whether it's valid to access the element at this position in + // the map. + bool is_valid() const; + }; + + natural_number_map(); + + void reserve(size_t count); + size_t size() const; + + void set(Key key, Value value); + void erase(Key key); + + bool contains(Key key) const; + bool try_get(Key key, Value &value) const; + + // Returns the value at the specified position. Prerequisite: position is + // valid. + Value get(const position &pos) const; + + // Finds the first element in the map, if any. Invalidated by changes in the + // map. + position find_first() const; + + // Finds the next element in the map after the specified position. + // Invalidated by changes in the map. + position find_next(const position &after_position) const; + + void clear(); + + private: + // Number of entries in the map. Not the same as size() of the + // _values_vector below. + size_t _size; + + // Array of values. The key is the index of the value. + std::vector _values_vector; + + // Values that are in the set have the corresponding bit index set + // to 1. + // + // Use a pointer here to allow for forward declaration of dynamic_bitset + // in public headers to avoid making boost a dependency for clients + // of DiskANN. + std::unique_ptr> _values_bitset; }; -} // namespace diskann +} // namespace diskann diff --git a/include/natural_number_set.h b/include/natural_number_set.h index eda3de5c3..ec5b827e6 100644 --- a/include/natural_number_set.h +++ b/include/natural_number_set.h @@ -8,7 +8,8 @@ #include "boost_dynamic_bitset_fwd.h" -namespace diskann { +namespace diskann +{ // A set of natural numbers (from 0 onwards). Made for scenario where the // pool of numbers is consecutive from zero to some max value and very // efficient methods for "add to set", "get any value from set", "is in set" @@ -19,31 +20,31 @@ namespace diskann { // Thread-safety: this class is not thread-safe in general. // Exception: multiple read-only operations (e.g. is_in_set, empty, size) are // safe on the object only if there are no writers to it in parallel. -template -class natural_number_set { - public: - static_assert(std::is_trivial::value, "Identifier must be a trivial type"); - - natural_number_set(); - - bool is_empty() const; - void reserve(size_t count); - void insert(T id); - T pop_any(); - void clear(); - size_t size() const; - bool is_in_set(T id) const; - - private: - // Values that are currently in set. - std::vector _values_vector; - - // Values that are in the set have the corresponding bit index set - // to 1. - // - // Use a pointer here to allow for forward declaration of dynamic_bitset - // in public headers to avoid making boost a dependency for clients - // of DiskANN. - std::unique_ptr> _values_bitset; +template class natural_number_set +{ + public: + static_assert(std::is_trivial::value, "Identifier must be a trivial type"); + + natural_number_set(); + + bool is_empty() const; + void reserve(size_t count); + void insert(T id); + T pop_any(); + void clear(); + size_t size() const; + bool is_in_set(T id) const; + + private: + // Values that are currently in set. + std::vector _values_vector; + + // Values that are in the set have the corresponding bit index set + // to 1. + // + // Use a pointer here to allow for forward declaration of dynamic_bitset + // in public headers to avoid making boost a dependency for clients + // of DiskANN. + std::unique_ptr> _values_bitset; }; -} // namespace diskann +} // namespace diskann diff --git a/include/neighbor.h b/include/neighbor.h index b6b66eb97..d7c0c25ed 100644 --- a/include/neighbor.h +++ b/include/neighbor.h @@ -8,106 +8,145 @@ #include #include "utils.h" -namespace diskann { +namespace diskann +{ -struct Neighbor { - unsigned id; - float distance; - bool expanded; +struct Neighbor +{ + unsigned id; + float distance; + bool expanded; - Neighbor() = default; + Neighbor() = default; - Neighbor(unsigned id, float distance) - : id{id}, distance{distance}, expanded(false) {} + Neighbor(unsigned id, float distance) : id{id}, distance{distance}, expanded(false) + { + } - inline bool operator<(const Neighbor &other) const { - return distance < other.distance || - (distance == other.distance && id < other.id); - } + inline bool operator<(const Neighbor &other) const + { + return distance < other.distance || (distance == other.distance && id < other.id); + } - inline bool operator==(const Neighbor &other) const { - return (id == other.id); - } + inline bool operator==(const Neighbor &other) const + { + return (id == other.id); + } }; // Invariant: after every `insert` and `closest_unexpanded()`, `_cur` points to // the first Neighbor which is unexpanded. -class NeighborPriorityQueue { - public: - NeighborPriorityQueue() : _size(0), _capacity(0), _cur(0) {} - - explicit NeighborPriorityQueue(size_t capacity) - : _size(0), _capacity(capacity), _cur(0), _data(capacity + 1) {} - - // Inserts the item ordered into the set up to the sets capacity. - // The item will be dropped if it is the same id as an exiting - // set item or it has a greated distance than the final - // item in the set. The set cursor that is used to pop() the - // next item will be set to the lowest index of an uncheck item - void insert(const Neighbor &nbr) { - if (_size == _capacity && _data[_size - 1] < nbr) { - return; +class NeighborPriorityQueue +{ + public: + NeighborPriorityQueue() : _size(0), _capacity(0), _cur(0) + { } - size_t lo = 0, hi = _size; - while (lo < hi) { - size_t mid = (lo + hi) >> 1; - if (nbr < _data[mid]) { - hi = mid; - // Make sure the same id isn't inserted into the set - } else if (_data[mid].id == nbr.id) { - return; - } else { - lo = mid + 1; - } + explicit NeighborPriorityQueue(size_t capacity) : _size(0), _capacity(capacity), _cur(0), _data(capacity + 1) + { } - if (lo < _capacity) { - std::memmove(&_data[lo + 1], &_data[lo], (_size - lo) * sizeof(Neighbor)); + // Inserts the item ordered into the set up to the sets capacity. + // The item will be dropped if it is the same id as an exiting + // set item or it has a greated distance than the final + // item in the set. The set cursor that is used to pop() the + // next item will be set to the lowest index of an uncheck item + void insert(const Neighbor &nbr) + { + if (_size == _capacity && _data[_size - 1] < nbr) + { + return; + } + + size_t lo = 0, hi = _size; + while (lo < hi) + { + size_t mid = (lo + hi) >> 1; + if (nbr < _data[mid]) + { + hi = mid; + // Make sure the same id isn't inserted into the set + } + else if (_data[mid].id == nbr.id) + { + return; + } + else + { + lo = mid + 1; + } + } + + if (lo < _capacity) + { + std::memmove(&_data[lo + 1], &_data[lo], (_size - lo) * sizeof(Neighbor)); + } + _data[lo] = {nbr.id, nbr.distance}; + if (_size < _capacity) + { + _size++; + } + if (lo < _cur) + { + _cur = lo; + } } - _data[lo] = {nbr.id, nbr.distance}; - if (_size < _capacity) { - _size++; - } - if (lo < _cur) { - _cur = lo; - } - } - Neighbor closest_unexpanded() { - _data[_cur].expanded = true; - size_t pre = _cur; - while (_cur < _size && _data[_cur].expanded) { - _cur++; + Neighbor closest_unexpanded() + { + _data[_cur].expanded = true; + size_t pre = _cur; + while (_cur < _size && _data[_cur].expanded) + { + _cur++; + } + return _data[pre]; } - return _data[pre]; - } - bool has_unexpanded_node() const { return _cur < _size; } + bool has_unexpanded_node() const + { + return _cur < _size; + } - size_t size() const { return _size; } + size_t size() const + { + return _size; + } - size_t capacity() const { return _capacity; } + size_t capacity() const + { + return _capacity; + } - void reserve(size_t capacity) { - if (capacity + 1 > _data.size()) { - _data.resize(capacity + 1); + void reserve(size_t capacity) + { + if (capacity + 1 > _data.size()) + { + _data.resize(capacity + 1); + } + _capacity = capacity; } - _capacity = capacity; - } - Neighbor &operator[](size_t i) { return _data[i]; } + Neighbor &operator[](size_t i) + { + return _data[i]; + } - Neighbor operator[](size_t i) const { return _data[i]; } + Neighbor operator[](size_t i) const + { + return _data[i]; + } - void clear() { - _size = 0; - _cur = 0; - } + void clear() + { + _size = 0; + _cur = 0; + } - private: - size_t _size, _capacity, _cur; - std::vector _data; + private: + size_t _size, _capacity, _cur; + std::vector _data; }; -} // namespace diskann +} // namespace diskann diff --git a/include/parameters.h b/include/parameters.h index 1a42c29dd..36f8a94e0 100644 --- a/include/parameters.h +++ b/include/parameters.h @@ -8,121 +8,117 @@ #include "defaults.h" -namespace diskann { +namespace diskann +{ class IndexWriteParameters { - public: - const uint32_t search_list_size; // L - const uint32_t max_degree; // R - const bool saturate_graph; - const uint32_t max_occlusion_size; // C - const float alpha; - const uint32_t num_rounds; - const uint32_t num_threads; - const uint32_t filter_list_size; // Lf - const uint32_t num_frozen_points; - - private: - IndexWriteParameters(const uint32_t search_list_size, - const uint32_t max_degree, const bool saturate_graph, - const uint32_t max_occlusion_size, const float alpha, - const uint32_t num_rounds, const uint32_t num_threads, - const uint32_t filter_list_size, - const uint32_t num_frozen_points) - : search_list_size(search_list_size), - max_degree(max_degree), - saturate_graph(saturate_graph), - max_occlusion_size(max_occlusion_size), - alpha(alpha), - num_rounds(num_rounds), - num_threads(num_threads), - filter_list_size(filter_list_size), - num_frozen_points(num_frozen_points) {} - - friend class IndexWriteParametersBuilder; + public: + const uint32_t search_list_size; // L + const uint32_t max_degree; // R + const bool saturate_graph; + const uint32_t max_occlusion_size; // C + const float alpha; + const uint32_t num_rounds; + const uint32_t num_threads; + const uint32_t filter_list_size; // Lf + const uint32_t num_frozen_points; + + private: + IndexWriteParameters(const uint32_t search_list_size, const uint32_t max_degree, const bool saturate_graph, + const uint32_t max_occlusion_size, const float alpha, const uint32_t num_rounds, + const uint32_t num_threads, const uint32_t filter_list_size, const uint32_t num_frozen_points) + : search_list_size(search_list_size), max_degree(max_degree), saturate_graph(saturate_graph), + max_occlusion_size(max_occlusion_size), alpha(alpha), num_rounds(num_rounds), num_threads(num_threads), + filter_list_size(filter_list_size), num_frozen_points(num_frozen_points) + { + } + + friend class IndexWriteParametersBuilder; }; -class IndexWriteParametersBuilder { - /** - * Fluent builder pattern to keep track of the 7 non-default properties - * and their order. The basic ctor was getting unwieldy. - */ - public: - IndexWriteParametersBuilder(const uint32_t search_list_size, // L - const uint32_t max_degree // R - ) - : _search_list_size(search_list_size), _max_degree(max_degree) {} - - IndexWriteParametersBuilder &with_max_occlusion_size( - const uint32_t max_occlusion_size) { - _max_occlusion_size = max_occlusion_size; - return *this; - } - - IndexWriteParametersBuilder &with_saturate_graph(const bool saturate_graph) { - _saturate_graph = saturate_graph; - return *this; - } - - IndexWriteParametersBuilder &with_alpha(const float alpha) { - _alpha = alpha; - return *this; - } - - IndexWriteParametersBuilder &with_num_rounds(const uint32_t num_rounds) { - _num_rounds = num_rounds; - return *this; - } - - IndexWriteParametersBuilder &with_num_threads(const uint32_t num_threads) { - _num_threads = num_threads; - return *this; - } - - IndexWriteParametersBuilder &with_filter_list_size( - const uint32_t filter_list_size) { - _filter_list_size = filter_list_size; - return *this; - } - - IndexWriteParametersBuilder &with_num_frozen_points( - const uint32_t num_frozen_points) { - _num_frozen_points = num_frozen_points; - return *this; - } - - IndexWriteParameters build() const { - return IndexWriteParameters(_search_list_size, _max_degree, _saturate_graph, - _max_occlusion_size, _alpha, _num_rounds, - _num_threads, _filter_list_size, - _num_frozen_points); - } - - IndexWriteParametersBuilder(const IndexWriteParameters &wp) - : _search_list_size(wp.search_list_size), - _max_degree(wp.max_degree), - _max_occlusion_size(wp.max_occlusion_size), - _saturate_graph(wp.saturate_graph), - _alpha(wp.alpha), - _num_rounds(wp.num_rounds), - _filter_list_size(wp.filter_list_size), - _num_frozen_points(wp.num_frozen_points) {} - IndexWriteParametersBuilder(const IndexWriteParametersBuilder &) = delete; - IndexWriteParametersBuilder &operator=(const IndexWriteParametersBuilder &) = - delete; - - private: - uint32_t _search_list_size{}; - uint32_t _max_degree{}; - uint32_t _max_occlusion_size{defaults::MAX_OCCLUSION_SIZE}; - bool _saturate_graph{defaults::SATURATE_GRAPH}; - float _alpha{defaults::ALPHA}; - uint32_t _num_rounds{defaults::NUM_ROUNDS}; - uint32_t _num_threads{defaults::NUM_THREADS}; - uint32_t _filter_list_size{defaults::FILTER_LIST_SIZE}; - uint32_t _num_frozen_points{defaults::NUM_FROZEN_POINTS}; +class IndexWriteParametersBuilder +{ + /** + * Fluent builder pattern to keep track of the 7 non-default properties + * and their order. The basic ctor was getting unwieldy. + */ + public: + IndexWriteParametersBuilder(const uint32_t search_list_size, // L + const uint32_t max_degree // R + ) + : _search_list_size(search_list_size), _max_degree(max_degree) + { + } + + IndexWriteParametersBuilder &with_max_occlusion_size(const uint32_t max_occlusion_size) + { + _max_occlusion_size = max_occlusion_size; + return *this; + } + + IndexWriteParametersBuilder &with_saturate_graph(const bool saturate_graph) + { + _saturate_graph = saturate_graph; + return *this; + } + + IndexWriteParametersBuilder &with_alpha(const float alpha) + { + _alpha = alpha; + return *this; + } + + IndexWriteParametersBuilder &with_num_rounds(const uint32_t num_rounds) + { + _num_rounds = num_rounds; + return *this; + } + + IndexWriteParametersBuilder &with_num_threads(const uint32_t num_threads) + { + _num_threads = num_threads; + return *this; + } + + IndexWriteParametersBuilder &with_filter_list_size(const uint32_t filter_list_size) + { + _filter_list_size = filter_list_size; + return *this; + } + + IndexWriteParametersBuilder &with_num_frozen_points(const uint32_t num_frozen_points) + { + _num_frozen_points = num_frozen_points; + return *this; + } + + IndexWriteParameters build() const + { + return IndexWriteParameters(_search_list_size, _max_degree, _saturate_graph, _max_occlusion_size, _alpha, + _num_rounds, _num_threads, _filter_list_size, _num_frozen_points); + } + + IndexWriteParametersBuilder(const IndexWriteParameters &wp) + : _search_list_size(wp.search_list_size), _max_degree(wp.max_degree), + _max_occlusion_size(wp.max_occlusion_size), _saturate_graph(wp.saturate_graph), _alpha(wp.alpha), + _num_rounds(wp.num_rounds), _filter_list_size(wp.filter_list_size), _num_frozen_points(wp.num_frozen_points) + { + } + IndexWriteParametersBuilder(const IndexWriteParametersBuilder &) = delete; + IndexWriteParametersBuilder &operator=(const IndexWriteParametersBuilder &) = delete; + + private: + uint32_t _search_list_size{}; + uint32_t _max_degree{}; + uint32_t _max_occlusion_size{defaults::MAX_OCCLUSION_SIZE}; + bool _saturate_graph{defaults::SATURATE_GRAPH}; + float _alpha{defaults::ALPHA}; + uint32_t _num_rounds{defaults::NUM_ROUNDS}; + uint32_t _num_threads{defaults::NUM_THREADS}; + uint32_t _filter_list_size{defaults::FILTER_LIST_SIZE}; + uint32_t _num_frozen_points{defaults::NUM_FROZEN_POINTS}; }; -} // namespace diskann +} // namespace diskann diff --git a/include/partition.h b/include/partition.h index a5af890be..c2c4c76ad 100644 --- a/include/partition.h +++ b/include/partition.h @@ -16,45 +16,34 @@ #include "windows_customizations.h" template -void gen_random_slice(const std::string base_file, - const std::string output_prefix, double sampling_rate); +void gen_random_slice(const std::string base_file, const std::string output_prefix, double sampling_rate); template -void gen_random_slice(const std::string data_file, double p_val, - float *&sampled_data, size_t &slice_size, size_t &ndims); +void gen_random_slice(const std::string data_file, double p_val, float *&sampled_data, size_t &slice_size, + size_t &ndims); template -void gen_random_slice(const T *inputdata, size_t npts, size_t ndims, - double p_val, float *&sampled_data, size_t &slice_size); +void gen_random_slice(const T *inputdata, size_t npts, size_t ndims, double p_val, float *&sampled_data, + size_t &slice_size); -int estimate_cluster_sizes(float *test_data_float, size_t num_test, - float *pivots, const size_t num_centers, - const size_t dim, const size_t k_base, - std::vector &cluster_sizes); +int estimate_cluster_sizes(float *test_data_float, size_t num_test, float *pivots, const size_t num_centers, + const size_t dim, const size_t k_base, std::vector &cluster_sizes); template -int shard_data_into_clusters(const std::string data_file, float *pivots, - const size_t num_centers, const size_t dim, +int shard_data_into_clusters(const std::string data_file, float *pivots, const size_t num_centers, const size_t dim, const size_t k_base, std::string prefix_path); template -int shard_data_into_clusters_only_ids(const std::string data_file, - float *pivots, const size_t num_centers, - const size_t dim, const size_t k_base, - std::string prefix_path); +int shard_data_into_clusters_only_ids(const std::string data_file, float *pivots, const size_t num_centers, + const size_t dim, const size_t k_base, std::string prefix_path); template -int retrieve_shard_data_from_ids(const std::string data_file, - std::string idmap_filename, - std::string data_filename); +int retrieve_shard_data_from_ids(const std::string data_file, std::string idmap_filename, std::string data_filename); template -int partition(const std::string data_file, const float sampling_rate, - size_t num_centers, size_t max_k_means_reps, +int partition(const std::string data_file, const float sampling_rate, size_t num_centers, size_t max_k_means_reps, const std::string prefix_path, size_t k_base); template -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); +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); diff --git a/include/percentile_stats.h b/include/percentile_stats.h index c7a70fe16..793257577 100644 --- a/include/percentile_stats.h +++ b/include/percentile_stats.h @@ -16,48 +16,50 @@ #include "distance.h" #include "parameters.h" -namespace diskann { -struct QueryStats { - float total_us = 0; // total time to process query in micros - float io_us = 0; // total time spent in IO - float cpu_us = 0; // total time spent in CPU - - unsigned n_4k = 0; // # of 4kB reads - unsigned n_8k = 0; // # of 8kB reads - unsigned n_12k = 0; // # of 12kB reads - unsigned n_ios = 0; // total # of IOs issued - unsigned read_size = 0; // total # of bytes read - unsigned n_cmps_saved = 0; // # cmps saved - unsigned n_cmps = 0; // # cmps - unsigned n_cache_hits = 0; // # cache_hits - unsigned n_hops = 0; // # search hops +namespace diskann +{ +struct QueryStats +{ + float total_us = 0; // total time to process query in micros + float io_us = 0; // total time spent in IO + float cpu_us = 0; // total time spent in CPU + + unsigned n_4k = 0; // # of 4kB reads + unsigned n_8k = 0; // # of 8kB reads + unsigned n_12k = 0; // # of 12kB reads + unsigned n_ios = 0; // total # of IOs issued + unsigned read_size = 0; // total # of bytes read + unsigned n_cmps_saved = 0; // # cmps saved + unsigned n_cmps = 0; // # cmps + unsigned n_cache_hits = 0; // # cache_hits + unsigned n_hops = 0; // # search hops }; template -inline T get_percentile_stats( - QueryStats *stats, uint64_t len, float percentile, - const std::function &member_fn) { - std::vector vals(len); - for (uint64_t i = 0; i < len; i++) { - vals[i] = member_fn(stats[i]); - } - - std::sort(vals.begin(), vals.end(), - [](const T &left, const T &right) { return left < right; }); - - auto retval = vals[(uint64_t)(percentile * len)]; - vals.clear(); - return retval; +inline T get_percentile_stats(QueryStats *stats, uint64_t len, float percentile, + const std::function &member_fn) +{ + std::vector vals(len); + for (uint64_t i = 0; i < len; i++) + { + vals[i] = member_fn(stats[i]); + } + + std::sort(vals.begin(), vals.end(), [](const T &left, const T &right) { return left < right; }); + + auto retval = vals[(uint64_t)(percentile * len)]; + vals.clear(); + return retval; } template -inline double get_mean_stats( - QueryStats *stats, uint64_t len, - const std::function &member_fn) { - double avg = 0; - for (uint64_t i = 0; i < len; i++) { - avg += (double)member_fn(stats[i]); - } - return avg / len; +inline double get_mean_stats(QueryStats *stats, uint64_t len, const std::function &member_fn) +{ + double avg = 0; + for (uint64_t i = 0; i < len; i++) + { + avg += (double)member_fn(stats[i]); + } + return avg / len; } -} // namespace diskann +} // namespace diskann diff --git a/include/pq.h b/include/pq.h index f26f566aa..acfa1b30a 100644 --- a/include/pq.h +++ b/include/pq.h @@ -12,134 +12,114 @@ #define MAX_PQ_TRAINING_SET_SIZE 256000 #define MAX_PQ_CHUNKS 512 -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(); +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); + 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); + void load_pq_centroid_bin(const char *pq_table_file, size_t num_chunks); #endif - uint32_t get_num_chunks(); + uint32_t get_num_chunks(); - void preprocess_query(float *query_vec); + void preprocess_query(float *query_vec); - // assumes pre-processed query - void populate_chunk_distances(const float *query_vec, float *dist_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 l2_distance(const float *query_vec, uint8_t *base_vec); - float inner_product(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); + // 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 populate_chunk_inner_products(const float *query_vec, float *dist_vec); }; -template -struct PQScratch { - float *aligned_pqtable_dist_scratch = - nullptr; // MUST BE AT LEAST [256 * NCHUNKS] - float *aligned_dist_scratch = nullptr; // MUST BE AT LEAST diskann MAX_DEGREE - uint8_t *aligned_pq_coord_scratch = - nullptr; // MUST BE AT LEAST [N_CHUNKS * MAX_DEGREE] - float *rotated_query = nullptr; - float *aligned_query_float = nullptr; - - PQScratch(size_t graph_degree, size_t aligned_dim) { - diskann::alloc_aligned( - (void **)&aligned_pq_coord_scratch, - (size_t)graph_degree * (size_t)MAX_PQ_CHUNKS * sizeof(uint8_t), 256); - diskann::alloc_aligned((void **)&aligned_pqtable_dist_scratch, - 256 * (size_t)MAX_PQ_CHUNKS * sizeof(float), 256); - diskann::alloc_aligned((void **)&aligned_dist_scratch, - (size_t)graph_degree * sizeof(float), 256); - diskann::alloc_aligned((void **)&aligned_query_float, - aligned_dim * sizeof(float), 8 * sizeof(float)); - diskann::alloc_aligned((void **)&rotated_query, aligned_dim * sizeof(float), - 8 * sizeof(float)); - - memset(aligned_query_float, 0, aligned_dim * sizeof(float)); - memset(rotated_query, 0, aligned_dim * sizeof(float)); - } - - void set(size_t dim, T *query, const float norm = 1.0f) { - for (size_t d = 0; d < dim; ++d) { - if (norm != 1.0f) - rotated_query[d] = aligned_query_float[d] = - static_cast(query[d]) / norm; - else - rotated_query[d] = aligned_query_float[d] = - static_cast(query[d]); +template struct PQScratch +{ + float *aligned_pqtable_dist_scratch = nullptr; // MUST BE AT LEAST [256 * NCHUNKS] + float *aligned_dist_scratch = nullptr; // MUST BE AT LEAST diskann MAX_DEGREE + uint8_t *aligned_pq_coord_scratch = nullptr; // MUST BE AT LEAST [N_CHUNKS * MAX_DEGREE] + float *rotated_query = nullptr; + float *aligned_query_float = nullptr; + + PQScratch(size_t graph_degree, size_t aligned_dim) + { + diskann::alloc_aligned((void **)&aligned_pq_coord_scratch, + (size_t)graph_degree * (size_t)MAX_PQ_CHUNKS * sizeof(uint8_t), 256); + diskann::alloc_aligned((void **)&aligned_pqtable_dist_scratch, 256 * (size_t)MAX_PQ_CHUNKS * sizeof(float), + 256); + diskann::alloc_aligned((void **)&aligned_dist_scratch, (size_t)graph_degree * sizeof(float), 256); + diskann::alloc_aligned((void **)&aligned_query_float, aligned_dim * sizeof(float), 8 * sizeof(float)); + diskann::alloc_aligned((void **)&rotated_query, aligned_dim * sizeof(float), 8 * sizeof(float)); + + memset(aligned_query_float, 0, aligned_dim * sizeof(float)); + memset(rotated_query, 0, aligned_dim * sizeof(float)); + } + + void set(size_t dim, T *query, const float norm = 1.0f) + { + for (size_t d = 0; d < dim; ++d) + { + if (norm != 1.0f) + rotated_query[d] = aligned_query_float[d] = static_cast(query[d]) / norm; + else + rotated_query[d] = aligned_query_float[d] = static_cast(query[d]); + } } - } }; -void aggregate_coords(const std::vector &ids, - const uint8_t *all_coords, const uint64_t ndims, - uint8_t *out); +void aggregate_coords(const std::vector &ids, const uint8_t *all_coords, const uint64_t ndims, uint8_t *out); -void pq_dist_lookup(const uint8_t *pq_ids, const size_t n_pts, - const size_t pq_nchunks, const float *pq_dists, +void pq_dist_lookup(const uint8_t *pq_ids, const size_t n_pts, const size_t pq_nchunks, const float *pq_dists, std::vector &dists_out); // Need to replace calls to these with calls to vector& based functions above -void aggregate_coords(const unsigned *ids, const uint64_t n_ids, - const uint8_t *all_coords, const uint64_t ndims, +void aggregate_coords(const unsigned *ids, const uint64_t n_ids, const uint8_t *all_coords, const uint64_t ndims, uint8_t *out); -void pq_dist_lookup(const uint8_t *pq_ids, const size_t n_pts, - const size_t pq_nchunks, const float *pq_dists, +void pq_dist_lookup(const uint8_t *pq_ids, const size_t n_pts, const size_t pq_nchunks, const float *pq_dists, float *dists_out); -DISKANN_DLLEXPORT int generate_pq_pivots( - const float *const 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, bool make_zero_mean = false); +DISKANN_DLLEXPORT int generate_pq_pivots(const float *const 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, bool make_zero_mean = false); -DISKANN_DLLEXPORT int generate_opq_pivots(const float *train_data, - size_t num_train, unsigned dim, - unsigned num_centers, - unsigned num_pq_chunks, - std::string opq_pivots_path, +DISKANN_DLLEXPORT int generate_opq_pivots(const float *train_data, size_t num_train, unsigned dim, unsigned num_centers, + unsigned num_pq_chunks, std::string opq_pivots_path, bool make_zero_mean = false); template -int generate_pq_data_from_pivots(const std::string &data_file, - unsigned num_centers, unsigned num_pq_chunks, - const std::string &pq_pivots_path, - const std::string &pq_compressed_vectors_path, +int generate_pq_data_from_pivots(const std::string &data_file, unsigned num_centers, unsigned num_pq_chunks, + const std::string &pq_pivots_path, const std::string &pq_compressed_vectors_path, bool use_opq = false); template -void generate_disk_quantized_data( - const std::string &data_file_to_use, const std::string &disk_pq_pivots_path, - const std::string &disk_pq_compressed_vectors_path, - const diskann::Metric compareMetric, const double p_val, - size_t &disk_pq_dims); +void generate_disk_quantized_data(const std::string &data_file_to_use, const std::string &disk_pq_pivots_path, + const std::string &disk_pq_compressed_vectors_path, + const diskann::Metric compareMetric, const double p_val, size_t &disk_pq_dims); template -void generate_quantized_data(const std::string &data_file_to_use, - const std::string &pq_pivots_path, - const std::string &pq_compressed_vectors_path, - const diskann::Metric compareMetric, - const double p_val, const uint64_t num_pq_chunks, - const bool use_opq, +void generate_quantized_data(const std::string &data_file_to_use, const std::string &pq_pivots_path, + const std::string &pq_compressed_vectors_path, const diskann::Metric compareMetric, + const double p_val, const uint64_t num_pq_chunks, const bool use_opq, const std::string &codebook_prefix = ""); -} // namespace diskann +} // namespace diskann diff --git a/include/pq_flash_index.h b/include/pq_flash_index.h index df5dc5a99..e11c0dc31 100644 --- a/include/pq_flash_index.h +++ b/include/pq_flash_index.h @@ -18,209 +18,191 @@ #define FULL_PRECISION_REORDER_MULTIPLIER 3 -namespace diskann { +namespace diskann +{ -template -class PQFlashIndex { - public: - DISKANN_DLLEXPORT PQFlashIndex(std::shared_ptr &fileReader, - diskann::Metric metric = diskann::Metric::L2); - DISKANN_DLLEXPORT ~PQFlashIndex(); +template class PQFlashIndex +{ + public: + DISKANN_DLLEXPORT PQFlashIndex(std::shared_ptr &fileReader, + diskann::Metric metric = diskann::Metric::L2); + DISKANN_DLLEXPORT ~PQFlashIndex(); #ifdef EXEC_ENV_OLS - DISKANN_DLLEXPORT int load(diskann::MemoryMappedFiles &files, - uint32_t num_threads, const char *index_prefix); + DISKANN_DLLEXPORT int load(diskann::MemoryMappedFiles &files, uint32_t num_threads, const char *index_prefix); #else - // load compressed data, and obtains the handle to the disk-resident index - DISKANN_DLLEXPORT int load(uint32_t num_threads, const char *index_prefix); + // load compressed data, and obtains the handle to the disk-resident index + DISKANN_DLLEXPORT int load(uint32_t num_threads, const char *index_prefix); #endif #ifdef EXEC_ENV_OLS - DISKANN_DLLEXPORT int load_from_separate_paths( - diskann::MemoryMappedFiles &files, uint32_t num_threads, - const char *index_filepath, const char *pivots_filepath, - const char *compressed_filepath); + DISKANN_DLLEXPORT int load_from_separate_paths(diskann::MemoryMappedFiles &files, uint32_t num_threads, + const char *index_filepath, const char *pivots_filepath, + const char *compressed_filepath); #else - DISKANN_DLLEXPORT int load_from_separate_paths( - uint32_t num_threads, const char *index_filepath, - const char *pivots_filepath, const char *compressed_filepath); + DISKANN_DLLEXPORT int load_from_separate_paths(uint32_t num_threads, const char *index_filepath, + const char *pivots_filepath, const char *compressed_filepath); #endif - DISKANN_DLLEXPORT void load_cache_list(std::vector &node_list); + DISKANN_DLLEXPORT void load_cache_list(std::vector &node_list); #ifdef EXEC_ENV_OLS - DISKANN_DLLEXPORT void generate_cache_list_from_sample_queries( - MemoryMappedFiles &files, std::string sample_bin, uint64_t l_search, - uint64_t beamwidth, uint64_t num_nodes_to_cache, uint32_t nthreads, - std::vector &node_list); + DISKANN_DLLEXPORT void generate_cache_list_from_sample_queries(MemoryMappedFiles &files, std::string sample_bin, + uint64_t l_search, uint64_t beamwidth, + uint64_t num_nodes_to_cache, uint32_t nthreads, + std::vector &node_list); #else - DISKANN_DLLEXPORT void generate_cache_list_from_sample_queries( - std::string sample_bin, uint64_t l_search, uint64_t beamwidth, - uint64_t num_nodes_to_cache, uint32_t num_threads, - std::vector &node_list); + DISKANN_DLLEXPORT void generate_cache_list_from_sample_queries(std::string sample_bin, uint64_t l_search, + uint64_t beamwidth, uint64_t num_nodes_to_cache, + uint32_t num_threads, + std::vector &node_list); #endif - DISKANN_DLLEXPORT void cache_bfs_levels(uint64_t num_nodes_to_cache, - std::vector &node_list, - const bool shuffle = false); - - DISKANN_DLLEXPORT void cached_beam_search( - const T *query, const uint64_t k_search, const uint64_t l_search, - uint64_t *res_ids, float *res_dists, const uint64_t beam_width, - const bool use_reorder_data = false, QueryStats *stats = nullptr); - - DISKANN_DLLEXPORT void cached_beam_search( - const T *query, const uint64_t k_search, const uint64_t l_search, - uint64_t *res_ids, float *res_dists, const uint64_t beam_width, - const bool use_filter, const LabelT &filter_label, - const bool use_reorder_data = false, QueryStats *stats = nullptr); - - DISKANN_DLLEXPORT void cached_beam_search( - const T *query, const uint64_t k_search, const uint64_t l_search, - uint64_t *res_ids, float *res_dists, const uint64_t beam_width, - const uint32_t io_limit, const bool use_reorder_data = false, - QueryStats *stats = nullptr); - - DISKANN_DLLEXPORT void cached_beam_search( - const T *query, const uint64_t k_search, const uint64_t l_search, - uint64_t *res_ids, float *res_dists, const uint64_t beam_width, - const bool use_filter, const LabelT &filter_label, - const uint32_t io_limit, const bool use_reorder_data = false, - QueryStats *stats = nullptr); - - DISKANN_DLLEXPORT LabelT get_converted_label(const std::string &filter_label); - - DISKANN_DLLEXPORT uint32_t range_search(const T *query1, const double range, - const uint64_t min_l_search, - const uint64_t max_l_search, - std::vector &indices, - std::vector &distances, - const uint64_t min_beam_width, - QueryStats *stats = nullptr); - - DISKANN_DLLEXPORT uint64_t get_data_dim(); - - std::shared_ptr &reader; - - DISKANN_DLLEXPORT diskann::Metric get_metric(); - - protected: - DISKANN_DLLEXPORT void use_medoids_data_as_centroids(); - DISKANN_DLLEXPORT void setup_thread_data(uint64_t nthreads, - uint64_t visited_reserve = 4096); - - DISKANN_DLLEXPORT void set_universal_label(const LabelT &label); - - private: - DISKANN_DLLEXPORT inline bool point_has_label(uint32_t point_id, - uint32_t label_id); - std::unordered_map load_label_map( - const std::string &map_file); - DISKANN_DLLEXPORT void parse_label_file(const std::string &map_file, - size_t &num_pts_labels); - DISKANN_DLLEXPORT void get_label_file_metadata(std::string map_file, - uint32_t &num_pts, - uint32_t &num_total_labels); - DISKANN_DLLEXPORT inline int32_t get_filter_number( - const LabelT &filter_label); - - // index info - // nhood of node `i` is in sector: [i / nnodes_per_sector] - // offset in sector: [(i % nnodes_per_sector) * max_node_len] - // nnbrs of node `i`: *(unsigned*) (buf) - // nbrs of node `i`: ((unsigned*)buf) + 1 - - uint64_t max_node_len = 0, nnodes_per_sector = 0, max_degree = 0; - - // Data used for searching with re-order vectors - uint64_t ndims_reorder_vecs = 0, reorder_data_start_sector = 0, - nvecs_per_sector = 0; - - diskann::Metric metric = diskann::Metric::L2; - - // used only for inner product search to re-scale the result value - // (due to the pre-processing of base during index build) - float max_base_norm = 0.0f; - - // data info - uint64_t num_points = 0; - uint64_t num_frozen_points = 0; - uint64_t frozen_location = 0; - uint64_t data_dim = 0; - uint64_t disk_data_dim = 0; // will be different from data_dim only if we use - // PQ for disk data (very large dimensionality) - uint64_t aligned_dim = 0; - uint64_t disk_bytes_per_point = 0; - - std::string disk_index_file; - std::vector> node_visit_counter; - - // PQ data - // n_chunks = # of chunks ndims is split into - // data: char * n_chunks - // chunk_size = chunk size of each dimension chunk - // pq_tables = float* [[2^8 * [chunk_size]] * n_chunks] - uint8_t *data = nullptr; - uint64_t n_chunks; - FixedChunkPQTable pq_table; - - // distance comparator - std::shared_ptr> dist_cmp; - std::shared_ptr> dist_cmp_float; - - // for very large datasets: we use PQ even for the disk resident index - bool use_disk_index_pq = false; - uint64_t disk_pq_n_chunks = 0; - FixedChunkPQTable disk_pq_table; - - // medoid/start info - - // graph has one entry point by default, - // we can optionally have multiple starting points - uint32_t *medoids = nullptr; - // defaults to 1 - size_t num_medoids; - // by default, it is empty. If there are multiple - // centroids, we pick the medoid corresponding to the - // closest centroid as the starting point of search - float *centroid_data = nullptr; - - // nhood_cache - unsigned *nhood_cache_buf = nullptr; - tsl::robin_map> nhood_cache; - - // coord_cache - T *coord_cache_buf = nullptr; - tsl::robin_map coord_cache; - - // thread-specific scratch - ConcurrentQueue *> thread_data; - uint64_t max_nthreads; - bool load_flag = false; - bool count_visited_nodes = false; - bool reorder_data_exists = false; - uint64_t reoreder_data_offset = 0; - - // filter support - uint32_t *_pts_to_label_offsets = nullptr; - uint32_t *_pts_to_labels = nullptr; - tsl::robin_set _labels; - std::unordered_map _filter_to_medoid_id; - bool _use_universal_label; - uint32_t _universal_filter_num; - std::vector _filter_list; - tsl::robin_set _dummy_pts; - tsl::robin_set _has_dummy_pts; - tsl::robin_map _dummy_to_real_map; - tsl::robin_map> _real_to_dummy_map; - std::unordered_map _label_map; + DISKANN_DLLEXPORT void cache_bfs_levels(uint64_t num_nodes_to_cache, std::vector &node_list, + const bool shuffle = false); + + DISKANN_DLLEXPORT void cached_beam_search(const T *query, const uint64_t k_search, const uint64_t l_search, + uint64_t *res_ids, float *res_dists, const uint64_t beam_width, + const bool use_reorder_data = false, QueryStats *stats = nullptr); + + DISKANN_DLLEXPORT void cached_beam_search(const T *query, const uint64_t k_search, const uint64_t l_search, + uint64_t *res_ids, float *res_dists, const uint64_t beam_width, + const bool use_filter, const LabelT &filter_label, + const bool use_reorder_data = false, QueryStats *stats = nullptr); + + DISKANN_DLLEXPORT void cached_beam_search(const T *query, const uint64_t k_search, const uint64_t l_search, + uint64_t *res_ids, float *res_dists, const uint64_t beam_width, + const uint32_t io_limit, const bool use_reorder_data = false, + QueryStats *stats = nullptr); + + DISKANN_DLLEXPORT void cached_beam_search(const T *query, const uint64_t k_search, const uint64_t l_search, + uint64_t *res_ids, float *res_dists, const uint64_t beam_width, + const bool use_filter, const LabelT &filter_label, + const uint32_t io_limit, const bool use_reorder_data = false, + QueryStats *stats = nullptr); + + DISKANN_DLLEXPORT LabelT get_converted_label(const std::string &filter_label); + + DISKANN_DLLEXPORT uint32_t range_search(const T *query1, const double range, const uint64_t min_l_search, + const uint64_t max_l_search, std::vector &indices, + std::vector &distances, const uint64_t min_beam_width, + QueryStats *stats = nullptr); + + DISKANN_DLLEXPORT uint64_t get_data_dim(); + + std::shared_ptr &reader; + + DISKANN_DLLEXPORT diskann::Metric get_metric(); + + protected: + DISKANN_DLLEXPORT void use_medoids_data_as_centroids(); + DISKANN_DLLEXPORT void setup_thread_data(uint64_t nthreads, uint64_t visited_reserve = 4096); + + DISKANN_DLLEXPORT void set_universal_label(const LabelT &label); + + private: + DISKANN_DLLEXPORT inline bool point_has_label(uint32_t point_id, uint32_t label_id); + std::unordered_map load_label_map(const std::string &map_file); + DISKANN_DLLEXPORT void parse_label_file(const std::string &map_file, size_t &num_pts_labels); + DISKANN_DLLEXPORT void get_label_file_metadata(std::string map_file, uint32_t &num_pts, uint32_t &num_total_labels); + DISKANN_DLLEXPORT inline int32_t get_filter_number(const LabelT &filter_label); + + // index info + // nhood of node `i` is in sector: [i / nnodes_per_sector] + // offset in sector: [(i % nnodes_per_sector) * max_node_len] + // nnbrs of node `i`: *(unsigned*) (buf) + // nbrs of node `i`: ((unsigned*)buf) + 1 + + uint64_t max_node_len = 0, nnodes_per_sector = 0, max_degree = 0; + + // Data used for searching with re-order vectors + uint64_t ndims_reorder_vecs = 0, reorder_data_start_sector = 0, nvecs_per_sector = 0; + + diskann::Metric metric = diskann::Metric::L2; + + // used only for inner product search to re-scale the result value + // (due to the pre-processing of base during index build) + float max_base_norm = 0.0f; + + // data info + uint64_t num_points = 0; + uint64_t num_frozen_points = 0; + uint64_t frozen_location = 0; + uint64_t data_dim = 0; + uint64_t disk_data_dim = 0; // will be different from data_dim only if we use + // PQ for disk data (very large dimensionality) + uint64_t aligned_dim = 0; + uint64_t disk_bytes_per_point = 0; + + std::string disk_index_file; + std::vector> node_visit_counter; + + // PQ data + // n_chunks = # of chunks ndims is split into + // data: char * n_chunks + // chunk_size = chunk size of each dimension chunk + // pq_tables = float* [[2^8 * [chunk_size]] * n_chunks] + uint8_t *data = nullptr; + uint64_t n_chunks; + FixedChunkPQTable pq_table; + + // distance comparator + std::shared_ptr> dist_cmp; + std::shared_ptr> dist_cmp_float; + + // for very large datasets: we use PQ even for the disk resident index + bool use_disk_index_pq = false; + uint64_t disk_pq_n_chunks = 0; + FixedChunkPQTable disk_pq_table; + + // medoid/start info + + // graph has one entry point by default, + // we can optionally have multiple starting points + uint32_t *medoids = nullptr; + // defaults to 1 + size_t num_medoids; + // by default, it is empty. If there are multiple + // centroids, we pick the medoid corresponding to the + // closest centroid as the starting point of search + float *centroid_data = nullptr; + + // nhood_cache + unsigned *nhood_cache_buf = nullptr; + tsl::robin_map> nhood_cache; + + // coord_cache + T *coord_cache_buf = nullptr; + tsl::robin_map coord_cache; + + // thread-specific scratch + ConcurrentQueue *> thread_data; + uint64_t max_nthreads; + bool load_flag = false; + bool count_visited_nodes = false; + bool reorder_data_exists = false; + uint64_t reoreder_data_offset = 0; + + // filter support + uint32_t *_pts_to_label_offsets = nullptr; + uint32_t *_pts_to_labels = nullptr; + tsl::robin_set _labels; + std::unordered_map _filter_to_medoid_id; + bool _use_universal_label; + uint32_t _universal_filter_num; + std::vector _filter_list; + tsl::robin_set _dummy_pts; + tsl::robin_set _has_dummy_pts; + tsl::robin_map _dummy_to_real_map; + tsl::robin_map> _real_to_dummy_map; + std::unordered_map _label_map; #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 - // on how big the header can be. - static const int HEADER_SIZE = SECTOR_LEN; - char *getHeaderBytes(); + // Set to a larger value than the actual header to accommodate + // any additions we make to the header. This is an outer limit + // on how big the header can be. + static const int HEADER_SIZE = SECTOR_LEN; + char *getHeaderBytes(); #endif }; -} // namespace diskann +} // namespace diskann diff --git a/include/restapi/common.h b/include/restapi/common.h index 640afd323..b8339635a 100644 --- a/include/restapi/common.h +++ b/include/restapi/common.h @@ -6,15 +6,13 @@ #include #include -namespace diskann { +namespace diskann +{ // Constants -static const std::string VECTOR_KEY = "query", K_KEY = "k", - INDICES_KEY = "indices", DISTANCES_KEY = "distances", - TAGS_KEY = "tags", QUERY_ID_KEY = "query_id", - ERROR_MESSAGE_KEY = "error", L_KEY = "Ls", - TIME_TAKEN_KEY = "time_taken_in_us", - PARTITION_KEY = "partition", +static const std::string VECTOR_KEY = "query", K_KEY = "k", INDICES_KEY = "indices", DISTANCES_KEY = "distances", + TAGS_KEY = "tags", QUERY_ID_KEY = "query_id", ERROR_MESSAGE_KEY = "error", L_KEY = "Ls", + TIME_TAKEN_KEY = "time_taken_in_us", PARTITION_KEY = "partition", UNKNOWN_ERROR = "unknown_error"; const unsigned int DEFAULT_L = 100; -} // namespace diskann \ No newline at end of file +} // namespace diskann \ No newline at end of file diff --git a/include/restapi/search_wrapper.h b/include/restapi/search_wrapper.h index 9a8c71f14..ebd067d8a 100644 --- a/include/restapi/search_wrapper.h +++ b/include/restapi/search_wrapper.h @@ -10,108 +10,131 @@ #include #include -namespace diskann { -class SearchResult { - public: - SearchResult(unsigned int K, unsigned int elapsed_time_in_ms, - const unsigned *const indices, const float *const distances, - const std::string *const tags = nullptr, - const unsigned *const partitions = nullptr); - - const std::vector &get_indices() const { return _indices; } - const std::vector &get_distances() const { return _distances; } - bool tags_enabled() const { return _tags_enabled; } - const std::vector &get_tags() const { return _tags; } - bool partitions_enabled() const { return _partitions_enabled; } - const std::vector &get_partitions() const { return _partitions; } - unsigned get_time() const { return _search_time_in_ms; } - - private: - unsigned int _K; - unsigned int _search_time_in_ms; - std::vector _indices; - std::vector _distances; - - bool _tags_enabled; - std::vector _tags; - - bool _partitions_enabled; - std::vector _partitions; +namespace diskann +{ +class SearchResult +{ + public: + SearchResult(unsigned int K, unsigned int elapsed_time_in_ms, const unsigned *const indices, + const float *const distances, const std::string *const tags = nullptr, + const unsigned *const partitions = nullptr); + + const std::vector &get_indices() const + { + return _indices; + } + const std::vector &get_distances() const + { + return _distances; + } + bool tags_enabled() const + { + return _tags_enabled; + } + const std::vector &get_tags() const + { + return _tags; + } + bool partitions_enabled() const + { + return _partitions_enabled; + } + const std::vector &get_partitions() const + { + return _partitions; + } + unsigned get_time() const + { + return _search_time_in_ms; + } + + private: + unsigned int _K; + unsigned int _search_time_in_ms; + std::vector _indices; + std::vector _distances; + + bool _tags_enabled; + std::vector _tags; + + bool _partitions_enabled; + std::vector _partitions; }; -class SearchNotImplementedException : public std::logic_error { - private: - std::string _errormsg; - - public: - SearchNotImplementedException(const char *type) - : std::logic_error("Not Implemented") { - _errormsg = "Search with data type "; - _errormsg += std::string(type); - _errormsg += " not implemented : "; - _errormsg += __FUNCTION__; - } - - virtual const char *what() const throw() { return _errormsg.c_str(); } +class SearchNotImplementedException : public std::logic_error +{ + private: + std::string _errormsg; + + public: + SearchNotImplementedException(const char *type) : std::logic_error("Not Implemented") + { + _errormsg = "Search with data type "; + _errormsg += std::string(type); + _errormsg += " not implemented : "; + _errormsg += __FUNCTION__; + } + + virtual const char *what() const throw() + { + return _errormsg.c_str(); + } }; -class BaseSearch { - public: - BaseSearch(const std::string &tagsFile = nullptr); - virtual SearchResult search(const float *query, const unsigned int dimensions, - const unsigned int K, const unsigned int Ls) { - throw SearchNotImplementedException("float"); - } - virtual SearchResult search(const int8_t *query, - const unsigned int dimensions, - const unsigned int K, const unsigned int Ls) { - throw SearchNotImplementedException("int8_t"); - } - - virtual SearchResult search(const uint8_t *query, - const unsigned int dimensions, - const unsigned int K, const unsigned int Ls) { - throw SearchNotImplementedException("uint8_t"); - } - - void lookup_tags(const unsigned K, const unsigned *indices, - std::string *ret_tags); - - protected: - bool _tags_enabled; - std::vector _tags_str; +class BaseSearch +{ + public: + BaseSearch(const std::string &tagsFile = nullptr); + virtual SearchResult search(const float *query, const unsigned int dimensions, const unsigned int K, + const unsigned int Ls) + { + throw SearchNotImplementedException("float"); + } + virtual SearchResult search(const int8_t *query, const unsigned int dimensions, const unsigned int K, + const unsigned int Ls) + { + throw SearchNotImplementedException("int8_t"); + } + + virtual SearchResult search(const uint8_t *query, const unsigned int dimensions, const unsigned int K, + const unsigned int Ls) + { + throw SearchNotImplementedException("uint8_t"); + } + + void lookup_tags(const unsigned K, const unsigned *indices, std::string *ret_tags); + + protected: + bool _tags_enabled; + std::vector _tags_str; }; -template -class InMemorySearch : public BaseSearch { - public: - InMemorySearch(const std::string &baseFile, const std::string &indexFile, - const std::string &tagsFile, Metric m, uint32_t num_threads, - uint32_t search_l); - virtual ~InMemorySearch(); +template class InMemorySearch : public BaseSearch +{ + public: + InMemorySearch(const std::string &baseFile, const std::string &indexFile, const std::string &tagsFile, Metric m, + uint32_t num_threads, uint32_t search_l); + virtual ~InMemorySearch(); - SearchResult search(const T *query, const unsigned int dimensions, - const unsigned int K, const unsigned int Ls); + SearchResult search(const T *query, const unsigned int dimensions, const unsigned int K, const unsigned int Ls); - private: - unsigned int _dimensions, _numPoints; - std::unique_ptr> _index; + private: + unsigned int _dimensions, _numPoints; + std::unique_ptr> _index; }; -template -class PQFlashSearch : public BaseSearch { - public: - PQFlashSearch(const std::string &indexPrefix, - const unsigned num_nodes_to_cache, const unsigned num_threads, - const std::string &tagsFile, Metric m); - virtual ~PQFlashSearch(); - - SearchResult search(const T *query, const unsigned int dimensions, - const unsigned int K, const unsigned int Ls); - - private: - unsigned int _dimensions, _numPoints; - std::unique_ptr> _index; - std::shared_ptr reader; +template class PQFlashSearch : public BaseSearch +{ + public: + PQFlashSearch(const std::string &indexPrefix, const unsigned num_nodes_to_cache, const unsigned num_threads, + const std::string &tagsFile, Metric m); + virtual ~PQFlashSearch(); + + SearchResult search(const T *query, const unsigned int dimensions, const unsigned int K, const unsigned int Ls); + + private: + unsigned int _dimensions, _numPoints; + std::unique_ptr> _index; + std::shared_ptr reader; }; -} // namespace diskann +} // namespace diskann diff --git a/include/restapi/server.h b/include/restapi/server.h index 2dcc0ebc1..1d75847a2 100644 --- a/include/restapi/server.h +++ b/include/restapi/server.h @@ -6,44 +6,40 @@ #include #include -namespace diskann { -class Server { - public: - Server(web::uri &url, - std::vector> &multi_searcher, - const std::string &typestring); - virtual ~Server(); - - pplx::task open(); - pplx::task close(); - - protected: - template - void handle_post(web::http::http_request message); - - template - web::json::value toJsonArray( - const std::vector &v, - std::function valConverter); - web::json::value prepareResponse(const int64_t &queryId, const int k); - - template - void parseJson(const utility::string_t &body, unsigned int &k, - int64_t &queryId, T *&queryVector, unsigned int &dimensions, - unsigned &Ls); - - web::json::value idsToJsonArray(const diskann::SearchResult &result); - web::json::value distancesToJsonArray(const diskann::SearchResult &result); - web::json::value tagsToJsonArray(const diskann::SearchResult &result); - web::json::value partitionsToJsonArray(const diskann::SearchResult &result); - - SearchResult aggregate_results( - const unsigned K, const std::vector &results); - - private: - bool _isDebug; - std::unique_ptr _listener; - const bool _multi_search; - std::vector> _multi_searcher; +namespace diskann +{ +class Server +{ + public: + Server(web::uri &url, std::vector> &multi_searcher, + const std::string &typestring); + virtual ~Server(); + + pplx::task open(); + pplx::task close(); + + protected: + template void handle_post(web::http::http_request message); + + template + web::json::value toJsonArray(const std::vector &v, std::function valConverter); + web::json::value prepareResponse(const int64_t &queryId, const int k); + + template + void parseJson(const utility::string_t &body, unsigned int &k, int64_t &queryId, T *&queryVector, + unsigned int &dimensions, unsigned &Ls); + + web::json::value idsToJsonArray(const diskann::SearchResult &result); + web::json::value distancesToJsonArray(const diskann::SearchResult &result); + web::json::value tagsToJsonArray(const diskann::SearchResult &result); + web::json::value partitionsToJsonArray(const diskann::SearchResult &result); + + SearchResult aggregate_results(const unsigned K, const std::vector &results); + + private: + bool _isDebug; + std::unique_ptr _listener; + const bool _multi_search; + std::vector> _multi_searcher; }; -} // namespace diskann +} // namespace diskann diff --git a/include/scratch.h b/include/scratch.h index 0875275d3..4ee345aa9 100644 --- a/include/scratch.h +++ b/include/scratch.h @@ -25,167 +25,209 @@ #define SECTOR_LEN (size_t)4096 #define MAX_N_SECTOR_READS 128 -namespace diskann { +namespace diskann +{ // // Scratch space for in-memory index based search // -template -class InMemQueryScratch { - public: - ~InMemQueryScratch(); - // REFACTOR TODO: move all parameters to a new class. - InMemQueryScratch(uint32_t search_l, uint32_t indexing_l, uint32_t r, - uint32_t maxc, size_t dim, size_t aligned_dim, - size_t alignment_factor, bool init_pq_scratch = false); - void resize_for_new_L(uint32_t new_search_l); - void clear(); - - inline uint32_t get_L() { return _L; } - inline uint32_t get_R() { return _R; } - inline uint32_t get_maxc() { return _maxc; } - inline T *aligned_query() { return _aligned_query; } - inline PQScratch *pq_scratch() { return _pq_scratch; } - inline std::vector &pool() { return _pool; } - inline NeighborPriorityQueue &best_l_nodes() { return _best_l_nodes; } - inline std::vector &occlude_factor() { return _occlude_factor; } - inline tsl::robin_set &inserted_into_pool_rs() { - return _inserted_into_pool_rs; - } - inline boost::dynamic_bitset<> &inserted_into_pool_bs() { - return *_inserted_into_pool_bs; - } - inline std::vector &id_scratch() { return _id_scratch; } - inline std::vector &dist_scratch() { return _dist_scratch; } - inline tsl::robin_set &expanded_nodes_set() { - return _expanded_nodes_set; - } - inline std::vector &expanded_nodes_vec() { - return _expanded_nghrs_vec; - } - inline std::vector &occlude_list_output() { - return _occlude_list_output; - } - - private: - uint32_t _L; - uint32_t _R; - uint32_t _maxc; - - T *_aligned_query = nullptr; - - PQScratch *_pq_scratch = nullptr; - - // _pool stores all neighbors explored from best_L_nodes. - // Usually around L+R, but could be higher. - // Initialized to 3L+R for some slack, expands as needed. - std::vector _pool; - - // _best_l_nodes is reserved for storing best L entries - // Underlying storage is L+1 to support inserts - NeighborPriorityQueue _best_l_nodes; - - // _occlude_factor.size() >= pool.size() in occlude_list function - // _pool is clipped to maxc in occlude_list before affecting _occlude_factor - // _occlude_factor is initialized to maxc size - std::vector _occlude_factor; - - // Capacity initialized to 20L - tsl::robin_set _inserted_into_pool_rs; - - // Use a pointer here to allow for forward declaration of dynamic_bitset - // in public headers to avoid making boost a dependency for clients - // of DiskANN. - boost::dynamic_bitset<> *_inserted_into_pool_bs; - - // _id_scratch.size() must be > R*GRAPH_SLACK_FACTOR for iterate_to_fp - std::vector _id_scratch; - - // _dist_scratch must be > R*GRAPH_SLACK_FACTOR for iterate_to_fp - // _dist_scratch should be at least the size of id_scratch - std::vector _dist_scratch; - - // Buffers used in process delete, capacity increases as needed - tsl::robin_set _expanded_nodes_set; - std::vector _expanded_nghrs_vec; - std::vector _occlude_list_output; +template class InMemQueryScratch +{ + public: + ~InMemQueryScratch(); + // REFACTOR TODO: move all parameters to a new class. + InMemQueryScratch(uint32_t search_l, uint32_t indexing_l, uint32_t r, uint32_t maxc, size_t dim, size_t aligned_dim, + size_t alignment_factor, bool init_pq_scratch = false); + void resize_for_new_L(uint32_t new_search_l); + void clear(); + + inline uint32_t get_L() + { + return _L; + } + inline uint32_t get_R() + { + return _R; + } + inline uint32_t get_maxc() + { + return _maxc; + } + inline T *aligned_query() + { + return _aligned_query; + } + inline PQScratch *pq_scratch() + { + return _pq_scratch; + } + inline std::vector &pool() + { + return _pool; + } + inline NeighborPriorityQueue &best_l_nodes() + { + return _best_l_nodes; + } + inline std::vector &occlude_factor() + { + return _occlude_factor; + } + inline tsl::robin_set &inserted_into_pool_rs() + { + return _inserted_into_pool_rs; + } + inline boost::dynamic_bitset<> &inserted_into_pool_bs() + { + return *_inserted_into_pool_bs; + } + inline std::vector &id_scratch() + { + return _id_scratch; + } + inline std::vector &dist_scratch() + { + return _dist_scratch; + } + inline tsl::robin_set &expanded_nodes_set() + { + return _expanded_nodes_set; + } + inline std::vector &expanded_nodes_vec() + { + return _expanded_nghrs_vec; + } + inline std::vector &occlude_list_output() + { + return _occlude_list_output; + } + + private: + uint32_t _L; + uint32_t _R; + uint32_t _maxc; + + T *_aligned_query = nullptr; + + PQScratch *_pq_scratch = nullptr; + + // _pool stores all neighbors explored from best_L_nodes. + // Usually around L+R, but could be higher. + // Initialized to 3L+R for some slack, expands as needed. + std::vector _pool; + + // _best_l_nodes is reserved for storing best L entries + // Underlying storage is L+1 to support inserts + NeighborPriorityQueue _best_l_nodes; + + // _occlude_factor.size() >= pool.size() in occlude_list function + // _pool is clipped to maxc in occlude_list before affecting _occlude_factor + // _occlude_factor is initialized to maxc size + std::vector _occlude_factor; + + // Capacity initialized to 20L + tsl::robin_set _inserted_into_pool_rs; + + // Use a pointer here to allow for forward declaration of dynamic_bitset + // in public headers to avoid making boost a dependency for clients + // of DiskANN. + boost::dynamic_bitset<> *_inserted_into_pool_bs; + + // _id_scratch.size() must be > R*GRAPH_SLACK_FACTOR for iterate_to_fp + std::vector _id_scratch; + + // _dist_scratch must be > R*GRAPH_SLACK_FACTOR for iterate_to_fp + // _dist_scratch should be at least the size of id_scratch + std::vector _dist_scratch; + + // Buffers used in process delete, capacity increases as needed + tsl::robin_set _expanded_nodes_set; + std::vector _expanded_nghrs_vec; + std::vector _occlude_list_output; }; // // Scratch space for SSD index based search // -template -class SSDQueryScratch { - public: - T *coord_scratch = nullptr; // MUST BE AT LEAST [MAX_N_CMPS * data_dim] - size_t coord_idx = 0; // index of next [data_dim] scratch to use +template class SSDQueryScratch +{ + public: + T *coord_scratch = nullptr; // MUST BE AT LEAST [MAX_N_CMPS * data_dim] + size_t coord_idx = 0; // index of next [data_dim] scratch to use - char *sector_scratch = - nullptr; // MUST BE AT LEAST [MAX_N_SECTOR_READS * SECTOR_LEN] - size_t sector_idx = 0; // index of next [SECTOR_LEN] scratch to use + char *sector_scratch = nullptr; // MUST BE AT LEAST [MAX_N_SECTOR_READS * SECTOR_LEN] + size_t sector_idx = 0; // index of next [SECTOR_LEN] scratch to use - T *aligned_query_T = nullptr; + T *aligned_query_T = nullptr; - PQScratch *_pq_scratch; + PQScratch *_pq_scratch; - tsl::robin_set visited; - NeighborPriorityQueue retset; - std::vector full_retset; + tsl::robin_set visited; + NeighborPriorityQueue retset; + std::vector full_retset; - SSDQueryScratch(size_t aligned_dim, size_t visited_reserve); - ~SSDQueryScratch(); + SSDQueryScratch(size_t aligned_dim, size_t visited_reserve); + ~SSDQueryScratch(); - void reset(); + void reset(); }; -template -class SSDThreadData { - public: - SSDQueryScratch scratch; - IOContext ctx; +template class SSDThreadData +{ + public: + SSDQueryScratch scratch; + IOContext ctx; - SSDThreadData(size_t aligned_dim, size_t visited_reserve); - void clear(); + SSDThreadData(size_t aligned_dim, size_t visited_reserve); + void clear(); }; // // Class to avoid the hassle of pushing and popping the query scratch. // -template -class ScratchStoreManager { - public: - ScratchStoreManager(ConcurrentQueue &query_scratch) - : _scratch_pool(query_scratch) { - _scratch = query_scratch.pop(); - while (_scratch == nullptr) { - query_scratch.wait_for_push_notify(); - _scratch = query_scratch.pop(); - } - } - T *scratch_space() { return _scratch; } - - ~ScratchStoreManager() { - _scratch->clear(); - _scratch_pool.push(_scratch); - _scratch_pool.push_notify_all(); - } - - void destroy() { - while (!_scratch_pool.empty()) { - auto scratch = _scratch_pool.pop(); - while (scratch == nullptr) { - _scratch_pool.wait_for_push_notify(); - scratch = _scratch_pool.pop(); - } - delete scratch; - } - } - - private: - T *_scratch; - ConcurrentQueue &_scratch_pool; - ScratchStoreManager(const ScratchStoreManager &); - ScratchStoreManager &operator=(const ScratchStoreManager &); +template class ScratchStoreManager +{ + public: + ScratchStoreManager(ConcurrentQueue &query_scratch) : _scratch_pool(query_scratch) + { + _scratch = query_scratch.pop(); + while (_scratch == nullptr) + { + query_scratch.wait_for_push_notify(); + _scratch = query_scratch.pop(); + } + } + T *scratch_space() + { + return _scratch; + } + + ~ScratchStoreManager() + { + _scratch->clear(); + _scratch_pool.push(_scratch); + _scratch_pool.push_notify_all(); + } + + void destroy() + { + while (!_scratch_pool.empty()) + { + auto scratch = _scratch_pool.pop(); + while (scratch == nullptr) + { + _scratch_pool.wait_for_push_notify(); + scratch = _scratch_pool.pop(); + } + delete scratch; + } + } + + private: + T *_scratch; + ConcurrentQueue &_scratch_pool; + ScratchStoreManager(const ScratchStoreManager &); + ScratchStoreManager &operator=(const ScratchStoreManager &); }; -} // namespace diskann +} // namespace diskann diff --git a/include/simd_utils.h b/include/simd_utils.h index fb94c35c2..4b0736998 100644 --- a/include/simd_utils.h +++ b/include/simd_utils.h @@ -9,97 +9,98 @@ #include #endif -namespace diskann { -static inline __m256 _mm256_mul_epi8(__m256i X) { - __m256i zero = _mm256_setzero_si256(); +namespace diskann +{ +static inline __m256 _mm256_mul_epi8(__m256i X) +{ + __m256i zero = _mm256_setzero_si256(); - __m256i sign_x = _mm256_cmpgt_epi8(zero, X); + __m256i sign_x = _mm256_cmpgt_epi8(zero, X); - __m256i xlo = _mm256_unpacklo_epi8(X, sign_x); - __m256i xhi = _mm256_unpackhi_epi8(X, sign_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))); + 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); +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))); + 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); +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))); + 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(); +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 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); + __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))); + 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_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 __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(); +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 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); + __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))); + 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 __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); +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 +} // namespace diskann diff --git a/include/timer.h b/include/timer.h index d187a4e3f..963927bd7 100644 --- a/include/timer.h +++ b/include/timer.h @@ -3,27 +3,37 @@ #include -namespace diskann { -class Timer { - typedef std::chrono::high_resolution_clock _clock; - std::chrono::time_point<_clock> check_point; +namespace diskann +{ +class Timer +{ + typedef std::chrono::high_resolution_clock _clock; + std::chrono::time_point<_clock> check_point; - public: - Timer() : check_point(_clock::now()) {} + public: + Timer() : check_point(_clock::now()) + { + } - void reset() { check_point = _clock::now(); } + void reset() + { + check_point = _clock::now(); + } - long long elapsed() const { - return std::chrono::duration_cast(_clock::now() - - check_point) - .count(); - } + long long elapsed() const + { + return std::chrono::duration_cast(_clock::now() - check_point).count(); + } - float elapsed_seconds() const { return (float)elapsed() / 1000000.0; } + float elapsed_seconds() const + { + return (float)elapsed() / 1000000.0; + } - std::string elapsed_seconds_for_step(const std::string &step) const { - return std::string("Time for ") + step + std::string(": ") + - std::to_string(elapsed_seconds()) + std::string(" seconds"); - } + std::string elapsed_seconds_for_step(const std::string &step) const + { + return std::string("Time for ") + step + std::string(": ") + std::to_string(elapsed_seconds()) + + std::string(" seconds"); + } }; -} // namespace diskann +} // namespace diskann diff --git a/include/types.h b/include/types.h index 9b2794474..ea04cd34d 100644 --- a/include/types.h +++ b/include/types.h @@ -6,6 +6,7 @@ #include #include -namespace diskann { +namespace diskann +{ typedef uint32_t location_t; -} // namespace diskann \ No newline at end of file +} // namespace diskann \ No newline at end of file diff --git a/include/utils.h b/include/utils.h index 0d4aa8722..1b9cce924 100644 --- a/include/utils.h +++ b/include/utils.h @@ -36,8 +36,7 @@ typedef int FileHandle; // taken from // https://github.com/Microsoft/BLAS-on-flash/blob/master/include/utils.h // round up X to the nearest multiple of Y -#define ROUND_UP(X, Y) \ - ((((uint64_t)(X) / (Y)) + ((uint64_t)(X) % (Y) != 0)) * (Y)) +#define ROUND_UP(X, Y) ((((uint64_t)(X) / (Y)) + ((uint64_t)(X) % (Y) != 0)) * (Y)) #define DIV_ROUND_UP(X, Y) (((uint64_t)(X) / (Y)) + ((uint64_t)(X) % (Y) != 0)) @@ -48,730 +47,750 @@ 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 +#define METADATA_SIZE \ + 4096 // all metadata of individual sub-component files is written in first + // 4KB for unified files #define BUFFER_SIZE_FOR_CACHED_IO (size_t)1024 * (size_t)1048576 #define PBSTR "||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||" #define PBWIDTH 60 -inline bool file_exists(const std::string &name, bool dirCheck = false) { - int val; +inline bool file_exists(const std::string &name, bool dirCheck = false) +{ + int val; #ifndef _WINDOWS - struct stat buffer; - val = stat(name.c_str(), &buffer); + struct stat buffer; + val = stat(name.c_str(), &buffer); #else - // It is the 21st century but Windows API still thinks in 32-bit terms. - // Turns out calling stat() on a file > 4GB results in errno = 132 - // (OVERFLOW). How silly is this!? So calling _stat64() - struct _stat64 buffer; - val = _stat64(name.c_str(), &buffer); + // It is the 21st century but Windows API still thinks in 32-bit terms. + // Turns out calling stat() on a file > 4GB results in errno = 132 + // (OVERFLOW). How silly is this!? So calling _stat64() + struct _stat64 buffer; + val = _stat64(name.c_str(), &buffer); #endif - if (val != 0) { - switch (errno) { - case EINVAL: - diskann::cout << "Invalid argument passed to stat()" << std::endl; - break; - case ENOENT: - // file is not existing, not an issue, so we won't cout anything. - break; - default: - diskann::cout << "Unexpected error in stat():" << errno << std::endl; - break; + if (val != 0) + { + switch (errno) + { + case EINVAL: + diskann::cout << "Invalid argument passed to stat()" << std::endl; + break; + case ENOENT: + // file is not existing, not an issue, so we won't cout anything. + 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; } - return false; - } else { - // the file entry exists. If reqd, check if this is a directory. - return dirCheck ? buffer.st_mode & S_IFDIR : true; - } } -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()) { - char buff[1024]; +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()) + { + char buff[1024]; #ifdef _WINDOWS - auto ret = std::to_string(strerror_s(buff, 1024, errno)); + auto ret = std::to_string(strerror_s(buff, 1024, errno)); #else - auto ret = std::string(strerror_r(errno, buff, 1024)); + auto ret = std::string(strerror_r(errno, buff, 1024)); #endif - auto message = std::string("Failed to open file") + filename + - " for write because " + buff + ", ret=" + ret; - diskann::cerr << message << std::endl; - throw diskann::ANNException(message, -1); - } + auto message = std::string("Failed to open file") + filename + " for write because " + buff + ", ret=" + ret; + diskann::cerr << message << std::endl; + throw diskann::ANNException(message, -1); + } } -inline size_t get_file_size(const std::string &fname) { - std::ifstream reader(fname, std::ios::binary | std::ios::ate); - if (!reader.fail() && reader.is_open()) { - size_t end_pos = reader.tellg(); - reader.close(); - return end_pos; - } else { - diskann::cerr << "Could not open file: " << fname << std::endl; - return 0; - } +inline size_t get_file_size(const std::string &fname) +{ + std::ifstream reader(fname, std::ios::binary | std::ios::ate); + if (!reader.fail() && reader.is_open()) + { + size_t 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; +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; } - return rc; - } else { - return 0; - } } -inline void convert_labels_string_to_int(const std::string &inFileName, - const std::string &outFileName, - const std::string &mapFileName, - const std::string &unv_label) { - std::unordered_map string_int_map; - std::ofstream label_writer(outFileName); - std::ifstream label_reader(inFileName); - if (unv_label != "") string_int_map[unv_label] = 0; - std::string line, token; - while (std::getline(label_reader, line)) { - std::istringstream new_iss(line); - std::vector lbls; - while (getline(new_iss, token, ',')) { - token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); - token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); - if (string_int_map.find(token) == string_int_map.end()) { - uint32_t nextId = (uint32_t)string_int_map.size() + 1; - string_int_map[token] = nextId; - } - lbls.push_back(string_int_map[token]); +inline void convert_labels_string_to_int(const std::string &inFileName, const std::string &outFileName, + const std::string &mapFileName, const std::string &unv_label) +{ + std::unordered_map string_int_map; + std::ofstream label_writer(outFileName); + std::ifstream label_reader(inFileName); + if (unv_label != "") + string_int_map[unv_label] = 0; + std::string line, token; + while (std::getline(label_reader, line)) + { + std::istringstream new_iss(line); + std::vector lbls; + while (getline(new_iss, token, ',')) + { + token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); + token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); + if (string_int_map.find(token) == string_int_map.end()) + { + uint32_t nextId = (uint32_t)string_int_map.size() + 1; + string_int_map[token] = nextId; + } + lbls.push_back(string_int_map[token]); + } + if (lbls.size() <= 0) + { + std::cout << "No label found"; + exit(-1); + } + for (size_t j = 0; j < lbls.size(); j++) + { + if (j != lbls.size() - 1) + label_writer << lbls[j] << ","; + else + label_writer << lbls[j] << std::endl; + } } - if (lbls.size() <= 0) { - std::cout << "No label found"; - exit(-1); - } - for (size_t j = 0; j < lbls.size(); j++) { - if (j != lbls.size() - 1) - label_writer << lbls[j] << ","; - else - label_writer << lbls[j] << std::endl; + label_writer.close(); + + std::ofstream map_writer(mapFileName); + for (auto mp : string_int_map) + { + map_writer << mp.first << "\t" << mp.second << std::endl; } - } - label_writer.close(); - - std::ofstream map_writer(mapFileName); - for (auto mp : string_int_map) { - map_writer << mp.first << "\t" << mp.second << std::endl; - } - map_writer.close(); + map_writer.close(); } #ifdef EXEC_ENV_OLS class AlignedFileReader; #endif -namespace diskann { +namespace diskann +{ static const size_t MAX_SIZE_OF_STREAMBUF = 2LL * 1024 * 1024 * 1024; -inline void print_error_and_terminate(std::stringstream &error_stream) { - diskann::cerr << error_stream.str() << std::endl; - throw diskann::ANNException(error_stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); +inline void print_error_and_terminate(std::stringstream &error_stream) +{ + diskann::cerr << error_stream.str() << std::endl; + throw diskann::ANNException(error_stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } -inline void report_memory_allocation_failure() { - std::stringstream stream; - stream << "Memory Allocation Failed."; - print_error_and_terminate(stream); +inline void report_memory_allocation_failure() +{ + std::stringstream stream; + stream << "Memory Allocation Failed."; + print_error_and_terminate(stream); } -inline void report_misalignment_of_requested_size(size_t align) { - std::stringstream stream; - stream << "Requested memory size is not a multiple of " << align - << ". Can not be allocated."; - print_error_and_terminate(stream); +inline void report_misalignment_of_requested_size(size_t align) +{ + std::stringstream stream; + stream << "Requested memory size is not a multiple of " << align << ". Can not be allocated."; + print_error_and_terminate(stream); } -inline void alloc_aligned(void **ptr, size_t size, size_t align) { - *ptr = nullptr; - if (IS_ALIGNED(size, align) == 0) - report_misalignment_of_requested_size(align); +inline void alloc_aligned(void **ptr, size_t size, size_t align) +{ + *ptr = nullptr; + if (IS_ALIGNED(size, align) == 0) + report_misalignment_of_requested_size(align); #ifndef _WINDOWS - *ptr = ::aligned_alloc(align, size); + *ptr = ::aligned_alloc(align, size); #else - *ptr = ::_aligned_malloc(size, align); // note the swapped arguments! + *ptr = ::_aligned_malloc(size, align); // note the swapped arguments! #endif - if (*ptr == nullptr) report_memory_allocation_failure(); + if (*ptr == nullptr) + report_memory_allocation_failure(); } -inline void realloc_aligned(void **ptr, size_t size, size_t align) { - if (IS_ALIGNED(size, align) == 0) - report_misalignment_of_requested_size(align); +inline void realloc_aligned(void **ptr, size_t size, size_t align) +{ + if (IS_ALIGNED(size, align) == 0) + report_misalignment_of_requested_size(align); #ifdef _WINDOWS - *ptr = ::_aligned_realloc(*ptr, size, align); + *ptr = ::_aligned_realloc(*ptr, size, align); #else - diskann::cerr << "No aligned realloc on GCC. Must malloc and mem_align, " - "left it out for now." - << std::endl; + diskann::cerr << "No aligned realloc on GCC. Must malloc and mem_align, " + "left it out for now." + << std::endl; #endif - if (*ptr == nullptr) report_memory_allocation_failure(); + if (*ptr == nullptr) + report_memory_allocation_failure(); } -inline void check_stop(std::string arnd) { - int brnd; - diskann::cout << arnd << std::endl; - std::cin >> brnd; +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; - } +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; + } #ifndef _WINDOWS - free(ptr); + free(ptr); #else - ::_aligned_free(ptr); + ::_aligned_free(ptr); #endif } -inline void GenRandom(std::mt19937 &rng, unsigned *addr, unsigned size, - unsigned N) { - for (unsigned i = 0; i < size; ++i) { - addr[i] = rng() % (N - size); - } +inline void GenRandom(std::mt19937 &rng, unsigned *addr, unsigned size, unsigned N) +{ + for (unsigned i = 0; i < size; ++i) + { + addr[i] = rng() % (N - size); + } - std::sort(addr, addr + size); - for (unsigned i = 1; i < size; ++i) { - if (addr[i] <= addr[i - 1]) { - addr[i] = addr[i - 1] + 1; + std::sort(addr, addr + size); + for (unsigned i = 1; i < size; ++i) + { + if (addr[i] <= addr[i - 1]) + { + addr[i] = addr[i - 1] + 1; + } + } + unsigned off = rng() % N; + for (unsigned i = 0; i < size; ++i) + { + addr[i] = (addr[i] + off) % N; } - } - unsigned off = rng() % N; - for (unsigned i = 0; i < size; ++i) { - addr[i] = (addr[i] + off) % N; - } } // get_bin_metadata functions START -inline void get_bin_metadata_impl(std::basic_istream &reader, - 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; - ncols = ncols_32; +inline void get_bin_metadata_impl(std::basic_istream &reader, 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; + ncols = ncols_32; } #ifdef EXEC_ENV_OLS -inline void get_bin_metadata(MemoryMappedFiles &files, - const std::string &bin_file, size_t &nrows, - 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, offset); - - 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; +inline void get_bin_metadata(MemoryMappedFiles &files, const std::string &bin_file, size_t &nrows, 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, offset); + + 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 offset = 0) { - std::ifstream reader(bin_file.c_str(), std::ios::binary); - get_bin_metadata_impl(reader, nrows, ncols, offset); +inline void get_bin_metadata(const std::string &bin_file, size_t &nrows, size_t &ncols, size_t offset = 0) +{ + std::ifstream reader(bin_file.c_str(), std::ios::binary); + get_bin_metadata_impl(reader, nrows, ncols, offset); } // get_bin_metadata functions END -template -inline std::string getValues(T *data, size_t num) { - std::stringstream stream; - stream << "["; - for (size_t i = 0; i < num; i++) { - stream << std::to_string(data[i]) << ","; - } - stream << "]" << std::endl; - - return stream.str(); +template inline std::string getValues(T *data, size_t num) +{ + std::stringstream stream; + stream << "["; + for (size_t i = 0; i < num; i++) + { + stream << std::to_string(data[i]) << ","; + } + stream << "]" << std::endl; + + return stream.str(); } // load_bin functions START template -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; +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; - dim = (unsigned)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; + dim = (unsigned)dim_i32; - std::cout << "Metadata: #pts = " << npts << ", #dims = " << dim << "..." - << std::endl; + std::cout << "Metadata: #pts = " << npts << ", #dims = " << dim << "..." << std::endl; - data = new T[npts * dim]; - reader.read((char *)data, npts * dim * sizeof(T)); + data = new T[npts * dim]; + reader.read((char *)data, npts * dim * sizeof(T)); } #ifdef EXEC_ENV_OLS template -inline void load_bin(MemoryMappedFiles &files, const std::string &bin_file, - 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 *)((char *)fc._content + offset); - t_npts = *(contentAsIntPtr); - t_dim = *(contentAsIntPtr + 1); - - npts = t_npts; - dim = t_dim; - - data = (T *)((char *)fc._content + offset + - 2 * sizeof(uint32_t)); // No need to copy! +inline void load_bin(MemoryMappedFiles &files, const std::string &bin_file, 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 *)((char *)fc._content + offset); + t_npts = *(contentAsIntPtr); + t_dim = *(contentAsIntPtr + 1); + + npts = t_npts; + dim = t_dim; + + data = (T *)((char *)fc._content + offset + 2 * sizeof(uint32_t)); // No need to copy! } -DISKANN_DLLEXPORT void get_bin_metadata(AlignedFileReader &reader, size_t &npts, - size_t &ndim, size_t offset = 0); +DISKANN_DLLEXPORT void get_bin_metadata(AlignedFileReader &reader, size_t &npts, size_t &ndim, size_t offset = 0); template -DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, T *&data, - size_t &npts, size_t &ndim, size_t offset = 0); +DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, T *&data, size_t &npts, size_t &ndim, size_t offset = 0); template -DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, - std::unique_ptr &data, size_t &npts, - size_t &ndim, size_t offset = 0); +DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, std::unique_ptr &data, size_t &npts, size_t &ndim, + size_t offset = 0); template -DISKANN_DLLEXPORT void copy_aligned_data_from_file(AlignedFileReader &reader, - T *&data, size_t &npts, - size_t &dim, - const size_t &rounded_dim, - size_t offset = 0); +DISKANN_DLLEXPORT void copy_aligned_data_from_file(AlignedFileReader &reader, T *&data, size_t &npts, size_t &dim, + const size_t &rounded_dim, size_t offset = 0); // Unlike load_bin, assumes that data is already allocated 'size' entries template -DISKANN_DLLEXPORT void read_array(AlignedFileReader &reader, T *data, - size_t size, size_t offset = 0); +DISKANN_DLLEXPORT void read_array(AlignedFileReader &reader, T *data, size_t size, size_t offset = 0); -template -DISKANN_DLLEXPORT void read_value(AlignedFileReader &reader, T &value, - size_t offset = 0); +template DISKANN_DLLEXPORT void read_value(AlignedFileReader &reader, T &value, size_t offset = 0); #endif template -inline void load_bin(const std::string &bin_file, T *&data, size_t &npts, - size_t &dim, size_t offset = 0) { - diskann::cout << "Reading bin file " << bin_file.c_str() << " ..." - << std::endl; - std::ifstream reader; - reader.exceptions(std::ifstream::failbit | std::ifstream::badbit); - - try { - diskann::cout << "Opening bin file " << bin_file.c_str() << "... " - << std::endl; - reader.open(bin_file, std::ios::binary | std::ios::ate); - reader.seekg(0); - load_bin_impl(reader, data, npts, dim, offset); - } catch (std::system_error &e) { - throw FileException(bin_file, e, __FUNCSIG__, __FILE__, __LINE__); - } - diskann::cout << "done." << std::endl; +inline void load_bin(const std::string &bin_file, T *&data, size_t &npts, size_t &dim, size_t offset = 0) +{ + diskann::cout << "Reading bin file " << bin_file.c_str() << " ..." << std::endl; + std::ifstream reader; + reader.exceptions(std::ifstream::failbit | std::ifstream::badbit); + + try + { + diskann::cout << "Opening bin file " << bin_file.c_str() << "... " << std::endl; + reader.open(bin_file, std::ios::binary | std::ios::ate); + reader.seekg(0); + load_bin_impl(reader, data, npts, dim, offset); + } + catch (std::system_error &e) + { + throw FileException(bin_file, e, __FUNCSIG__, __FILE__, __LINE__); + } + diskann::cout << "done." << std::endl; } -inline void wait_for_keystroke() { - int a; - std::cout << "Press any number to continue.." << std::endl; - std::cin >> a; +inline void wait_for_keystroke() +{ + int a; + std::cout << "Press any number to continue.." << std::endl; + std::cin >> a; } // load_bin functions END -inline void load_truthset(const std::string &bin_file, uint32_t *&ids, - float *&dists, size_t &npts, size_t &dim) { - size_t read_blk_size = 64 * 1024 * 1024; - cached_ifstream reader(bin_file, read_blk_size); - diskann::cout << "Reading truthset file " << bin_file.c_str() << " ..." - << std::endl; - size_t actual_file_size = reader.get_file_size(); +inline void load_truthset(const std::string &bin_file, uint32_t *&ids, float *&dists, size_t &npts, size_t &dim) +{ + size_t read_blk_size = 64 * 1024 * 1024; + cached_ifstream reader(bin_file, read_blk_size); + diskann::cout << "Reading truthset file " << bin_file.c_str() << " ..." << std::endl; + size_t actual_file_size = reader.get_file_size(); + + 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; + + diskann::cout << "Metadata: #pts = " << npts << ", #dims = " << dim << "... " << std::endl; + + int truthset_type = -1; // 1 means truthset has ids and distances, 2 means + // only ids, -1 is error + size_t expected_file_size_with_dists = 2 * npts * dim * sizeof(uint32_t) + 2 * sizeof(uint32_t); + + if (actual_file_size == expected_file_size_with_dists) + truthset_type = 1; + + size_t expected_file_size_just_ids = npts * dim * sizeof(uint32_t) + 2 * sizeof(uint32_t); + + if (actual_file_size == expected_file_size_just_ids) + truthset_type = 2; + + if (truthset_type == -1) + { + std::stringstream stream; + stream << "Error. File size mismatch. File should have bin format, with " + "npts followed by ngt followed by npts*ngt ids and optionally " + "followed by npts*ngt distance values; actual size: " + << actual_file_size << ", expected: " << expected_file_size_with_dists << " or " + << expected_file_size_just_ids; + diskann::cout << stream.str(); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } - 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; + ids = new uint32_t[npts * dim]; + reader.read((char *)ids, npts * dim * sizeof(uint32_t)); - diskann::cout << "Metadata: #pts = " << npts << ", #dims = " << dim << "... " - << std::endl; + if (truthset_type == 1) + { + dists = new float[npts * dim]; + reader.read((char *)dists, npts * dim * sizeof(float)); + } +} - int truthset_type = -1; // 1 means truthset has ids and distances, 2 means - // only ids, -1 is error - size_t expected_file_size_with_dists = - 2 * npts * dim * sizeof(uint32_t) + 2 * sizeof(uint32_t); +inline void prune_truthset_for_range(const std::string &bin_file, float range, + std::vector> &groundtruth, size_t &npts) +{ + size_t read_blk_size = 64 * 1024 * 1024; + cached_ifstream reader(bin_file, read_blk_size); + diskann::cout << "Reading truthset file " << bin_file.c_str() << "... " << std::endl; + size_t actual_file_size = reader.get_file_size(); + + int npts_i32, dim_i32; + reader.read((char *)&npts_i32, sizeof(int)); + reader.read((char *)&dim_i32, sizeof(int)); + npts = (unsigned)npts_i32; + uint64_t dim = (unsigned)dim_i32; + uint32_t *ids; + float *dists; + + diskann::cout << "Metadata: #pts = " << npts << ", #dims = " << dim << "... " << std::endl; + + int truthset_type = -1; // 1 means truthset has ids and distances, 2 means + // only ids, -1 is error + size_t expected_file_size_with_dists = 2 * npts * dim * sizeof(uint32_t) + 2 * sizeof(uint32_t); + + if (actual_file_size == expected_file_size_with_dists) + truthset_type = 1; + + if (truthset_type == -1) + { + std::stringstream stream; + stream << "Error. File size mismatch. File should have bin format, with " + "npts followed by ngt followed by npts*ngt ids and optionally " + "followed by npts*ngt distance values; actual size: " + << actual_file_size << ", expected: " << expected_file_size_with_dists; + diskann::cout << stream.str(); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } - if (actual_file_size == expected_file_size_with_dists) truthset_type = 1; + ids = new uint32_t[npts * dim]; + reader.read((char *)ids, npts * dim * sizeof(uint32_t)); - size_t expected_file_size_just_ids = - npts * dim * sizeof(uint32_t) + 2 * sizeof(uint32_t); + if (truthset_type == 1) + { + dists = new float[npts * dim]; + reader.read((char *)dists, npts * dim * sizeof(float)); + } + float min_dist = std::numeric_limits::max(); + float max_dist = 0; + groundtruth.resize(npts); + for (uint32_t i = 0; i < npts; i++) + { + groundtruth[i].clear(); + for (uint32_t j = 0; j < dim; j++) + { + if (dists[i * dim + j] <= range) + { + groundtruth[i].emplace_back(ids[i * dim + j]); + } + min_dist = min_dist > dists[i * dim + j] ? dists[i * dim + j] : min_dist; + max_dist = max_dist < dists[i * dim + j] ? dists[i * dim + j] : max_dist; + } + // std::cout<> &groundtruth, + uint64_t >_num) +{ + size_t read_blk_size = 64 * 1024 * 1024; + cached_ifstream reader(bin_file, read_blk_size); + diskann::cout << "Reading truthset file " << bin_file.c_str() << "... " << std::flush; + size_t actual_file_size = reader.get_file_size(); - if (truthset_type == -1) { - std::stringstream stream; - stream << "Error. File size mismatch. File should have bin format, with " - "npts followed by ngt followed by npts*ngt ids and optionally " - "followed by npts*ngt distance values; actual size: " - << actual_file_size - << ", expected: " << expected_file_size_with_dists << " or " - << expected_file_size_just_ids; - diskann::cout << stream.str(); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); - } - - ids = new uint32_t[npts * dim]; - reader.read((char *)ids, npts * dim * sizeof(uint32_t)); - - if (truthset_type == 1) { - dists = new float[npts * dim]; - reader.read((char *)dists, npts * dim * sizeof(float)); - } -} + int nptsuint32_t, totaluint32_t; + reader.read((char *)&nptsuint32_t, sizeof(int)); + reader.read((char *)&totaluint32_t, sizeof(int)); -inline void prune_truthset_for_range( - const std::string &bin_file, float range, - std::vector> &groundtruth, size_t &npts) { - size_t read_blk_size = 64 * 1024 * 1024; - cached_ifstream reader(bin_file, read_blk_size); - diskann::cout << "Reading truthset file " << bin_file.c_str() << "... " - << std::endl; - size_t actual_file_size = reader.get_file_size(); - - int npts_i32, dim_i32; - reader.read((char *)&npts_i32, sizeof(int)); - reader.read((char *)&dim_i32, sizeof(int)); - npts = (unsigned)npts_i32; - uint64_t dim = (unsigned)dim_i32; - uint32_t *ids; - float *dists; - - diskann::cout << "Metadata: #pts = " << npts << ", #dims = " << dim << "... " - << std::endl; - - int truthset_type = -1; // 1 means truthset has ids and distances, 2 means - // only ids, -1 is error - size_t expected_file_size_with_dists = - 2 * npts * dim * sizeof(uint32_t) + 2 * sizeof(uint32_t); - - if (actual_file_size == expected_file_size_with_dists) truthset_type = 1; - - if (truthset_type == -1) { - std::stringstream stream; - stream << "Error. File size mismatch. File should have bin format, with " - "npts followed by ngt followed by npts*ngt ids and optionally " - "followed by npts*ngt distance values; actual size: " - << actual_file_size - << ", expected: " << expected_file_size_with_dists; - diskann::cout << stream.str(); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); - } - - ids = new uint32_t[npts * dim]; - reader.read((char *)ids, npts * dim * sizeof(uint32_t)); - - if (truthset_type == 1) { - dists = new float[npts * dim]; - reader.read((char *)dists, npts * dim * sizeof(float)); - } - float min_dist = std::numeric_limits::max(); - float max_dist = 0; - groundtruth.resize(npts); - for (uint32_t i = 0; i < npts; i++) { - groundtruth[i].clear(); - for (uint32_t j = 0; j < dim; j++) { - if (dists[i * dim + j] <= range) { - groundtruth[i].emplace_back(ids[i * dim + j]); - } - min_dist = min_dist > dists[i * dim + j] ? dists[i * dim + j] : min_dist; - max_dist = max_dist < dists[i * dim + j] ? dists[i * dim + j] : max_dist; - } - // std::cout<> &groundtruth, - uint64_t >_num) { - size_t read_blk_size = 64 * 1024 * 1024; - cached_ifstream reader(bin_file, read_blk_size); - diskann::cout << "Reading truthset file " << bin_file.c_str() << "... " - << std::flush; - size_t actual_file_size = reader.get_file_size(); + diskann::cout << "Metadata: #pts = " << gt_num << ", #total_results = " << total_res << "..." << std::endl; - int nptsuint32_t, totaluint32_t; - reader.read((char *)&nptsuint32_t, sizeof(int)); - reader.read((char *)&totaluint32_t, sizeof(int)); + size_t expected_file_size = 2 * sizeof(uint32_t) + gt_num * sizeof(uint32_t) + total_res * sizeof(uint32_t); - gt_num = (uint64_t)nptsuint32_t; - uint64_t total_res = (uint64_t)totaluint32_t; + if (actual_file_size != expected_file_size) + { + std::stringstream stream; + stream << "Error. File size mismatch in range truthset. actual size: " << actual_file_size + << ", expected: " << expected_file_size; + diskann::cout << stream.str(); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } + groundtruth.clear(); + groundtruth.resize(gt_num); + std::vector gt_count(gt_num); - diskann::cout << "Metadata: #pts = " << gt_num - << ", #total_results = " << total_res << "..." << std::endl; + reader.read((char *)gt_count.data(), sizeof(uint32_t) * gt_num); - size_t expected_file_size = 2 * sizeof(uint32_t) + gt_num * sizeof(uint32_t) + - total_res * sizeof(uint32_t); + std::vector gt_stats(gt_count); + std::sort(gt_stats.begin(), gt_stats.end()); - if (actual_file_size != expected_file_size) { - std::stringstream stream; - stream << "Error. File size mismatch in range truthset. actual size: " - << actual_file_size << ", expected: " << expected_file_size; - diskann::cout << stream.str(); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); - } - groundtruth.clear(); - groundtruth.resize(gt_num); - std::vector gt_count(gt_num); - - reader.read((char *)gt_count.data(), sizeof(uint32_t) * gt_num); - - std::vector gt_stats(gt_count); - std::sort(gt_stats.begin(), gt_stats.end()); - - std::cout << "GT count percentiles:" << std::endl; - for (uint32_t p = 0; p < 100; p += 5) - std::cout << "percentile " << p << ": " - << gt_stats[static_cast(std::floor((p / 100.0) * gt_num))] - << std::endl; - std::cout << "percentile 100" - << ": " << gt_stats[gt_num - 1] << std::endl; - - for (uint32_t i = 0; i < gt_num; i++) { - groundtruth[i].clear(); - groundtruth[i].resize(gt_count[i]); - if (gt_count[i] != 0) - reader.read((char *)groundtruth[i].data(), - sizeof(uint32_t) * gt_count[i]); - } + std::cout << "GT count percentiles:" << std::endl; + for (uint32_t p = 0; p < 100; p += 5) + std::cout << "percentile " << p << ": " << gt_stats[static_cast(std::floor((p / 100.0) * gt_num))] + << std::endl; + std::cout << "percentile 100" + << ": " << gt_stats[gt_num - 1] << std::endl; + + for (uint32_t i = 0; i < gt_num; i++) + { + groundtruth[i].clear(); + groundtruth[i].resize(gt_count[i]); + if (gt_count[i] != 0) + reader.read((char *)groundtruth[i].data(), sizeof(uint32_t) * gt_count[i]); + } } #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, - size_t offset = 0) { - T *ptr; - load_bin(files, bin_file, ptr, npts, dim, offset); - data.reset(ptr); +inline void load_bin(MemoryMappedFiles &files, const std::string &bin_file, std::unique_ptr &data, size_t &npts, + size_t &dim, size_t offset = 0) +{ + T *ptr; + load_bin(files, bin_file, ptr, npts, dim, offset); + data.reset(ptr); } #endif -inline void copy_file(std::string in_file, std::string out_file) { - std::ifstream source(in_file, std::ios::binary); - std::ofstream dest(out_file, std::ios::binary); +inline void copy_file(std::string in_file, std::string out_file) +{ + std::ifstream source(in_file, std::ios::binary); + std::ofstream dest(out_file, std::ios::binary); - std::istreambuf_iterator begin_source(source); - std::istreambuf_iterator end_source; - std::ostreambuf_iterator begin_dest(dest); - std::copy(begin_source, end_source, begin_dest); + std::istreambuf_iterator begin_source(source); + std::istreambuf_iterator end_source; + std::ostreambuf_iterator begin_dest(dest); + std::copy(begin_source, end_source, begin_dest); - source.close(); - dest.close(); + source.close(); + dest.close(); } -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); -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 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 double calculate_range_search_recall( - unsigned num_queries, std::vector> &groundtruth, - std::vector> &our_results); +DISKANN_DLLEXPORT double calculate_range_search_recall(unsigned num_queries, + std::vector> &groundtruth, + std::vector> &our_results); template -inline void load_bin(const std::string &bin_file, std::unique_ptr &data, - size_t &npts, size_t &dim, size_t offset = 0) { - T *ptr; - load_bin(bin_file, ptr, npts, dim, offset); - data.reset(ptr); +inline void load_bin(const std::string &bin_file, std::unique_ptr &data, size_t &npts, size_t &dim, + size_t offset = 0) +{ + T *ptr; + load_bin(bin_file, ptr, npts, dim, offset); + data.reset(ptr); } -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()) { - char buff[1024]; +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()) + { + char buff[1024]; #ifdef _WINDOWS - auto ret = std::to_string(strerror_s(buff, 1024, errno)); + auto ret = std::to_string(strerror_s(buff, 1024, errno)); #else - auto ret = std::string(strerror_r(errno, buff, 1024)); + auto ret = std::string(strerror_r(errno, buff, 1024)); #endif - std::string error_message = std::string("Failed to open file") + filename + - " for write because " + buff + ", ret=" + ret; - diskann::cerr << error_message << std::endl; - throw diskann::ANNException(error_message, -1); - } + std::string error_message = + std::string("Failed to open file") + filename + " for write because " + buff + ", ret=" + ret; + diskann::cerr << error_message << std::endl; + throw diskann::ANNException(error_message, -1); + } } template -inline size_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; - 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 = " << bytes_written << "B" << std::endl; - - writer.write((char *)data, npts * ndims * sizeof(T)); - writer.close(); - diskann::cout << "Finished writing bin." << std::endl; - return bytes_written; +inline size_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; + 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 = " << bytes_written << "B" + << std::endl; + + writer.write((char *)data, npts * ndims * sizeof(T)); + writer.close(); + diskann::cout << "Finished writing bin." << std::endl; + return bytes_written; } -inline void print_progress(double percentage) { - int val = (int)(percentage * 100); - int lpad = (int)(percentage * PBWIDTH); - int rpad = PBWIDTH - lpad; - printf("\r%3d%% [%.*s%*s]", val, lpad, PBSTR, rpad, ""); - fflush(stdout); +inline void print_progress(double percentage) +{ + int val = (int)(percentage * 100); + int lpad = (int)(percentage * PBWIDTH); + int rpad = PBWIDTH - lpad; + printf("\r%3d%% [%.*s%*s]", val, lpad, PBSTR, rpad, ""); + fflush(stdout); } // load_aligned_bin functions START template -inline void load_aligned_bin_impl(std::basic_istream &reader, - size_t actual_file_size, T *&data, - size_t &npts, size_t &dim, - size_t &rounded_dim) { - 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; - - 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 of " << allocSize << " bytes... " - << std::flush; - alloc_aligned(((void **)&data), allocSize, 8 * sizeof(T)); - diskann::cout << "done. Copying data to mem_aligned buffer..." << 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; +inline void load_aligned_bin_impl(std::basic_istream &reader, size_t actual_file_size, T *&data, size_t &npts, + size_t &dim, size_t &rounded_dim) +{ + 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; + + 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 of " << allocSize << " bytes... " << std::flush; + alloc_aligned(((void **)&data), allocSize, 8 * sizeof(T)); + diskann::cout << "done. Copying data to mem_aligned buffer..." << 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; } #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) { - try { - diskann::cout << "Opening 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 actual_file_size = fc._size; - load_aligned_bin_impl(reader, actual_file_size, data, npts, dim, - rounded_dim); - } catch (std::system_error &e) { - throw FileException(bin_file, e, __FUNCSIG__, __FILE__, __LINE__); - } +inline void load_aligned_bin(MemoryMappedFiles &files, const std::string &bin_file, T *&data, size_t &npts, size_t &dim, + size_t &rounded_dim) +{ + try + { + diskann::cout << "Opening 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 actual_file_size = fc._size; + load_aligned_bin_impl(reader, actual_file_size, data, npts, dim, rounded_dim); + } + catch (std::system_error &e) + { + throw FileException(bin_file, e, __FUNCSIG__, __FILE__, __LINE__); + } } #endif template -inline void load_aligned_bin(const std::string &bin_file, T *&data, - size_t &npts, size_t &dim, size_t &rounded_dim) { - std::ifstream reader; - reader.exceptions(std::ifstream::failbit | std::ifstream::badbit); - - try { - diskann::cout << "Reading (with alignment) bin file " << bin_file << " ..." - << std::flush; - reader.open(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); - } catch (std::system_error &e) { - throw FileException(bin_file, e, __FUNCSIG__, __FILE__, __LINE__); - } +inline void load_aligned_bin(const std::string &bin_file, T *&data, size_t &npts, size_t &dim, size_t &rounded_dim) +{ + std::ifstream reader; + reader.exceptions(std::ifstream::failbit | std::ifstream::badbit); + + try + { + diskann::cout << "Reading (with alignment) bin file " << bin_file << " ..." << std::flush; + reader.open(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); + } + catch (std::system_error &e) + { + throw FileException(bin_file, e, __FUNCSIG__, __FILE__, __LINE__); + } } template -void convert_types(const InType *srcmat, OutType *destmat, size_t npts, - size_t dim) { +void convert_types(const InType *srcmat, OutType *destmat, size_t npts, size_t dim) +{ #pragma omp parallel for schedule(static, 65536) - for (int64_t i = 0; i < (int64_t)npts; i++) { - for (uint64_t j = 0; j < dim; j++) { - destmat[i * dim + j] = (OutType)srcmat[i * dim + j]; + for (int64_t i = 0; i < (int64_t)npts; i++) + { + for (uint64_t j = 0; j < dim; j++) + { + destmat[i * dim + j] = (OutType)srcmat[i * dim + j]; + } } - } } // this function will take in_file of n*d dimensions and save the output as a @@ -782,282 +801,295 @@ void convert_types(const InType *srcmat, OutType *destmat, size_t npts, // from MIPS to L2 search from "On Symmetric and Asymmetric LSHs for Inner // Product Search" by Neyshabur and Srebro -template -float prepare_base_for_inner_products(const std::string in_file, - const std::string out_file) { - std::cout << "Pre-processing base file by adding extra coordinate" - << std::endl; - std::ifstream in_reader(in_file.c_str(), std::ios::binary); - std::ofstream out_writer(out_file.c_str(), std::ios::binary); - uint64_t npts, in_dims, out_dims; - float max_norm = 0; - - uint32_t npts32, dims32; - in_reader.read((char *)&npts32, sizeof(uint32_t)); - in_reader.read((char *)&dims32, sizeof(uint32_t)); - - npts = npts32; - in_dims = dims32; - out_dims = in_dims + 1; - uint32_t outdims32 = (uint32_t)out_dims; - - out_writer.write((char *)&npts32, sizeof(uint32_t)); - out_writer.write((char *)&outdims32, sizeof(uint32_t)); - - size_t BLOCK_SIZE = 100000; - size_t block_size = npts <= BLOCK_SIZE ? npts : BLOCK_SIZE; - std::unique_ptr in_block_data = - std::make_unique(block_size * in_dims); - std::unique_ptr out_block_data = - std::make_unique(block_size * out_dims); - - std::memset(out_block_data.get(), 0, sizeof(float) * block_size * out_dims); - uint64_t num_blocks = DIV_ROUND_UP(npts, block_size); - - std::vector norms(npts, 0); - - for (uint64_t b = 0; b < num_blocks; b++) { - uint64_t start_id = b * block_size; - uint64_t end_id = (b + 1) * block_size < npts ? (b + 1) * block_size : npts; - uint64_t block_pts = end_id - start_id; - in_reader.read((char *)in_block_data.get(), - block_pts * in_dims * sizeof(T)); - for (uint64_t p = 0; p < block_pts; p++) { - for (uint64_t j = 0; j < in_dims; j++) { - norms[start_id + p] += - in_block_data[p * in_dims + j] * in_block_data[p * in_dims + j]; - } - max_norm = - max_norm > norms[start_id + p] ? max_norm : norms[start_id + p]; +template float prepare_base_for_inner_products(const std::string in_file, const std::string out_file) +{ + std::cout << "Pre-processing base file by adding extra coordinate" << std::endl; + std::ifstream in_reader(in_file.c_str(), std::ios::binary); + std::ofstream out_writer(out_file.c_str(), std::ios::binary); + uint64_t npts, in_dims, out_dims; + float max_norm = 0; + + uint32_t npts32, dims32; + in_reader.read((char *)&npts32, sizeof(uint32_t)); + in_reader.read((char *)&dims32, sizeof(uint32_t)); + + npts = npts32; + in_dims = dims32; + out_dims = in_dims + 1; + uint32_t outdims32 = (uint32_t)out_dims; + + out_writer.write((char *)&npts32, sizeof(uint32_t)); + out_writer.write((char *)&outdims32, sizeof(uint32_t)); + + size_t BLOCK_SIZE = 100000; + size_t block_size = npts <= BLOCK_SIZE ? npts : BLOCK_SIZE; + std::unique_ptr in_block_data = std::make_unique(block_size * in_dims); + std::unique_ptr out_block_data = std::make_unique(block_size * out_dims); + + std::memset(out_block_data.get(), 0, sizeof(float) * block_size * out_dims); + uint64_t num_blocks = DIV_ROUND_UP(npts, block_size); + + std::vector norms(npts, 0); + + for (uint64_t b = 0; b < num_blocks; b++) + { + uint64_t start_id = b * block_size; + uint64_t end_id = (b + 1) * block_size < npts ? (b + 1) * block_size : npts; + uint64_t block_pts = end_id - start_id; + in_reader.read((char *)in_block_data.get(), block_pts * in_dims * sizeof(T)); + for (uint64_t p = 0; p < block_pts; p++) + { + for (uint64_t j = 0; j < in_dims; j++) + { + norms[start_id + p] += in_block_data[p * in_dims + j] * in_block_data[p * in_dims + j]; + } + max_norm = max_norm > norms[start_id + p] ? max_norm : norms[start_id + p]; + } } - } - - max_norm = std::sqrt(max_norm); - - in_reader.seekg(2 * sizeof(uint32_t), std::ios::beg); - for (uint64_t b = 0; b < num_blocks; b++) { - uint64_t start_id = b * block_size; - uint64_t end_id = (b + 1) * block_size < npts ? (b + 1) * block_size : npts; - uint64_t block_pts = end_id - start_id; - in_reader.read((char *)in_block_data.get(), - block_pts * in_dims * sizeof(T)); - for (uint64_t p = 0; p < block_pts; p++) { - for (uint64_t j = 0; j < in_dims; j++) { - out_block_data[p * out_dims + j] = - in_block_data[p * in_dims + j] / max_norm; - } - float res = 1 - (norms[start_id + p] / (max_norm * max_norm)); - res = res <= 0 ? 0 : std::sqrt(res); - out_block_data[p * out_dims + out_dims - 1] = res; + + max_norm = std::sqrt(max_norm); + + in_reader.seekg(2 * sizeof(uint32_t), std::ios::beg); + for (uint64_t b = 0; b < num_blocks; b++) + { + uint64_t start_id = b * block_size; + uint64_t end_id = (b + 1) * block_size < npts ? (b + 1) * block_size : npts; + uint64_t block_pts = end_id - start_id; + in_reader.read((char *)in_block_data.get(), block_pts * in_dims * sizeof(T)); + for (uint64_t p = 0; p < block_pts; p++) + { + for (uint64_t j = 0; j < in_dims; j++) + { + out_block_data[p * out_dims + j] = in_block_data[p * in_dims + j] / max_norm; + } + float res = 1 - (norms[start_id + p] / (max_norm * max_norm)); + res = res <= 0 ? 0 : std::sqrt(res); + out_block_data[p * out_dims + out_dims - 1] = res; + } + out_writer.write((char *)out_block_data.get(), block_pts * out_dims * sizeof(float)); } - out_writer.write((char *)out_block_data.get(), - block_pts * out_dims * sizeof(float)); - } - out_writer.close(); - return max_norm; + out_writer.close(); + return max_norm; } // 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); +template void save_Tvecs(const char *filename, T *data, size_t npts, size_t ndims) +{ + std::string fname(filename); - // create cached ofstream with 64MB cache - cached_ofstream writer(fname, 64 * 1048576); + // create cached ofstream with 64MB cache + cached_ofstream writer(fname, 64 * 1048576); - unsigned dims_u32 = (unsigned)ndims; + unsigned dims_u32 = (unsigned)ndims; - // start writing - for (size_t i = 0; i < npts; i++) { - // write dims in u32 - writer.write((char *)&dims_u32, sizeof(unsigned)); + // start writing + for (size_t i = 0; i < npts; i++) + { + // write dims in u32 + writer.write((char *)&dims_u32, sizeof(unsigned)); - // get cur point in data - T *cur_pt = data + i * ndims; - writer.write((char *)cur_pt, ndims * sizeof(T)); - } + // get cur point in data + T *cur_pt = data + i * ndims; + writer.write((char *)cur_pt, ndims * sizeof(T)); + } } template -inline size_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; - size_t 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; +inline size_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; + size_t 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; } template -inline void copy_aligned_data_from_file(const char *bin_file, T *&data, - size_t &npts, size_t &dim, - const size_t &rounded_dim, - size_t offset = 0) { - if (data == nullptr) { - diskann::cerr << "Memory was not allocated for " << data - << " before calling the load function. Exiting..." - << std::endl; - throw diskann::ANNException( - "Null pointer passed to copy_aligned_data_from_file function", -1, - __FUNCSIG__, __FILE__, __LINE__); - } - std::ifstream reader; - reader.exceptions(std::ios::badbit | std::ios::failbit); - reader.open(bin_file, std::ios::binary); - reader.seekg(offset, reader.beg); - - 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; - - 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)); - } +inline void copy_aligned_data_from_file(const char *bin_file, T *&data, size_t &npts, size_t &dim, + const size_t &rounded_dim, size_t offset = 0) +{ + if (data == nullptr) + { + diskann::cerr << "Memory was not allocated for " << data << " before calling the load function. Exiting..." + << std::endl; + throw diskann::ANNException("Null pointer passed to copy_aligned_data_from_file function", -1, __FUNCSIG__, + __FILE__, __LINE__); + } + std::ifstream reader; + reader.exceptions(std::ios::badbit | std::ios::failbit); + reader.open(bin_file, std::ios::binary); + reader.seekg(offset, reader.beg); + + 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; + + 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)); + } } // NOTE :: good efficiency when total_vec_size is integral multiple of 64 -inline void prefetch_vector(const char *vec, size_t vecsize) { - size_t max_prefetch_size = (vecsize / 64) * 64; - for (size_t d = 0; d < max_prefetch_size; d += 64) - _mm_prefetch((const char *)vec + d, _MM_HINT_T0); +inline void prefetch_vector(const char *vec, size_t vecsize) +{ + size_t max_prefetch_size = (vecsize / 64) * 64; + for (size_t d = 0; d < max_prefetch_size; d += 64) + _mm_prefetch((const char *)vec + d, _MM_HINT_T0); } // NOTE :: good efficiency when total_vec_size is integral multiple of 64 -inline void prefetch_vector_l2(const char *vec, size_t vecsize) { - size_t max_prefetch_size = (vecsize / 64) * 64; - for (size_t d = 0; d < max_prefetch_size; d += 64) - _mm_prefetch((const char *)vec + d, _MM_HINT_T1); +inline void prefetch_vector_l2(const char *vec, size_t vecsize) +{ + size_t max_prefetch_size = (vecsize / 64) * 64; + for (size_t d = 0; d < max_prefetch_size; d += 64) + _mm_prefetch((const char *)vec + d, _MM_HINT_T1); } // NOTE: Implementation in utils.cpp. -void block_convert(std::ofstream &writr, std::ifstream &readr, float *read_buf, - uint64_t npts, uint64_t ndims); +void block_convert(std::ofstream &writr, std::ifstream &readr, float *read_buf, uint64_t npts, uint64_t ndims); -DISKANN_DLLEXPORT void normalize_data_file(const std::string &inFileName, - const std::string &outFileName); +DISKANN_DLLEXPORT void normalize_data_file(const std::string &inFileName, const std::string &outFileName); -}; // namespace diskann +}; // namespace diskann -struct PivotContainer { - PivotContainer() = default; +struct PivotContainer +{ + PivotContainer() = default; - PivotContainer(size_t pivo_id, float pivo_dist) - : piv_id{pivo_id}, piv_dist{pivo_dist} {} + PivotContainer(size_t pivo_id, float pivo_dist) : piv_id{pivo_id}, piv_dist{pivo_dist} + { + } - bool operator<(const PivotContainer &p) const { - return p.piv_dist < piv_dist; - } + bool operator<(const PivotContainer &p) const + { + return p.piv_dist < piv_dist; + } - bool operator>(const PivotContainer &p) const { - return p.piv_dist > piv_dist; - } + bool operator>(const PivotContainer &p) const + { + return p.piv_dist > piv_dist; + } - size_t piv_id; - float piv_dist; + size_t piv_id; + float piv_dist; }; -inline bool validate_index_file_size(std::ifstream &in) { - if (!in.is_open()) - throw diskann::ANNException( - "Index file size check called on unopened file stream", -1, __FUNCSIG__, - __FILE__, __LINE__); - in.seekg(0, in.end); - size_t actual_file_size = in.tellg(); - in.seekg(0, in.beg); - size_t expected_file_size; - in.read((char *)&expected_file_size, sizeof(uint64_t)); - in.seekg(0, in.beg); - if (actual_file_size != expected_file_size) { - diskann::cerr << "Index file size error. Expected size (metadata): " - << expected_file_size - << ", actual file size : " << actual_file_size << "." - << std::endl; - return false; - } - return true; +inline bool validate_index_file_size(std::ifstream &in) +{ + if (!in.is_open()) + throw diskann::ANNException("Index file size check called on unopened file stream", -1, __FUNCSIG__, __FILE__, + __LINE__); + in.seekg(0, in.end); + size_t actual_file_size = in.tellg(); + in.seekg(0, in.beg); + size_t expected_file_size; + in.read((char *)&expected_file_size, sizeof(uint64_t)); + in.seekg(0, in.beg); + if (actual_file_size != expected_file_size) + { + diskann::cerr << "Index file size error. Expected size (metadata): " << expected_file_size + << ", actual file size : " << actual_file_size << "." << std::endl; + return false; + } + return true; } -template -inline float get_norm(T *arr, const size_t dim) { - float sum = 0.0f; - for (uint32_t i = 0; i < dim; i++) { - sum += arr[i] * arr[i]; - } - return sqrt(sum); +template inline float get_norm(T *arr, const size_t dim) +{ + float sum = 0.0f; + for (uint32_t i = 0; i < dim; i++) + { + sum += arr[i] * arr[i]; + } + return sqrt(sum); } // This function is valid only for float data type. -template -inline void normalize(T *arr, const size_t dim) { - float norm = get_norm(arr, dim); - for (uint32_t i = 0; i < dim; i++) { - arr[i] = (T)(arr[i] / norm); - } +template inline void normalize(T *arr, const size_t dim) +{ + float norm = get_norm(arr, dim); + for (uint32_t i = 0; i < dim; i++) + { + arr[i] = (T)(arr[i] / norm); + } } -inline std::vector read_file_to_vector_of_strings( - const std::string &filename, bool unique = false) { - std::vector result; - std::set elementSet; - if (filename != "") { - std::ifstream file(filename); - if (file.fail()) { - throw diskann::ANNException( - std::string("Failed to open file ") + filename, -1); +inline std::vector read_file_to_vector_of_strings(const std::string &filename, bool unique = false) +{ + std::vector result; + std::set elementSet; + if (filename != "") + { + std::ifstream file(filename); + if (file.fail()) + { + throw diskann::ANNException(std::string("Failed to open file ") + filename, -1); + } + std::string line; + while (std::getline(file, line)) + { + if (line.empty()) + { + break; + } + if (line.find(',') != std::string::npos) + { + std::cerr << "Every query must have exactly one filter" << std::endl; + exit(-1); + } + if (!line.empty() && (line.back() == '\r' || line.back() == '\n')) + { + line.erase(line.size() - 1); + } + if (!elementSet.count(line)) + { + result.push_back(line); + } + if (unique) + { + elementSet.insert(line); + } + } + file.close(); } - std::string line; - while (std::getline(file, line)) { - if (line.empty()) { - break; - } - if (line.find(',') != std::string::npos) { - std::cerr << "Every query must have exactly one filter" << std::endl; - exit(-1); - } - if (!line.empty() && (line.back() == '\r' || line.back() == '\n')) { - line.erase(line.size() - 1); - } - if (!elementSet.count(line)) { - result.push_back(line); - } - if (unique) { - elementSet.insert(line); - } + else + { + throw diskann::ANNException(std::string("Failed to open file. filename can not be blank"), -1); } - file.close(); - } else { - throw diskann::ANNException( - std::string("Failed to open file. filename can not be blank"), -1); - } - return result; + return result; } -inline void clean_up_artifacts(tsl::robin_set paths_to_clean, - tsl::robin_set path_suffixes) { - try { - for (const auto &path : paths_to_clean) { - for (const auto &suffix : path_suffixes) { - std::string curr_path_to_clean(path + "_" + suffix); - if (std::remove(curr_path_to_clean.c_str()) != 0) - diskann::cout << "Warning: Unable to remove file :" - << curr_path_to_clean << std::endl; - } +inline void clean_up_artifacts(tsl::robin_set paths_to_clean, tsl::robin_set path_suffixes) +{ + try + { + for (const auto &path : paths_to_clean) + { + for (const auto &suffix : path_suffixes) + { + std::string curr_path_to_clean(path + "_" + suffix); + if (std::remove(curr_path_to_clean.c_str()) != 0) + diskann::cout << "Warning: Unable to remove file :" << curr_path_to_clean << std::endl; + } + } + diskann::cout << "Cleaned all artifacts" << std::endl; + } + catch (const std::exception &e) + { + diskann::cout << "Warning: Unable to clean all artifacts" << std::endl; } - diskann::cout << "Cleaned all artifacts" << std::endl; - } catch (const std::exception &e) { - diskann::cout << "Warning: Unable to clean all artifacts" << std::endl; - } } #ifdef _WINDOWS @@ -1067,53 +1099,57 @@ inline void clean_up_artifacts(tsl::robin_set paths_to_clean, extern bool AvxSupportedCPU; extern bool Avx2SupportedCPU; -inline size_t getMemoryUsage() { - PROCESS_MEMORY_COUNTERS_EX pmc; - GetProcessMemoryInfo(GetCurrentProcess(), (PROCESS_MEMORY_COUNTERS *)&pmc, - sizeof(pmc)); - return pmc.PrivateUsage; +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 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.0 * 1024 * 1024) - << "GB Working set size: " - << counters.WorkingSetSize * 1.0 / (1024.0 * 1024 * 1024) - << "GB Private bytes " - << counters.PagefileUsage * 1.0 / (1024 * 1024 * 1024) << "GB]" - << std::endl; +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.0 * 1024 * 1024) + << "GB Working set size: " << counters.WorkingSetSize * 1.0 / (1024.0 * 1024 * 1024) + << "GB Private bytes " << counters.PagefileUsage * 1.0 / (1024 * 1024 * 1024) << "GB]" << std::endl; } #else // need to check and change this -inline bool avx2Supported() { return true; } -inline void printProcessMemory(const char *) {} +inline bool avx2Supported() +{ + return true; +} +inline void printProcessMemory(const char *) +{ +} -inline size_t -getMemoryUsage() { // for non-windows, we have not implemented this function - return 0; +inline size_t getMemoryUsage() +{ // for non-windows, we have not implemented this function + return 0; } #endif diff --git a/include/windows_aligned_file_reader.h b/include/windows_aligned_file_reader.h index e7092a973..0d9a3173c 100644 --- a/include/windows_aligned_file_reader.h +++ b/include/windows_aligned_file_reader.h @@ -17,39 +17,41 @@ #include "utils.h" #include "windows_customizations.h" -class WindowsAlignedFileReader : public AlignedFileReader { - private: +class WindowsAlignedFileReader : public AlignedFileReader +{ + private: #ifdef UNICODE - std::wstring m_filename; + std::wstring m_filename; #else - std::string m_filename; + std::string m_filename; #endif - protected: - // virtual IOContext createContext(); - - public: - DISKANN_DLLEXPORT WindowsAlignedFileReader(){}; - DISKANN_DLLEXPORT virtual ~WindowsAlignedFileReader(){}; - - // Open & close ops - // Blocking calls - DISKANN_DLLEXPORT virtual void open(const std::string &fname) override; - DISKANN_DLLEXPORT virtual void close() override; - - DISKANN_DLLEXPORT virtual void register_thread() override; - DISKANN_DLLEXPORT virtual void deregister_thread() override { - // TODO: Needs implementation. - } - DISKANN_DLLEXPORT virtual void deregister_all_threads() override { - // TODO: Needs implementation. - } - DISKANN_DLLEXPORT virtual IOContext &get_ctx() override; - - // 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) override; + protected: + // virtual IOContext createContext(); + + public: + DISKANN_DLLEXPORT WindowsAlignedFileReader(){}; + DISKANN_DLLEXPORT virtual ~WindowsAlignedFileReader(){}; + + // Open & close ops + // Blocking calls + DISKANN_DLLEXPORT virtual void open(const std::string &fname) override; + DISKANN_DLLEXPORT virtual void close() override; + + DISKANN_DLLEXPORT virtual void register_thread() override; + DISKANN_DLLEXPORT virtual void deregister_thread() override + { + // TODO: Needs implementation. + } + DISKANN_DLLEXPORT virtual void deregister_all_threads() override + { + // TODO: Needs implementation. + } + DISKANN_DLLEXPORT virtual IOContext &get_ctx() override; + + // 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) override; }; -#endif // USE_BING_INFRA -#endif //_WINDOWS +#endif // USE_BING_INFRA +#endif //_WINDOWS diff --git a/include/windows_slim_lock.h b/include/windows_slim_lock.h index 8f37951a2..5d0d65508 100644 --- a/include/windows_slim_lock.h +++ b/include/windows_slim_lock.h @@ -6,7 +6,8 @@ #endif #include "Windows.h" -namespace diskann { +namespace diskann +{ // A thin C++ wrapper around Windows exclusive functionality of Windows // SlimReaderWriterLock. // @@ -17,42 +18,55 @@ namespace diskann { // // Full documentation can be found at. // https://msdn.microsoft.com/en-us/library/windows/desktop/aa904937(v=vs.85).aspx -class windows_exclusive_slim_lock { - public: - windows_exclusive_slim_lock() : _lock(SRWLOCK_INIT) {} +class windows_exclusive_slim_lock +{ + public: + windows_exclusive_slim_lock() : _lock(SRWLOCK_INIT) + { + } - // The lock is non-copyable. This also disables move constructor/operator=. - windows_exclusive_slim_lock(const windows_exclusive_slim_lock &) = delete; - windows_exclusive_slim_lock &operator=(const windows_exclusive_slim_lock &) = - delete; + // The lock is non-copyable. This also disables move constructor/operator=. + windows_exclusive_slim_lock(const windows_exclusive_slim_lock &) = delete; + windows_exclusive_slim_lock &operator=(const windows_exclusive_slim_lock &) = delete; - void lock() { return AcquireSRWLockExclusive(&_lock); } + void lock() + { + return AcquireSRWLockExclusive(&_lock); + } - bool try_lock() { return TryAcquireSRWLockExclusive(&_lock) != FALSE; } + bool try_lock() + { + return TryAcquireSRWLockExclusive(&_lock) != FALSE; + } - void unlock() { return ReleaseSRWLockExclusive(&_lock); } + void unlock() + { + return ReleaseSRWLockExclusive(&_lock); + } - private: - SRWLOCK _lock; + private: + SRWLOCK _lock; }; // An exclusive lock over a SlimReaderWriterLock. -class windows_exclusive_slim_lock_guard { - public: - windows_exclusive_slim_lock_guard(windows_exclusive_slim_lock &p_lock) - : _lock(p_lock) { - _lock.lock(); - } - - // The lock is non-copyable. This also disables move constructor/operator=. - windows_exclusive_slim_lock_guard(const windows_exclusive_slim_lock_guard &) = - delete; - windows_exclusive_slim_lock_guard &operator=( - const windows_exclusive_slim_lock_guard &) = delete; - - ~windows_exclusive_slim_lock_guard() { _lock.unlock(); } - - private: - windows_exclusive_slim_lock &_lock; +class windows_exclusive_slim_lock_guard +{ + public: + windows_exclusive_slim_lock_guard(windows_exclusive_slim_lock &p_lock) : _lock(p_lock) + { + _lock.lock(); + } + + // The lock is non-copyable. This also disables move constructor/operator=. + windows_exclusive_slim_lock_guard(const windows_exclusive_slim_lock_guard &) = delete; + windows_exclusive_slim_lock_guard &operator=(const windows_exclusive_slim_lock_guard &) = delete; + + ~windows_exclusive_slim_lock_guard() + { + _lock.unlock(); + } + + private: + windows_exclusive_slim_lock &_lock; }; -} // namespace diskann +} // namespace diskann diff --git a/python/src/diskann_bindings.cpp b/python/src/diskann_bindings.cpp index c7b63233d..9b6b97bc0 100644 --- a/python/src/diskann_bindings.cpp +++ b/python/src/diskann_bindings.cpp @@ -30,458 +30,418 @@ PYBIND11_MAKE_OPAQUE(std::vector); namespace py = pybind11; using namespace diskann; -template -struct DiskANNIndex { - PQFlashIndex *pq_flash_index; - std::shared_ptr reader; +template struct DiskANNIndex +{ + PQFlashIndex *pq_flash_index; + std::shared_ptr reader; - DiskANNIndex(diskann::Metric metric) { + DiskANNIndex(diskann::Metric metric) + { #ifdef _WINDOWS - reader = std::make_shared(); + reader = std::make_shared(); #else - reader = std::make_shared(); + reader = std::make_shared(); #endif - pq_flash_index = new PQFlashIndex(reader, metric); - } - - ~DiskANNIndex() { delete pq_flash_index; } + pq_flash_index = new PQFlashIndex(reader, metric); + } - auto get_metric() { return pq_flash_index->get_metric(); } + ~DiskANNIndex() + { + delete pq_flash_index; + } - void cache_bfs_levels(size_t num_nodes_to_cache) { - std::vector node_list; - pq_flash_index->cache_bfs_levels(num_nodes_to_cache, node_list); - pq_flash_index->load_cache_list(node_list); - } + auto get_metric() + { + return pq_flash_index->get_metric(); + } - void cache_sample_paths(size_t num_nodes_to_cache, - const std::string &warmup_query_file, - uint32_t num_threads) { - if (!file_exists(warmup_query_file)) { - return; + void cache_bfs_levels(size_t num_nodes_to_cache) + { + std::vector node_list; + pq_flash_index->cache_bfs_levels(num_nodes_to_cache, node_list); + pq_flash_index->load_cache_list(node_list); } - std::vector node_list; - pq_flash_index->generate_cache_list_from_sample_queries( - warmup_query_file, 15, 4, num_nodes_to_cache, num_threads, node_list); - pq_flash_index->load_cache_list(node_list); - } - - int load_index(const std::string &index_path_prefix, const int num_threads, - const size_t num_nodes_to_cache, int cache_mechanism) { - int load_success = - pq_flash_index->load(num_threads, index_path_prefix.c_str()); - if (load_success != 0) { - throw std::runtime_error("load_index failed."); + void cache_sample_paths(size_t num_nodes_to_cache, const std::string &warmup_query_file, uint32_t num_threads) + { + if (!file_exists(warmup_query_file)) + { + return; + } + + std::vector node_list; + pq_flash_index->generate_cache_list_from_sample_queries(warmup_query_file, 15, 4, num_nodes_to_cache, + num_threads, node_list); + pq_flash_index->load_cache_list(node_list); } - if (cache_mechanism == 0) { - // Nothing to do - } else if (cache_mechanism == 1) { - std::string sample_file = - index_path_prefix + std::string("_sample_data.bin"); - cache_sample_paths(num_nodes_to_cache, sample_file, num_threads); - } else if (cache_mechanism == 2) { - cache_bfs_levels(num_nodes_to_cache); + + int load_index(const std::string &index_path_prefix, const int num_threads, const size_t num_nodes_to_cache, + int cache_mechanism) + { + int load_success = pq_flash_index->load(num_threads, index_path_prefix.c_str()); + if (load_success != 0) + { + throw std::runtime_error("load_index failed."); + } + if (cache_mechanism == 0) + { + // Nothing to do + } + else if (cache_mechanism == 1) + { + std::string sample_file = index_path_prefix + std::string("_sample_data.bin"); + cache_sample_paths(num_nodes_to_cache, sample_file, num_threads); + } + else if (cache_mechanism == 2) + { + cache_bfs_levels(num_nodes_to_cache); + } + return 0; } - return 0; - } - auto search(py::array_t &query, - const uint64_t knn, const uint64_t l_search, - const uint64_t beam_width) { - py::array_t ids(knn); - py::array_t dists(knn); + auto search(py::array_t &query, const uint64_t knn, + const uint64_t l_search, const uint64_t beam_width) + { + py::array_t ids(knn); + py::array_t dists(knn); - std::vector u32_ids(knn); - std::vector u64_ids(knn); - QueryStats stats; + std::vector u32_ids(knn); + std::vector u64_ids(knn); + QueryStats stats; - pq_flash_index->cached_beam_search(query.data(), knn, l_search, - u64_ids.data(), dists.mutable_data(), - beam_width, false, &stats); + pq_flash_index->cached_beam_search(query.data(), knn, l_search, u64_ids.data(), dists.mutable_data(), + beam_width, false, &stats); - auto r = ids.mutable_unchecked<1>(); - for (uint64_t i = 0; i < knn; ++i) r(i) = (unsigned)u64_ids[i]; + auto r = ids.mutable_unchecked<1>(); + for (uint64_t i = 0; i < knn; ++i) + r(i) = (unsigned)u64_ids[i]; - return std::make_pair(ids, dists); - } + return std::make_pair(ids, dists); + } - auto batch_search( - py::array_t &queries, - const uint64_t num_queries, const uint64_t knn, const uint64_t l_search, - const uint64_t beam_width, const int num_threads) { - py::array_t ids({num_queries, knn}); - py::array_t dists({num_queries, knn}); + auto batch_search(py::array_t &queries, const uint64_t num_queries, + const uint64_t knn, const uint64_t l_search, const uint64_t beam_width, const int num_threads) + { + py::array_t ids({num_queries, knn}); + py::array_t dists({num_queries, knn}); - omp_set_num_threads(num_threads); + omp_set_num_threads(num_threads); - std::vector u64_ids(knn * num_queries); + std::vector u64_ids(knn * num_queries); #pragma omp parallel for schedule(dynamic, 1) - for (int64_t i = 0; i < (int64_t)num_queries; i++) { - pq_flash_index->cached_beam_search(queries.data(i), knn, l_search, - u64_ids.data() + i * knn, - dists.mutable_data(i), beam_width); + for (int64_t i = 0; i < (int64_t)num_queries; i++) + { + pq_flash_index->cached_beam_search(queries.data(i), knn, l_search, u64_ids.data() + i * knn, + dists.mutable_data(i), beam_width); + } + + auto r = ids.mutable_unchecked(); + for (uint64_t i = 0; i < num_queries; ++i) + for (uint64_t j = 0; j < knn; ++j) + r(i, j) = (unsigned)u64_ids[i * knn + j]; + + return std::make_pair(ids, dists); } - - auto r = ids.mutable_unchecked(); - for (uint64_t i = 0; i < num_queries; ++i) - for (uint64_t j = 0; j < knn; ++j) - r(i, j) = (unsigned)u64_ids[i * knn + j]; - - return std::make_pair(ids, dists); - } }; typedef uint32_t IdT; typedef uint32_t filterT; -template -struct DynamicInMemIndex { - Index *_index; - const IndexWriteParameters write_params; - - DynamicInMemIndex(Metric m, const size_t dim, const size_t max_points, - const IndexWriteParameters &index_parameters, - const uint32_t initial_search_list_size, - const uint32_t search_threads, - const bool concurrent_consolidate) - : write_params(index_parameters) { - _index = new Index( - m, dim, max_points, - true, // dynamic_index - index_parameters, // used for insert - initial_search_list_size, // used to prepare the scratch space for - // searching. can / may be expanded if the - // search asks for a larger L. - search_threads, // also used for the scratch space - true, // enable_tags - concurrent_consolidate, - false, // pq_dist_build - 0, // num_pq_chunks - false); // use_opq = false - } - - ~DynamicInMemIndex() { delete _index; } - - int insert(py::array_t &vector, - const IdT id) { - return _index->insert_point(vector.data(), id); - } - - int mark_deleted(const IdT id) { return _index->lazy_delete(id); } - - auto search(py::array_t &query, - const uint64_t knn, const uint64_t l_search) { - py::array_t ids(knn); - py::array_t dists(knn); - std::vector empty_vector; - _index->search_with_tags(query.data(), knn, l_search, ids.mutable_data(), - dists.mutable_data(), empty_vector); - return std::make_pair(ids, dists); - } - - auto batch_search( - py::array_t &queries, - const uint64_t num_queries, const uint64_t knn, const uint64_t l_search, - const int num_threads) { - py::array_t ids({num_queries, knn}); - py::array_t dists({num_queries, knn}); - std::vector empty_vector; - - omp_set_num_threads(num_threads); +template struct DynamicInMemIndex +{ + Index *_index; + const IndexWriteParameters write_params; + + DynamicInMemIndex(Metric m, const size_t dim, const size_t max_points, const IndexWriteParameters &index_parameters, + const uint32_t initial_search_list_size, const uint32_t search_threads, + const bool concurrent_consolidate) + : write_params(index_parameters) + { + _index = new Index(m, dim, max_points, + true, // dynamic_index + index_parameters, // used for insert + initial_search_list_size, // used to prepare the scratch space for + // searching. can / may be expanded if the + // search asks for a larger L. + search_threads, // also used for the scratch space + true, // enable_tags + concurrent_consolidate, + false, // pq_dist_build + 0, // num_pq_chunks + false); // use_opq = false + } -#pragma omp parallel for schedule(dynamic, 1) - for (int64_t i = 0; i < (int64_t)num_queries; i++) { - _index->search_with_tags(queries.data(i), knn, l_search, - ids.mutable_data(i), dists.mutable_data(i), - empty_vector); + ~DynamicInMemIndex() + { + delete _index; } - return std::make_pair(ids, dists); - } + int insert(py::array_t &vector, const IdT id) + { + return _index->insert_point(vector.data(), id); + } - auto consolidate_delete() { - return _index->consolidate_deletes(write_params); - } -}; + int mark_deleted(const IdT id) + { + return _index->lazy_delete(id); + } + + auto search(py::array_t &query, const uint64_t knn, + const uint64_t l_search) + { + py::array_t ids(knn); + py::array_t dists(knn); + std::vector empty_vector; + _index->search_with_tags(query.data(), knn, l_search, ids.mutable_data(), dists.mutable_data(), empty_vector); + return std::make_pair(ids, dists); + } + + auto batch_search(py::array_t &queries, const uint64_t num_queries, + const uint64_t knn, const uint64_t l_search, const int num_threads) + { + py::array_t ids({num_queries, knn}); + py::array_t dists({num_queries, knn}); + std::vector empty_vector; -template -struct StaticInMemIndex { - Index *_index; - - StaticInMemIndex(Metric m, const std::string &data_path, - IndexWriteParameters &index_parameters) { - size_t ndims, npoints; - diskann::get_bin_metadata(data_path, npoints, ndims); - _index = new Index(m, ndims, npoints, - false, // not a dynamic_index - false, // no enable_tags/ids - false, // no concurrent_consolidate, - false, // pq_dist_build - 0, // num_pq_chunks - false, // use_opq = false - 0); // num_frozen_pts = 0 - _index->build(data_path.c_str(), npoints, index_parameters); - } - - ~StaticInMemIndex() { delete _index; } - - auto search(py::array_t &query, - const uint64_t knn, const uint64_t l_search) { - py::array_t ids(knn); - py::array_t dists(knn); - std::vector empty_vector; - _index->search(query.data(), knn, l_search, ids.mutable_data(), - dists.mutable_data()); - return std::make_pair(ids, dists); - } - - auto batch_search( - py::array_t &queries, - const uint64_t num_queries, const uint64_t knn, const uint64_t l_search, - const int num_threads) { - py::array_t ids({num_queries, knn}); - py::array_t dists({num_queries, knn}); - std::vector empty_vector; - - omp_set_num_threads(num_threads); + omp_set_num_threads(num_threads); #pragma omp parallel for schedule(dynamic, 1) - for (int64_t i = 0; i < (int64_t)num_queries; i++) { - _index->search(queries.data(i), knn, l_search, ids.mutable_data(i), - dists.mutable_data(i)); + for (int64_t i = 0; i < (int64_t)num_queries; i++) + { + _index->search_with_tags(queries.data(i), knn, l_search, ids.mutable_data(i), dists.mutable_data(i), + empty_vector); + } + + return std::make_pair(ids, dists); + } + + auto consolidate_delete() + { + return _index->consolidate_deletes(write_params); + } +}; + +template struct StaticInMemIndex +{ + Index *_index; + + StaticInMemIndex(Metric m, const std::string &data_path, IndexWriteParameters &index_parameters) + { + size_t ndims, npoints; + diskann::get_bin_metadata(data_path, npoints, ndims); + _index = new Index(m, ndims, npoints, + false, // not a dynamic_index + false, // no enable_tags/ids + false, // no concurrent_consolidate, + false, // pq_dist_build + 0, // num_pq_chunks + false, // use_opq = false + 0); // num_frozen_pts = 0 + _index->build(data_path.c_str(), npoints, index_parameters); + } + + ~StaticInMemIndex() + { + delete _index; } - return std::make_pair(ids, dists); - } + auto search(py::array_t &query, const uint64_t knn, + const uint64_t l_search) + { + py::array_t ids(knn); + py::array_t dists(knn); + std::vector empty_vector; + _index->search(query.data(), knn, l_search, ids.mutable_data(), dists.mutable_data()); + return std::make_pair(ids, dists); + } + + auto batch_search(py::array_t &queries, const uint64_t num_queries, + const uint64_t knn, const uint64_t l_search, const int num_threads) + { + py::array_t ids({num_queries, knn}); + py::array_t dists({num_queries, knn}); + std::vector empty_vector; + + omp_set_num_threads(num_threads); + +#pragma omp parallel for schedule(dynamic, 1) + for (int64_t i = 0; i < (int64_t)num_queries; i++) + { + _index->search(queries.data(i), knn, l_search, ids.mutable_data(i), dists.mutable_data(i)); + } + + return std::make_pair(ids, dists); + } }; -PYBIND11_MODULE(_diskannpy, m) { - m.doc() = "DiskANN Python Bindings"; +PYBIND11_MODULE(_diskannpy, m) +{ + m.doc() = "DiskANN Python Bindings"; #ifdef VERSION_INFO - m.attr("__version__") = VERSION_INFO; + m.attr("__version__") = VERSION_INFO; #else - m.attr("__version__") = "dev"; + m.attr("__version__") = "dev"; #endif - py::enum_(m, "Metric") - .value("L2", Metric::L2) - .value("INNER_PRODUCT", Metric::INNER_PRODUCT) - .export_values(); - - py::class_>(m, "DiskANNStaticInMemFloatIndex") - .def(py::init([](diskann::Metric metric, const std::string &data_path, - IndexWriteParameters &index_parameters) { - return std::unique_ptr>( - new StaticInMemIndex(metric, data_path, index_parameters)); - })) - .def("search", &StaticInMemIndex::search, py::arg("query"), - py::arg("knn"), py::arg("l_search")) - .def("batch_search", &StaticInMemIndex::batch_search, - py::arg("queries"), py::arg("num_queries"), py::arg("knn"), - py::arg("l_search"), py::arg("num_threads")); - - py::class_>(m, "DiskANNStaticInMemInt8Index") - .def(py::init([](diskann::Metric metric, const std::string &data_path, - IndexWriteParameters &index_parameters) { - return std::unique_ptr>( - new StaticInMemIndex(metric, data_path, index_parameters)); - })) - .def("search", &StaticInMemIndex::search, py::arg("query"), - py::arg("knn"), py::arg("l_search")) - .def("batch_search", &StaticInMemIndex::batch_search, - py::arg("queries"), py::arg("num_queries"), py::arg("knn"), - py::arg("l_search"), py::arg("num_threads")); - - py::class_>(m, "DiskANNStaticInMemUint8Index") - .def(py::init([](diskann::Metric metric, const std::string &data_path, - IndexWriteParameters &index_parameters) { - return std::unique_ptr>( - new StaticInMemIndex(metric, data_path, index_parameters)); - })) - .def("search", &StaticInMemIndex::search, py::arg("query"), - py::arg("knn"), py::arg("l_search")) - .def("batch_search", &StaticInMemIndex::batch_search, - py::arg("queries"), py::arg("num_queries"), py::arg("knn"), - py::arg("l_search"), py::arg("num_threads")); - - py::class_>(m, "DiskANNDynamicInMemFloatIndex") - .def(py::init( - [](diskann::Metric metric, const size_t dim, const size_t max_points, - const IndexWriteParameters &index_parameters, - const uint32_t initial_search_list_size, - const uint32_t search_threads, const bool concurrent_consolidate) { + py::enum_(m, "Metric") + .value("L2", Metric::L2) + .value("INNER_PRODUCT", Metric::INNER_PRODUCT) + .export_values(); + + py::class_>(m, "DiskANNStaticInMemFloatIndex") + .def(py::init([](diskann::Metric metric, const std::string &data_path, IndexWriteParameters &index_parameters) { + return std::unique_ptr>( + new StaticInMemIndex(metric, data_path, index_parameters)); + })) + .def("search", &StaticInMemIndex::search, py::arg("query"), py::arg("knn"), py::arg("l_search")) + .def("batch_search", &StaticInMemIndex::batch_search, py::arg("queries"), py::arg("num_queries"), + py::arg("knn"), py::arg("l_search"), py::arg("num_threads")); + + py::class_>(m, "DiskANNStaticInMemInt8Index") + .def(py::init([](diskann::Metric metric, const std::string &data_path, IndexWriteParameters &index_parameters) { + return std::unique_ptr>( + new StaticInMemIndex(metric, data_path, index_parameters)); + })) + .def("search", &StaticInMemIndex::search, py::arg("query"), py::arg("knn"), py::arg("l_search")) + .def("batch_search", &StaticInMemIndex::batch_search, py::arg("queries"), py::arg("num_queries"), + py::arg("knn"), py::arg("l_search"), py::arg("num_threads")); + + py::class_>(m, "DiskANNStaticInMemUint8Index") + .def(py::init([](diskann::Metric metric, const std::string &data_path, IndexWriteParameters &index_parameters) { + return std::unique_ptr>( + new StaticInMemIndex(metric, data_path, index_parameters)); + })) + .def("search", &StaticInMemIndex::search, py::arg("query"), py::arg("knn"), py::arg("l_search")) + .def("batch_search", &StaticInMemIndex::batch_search, py::arg("queries"), py::arg("num_queries"), + py::arg("knn"), py::arg("l_search"), py::arg("num_threads")); + + py::class_>(m, "DiskANNDynamicInMemFloatIndex") + .def(py::init([](diskann::Metric metric, const size_t dim, const size_t max_points, + const IndexWriteParameters &index_parameters, const uint32_t initial_search_list_size, + const uint32_t search_threads, const bool concurrent_consolidate) { return std::unique_ptr>( - new DynamicInMemIndex( - metric, dim, max_points, index_parameters, - initial_search_list_size, search_threads, - concurrent_consolidate)); - })) - .def("search", &DynamicInMemIndex::search, py::arg("query"), - py::arg("knn"), py::arg("l_search")) - .def("batch_search", &DynamicInMemIndex::batch_search, - py::arg("queries"), py::arg("num_queries"), py::arg("knn"), - py::arg("l_search"), py::arg("num_threads")) - .def("insert", &DynamicInMemIndex::insert, py::arg("vector"), - py::arg("id")) - .def("mark_deleted", &DynamicInMemIndex::mark_deleted, - py::arg("id")) - .def("consolidate_delete", &DynamicInMemIndex::consolidate_delete); - - py::class_>(m, "DiskANNDynamicInMemInt8Index") - .def(py::init( - [](diskann::Metric metric, const size_t dim, const size_t max_points, - const IndexWriteParameters &index_parameters, - const uint32_t initial_search_list_size, - const uint32_t search_threads, const bool concurrent_consolidate) { + new DynamicInMemIndex(metric, dim, max_points, index_parameters, initial_search_list_size, + search_threads, concurrent_consolidate)); + })) + .def("search", &DynamicInMemIndex::search, py::arg("query"), py::arg("knn"), py::arg("l_search")) + .def("batch_search", &DynamicInMemIndex::batch_search, py::arg("queries"), py::arg("num_queries"), + py::arg("knn"), py::arg("l_search"), py::arg("num_threads")) + .def("insert", &DynamicInMemIndex::insert, py::arg("vector"), py::arg("id")) + .def("mark_deleted", &DynamicInMemIndex::mark_deleted, py::arg("id")) + .def("consolidate_delete", &DynamicInMemIndex::consolidate_delete); + + py::class_>(m, "DiskANNDynamicInMemInt8Index") + .def(py::init([](diskann::Metric metric, const size_t dim, const size_t max_points, + const IndexWriteParameters &index_parameters, const uint32_t initial_search_list_size, + const uint32_t search_threads, const bool concurrent_consolidate) { return std::unique_ptr>( - new DynamicInMemIndex( - metric, dim, max_points, index_parameters, - initial_search_list_size, search_threads, - concurrent_consolidate)); - })) - .def("search", &DynamicInMemIndex::search, py::arg("query"), - py::arg("knn"), py::arg("l_search")) - .def("batch_search", &DynamicInMemIndex::batch_search, - py::arg("queries"), py::arg("num_queries"), py::arg("knn"), - py::arg("l_search"), py::arg("num_threads")) - .def("insert", &DynamicInMemIndex::insert, py::arg("vector"), - py::arg("id")) - .def("mark_deleted", &DynamicInMemIndex::mark_deleted, - py::arg("id")) - .def("consolidate_delete", - &DynamicInMemIndex::consolidate_delete); - - py::class_>(m, "DiskANNDynamicInMemUint8Index") - .def(py::init( - [](diskann::Metric metric, const size_t dim, const size_t max_points, - const IndexWriteParameters &index_parameters, - const uint32_t initial_search_list_size, - const uint32_t search_threads, const bool concurrent_consolidate) { + new DynamicInMemIndex(metric, dim, max_points, index_parameters, initial_search_list_size, + search_threads, concurrent_consolidate)); + })) + .def("search", &DynamicInMemIndex::search, py::arg("query"), py::arg("knn"), py::arg("l_search")) + .def("batch_search", &DynamicInMemIndex::batch_search, py::arg("queries"), py::arg("num_queries"), + py::arg("knn"), py::arg("l_search"), py::arg("num_threads")) + .def("insert", &DynamicInMemIndex::insert, py::arg("vector"), py::arg("id")) + .def("mark_deleted", &DynamicInMemIndex::mark_deleted, py::arg("id")) + .def("consolidate_delete", &DynamicInMemIndex::consolidate_delete); + + py::class_>(m, "DiskANNDynamicInMemUint8Index") + .def(py::init([](diskann::Metric metric, const size_t dim, const size_t max_points, + const IndexWriteParameters &index_parameters, const uint32_t initial_search_list_size, + const uint32_t search_threads, const bool concurrent_consolidate) { return std::unique_ptr>( - new DynamicInMemIndex( - metric, dim, max_points, index_parameters, - initial_search_list_size, search_threads, - concurrent_consolidate)); - })) - .def("search", &DynamicInMemIndex::search, py::arg("query"), - py::arg("knn"), py::arg("l_search")) - .def("batch_search", &DynamicInMemIndex::batch_search, - py::arg("queries"), py::arg("num_queries"), py::arg("knn"), - py::arg("l_search"), py::arg("num_threads")) - .def("insert", &DynamicInMemIndex::insert, py::arg("vector"), - py::arg("id")) - .def("mark_deleted", &DynamicInMemIndex::mark_deleted, - py::arg("id")) - .def("consolidate_delete", - &DynamicInMemIndex::consolidate_delete); - - py::class_>(m, "DiskANNFloatIndex") - .def(py::init([](diskann::Metric metric) { - return std::unique_ptr>( - new DiskANNIndex(metric)); - })) - .def("cache_bfs_levels", &DiskANNIndex::cache_bfs_levels, - py::arg("num_nodes_to_cache")) - .def("load_index", &DiskANNIndex::load_index, - py::arg("index_path_prefix"), py::arg("num_threads"), - py::arg("num_nodes_to_cache"), py::arg("cache_mechanism") = 1) - .def("search", &DiskANNIndex::search, py::arg("query"), - py::arg("knn"), py::arg("l_search"), py::arg("beam_width")) - .def("batch_search", &DiskANNIndex::batch_search, - py::arg("queries"), py::arg("num_queries"), py::arg("knn"), - py::arg("l_search"), py::arg("beam_width"), py::arg("num_threads")) - .def( - "build", - [](DiskANNIndex &self, const char *data_file_path, - const char *index_prefix_path, unsigned R, unsigned L, - double final_index_ram_limit, double indexing_ram_budget, - unsigned num_threads, unsigned pq_disk_bytes) { - std::string params = std::to_string(R) + " " + std::to_string(L) + - " " + std::to_string(final_index_ram_limit) + - " " + std::to_string(indexing_ram_budget) + - " " + std::to_string(num_threads); - if (pq_disk_bytes > 0) { - params = params + " " + std::to_string(pq_disk_bytes); - } - diskann::build_disk_index(data_file_path, index_prefix_path, - params.c_str(), self.get_metric()); - }, - py::arg("data_file_path"), py::arg("index_prefix_path"), py::arg("R"), - py::arg("L"), py::arg("final_index_ram_limit"), - py::arg("indexing_ram_limit"), py::arg("num_threads"), - py::arg("pq_disk_bytes") = 0); - - py::class_>(m, "DiskANNInt8Index") - .def(py::init([](diskann::Metric metric) { - return std::unique_ptr>( - new DiskANNIndex(metric)); - })) - .def("cache_bfs_levels", &DiskANNIndex::cache_bfs_levels, - py::arg("num_nodes_to_cache")) - .def("load_index", &DiskANNIndex::load_index, - py::arg("index_path_prefix"), py::arg("num_threads"), - py::arg("num_nodes_to_cache"), py::arg("cache_mechanism") = 1) - .def("search", &DiskANNIndex::search, py::arg("query"), - py::arg("knn"), py::arg("l_search"), py::arg("beam_width")) - .def("batch_search", &DiskANNIndex::batch_search, - py::arg("queries"), py::arg("num_queries"), py::arg("knn"), - py::arg("l_search"), py::arg("beam_width"), py::arg("num_threads")) - .def( - "build", - [](DiskANNIndex &self, const char *data_file_path, - const char *index_prefix_path, unsigned R, unsigned L, - double final_index_ram_limit, double indexing_ram_budget, - unsigned num_threads, unsigned pq_disk_bytes) { - std::string params = std::to_string(R) + " " + std::to_string(L) + - " " + std::to_string(final_index_ram_limit) + - " " + std::to_string(indexing_ram_budget) + - " " + std::to_string(num_threads); - if (pq_disk_bytes > 0) - params = params + " " + std::to_string(pq_disk_bytes); - diskann::build_disk_index(data_file_path, index_prefix_path, - params.c_str(), - self.get_metric()); - }, - py::arg("data_file_path"), py::arg("index_prefix_path"), py::arg("R"), - py::arg("L"), py::arg("final_index_ram_limit"), - py::arg("indexing_ram_limit"), py::arg("num_threads"), - py::arg("pq_disk_bytes") = 0); - - py::class_>(m, "DiskANNUInt8Index") - .def(py::init([](diskann::Metric metric) { - return std::unique_ptr>( - new DiskANNIndex(metric)); - })) - .def("cache_bfs_levels", &DiskANNIndex::cache_bfs_levels, - py::arg("num_nodes_to_cache")) - .def("load_index", &DiskANNIndex::load_index, - py::arg("index_path_prefix"), py::arg("num_threads"), - py::arg("num_nodes_to_cache"), py::arg("cache_mechanism") = 1) - .def("search", &DiskANNIndex::search, py::arg("query"), - py::arg("knn"), py::arg("l_search"), py::arg("beam_width")) - .def("batch_search", &DiskANNIndex::batch_search, - py::arg("queries"), py::arg("num_queries"), py::arg("knn"), - py::arg("l_search"), py::arg("beam_width"), py::arg("num_threads")) - .def( - "build", - [](DiskANNIndex &self, const char *data_file_path, - const char *index_prefix_path, unsigned R, unsigned L, - double final_index_ram_limit, double indexing_ram_budget, - unsigned num_threads, unsigned pq_disk_bytes) { - std::string params = std::to_string(R) + " " + std::to_string(L) + - " " + std::to_string(final_index_ram_limit) + - " " + std::to_string(indexing_ram_budget) + - " " + std::to_string(num_threads); - if (pq_disk_bytes > 0) - params = params + " " + std::to_string(pq_disk_bytes); - diskann::build_disk_index( - data_file_path, index_prefix_path, params.c_str(), - self.get_metric()); - }, - py::arg("data_file_path"), py::arg("index_prefix_path"), py::arg("R"), - py::arg("L"), py::arg("final_index_ram_limit"), - py::arg("indexing_ram_limit"), py::arg("num_threads"), - py::arg("pq_disk_bytes") = 0); + new DynamicInMemIndex(metric, dim, max_points, index_parameters, initial_search_list_size, + search_threads, concurrent_consolidate)); + })) + .def("search", &DynamicInMemIndex::search, py::arg("query"), py::arg("knn"), py::arg("l_search")) + .def("batch_search", &DynamicInMemIndex::batch_search, py::arg("queries"), py::arg("num_queries"), + py::arg("knn"), py::arg("l_search"), py::arg("num_threads")) + .def("insert", &DynamicInMemIndex::insert, py::arg("vector"), py::arg("id")) + .def("mark_deleted", &DynamicInMemIndex::mark_deleted, py::arg("id")) + .def("consolidate_delete", &DynamicInMemIndex::consolidate_delete); + + py::class_>(m, "DiskANNFloatIndex") + .def(py::init([](diskann::Metric metric) { + return std::unique_ptr>(new DiskANNIndex(metric)); + })) + .def("cache_bfs_levels", &DiskANNIndex::cache_bfs_levels, py::arg("num_nodes_to_cache")) + .def("load_index", &DiskANNIndex::load_index, py::arg("index_path_prefix"), py::arg("num_threads"), + py::arg("num_nodes_to_cache"), py::arg("cache_mechanism") = 1) + .def("search", &DiskANNIndex::search, py::arg("query"), py::arg("knn"), py::arg("l_search"), + py::arg("beam_width")) + .def("batch_search", &DiskANNIndex::batch_search, py::arg("queries"), py::arg("num_queries"), + py::arg("knn"), py::arg("l_search"), py::arg("beam_width"), py::arg("num_threads")) + .def( + "build", + [](DiskANNIndex &self, const char *data_file_path, const char *index_prefix_path, unsigned R, + unsigned L, double final_index_ram_limit, double indexing_ram_budget, unsigned num_threads, + unsigned pq_disk_bytes) { + std::string params = std::to_string(R) + " " + std::to_string(L) + " " + + std::to_string(final_index_ram_limit) + " " + std::to_string(indexing_ram_budget) + + " " + std::to_string(num_threads); + if (pq_disk_bytes > 0) + { + params = params + " " + std::to_string(pq_disk_bytes); + } + diskann::build_disk_index(data_file_path, index_prefix_path, params.c_str(), self.get_metric()); + }, + py::arg("data_file_path"), py::arg("index_prefix_path"), py::arg("R"), py::arg("L"), + py::arg("final_index_ram_limit"), py::arg("indexing_ram_limit"), py::arg("num_threads"), + py::arg("pq_disk_bytes") = 0); + + py::class_>(m, "DiskANNInt8Index") + .def(py::init([](diskann::Metric metric) { + return std::unique_ptr>(new DiskANNIndex(metric)); + })) + .def("cache_bfs_levels", &DiskANNIndex::cache_bfs_levels, py::arg("num_nodes_to_cache")) + .def("load_index", &DiskANNIndex::load_index, py::arg("index_path_prefix"), py::arg("num_threads"), + py::arg("num_nodes_to_cache"), py::arg("cache_mechanism") = 1) + .def("search", &DiskANNIndex::search, py::arg("query"), py::arg("knn"), py::arg("l_search"), + py::arg("beam_width")) + .def("batch_search", &DiskANNIndex::batch_search, py::arg("queries"), py::arg("num_queries"), + py::arg("knn"), py::arg("l_search"), py::arg("beam_width"), py::arg("num_threads")) + .def( + "build", + [](DiskANNIndex &self, const char *data_file_path, const char *index_prefix_path, unsigned R, + unsigned L, double final_index_ram_limit, double indexing_ram_budget, unsigned num_threads, + unsigned pq_disk_bytes) { + std::string params = std::to_string(R) + " " + std::to_string(L) + " " + + std::to_string(final_index_ram_limit) + " " + std::to_string(indexing_ram_budget) + + " " + std::to_string(num_threads); + if (pq_disk_bytes > 0) + params = params + " " + std::to_string(pq_disk_bytes); + diskann::build_disk_index(data_file_path, index_prefix_path, params.c_str(), self.get_metric()); + }, + py::arg("data_file_path"), py::arg("index_prefix_path"), py::arg("R"), py::arg("L"), + py::arg("final_index_ram_limit"), py::arg("indexing_ram_limit"), py::arg("num_threads"), + py::arg("pq_disk_bytes") = 0); + + py::class_>(m, "DiskANNUInt8Index") + .def(py::init([](diskann::Metric metric) { + return std::unique_ptr>(new DiskANNIndex(metric)); + })) + .def("cache_bfs_levels", &DiskANNIndex::cache_bfs_levels, py::arg("num_nodes_to_cache")) + .def("load_index", &DiskANNIndex::load_index, py::arg("index_path_prefix"), py::arg("num_threads"), + py::arg("num_nodes_to_cache"), py::arg("cache_mechanism") = 1) + .def("search", &DiskANNIndex::search, py::arg("query"), py::arg("knn"), py::arg("l_search"), + py::arg("beam_width")) + .def("batch_search", &DiskANNIndex::batch_search, py::arg("queries"), py::arg("num_queries"), + py::arg("knn"), py::arg("l_search"), py::arg("beam_width"), py::arg("num_threads")) + .def( + "build", + [](DiskANNIndex &self, const char *data_file_path, const char *index_prefix_path, unsigned R, + unsigned L, double final_index_ram_limit, double indexing_ram_budget, unsigned num_threads, + unsigned pq_disk_bytes) { + std::string params = std::to_string(R) + " " + std::to_string(L) + " " + + std::to_string(final_index_ram_limit) + " " + std::to_string(indexing_ram_budget) + + " " + std::to_string(num_threads); + if (pq_disk_bytes > 0) + params = params + " " + std::to_string(pq_disk_bytes); + diskann::build_disk_index(data_file_path, index_prefix_path, params.c_str(), + self.get_metric()); + }, + py::arg("data_file_path"), py::arg("index_prefix_path"), py::arg("R"), py::arg("L"), + py::arg("final_index_ram_limit"), py::arg("indexing_ram_limit"), py::arg("num_threads"), + py::arg("pq_disk_bytes") = 0); } diff --git a/src/abstract_data_store.cpp b/src/abstract_data_store.cpp index e154a9cf4..a980bd545 100644 --- a/src/abstract_data_store.cpp +++ b/src/abstract_data_store.cpp @@ -5,35 +5,42 @@ #include "abstract_data_store.h" -namespace diskann { +namespace diskann +{ template -AbstractDataStore::AbstractDataStore(const location_t capacity, - const size_t dim) - : _capacity(capacity), _dim(dim) {} +AbstractDataStore::AbstractDataStore(const location_t capacity, const size_t dim) + : _capacity(capacity), _dim(dim) +{ +} -template -location_t AbstractDataStore::capacity() const { - return _capacity; +template location_t AbstractDataStore::capacity() const +{ + return _capacity; } -template -size_t AbstractDataStore::get_dims() const { - return _dim; +template size_t AbstractDataStore::get_dims() const +{ + return _dim; } -template -location_t AbstractDataStore::resize(const location_t new_num_points) { - if (new_num_points > _capacity) { - return expand(new_num_points); - } else if (new_num_points < _capacity) { - return shrink(new_num_points); - } else { - return _capacity; - } +template location_t AbstractDataStore::resize(const location_t new_num_points) +{ + if (new_num_points > _capacity) + { + return expand(new_num_points); + } + else if (new_num_points < _capacity) + { + return shrink(new_num_points); + } + else + { + return _capacity; + } } template DISKANN_DLLEXPORT class AbstractDataStore; template DISKANN_DLLEXPORT class AbstractDataStore; template DISKANN_DLLEXPORT class AbstractDataStore; -} // namespace diskann +} // namespace diskann diff --git a/src/ann_exception.cpp b/src/ann_exception.cpp index 36753305a..ba55e3655 100644 --- a/src/ann_exception.cpp +++ b/src/ann_exception.cpp @@ -5,33 +5,32 @@ #include #include -namespace diskann { +namespace diskann +{ ANNException::ANNException(const std::string &message, int errorCode) - : std::runtime_error(message), _errorCode(errorCode) {} + : std::runtime_error(message), _errorCode(errorCode) +{ +} -std::string package_string(const std::string &item_name, - const std::string &item_val) { - return std::string("[") + item_name + ": " + std::string(item_val) + - std::string("]"); +std::string package_string(const std::string &item_name, const std::string &item_val) +{ + return std::string("[") + item_name + ": " + std::string(item_val) + std::string("]"); } -ANNException::ANNException(const std::string &message, int errorCode, - const std::string &funcSig, +ANNException::ANNException(const std::string &message, int errorCode, const std::string &funcSig, const std::string &fileName, uint32_t lineNum) - : ANNException( - package_string(std::string("FUNC"), funcSig) + - package_string(std::string("FILE"), fileName) + - package_string(std::string("LINE"), std::to_string(lineNum)) + - " " + message, - errorCode) {} + : ANNException(package_string(std::string("FUNC"), funcSig) + package_string(std::string("FILE"), fileName) + + package_string(std::string("LINE"), std::to_string(lineNum)) + " " + message, + errorCode) +{ +} -FileException::FileException(const std::string &filename, std::system_error &e, - const std::string &funcSig, +FileException::FileException(const std::string &filename, std::system_error &e, const std::string &funcSig, const std::string &fileName, uint32_t lineNum) - : ANNException(std::string(" While opening file \'") + filename + - std::string("\', error code: ") + - std::to_string(e.code().value()) + " " + - e.code().message(), - e.code().value(), funcSig, fileName, lineNum) {} + : ANNException(std::string(" While opening file \'") + filename + std::string("\', error code: ") + + std::to_string(e.code().value()) + " " + e.code().message(), + e.code().value(), funcSig, fileName, lineNum) +{ +} -} // namespace diskann \ No newline at end of file +} // namespace diskann \ No newline at end of file diff --git a/src/disk_utils.cpp b/src/disk_utils.cpp index a22ebe1fb..775ad5441 100644 --- a/src/disk_utils.cpp +++ b/src/disk_utils.cpp @@ -3,8 +3,7 @@ #include "common_includes.h" -#if defined(RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS) && \ - defined(DISKANN_BUILD) +#if defined(RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS) && defined(DISKANN_BUILD) #include "gperftools/malloc_extension.h" #endif @@ -20,456 +19,464 @@ #include "timer.h" #include "tsl/robin_set.h" -namespace diskann { - -void add_new_file_to_single_index(std::string index_file, - std::string new_file) { - std::unique_ptr metadata; - uint64_t nr, nc; - diskann::load_bin(index_file, metadata, nr, nc); - 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]; - size_t read_blk_size = 64 * 1024 * 1024; - cached_ofstream writer(index_file, read_blk_size); - size_t 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 (uint64_t 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 new_meta; - for (uint64_t i = 0; i < nr; i++) new_meta.push_back(metadata[i]); - new_meta.push_back(metadata[nr - 1] + fsize); - - diskann::save_bin(index_file, new_meta.data(), new_meta.size(), 1); +namespace diskann +{ + +void add_new_file_to_single_index(std::string index_file, std::string new_file) +{ + std::unique_ptr metadata; + uint64_t nr, nc; + diskann::load_bin(index_file, metadata, nr, nc); + 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]; + size_t read_blk_size = 64 * 1024 * 1024; + cached_ofstream writer(index_file, read_blk_size); + size_t 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 (uint64_t 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 new_meta; + for (uint64_t i = 0; i < nr; i++) + new_meta.push_back(metadata[i]); + new_meta.push_back(metadata[nr - 1] + fsize); + + diskann::save_bin(index_file, new_meta.data(), new_meta.size(), 1); } -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 = search_ram_budget - SPACE_FOR_CACHED_NODES_IN_GB; - } - return final_index_ram_limit * 1024 * 1024 * 1024; +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 = 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); +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, - const std::vector ¶m_list) { - size_t num_pq_chunks = (size_t)(std::floor)( - uint64_t(final_index_ram_limit / (double)points_num)); - diskann::cout << "Calculated num_pq_chunks :" << num_pq_chunks << std::endl; - if (param_list.size() >= 6) { - float compress_ratio = (float)atof(param_list[5].c_str()); - if (compress_ratio > 0 && compress_ratio <= 1) { - size_t chunks_by_cr = (size_t)(std::floor)(compress_ratio * dim); - - if (chunks_by_cr > 0 && chunks_by_cr < num_pq_chunks) { - diskann::cout << "Compress ratio:" << compress_ratio - << " new #pq_chunks:" << chunks_by_cr << std::endl; - num_pq_chunks = chunks_by_cr; - } else { - diskann::cout << "Compress ratio: " << compress_ratio - << " #new pq_chunks: " << chunks_by_cr - << " is either zero or greater than num_pq_chunks: " - << num_pq_chunks << ". num_pq_chunks is unchanged. " - << std::endl; - } - } else { - diskann::cerr << "Compression ratio: " << compress_ratio - << " should be in (0,1]" << std::endl; +size_t calculate_num_pq_chunks(double final_index_ram_limit, size_t points_num, uint32_t dim, + const std::vector ¶m_list) +{ + size_t num_pq_chunks = (size_t)(std::floor)(uint64_t(final_index_ram_limit / (double)points_num)); + diskann::cout << "Calculated num_pq_chunks :" << num_pq_chunks << std::endl; + if (param_list.size() >= 6) + { + float compress_ratio = (float)atof(param_list[5].c_str()); + if (compress_ratio > 0 && compress_ratio <= 1) + { + size_t chunks_by_cr = (size_t)(std::floor)(compress_ratio * dim); + + if (chunks_by_cr > 0 && chunks_by_cr < num_pq_chunks) + { + diskann::cout << "Compress ratio:" << compress_ratio << " new #pq_chunks:" << chunks_by_cr << std::endl; + num_pq_chunks = chunks_by_cr; + } + else + { + diskann::cout << "Compress ratio: " << compress_ratio << " #new pq_chunks: " << chunks_by_cr + << " is either zero or greater than num_pq_chunks: " << num_pq_chunks + << ". num_pq_chunks is unchanged. " << std::endl; + } + } + else + { + diskann::cerr << "Compression ratio: " << compress_ratio << " should be in (0,1]" << 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; + 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; + diskann::cout << "Compressing " << dim << "-dimensional data into " << num_pq_chunks << " bytes per vector." + << std::endl; + return num_pq_chunks; } -template -T *generateRandomWarmup(uint64_t warmup_num, uint64_t warmup_dim, - uint64_t warmup_aligned_dim) { - T *warmup = nullptr; - warmup_num = 100000; - diskann::cout << "Generating random warmup file with dim " << warmup_dim - << " and aligned dim " << warmup_aligned_dim << std::flush; - diskann::alloc_aligned(((void **)&warmup), - warmup_num * warmup_aligned_dim * sizeof(T), - 8 * sizeof(T)); - std::memset(warmup, 0, warmup_num * warmup_aligned_dim * sizeof(T)); - std::random_device rd; - std::mt19937 gen(rd()); - std::uniform_int_distribution<> dis(-128, 127); - for (uint32_t i = 0; i < warmup_num; i++) { - for (uint32_t d = 0; d < warmup_dim; d++) { - warmup[i * warmup_aligned_dim + d] = (T)dis(gen); +template T *generateRandomWarmup(uint64_t warmup_num, uint64_t warmup_dim, uint64_t warmup_aligned_dim) +{ + T *warmup = nullptr; + warmup_num = 100000; + diskann::cout << "Generating random warmup file with dim " << warmup_dim << " and aligned dim " + << warmup_aligned_dim << std::flush; + diskann::alloc_aligned(((void **)&warmup), warmup_num * warmup_aligned_dim * sizeof(T), 8 * sizeof(T)); + std::memset(warmup, 0, warmup_num * warmup_aligned_dim * sizeof(T)); + std::random_device rd; + std::mt19937 gen(rd()); + std::uniform_int_distribution<> dis(-128, 127); + for (uint32_t i = 0; i < warmup_num; i++) + { + for (uint32_t d = 0; d < warmup_dim; d++) + { + warmup[i * warmup_aligned_dim + d] = (T)dis(gen); + } } - } - diskann::cout << "..done" << std::endl; - return warmup; + diskann::cout << "..done" << std::endl; + return warmup; } #ifdef EXEC_ENV_OLS template -T *load_warmup(MemoryMappedFiles &files, const std::string &cache_warmup_file, - uint64_t &warmup_num, uint64_t warmup_dim, - uint64_t warmup_aligned_dim) { - T *warmup = nullptr; - uint64_t file_dim, file_aligned_dim; - - 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 = " << file_dim - << " file_aligned_dim: " << file_aligned_dim - << " index_dim: " << warmup_dim - << " index_aligned_dim: " << warmup_aligned_dim << std::endl; - diskann::cerr << stream.str(); - throw diskann::ANNException(stream.str(), -1); +T *load_warmup(MemoryMappedFiles &files, const std::string &cache_warmup_file, uint64_t &warmup_num, + uint64_t warmup_dim, uint64_t warmup_aligned_dim) +{ + T *warmup = nullptr; + uint64_t file_dim, file_aligned_dim; + + 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 = " << file_dim + << " file_aligned_dim: " << file_aligned_dim << " index_dim: " << warmup_dim + << " index_aligned_dim: " << warmup_aligned_dim << std::endl; + diskann::cerr << stream.str(); + throw diskann::ANNException(stream.str(), -1); + } } - } else { - warmup = - generateRandomWarmup(warmup_num, warmup_dim, warmup_aligned_dim); - } - return warmup; + else + { + warmup = generateRandomWarmup(warmup_num, warmup_dim, warmup_aligned_dim); + } + return warmup; } #endif template -T *load_warmup(const std::string &cache_warmup_file, uint64_t &warmup_num, - uint64_t warmup_dim, uint64_t warmup_aligned_dim) { - T *warmup = nullptr; - uint64_t file_dim, file_aligned_dim; - - if (file_exists(cache_warmup_file)) { - diskann::load_aligned_bin(cache_warmup_file, warmup, warmup_num, - file_dim, file_aligned_dim); - if (file_dim != warmup_dim || file_aligned_dim != warmup_aligned_dim) { - std::stringstream stream; - stream << "Mismatched dimensions in sample file. file_dim = " << file_dim - << " file_aligned_dim: " << file_aligned_dim - << " index_dim: " << warmup_dim - << " index_aligned_dim: " << warmup_aligned_dim << std::endl; - throw diskann::ANNException(stream.str(), -1); +T *load_warmup(const std::string &cache_warmup_file, uint64_t &warmup_num, uint64_t warmup_dim, + uint64_t warmup_aligned_dim) +{ + T *warmup = nullptr; + uint64_t file_dim, file_aligned_dim; + + if (file_exists(cache_warmup_file)) + { + diskann::load_aligned_bin(cache_warmup_file, warmup, warmup_num, file_dim, file_aligned_dim); + if (file_dim != warmup_dim || file_aligned_dim != warmup_aligned_dim) + { + std::stringstream stream; + stream << "Mismatched dimensions in sample file. file_dim = " << file_dim + << " file_aligned_dim: " << file_aligned_dim << " index_dim: " << warmup_dim + << " index_aligned_dim: " << warmup_aligned_dim << std::endl; + throw diskann::ANNException(stream.str(), -1); + } + } + else + { + warmup = generateRandomWarmup(warmup_num, warmup_dim, warmup_aligned_dim); } - } else { - warmup = - generateRandomWarmup(warmup_num, warmup_dim, warmup_aligned_dim); - } - return warmup; + return warmup; } /*************************************************** Support for Merging Many Vamana Indices ***************************************************/ -void read_idmap(const std::string &fname, std::vector &ivecs) { - uint32_t npts32, dim; - size_t actual_file_size = get_file_size(fname); - 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)) { - std::stringstream stream; - stream << "Error reading idmap file. Check if the file is bin file with " - "1 dimensional data. Actual: " - << actual_file_size - << ", expected: " << (size_t)npts32 + 2 * sizeof(uint32_t) - << std::endl; - - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); - } - ivecs.resize(npts32); - reader.read((char *)ivecs.data(), ((size_t)npts32) * sizeof(uint32_t)); - reader.close(); +void read_idmap(const std::string &fname, std::vector &ivecs) +{ + uint32_t npts32, dim; + size_t actual_file_size = get_file_size(fname); + 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)) + { + std::stringstream stream; + stream << "Error reading idmap file. Check if the file is bin file with " + "1 dimensional data. Actual: " + << actual_file_size << ", expected: " << (size_t)npts32 + 2 * sizeof(uint32_t) << std::endl; + + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } + ivecs.resize(npts32); + reader.read((char *)ivecs.data(), ((size_t)npts32) * sizeof(uint32_t)); + reader.close(); } -int merge_shards(const std::string &vamana_prefix, - const std::string &vamana_suffix, - const std::string &idmaps_prefix, - const std::string &idmaps_suffix, const uint64_t nshards, - uint32_t max_degree, const std::string &output_vamana, - const std::string &medoids_file, bool use_filters, - const std::string &labels_to_medoids_file) { - // Read ID maps - std::vector vamana_names(nshards); - std::vector> idmaps(nshards); - for (uint64_t shard = 0; shard < nshards; shard++) { - vamana_names[shard] = vamana_prefix + std::to_string(shard) + vamana_suffix; - read_idmap(idmaps_prefix + std::to_string(shard) + idmaps_suffix, - idmaps[shard]); - } - - // find max node id - size_t nnodes = 0; - size_t nelems = 0; - for (auto &idmap : idmaps) { - for (auto &id : idmap) { - nnodes = std::max(nnodes, (size_t)id); +int merge_shards(const std::string &vamana_prefix, const std::string &vamana_suffix, const std::string &idmaps_prefix, + const std::string &idmaps_suffix, const uint64_t nshards, uint32_t max_degree, + const std::string &output_vamana, const std::string &medoids_file, bool use_filters, + const std::string &labels_to_medoids_file) +{ + // Read ID maps + std::vector vamana_names(nshards); + std::vector> idmaps(nshards); + for (uint64_t shard = 0; shard < nshards; shard++) + { + vamana_names[shard] = vamana_prefix + std::to_string(shard) + vamana_suffix; + read_idmap(idmaps_prefix + std::to_string(shard) + idmaps_suffix, idmaps[shard]); } - nelems += idmap.size(); - } - nnodes++; - diskann::cout << "# nodes: " << nnodes << ", max. degree: " << max_degree - << std::endl; - - // compute inverse map: node -> shards - std::vector> node_shard; - node_shard.reserve(nelems); - for (size_t shard = 0; shard < nshards; shard++) { - diskann::cout << "Creating inverse map -- shard #" << shard << std::endl; - for (size_t idx = 0; idx < idmaps[shard].size(); idx++) { - size_t node_id = idmaps[shard][idx]; - node_shard.push_back(std::make_pair((uint32_t)node_id, (uint32_t)shard)); + + // find max node id + size_t nnodes = 0; + size_t nelems = 0; + for (auto &idmap : idmaps) + { + for (auto &id : idmap) + { + nnodes = std::max(nnodes, (size_t)id); + } + nelems += idmap.size(); } - } - 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; - - // will merge all the labels to medoids files of each shard into one - // combined file - if (use_filters) { - std::unordered_map> global_label_to_medoids; - - for (size_t i = 0; i < nshards; i++) { - std::ifstream mapping_reader; - std::string map_file = vamana_names[i] + "_labels_to_medoids.txt"; - mapping_reader.open(map_file); - - std::string line, token; - uint32_t line_cnt = 0; - - while (std::getline(mapping_reader, line)) { - std::istringstream iss(line); - uint32_t cnt = 0; - uint32_t medoid = 0; - uint32_t label = 0; - while (std::getline(iss, token, ',')) { - token.erase(std::remove(token.begin(), token.end(), '\n'), - token.end()); - token.erase(std::remove(token.begin(), token.end(), '\r'), - token.end()); - - uint32_t token_as_num = std::stoul(token); - - if (cnt == 0) - label = token_as_num; - else - medoid = token_as_num; - cnt++; + nnodes++; + diskann::cout << "# nodes: " << nnodes << ", max. degree: " << max_degree << std::endl; + + // compute inverse map: node -> shards + std::vector> node_shard; + node_shard.reserve(nelems); + for (size_t shard = 0; shard < nshards; shard++) + { + diskann::cout << "Creating inverse map -- shard #" << shard << std::endl; + for (size_t idx = 0; idx < idmaps[shard].size(); idx++) + { + size_t node_id = idmaps[shard][idx]; + node_shard.push_back(std::make_pair((uint32_t)node_id, (uint32_t)shard)); } - global_label_to_medoids[label].push_back(idmaps[i][medoid]); - line_cnt++; - } - mapping_reader.close(); } + 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; + + // will merge all the labels to medoids files of each shard into one + // combined file + if (use_filters) + { + std::unordered_map> global_label_to_medoids; + + for (size_t i = 0; i < nshards; i++) + { + std::ifstream mapping_reader; + std::string map_file = vamana_names[i] + "_labels_to_medoids.txt"; + mapping_reader.open(map_file); + + std::string line, token; + uint32_t line_cnt = 0; + + while (std::getline(mapping_reader, line)) + { + std::istringstream iss(line); + uint32_t cnt = 0; + uint32_t medoid = 0; + uint32_t label = 0; + while (std::getline(iss, token, ',')) + { + token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); + token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); + + uint32_t token_as_num = std::stoul(token); + + if (cnt == 0) + label = token_as_num; + else + medoid = token_as_num; + cnt++; + } + global_label_to_medoids[label].push_back(idmaps[i][medoid]); + line_cnt++; + } + mapping_reader.close(); + } - std::ofstream mapping_writer(labels_to_medoids_file); - assert(mapping_writer.is_open()); - for (auto iter : global_label_to_medoids) { - mapping_writer << iter.first << ", "; - auto &vec = iter.second; - for (uint32_t idx = 0; idx < vec.size() - 1; idx++) { - mapping_writer << vec[idx] << ", "; - } - mapping_writer << vec[vec.size() - 1] << std::endl; + std::ofstream mapping_writer(labels_to_medoids_file); + assert(mapping_writer.is_open()); + for (auto iter : global_label_to_medoids) + { + mapping_writer << iter.first << ", "; + auto &vec = iter.second; + for (uint32_t idx = 0; idx < vec.size() - 1; idx++) + { + mapping_writer << vec[idx] << ", "; + } + mapping_writer << vec[vec.size() - 1] << std::endl; + } + mapping_writer.close(); + } + + // create cached vamana readers + std::vector vamana_readers(nshards); + for (size_t i = 0; i < nshards; i++) + { + vamana_readers[i].open(vamana_names[i], BUFFER_SIZE_FOR_CACHED_IO); + size_t expected_file_size; + vamana_readers[i].read((char *)&expected_file_size, sizeof(uint64_t)); } - mapping_writer.close(); - } - - // create cached vamana readers - std::vector vamana_readers(nshards); - for (size_t i = 0; i < nshards; i++) { - vamana_readers[i].open(vamana_names[i], BUFFER_SIZE_FOR_CACHED_IO); - size_t expected_file_size; - vamana_readers[i].read((char *)&expected_file_size, sizeof(uint64_t)); - } - - size_t vamana_metadata_size = - sizeof(uint64_t) + sizeof(uint32_t) + sizeof(uint32_t) + - sizeof(uint64_t); // expected file size + max degree + - // medoid_id + frozen_point info - - // create cached vamana writers - cached_ofstream merged_vamana_writer(output_vamana, - BUFFER_SIZE_FOR_CACHED_IO); - - size_t merged_index_size = - vamana_metadata_size; // we initialize the size of the merged index to - // the metadata size - size_t merged_index_frozen = 0; - merged_vamana_writer.write( - (char *)&merged_index_size, - sizeof(uint64_t)); // we will overwrite the index size at the end - - uint32_t output_width = max_degree; - uint32_t max_input_width = 0; - // read width from each vamana to advance buffer by sizeof(uint32_t) bytes - for (auto &reader : vamana_readers) { - uint32_t input_width; - reader.read((char *)&input_width, sizeof(uint32_t)); - max_input_width = - input_width > max_input_width ? input_width : max_input_width; - } - - diskann::cout << "Max input width: " << max_input_width - << ", output width: " << output_width << std::endl; - - merged_vamana_writer.write((char *)&output_width, sizeof(uint32_t)); - std::ofstream medoid_writer(medoids_file.c_str(), std::ios::binary); - uint32_t nshards_u32 = (uint32_t)nshards; - uint32_t one_val = 1; - medoid_writer.write((char *)&nshards_u32, sizeof(uint32_t)); - medoid_writer.write((char *)&one_val, sizeof(uint32_t)); - - uint64_t vamana_index_frozen = - 0; // as of now the functionality to merge many overlapping vamana - // indices is supported only for bulk indices without frozen point. - // Hence the final index will also not have any frozen points. - for (uint64_t shard = 0; shard < nshards; shard++) { - uint32_t medoid; - // read medoid - vamana_readers[shard].read((char *)&medoid, sizeof(uint32_t)); - vamana_readers[shard].read((char *)&vamana_index_frozen, sizeof(uint64_t)); - assert(vamana_index_frozen == false); - // rename medoid - medoid = idmaps[shard][medoid]; - - medoid_writer.write((char *)&medoid, sizeof(uint32_t)); - // write renamed medoid - if (shard == (nshards - 1)) //--> uncomment if running hierarchical - merged_vamana_writer.write((char *)&medoid, sizeof(uint32_t)); - } - merged_vamana_writer.write((char *)&merged_index_frozen, sizeof(uint64_t)); - medoid_writer.close(); - - diskann::cout << "Starting merge" << std::endl; - - // Gopal. random_shuffle() is deprecated. - std::random_device rng; - std::mt19937 urng(rng()); - - std::vector nhood_set(nnodes, 0); - std::vector final_nhood; - - uint32_t nnbrs = 0, shard_nnbrs = 0; - uint32_t cur_id = 0; - for (const auto &id_shard : node_shard) { - uint32_t node_id = id_shard.first; - uint32_t shard_id = id_shard.second; - if (cur_id < node_id) { - // Gopal. random_shuffle() is deprecated. - std::shuffle(final_nhood.begin(), final_nhood.end(), urng); - nnbrs = (uint32_t)(std::min)(final_nhood.size(), (uint64_t)max_degree); - // write into merged ofstream - merged_vamana_writer.write((char *)&nnbrs, sizeof(uint32_t)); - merged_vamana_writer.write((char *)final_nhood.data(), - nnbrs * sizeof(uint32_t)); - merged_index_size += (sizeof(uint32_t) + nnbrs * sizeof(uint32_t)); - if (cur_id % 499999 == 1) { - diskann::cout << "." << std::flush; - } - cur_id = node_id; - nnbrs = 0; - for (auto &p : final_nhood) nhood_set[p] = 0; - final_nhood.clear(); + + size_t vamana_metadata_size = + sizeof(uint64_t) + sizeof(uint32_t) + sizeof(uint32_t) + sizeof(uint64_t); // expected file size + max degree + + // medoid_id + frozen_point info + + // create cached vamana writers + cached_ofstream merged_vamana_writer(output_vamana, BUFFER_SIZE_FOR_CACHED_IO); + + size_t merged_index_size = vamana_metadata_size; // we initialize the size of the merged index to + // the metadata size + size_t merged_index_frozen = 0; + merged_vamana_writer.write((char *)&merged_index_size, + sizeof(uint64_t)); // we will overwrite the index size at the end + + uint32_t output_width = max_degree; + uint32_t max_input_width = 0; + // read width from each vamana to advance buffer by sizeof(uint32_t) bytes + for (auto &reader : vamana_readers) + { + uint32_t input_width; + reader.read((char *)&input_width, sizeof(uint32_t)); + max_input_width = input_width > max_input_width ? input_width : max_input_width; } - // read from shard_id ifstream - vamana_readers[shard_id].read((char *)&shard_nnbrs, sizeof(uint32_t)); - if (shard_nnbrs == 0) { - diskann::cout << "WARNING: shard #" << shard_id << ", node_id " << node_id - << " has 0 nbrs" << std::endl; + diskann::cout << "Max input width: " << max_input_width << ", output width: " << output_width << std::endl; + + merged_vamana_writer.write((char *)&output_width, sizeof(uint32_t)); + std::ofstream medoid_writer(medoids_file.c_str(), std::ios::binary); + uint32_t nshards_u32 = (uint32_t)nshards; + uint32_t one_val = 1; + medoid_writer.write((char *)&nshards_u32, sizeof(uint32_t)); + medoid_writer.write((char *)&one_val, sizeof(uint32_t)); + + uint64_t vamana_index_frozen = 0; // as of now the functionality to merge many overlapping vamana + // indices is supported only for bulk indices without frozen point. + // Hence the final index will also not have any frozen points. + for (uint64_t shard = 0; shard < nshards; shard++) + { + uint32_t medoid; + // read medoid + vamana_readers[shard].read((char *)&medoid, sizeof(uint32_t)); + vamana_readers[shard].read((char *)&vamana_index_frozen, sizeof(uint64_t)); + assert(vamana_index_frozen == false); + // rename medoid + medoid = idmaps[shard][medoid]; + + medoid_writer.write((char *)&medoid, sizeof(uint32_t)); + // write renamed medoid + if (shard == (nshards - 1)) //--> uncomment if running hierarchical + merged_vamana_writer.write((char *)&medoid, sizeof(uint32_t)); } + merged_vamana_writer.write((char *)&merged_index_frozen, sizeof(uint64_t)); + medoid_writer.close(); + + diskann::cout << "Starting merge" << std::endl; + + // Gopal. random_shuffle() is deprecated. + std::random_device rng; + std::mt19937 urng(rng()); + + std::vector nhood_set(nnodes, 0); + std::vector final_nhood; + + uint32_t nnbrs = 0, shard_nnbrs = 0; + uint32_t cur_id = 0; + for (const auto &id_shard : node_shard) + { + uint32_t node_id = id_shard.first; + uint32_t shard_id = id_shard.second; + if (cur_id < node_id) + { + // Gopal. random_shuffle() is deprecated. + std::shuffle(final_nhood.begin(), final_nhood.end(), urng); + nnbrs = (uint32_t)(std::min)(final_nhood.size(), (uint64_t)max_degree); + // write into merged ofstream + merged_vamana_writer.write((char *)&nnbrs, sizeof(uint32_t)); + merged_vamana_writer.write((char *)final_nhood.data(), nnbrs * sizeof(uint32_t)); + merged_index_size += (sizeof(uint32_t) + nnbrs * sizeof(uint32_t)); + if (cur_id % 499999 == 1) + { + diskann::cout << "." << std::flush; + } + cur_id = node_id; + nnbrs = 0; + for (auto &p : final_nhood) + nhood_set[p] = 0; + final_nhood.clear(); + } + // read from shard_id ifstream + vamana_readers[shard_id].read((char *)&shard_nnbrs, sizeof(uint32_t)); + + if (shard_nnbrs == 0) + { + diskann::cout << "WARNING: shard #" << shard_id << ", node_id " << node_id << " has 0 nbrs" << std::endl; + } - std::vector shard_nhood(shard_nnbrs); - if (shard_nnbrs > 0) - vamana_readers[shard_id].read((char *)shard_nhood.data(), - shard_nnbrs * sizeof(uint32_t)); - // rename nodes - for (uint64_t j = 0; j < shard_nnbrs; j++) { - if (nhood_set[idmaps[shard_id][shard_nhood[j]]] == 0) { - nhood_set[idmaps[shard_id][shard_nhood[j]]] = 1; - final_nhood.emplace_back(idmaps[shard_id][shard_nhood[j]]); - } + std::vector shard_nhood(shard_nnbrs); + if (shard_nnbrs > 0) + vamana_readers[shard_id].read((char *)shard_nhood.data(), shard_nnbrs * sizeof(uint32_t)); + // rename nodes + for (uint64_t j = 0; j < shard_nnbrs; j++) + { + if (nhood_set[idmaps[shard_id][shard_nhood[j]]] == 0) + { + nhood_set[idmaps[shard_id][shard_nhood[j]]] = 1; + final_nhood.emplace_back(idmaps[shard_id][shard_nhood[j]]); + } + } } - } - - // Gopal. random_shuffle() is deprecated. - std::shuffle(final_nhood.begin(), final_nhood.end(), urng); - nnbrs = (uint32_t)(std::min)(final_nhood.size(), (uint64_t)max_degree); - // write into merged ofstream - merged_vamana_writer.write((char *)&nnbrs, sizeof(uint32_t)); - if (nnbrs > 0) { - merged_vamana_writer.write((char *)final_nhood.data(), - nnbrs * sizeof(uint32_t)); - } - merged_index_size += (sizeof(uint32_t) + nnbrs * sizeof(uint32_t)); - for (auto &p : final_nhood) nhood_set[p] = 0; - final_nhood.clear(); - - diskann::cout << "Expected size: " << merged_index_size << std::endl; - - merged_vamana_writer.reset(); - merged_vamana_writer.write((char *)&merged_index_size, sizeof(uint64_t)); - - diskann::cout << "Finished merge" << std::endl; - return 0; + + // Gopal. random_shuffle() is deprecated. + std::shuffle(final_nhood.begin(), final_nhood.end(), urng); + nnbrs = (uint32_t)(std::min)(final_nhood.size(), (uint64_t)max_degree); + // write into merged ofstream + merged_vamana_writer.write((char *)&nnbrs, sizeof(uint32_t)); + if (nnbrs > 0) + { + merged_vamana_writer.write((char *)final_nhood.data(), nnbrs * sizeof(uint32_t)); + } + merged_index_size += (sizeof(uint32_t) + nnbrs * sizeof(uint32_t)); + for (auto &p : final_nhood) + nhood_set[p] = 0; + final_nhood.clear(); + + diskann::cout << "Expected size: " << merged_index_size << std::endl; + + merged_vamana_writer.reset(); + merged_vamana_writer.write((char *)&merged_index_size, sizeof(uint64_t)); + + diskann::cout << "Finished merge" << std::endl; + return 0; } // TODO: Make this a streaming implementation to avoid exceeding the memory @@ -481,281 +488,277 @@ int merge_shards(const std::string &vamana_prefix, the new nodes at the end. The dummy map contains the real graph id of the new nodes added to the graph */ template -void breakup_dense_points(const std::string data_file, - const std::string labels_file, uint32_t density, - const std::string out_data_file, - const std::string out_labels_file, - const std::string out_metadata_file) { - std::string token, line; - std::ifstream labels_stream(labels_file); - T *data; - uint64_t npts, ndims; - diskann::load_bin(data_file, data, npts, ndims); - - std::unordered_map dummy_pt_ids; - uint32_t next_dummy_id = (uint32_t)npts; - - uint32_t point_cnt = 0; - - std::vector> labels_per_point; - labels_per_point.resize(npts); - - uint32_t dense_pts = 0; - if (labels_stream.is_open()) { - while (getline(labels_stream, line)) { - std::stringstream iss(line); - uint32_t lbl_cnt = 0; - uint32_t label_host = point_cnt; - while (getline(iss, token, ',')) { - if (lbl_cnt == density) { - if (label_host == point_cnt) dense_pts++; - label_host = next_dummy_id; - labels_per_point.resize(next_dummy_id + 1); - dummy_pt_ids[next_dummy_id] = (uint32_t)point_cnt; - next_dummy_id++; - lbl_cnt = 0; +void breakup_dense_points(const std::string data_file, const std::string labels_file, uint32_t density, + const std::string out_data_file, const std::string out_labels_file, + const std::string out_metadata_file) +{ + std::string token, line; + std::ifstream labels_stream(labels_file); + T *data; + uint64_t npts, ndims; + diskann::load_bin(data_file, data, npts, ndims); + + std::unordered_map dummy_pt_ids; + uint32_t next_dummy_id = (uint32_t)npts; + + uint32_t point_cnt = 0; + + std::vector> labels_per_point; + labels_per_point.resize(npts); + + uint32_t dense_pts = 0; + if (labels_stream.is_open()) + { + while (getline(labels_stream, line)) + { + std::stringstream iss(line); + uint32_t lbl_cnt = 0; + uint32_t label_host = point_cnt; + while (getline(iss, token, ',')) + { + if (lbl_cnt == density) + { + if (label_host == point_cnt) + dense_pts++; + label_host = next_dummy_id; + labels_per_point.resize(next_dummy_id + 1); + dummy_pt_ids[next_dummy_id] = (uint32_t)point_cnt; + next_dummy_id++; + lbl_cnt = 0; + } + token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); + token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); + uint32_t token_as_num = std::stoul(token); + labels_per_point[label_host].push_back(token_as_num); + lbl_cnt++; + } + point_cnt++; } - token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); - token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); - uint32_t token_as_num = std::stoul(token); - labels_per_point[label_host].push_back(token_as_num); - lbl_cnt++; - } - point_cnt++; } - } - diskann::cout << "fraction of dense points with >= " << density - << " labels = " << (float)dense_pts / (float)npts << std::endl; - - if (labels_per_point.size() != 0) { - diskann::cout << labels_per_point.size() << " is the new number of points" + diskann::cout << "fraction of dense points with >= " << density << " labels = " << (float)dense_pts / (float)npts << std::endl; - std::ofstream label_writer(out_labels_file); - assert(label_writer.is_open()); - for (uint32_t i = 0; i < labels_per_point.size(); i++) { - for (uint32_t j = 0; j < (labels_per_point[i].size() - 1); j++) { - label_writer << labels_per_point[i][j] << ","; - } - if (labels_per_point[i].size() != 0) - label_writer << labels_per_point[i][labels_per_point[i].size() - 1]; - label_writer << std::endl; + + if (labels_per_point.size() != 0) + { + diskann::cout << labels_per_point.size() << " is the new number of points" << std::endl; + std::ofstream label_writer(out_labels_file); + assert(label_writer.is_open()); + for (uint32_t i = 0; i < labels_per_point.size(); i++) + { + for (uint32_t j = 0; j < (labels_per_point[i].size() - 1); j++) + { + label_writer << labels_per_point[i][j] << ","; + } + if (labels_per_point[i].size() != 0) + label_writer << labels_per_point[i][labels_per_point[i].size() - 1]; + label_writer << std::endl; + } + label_writer.close(); } - label_writer.close(); - } - - if (dummy_pt_ids.size() != 0) { - diskann::cout << dummy_pt_ids.size() - << " is the number of dummy points created" << std::endl; - data = (T *)std::realloc((void *)data, - labels_per_point.size() * ndims * sizeof(T)); - std::ofstream dummy_writer(out_metadata_file); - assert(dummy_writer.is_open()); - for (auto i = dummy_pt_ids.begin(); i != dummy_pt_ids.end(); i++) { - dummy_writer << i->first << "," << i->second << std::endl; - std::memcpy(data + i->first * ndims, data + i->second * ndims, - ndims * sizeof(T)); + + if (dummy_pt_ids.size() != 0) + { + diskann::cout << dummy_pt_ids.size() << " is the number of dummy points created" << std::endl; + data = (T *)std::realloc((void *)data, labels_per_point.size() * ndims * sizeof(T)); + std::ofstream dummy_writer(out_metadata_file); + assert(dummy_writer.is_open()); + for (auto i = dummy_pt_ids.begin(); i != dummy_pt_ids.end(); i++) + { + dummy_writer << i->first << "," << i->second << std::endl; + std::memcpy(data + i->first * ndims, data + i->second * ndims, ndims * sizeof(T)); + } + dummy_writer.close(); } - dummy_writer.close(); - } - diskann::save_bin(out_data_file, data, labels_per_point.size(), ndims); + diskann::save_bin(out_data_file, data, labels_per_point.size(), ndims); } -void extract_shard_labels( - const std::string &in_label_file, const std::string &shard_ids_bin, - const std::string &shard_label_file) { // assumes ith row is for ith - // point in labels file - diskann::cout << "Extracting labels for shard" << std::endl; - - uint32_t *ids = nullptr; - uint64_t num_ids, tmp_dim; - diskann::load_bin(shard_ids_bin, ids, num_ids, tmp_dim); - - uint32_t counter = 0, shard_counter = 0; - std::string cur_line; - - std::ifstream label_reader(in_label_file); - std::ofstream label_writer(shard_label_file); - assert(label_reader.is_open()); - assert(label_reader.is_open()); - if (label_reader && label_writer) { - while (std::getline(label_reader, cur_line)) { - if (shard_counter >= num_ids) { - break; - } - if (counter == ids[shard_counter]) { - label_writer << cur_line << "\n"; - shard_counter++; - } - counter++; +void extract_shard_labels(const std::string &in_label_file, const std::string &shard_ids_bin, + const std::string &shard_label_file) +{ // assumes ith row is for ith + // point in labels file + diskann::cout << "Extracting labels for shard" << std::endl; + + uint32_t *ids = nullptr; + uint64_t num_ids, tmp_dim; + diskann::load_bin(shard_ids_bin, ids, num_ids, tmp_dim); + + uint32_t counter = 0, shard_counter = 0; + std::string cur_line; + + std::ifstream label_reader(in_label_file); + std::ofstream label_writer(shard_label_file); + assert(label_reader.is_open()); + assert(label_reader.is_open()); + if (label_reader && label_writer) + { + while (std::getline(label_reader, cur_line)) + { + if (shard_counter >= num_ids) + { + break; + } + if (counter == ids[shard_counter]) + { + label_writer << cur_line << "\n"; + shard_counter++; + } + counter++; + } } - } - if (ids != nullptr) delete[] ids; + if (ids != nullptr) + delete[] ids; } template -int build_merged_vamana_index( - std::string base_file, diskann::Metric compareMetric, uint32_t L, - uint32_t R, double sampling_rate, double ram_budget, - std::string mem_index_path, std::string medoids_file, - std::string centroids_file, size_t build_pq_bytes, bool use_opq, - uint32_t num_threads, bool use_filters, const std::string &label_file, - const std::string &labels_to_medoids_file, - const std::string &universal_label, const uint32_t Lf) { - size_t base_num, base_dim; - diskann::get_bin_metadata(base_file, base_num, base_dim); - - double full_index_ram = estimate_ram_usage(base_num, base_dim, sizeof(T), R); - - // TODO: Make this honest when there is filter support - if (full_index_ram < ram_budget * 1024 * 1024 * 1024) { - diskann::cout << "Full index fits in RAM budget, should consume at most " - << full_index_ram / (1024 * 1024 * 1024) - << "GiBs, so building in one shot" << std::endl; - - diskann::IndexWriteParameters paras = - diskann::IndexWriteParametersBuilder(L, R) - .with_filter_list_size(Lf) - .with_saturate_graph(!use_filters) - .with_num_threads(num_threads) - .build(); - using TagT = uint32_t; - diskann::Index _index( - compareMetric, base_dim, base_num, false, false, false, - build_pq_bytes > 0, build_pq_bytes, use_opq); - if (!use_filters) - _index.build(base_file.c_str(), base_num, paras); - else { - if (universal_label != "") { // indicates no universal label - LabelT unv_label_as_num = 0; - _index.set_universal_label(unv_label_as_num); - } - _index.build_filtered_index(base_file.c_str(), label_file, base_num, - paras); - } - _index.save(mem_index_path.c_str()); - - if (use_filters) { - // need to copy the labels_to_medoids file to the specified input - // file - std::remove(labels_to_medoids_file.c_str()); - std::string mem_labels_to_medoid_file = - mem_index_path + "_labels_to_medoids.txt"; - copy_file(mem_labels_to_medoid_file, labels_to_medoids_file); - std::remove(mem_labels_to_medoid_file.c_str()); - } +int build_merged_vamana_index(std::string base_file, diskann::Metric compareMetric, uint32_t L, uint32_t R, + double sampling_rate, double ram_budget, std::string mem_index_path, + std::string medoids_file, std::string centroids_file, size_t build_pq_bytes, bool use_opq, + uint32_t num_threads, bool use_filters, const std::string &label_file, + const std::string &labels_to_medoids_file, const std::string &universal_label, + const uint32_t Lf) +{ + size_t base_num, base_dim; + diskann::get_bin_metadata(base_file, base_num, base_dim); + + double full_index_ram = estimate_ram_usage(base_num, base_dim, sizeof(T), R); + + // TODO: Make this honest when there is filter support + if (full_index_ram < ram_budget * 1024 * 1024 * 1024) + { + diskann::cout << "Full index fits in RAM budget, should consume at most " + << full_index_ram / (1024 * 1024 * 1024) << "GiBs, so building in one shot" << std::endl; + + diskann::IndexWriteParameters paras = diskann::IndexWriteParametersBuilder(L, R) + .with_filter_list_size(Lf) + .with_saturate_graph(!use_filters) + .with_num_threads(num_threads) + .build(); + using TagT = uint32_t; + diskann::Index _index(compareMetric, base_dim, base_num, false, false, false, + build_pq_bytes > 0, build_pq_bytes, use_opq); + if (!use_filters) + _index.build(base_file.c_str(), base_num, paras); + else + { + if (universal_label != "") + { // indicates no universal label + LabelT unv_label_as_num = 0; + _index.set_universal_label(unv_label_as_num); + } + _index.build_filtered_index(base_file.c_str(), label_file, base_num, paras); + } + _index.save(mem_index_path.c_str()); + + if (use_filters) + { + // need to copy the labels_to_medoids file to the specified input + // file + std::remove(labels_to_medoids_file.c_str()); + std::string mem_labels_to_medoid_file = mem_index_path + "_labels_to_medoids.txt"; + copy_file(mem_labels_to_medoid_file, labels_to_medoids_file); + std::remove(mem_labels_to_medoid_file.c_str()); + } - std::remove(medoids_file.c_str()); - std::remove(centroids_file.c_str()); - return 0; - } - - // where the universal label is to be saved in the final graph - std::string final_index_universal_label_file = - mem_index_path + "_universal_label.txt"; - - std::string merged_index_prefix = mem_index_path + "_tempFiles"; - - Timer timer; - int num_parts = partition_with_ram_budget( - base_file, sampling_rate, ram_budget, 2 * R / 3, merged_index_prefix, 2); - diskann::cout << timer.elapsed_seconds_for_step("partitioning data") - << std::endl; - - std::string cur_centroid_filepath = merged_index_prefix + "_centroids.bin"; - std::rename(cur_centroid_filepath.c_str(), centroids_file.c_str()); - - timer.reset(); - for (int p = 0; p < num_parts; p++) { - std::string shard_base_file = - merged_index_prefix + "_subshard-" + std::to_string(p) + ".bin"; - - std::string shard_ids_file = merged_index_prefix + "_subshard-" + - std::to_string(p) + "_ids_uint32.bin"; - - std::string shard_labels_file = - merged_index_prefix + "_subshard-" + std::to_string(p) + "_labels.txt"; - - retrieve_shard_data_from_ids(base_file, shard_ids_file, shard_base_file); - - std::string shard_index_file = - merged_index_prefix + "_subshard-" + std::to_string(p) + "_mem.index"; - - diskann::IndexWriteParameters paras = - diskann::IndexWriteParametersBuilder(L, (2 * R / 3)) - .with_filter_list_size(Lf) - .build(); - - uint64_t shard_base_dim, shard_base_pts; - get_bin_metadata(shard_base_file, shard_base_pts, shard_base_dim); - diskann::Index _index(compareMetric, shard_base_dim, shard_base_pts, - false, false, false, build_pq_bytes > 0, - build_pq_bytes, use_opq); - if (!use_filters) { - _index.build(shard_base_file.c_str(), shard_base_pts, paras); - } else { - diskann::extract_shard_labels(label_file, shard_ids_file, - shard_labels_file); - if (universal_label != "") { // indicates no universal label - LabelT unv_label_as_num = 0; - _index.set_universal_label(unv_label_as_num); - } - _index.build_filtered_index(shard_base_file.c_str(), shard_labels_file, - shard_base_pts, paras); - } - _index.save(shard_index_file.c_str()); - // copy universal label file from first shard to the final destination - // index, since all shards anyway share the universal label - if (p == 0) { - std::string shard_universal_label_file = - shard_index_file + "_universal_label.txt"; - if (universal_label != "") { - copy_file(shard_universal_label_file, final_index_universal_label_file); - } + std::remove(medoids_file.c_str()); + std::remove(centroids_file.c_str()); + return 0; } - std::remove(shard_base_file.c_str()); - } - diskann::cout << timer.elapsed_seconds_for_step("building indices on shards") - << std::endl; - - timer.reset(); - diskann::merge_shards(merged_index_prefix + "_subshard-", "_mem.index", - merged_index_prefix + "_subshard-", "_ids_uint32.bin", - num_parts, R, mem_index_path, medoids_file, use_filters, - labels_to_medoids_file); - diskann::cout << timer.elapsed_seconds_for_step("merging indices") - << std::endl; - - // delete tempFiles - for (int p = 0; p < num_parts; p++) { - std::string shard_base_file = - merged_index_prefix + "_subshard-" + std::to_string(p) + ".bin"; - std::string shard_id_file = merged_index_prefix + "_subshard-" + - std::to_string(p) + "_ids_uint32.bin"; - std::string shard_labels_file = - merged_index_prefix + "_subshard-" + std::to_string(p) + "_labels.txt"; - std::string shard_index_file = - merged_index_prefix + "_subshard-" + std::to_string(p) + "_mem.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()); - if (use_filters) { - std::string shard_index_label_file = shard_index_file + "_labels.txt"; - std::string shard_index_univ_label_file = - shard_index_file + "_universal_label.txt"; - std::string shard_index_label_map_file = - shard_index_file + "_labels_to_medoids.txt"; - std::remove(shard_labels_file.c_str()); - std::remove(shard_index_label_file.c_str()); - std::remove(shard_index_label_map_file.c_str()); - std::remove(shard_index_univ_label_file.c_str()); + // where the universal label is to be saved in the final graph + std::string final_index_universal_label_file = mem_index_path + "_universal_label.txt"; + + std::string merged_index_prefix = mem_index_path + "_tempFiles"; + + Timer timer; + int num_parts = + partition_with_ram_budget(base_file, sampling_rate, ram_budget, 2 * R / 3, merged_index_prefix, 2); + diskann::cout << timer.elapsed_seconds_for_step("partitioning data") << std::endl; + + std::string cur_centroid_filepath = merged_index_prefix + "_centroids.bin"; + std::rename(cur_centroid_filepath.c_str(), centroids_file.c_str()); + + timer.reset(); + for (int p = 0; p < num_parts; p++) + { + std::string shard_base_file = merged_index_prefix + "_subshard-" + std::to_string(p) + ".bin"; + + std::string shard_ids_file = merged_index_prefix + "_subshard-" + std::to_string(p) + "_ids_uint32.bin"; + + std::string shard_labels_file = merged_index_prefix + "_subshard-" + std::to_string(p) + "_labels.txt"; + + retrieve_shard_data_from_ids(base_file, shard_ids_file, shard_base_file); + + std::string shard_index_file = merged_index_prefix + "_subshard-" + std::to_string(p) + "_mem.index"; + + diskann::IndexWriteParameters paras = + diskann::IndexWriteParametersBuilder(L, (2 * R / 3)).with_filter_list_size(Lf).build(); + + uint64_t shard_base_dim, shard_base_pts; + get_bin_metadata(shard_base_file, shard_base_pts, shard_base_dim); + diskann::Index _index(compareMetric, shard_base_dim, shard_base_pts, false, false, false, build_pq_bytes > 0, + build_pq_bytes, use_opq); + if (!use_filters) + { + _index.build(shard_base_file.c_str(), shard_base_pts, paras); + } + else + { + diskann::extract_shard_labels(label_file, shard_ids_file, shard_labels_file); + if (universal_label != "") + { // indicates no universal label + LabelT unv_label_as_num = 0; + _index.set_universal_label(unv_label_as_num); + } + _index.build_filtered_index(shard_base_file.c_str(), shard_labels_file, shard_base_pts, paras); + } + _index.save(shard_index_file.c_str()); + // copy universal label file from first shard to the final destination + // index, since all shards anyway share the universal label + if (p == 0) + { + std::string shard_universal_label_file = shard_index_file + "_universal_label.txt"; + if (universal_label != "") + { + copy_file(shard_universal_label_file, final_index_universal_label_file); + } + } + + std::remove(shard_base_file.c_str()); } - } - return 0; + diskann::cout << timer.elapsed_seconds_for_step("building indices on shards") << std::endl; + + timer.reset(); + diskann::merge_shards(merged_index_prefix + "_subshard-", "_mem.index", merged_index_prefix + "_subshard-", + "_ids_uint32.bin", num_parts, R, mem_index_path, medoids_file, use_filters, + labels_to_medoids_file); + diskann::cout << timer.elapsed_seconds_for_step("merging indices") << std::endl; + + // delete tempFiles + for (int p = 0; p < num_parts; p++) + { + std::string shard_base_file = merged_index_prefix + "_subshard-" + std::to_string(p) + ".bin"; + std::string shard_id_file = merged_index_prefix + "_subshard-" + std::to_string(p) + "_ids_uint32.bin"; + std::string shard_labels_file = merged_index_prefix + "_subshard-" + std::to_string(p) + "_labels.txt"; + std::string shard_index_file = merged_index_prefix + "_subshard-" + std::to_string(p) + "_mem.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()); + if (use_filters) + { + std::string shard_index_label_file = shard_index_file + "_labels.txt"; + std::string shard_index_univ_label_file = shard_index_file + "_universal_label.txt"; + std::string shard_index_label_map_file = shard_index_file + "_labels_to_medoids.txt"; + std::remove(shard_labels_file.c_str()); + std::remove(shard_index_label_file.c_str()); + std::remove(shard_index_label_map_file.c_str()); + std::remove(shard_index_univ_label_file.c_str()); + } + } + return 0; } // General purpose support for DiskANN interface @@ -763,690 +766,645 @@ int build_merged_vamana_index( // 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, uint64_t tuning_sample_num, - uint64_t 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(); +uint32_t optimize_beamwidth(std::unique_ptr> &pFlashIndex, T *tuning_sample, + uint64_t tuning_sample_num, uint64_t 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 (int64_t 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, false, - stats + i); - } - 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; + for (int64_t 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, false, stats + i); + } + 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()); - delete[] stats; - } - return best_bw; + 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, - const std::string reorder_data_file) { - uint32_t npts, ndims; - - // amount to read or write in one shot - size_t read_blk_size = 64 * 1024 * 1024; - size_t write_blk_size = read_blk_size; - cached_ifstream base_reader(base_file, read_blk_size); - base_reader.read((char *)&npts, sizeof(uint32_t)); - base_reader.read((char *)&ndims, sizeof(uint32_t)); - - size_t npts_64, ndims_64; - npts_64 = npts; - ndims_64 = ndims; - - // Check if we need to append data for re-ordering - bool append_reorder_data = false; - std::ifstream reorder_data_reader; - - uint32_t npts_reorder_file = 0, ndims_reorder_file = 0; - if (reorder_data_file != std::string("")) { - append_reorder_data = true; - size_t reorder_data_file_size = get_file_size(reorder_data_file); - reorder_data_reader.exceptions(std::ofstream::failbit | - std::ofstream::badbit); - - try { - reorder_data_reader.open(reorder_data_file, std::ios::binary); - reorder_data_reader.read((char *)&npts_reorder_file, sizeof(uint32_t)); - reorder_data_reader.read((char *)&ndims_reorder_file, sizeof(uint32_t)); - if (npts_reorder_file != npts) - throw ANNException( - "Mismatch in num_points between reorder " - "data file and base file", - -1, __FUNCSIG__, __FILE__, __LINE__); - if (reorder_data_file_size != 8 + sizeof(float) * - (size_t)npts_reorder_file * - (size_t)ndims_reorder_file) - throw ANNException("Discrepancy in reorder data file size ", -1, - __FUNCSIG__, __FILE__, __LINE__); - } catch (std::system_error &e) { - throw FileException(reorder_data_file, e, __FUNCSIG__, __FILE__, - __LINE__); +void create_disk_layout(const std::string base_file, const std::string mem_index_file, const std::string output_file, + const std::string reorder_data_file) +{ + uint32_t npts, ndims; + + // amount to read or write in one shot + size_t read_blk_size = 64 * 1024 * 1024; + size_t write_blk_size = read_blk_size; + cached_ifstream base_reader(base_file, read_blk_size); + base_reader.read((char *)&npts, sizeof(uint32_t)); + base_reader.read((char *)&ndims, sizeof(uint32_t)); + + size_t npts_64, ndims_64; + npts_64 = npts; + ndims_64 = ndims; + + // Check if we need to append data for re-ordering + bool append_reorder_data = false; + std::ifstream reorder_data_reader; + + uint32_t npts_reorder_file = 0, ndims_reorder_file = 0; + if (reorder_data_file != std::string("")) + { + append_reorder_data = true; + size_t reorder_data_file_size = get_file_size(reorder_data_file); + reorder_data_reader.exceptions(std::ofstream::failbit | std::ofstream::badbit); + + try + { + reorder_data_reader.open(reorder_data_file, std::ios::binary); + reorder_data_reader.read((char *)&npts_reorder_file, sizeof(uint32_t)); + reorder_data_reader.read((char *)&ndims_reorder_file, sizeof(uint32_t)); + if (npts_reorder_file != npts) + throw ANNException("Mismatch in num_points between reorder " + "data file and base file", + -1, __FUNCSIG__, __FILE__, __LINE__); + if (reorder_data_file_size != 8 + sizeof(float) * (size_t)npts_reorder_file * (size_t)ndims_reorder_file) + throw ANNException("Discrepancy in reorder data file size ", -1, __FUNCSIG__, __FILE__, __LINE__); + } + catch (std::system_error &e) + { + throw FileException(reorder_data_file, e, __FUNCSIG__, __FILE__, __LINE__); + } } - } - - // create cached reader + writer - size_t actual_file_size = get_file_size(mem_index_file); - diskann::cout << "Vamana index file size=" << actual_file_size << std::endl; - std::ifstream vamana_reader(mem_index_file, std::ios::binary); - cached_ofstream diskann_writer(output_file, write_blk_size); - - // metadata: width, medoid - uint32_t 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__); - } - uint64_t vamana_frozen_num = false, vamana_frozen_loc = 0; - - vamana_reader.read((char *)&width_u32, sizeof(uint32_t)); - vamana_reader.read((char *)&medoid_u32, sizeof(uint32_t)); - vamana_reader.read((char *)&vamana_frozen_num, sizeof(uint64_t)); - // compute - uint64_t medoid, max_node_len, nnodes_per_sector; - npts_64 = (uint64_t)npts; - medoid = (uint64_t)medoid_u32; - if (vamana_frozen_num == 1) vamana_frozen_loc = medoid; - max_node_len = - (((uint64_t)width_u32 + 1) * sizeof(uint32_t)) + (ndims_64 * sizeof(T)); - nnodes_per_sector = SECTOR_LEN / max_node_len; - - diskann::cout << "medoid: " << medoid << "B" << std::endl; - diskann::cout << "max_node_len: " << max_node_len << "B" << std::endl; - diskann::cout << "nnodes_per_sector: " << nnodes_per_sector << "B" - << 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); - uint32_t &nnbrs = *(uint32_t *)(node_buf.get() + ndims_64 * sizeof(T)); - uint32_t *nhood_buf = - (uint32_t *)(node_buf.get() + (ndims_64 * sizeof(T)) + sizeof(uint32_t)); - - // number of sectors (1 for meta data) - uint64_t n_sectors = ROUND_UP(npts_64, nnodes_per_sector) / nnodes_per_sector; - uint64_t n_reorder_sectors = 0; - uint64_t n_data_nodes_per_sector = 0; - - if (append_reorder_data) { - n_data_nodes_per_sector = SECTOR_LEN / (ndims_reorder_file * sizeof(float)); - n_reorder_sectors = - ROUND_UP(npts_64, n_data_nodes_per_sector) / n_data_nodes_per_sector; - } - uint64_t disk_index_file_size = - (n_sectors + n_reorder_sectors + 1) * SECTOR_LEN; - - std::vector 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((uint64_t)append_reorder_data); - if (append_reorder_data) { - output_file_meta.push_back(n_sectors + 1); - output_file_meta.push_back(ndims_reorder_file); - output_file_meta.push_back(n_data_nodes_per_sector); - } - output_file_meta.push_back(disk_index_file_size); - - diskann_writer.write(sector_buf.get(), SECTOR_LEN); - - std::unique_ptr cur_node_coords = std::make_unique(ndims_64); - diskann::cout << "# sectors: " << n_sectors << std::endl; - uint64_t cur_node_id = 0; - for (uint64_t sector = 0; sector < n_sectors; sector++) { - if (sector % 100000 == 0) { - diskann::cout << "Sector #" << sector << "written" << std::endl; + + // create cached reader + writer + size_t actual_file_size = get_file_size(mem_index_file); + diskann::cout << "Vamana index file size=" << actual_file_size << std::endl; + std::ifstream vamana_reader(mem_index_file, std::ios::binary); + cached_ofstream diskann_writer(output_file, write_blk_size); + + // metadata: width, medoid + uint32_t 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__); } - memset(sector_buf.get(), 0, SECTOR_LEN); - for (uint64_t sector_node_id = 0; - sector_node_id < nnodes_per_sector && cur_node_id < npts_64; - sector_node_id++) { - memset(node_buf.get(), 0, max_node_len); - // read cur node's nnbrs - vamana_reader.read((char *)&nnbrs, sizeof(uint32_t)); - - // sanity checks on nnbrs - assert(nnbrs > 0); - assert(nnbrs <= width_u32); - - // read node's nhood - vamana_reader.read((char *)nhood_buf, - (std::min)(nnbrs, width_u32) * sizeof(uint32_t)); - if (nnbrs > width_u32) { - vamana_reader.seekg((nnbrs - width_u32) * sizeof(uint32_t), - vamana_reader.cur); - } - - // write coords of node first - // T *node_coords = data + ((uint64_t) ndims_64 * cur_node_id); - 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 - *(uint32_t *)(node_buf.get() + ndims_64 * sizeof(T)) = - (std::min)(nnbrs, width_u32); - - // write nhood next - memcpy(node_buf.get() + ndims_64 * sizeof(T) + sizeof(uint32_t), - nhood_buf, (std::min)(nnbrs, width_u32) * sizeof(uint32_t)); - - // get offset into sector_buf - char *sector_node_buf = - sector_buf.get() + (sector_node_id * max_node_len); - - // copy node buf into sector_node_buf - memcpy(sector_node_buf, node_buf.get(), max_node_len); - cur_node_id++; + uint64_t vamana_frozen_num = false, vamana_frozen_loc = 0; + + vamana_reader.read((char *)&width_u32, sizeof(uint32_t)); + vamana_reader.read((char *)&medoid_u32, sizeof(uint32_t)); + vamana_reader.read((char *)&vamana_frozen_num, sizeof(uint64_t)); + // compute + uint64_t medoid, max_node_len, nnodes_per_sector; + npts_64 = (uint64_t)npts; + medoid = (uint64_t)medoid_u32; + if (vamana_frozen_num == 1) + vamana_frozen_loc = medoid; + max_node_len = (((uint64_t)width_u32 + 1) * sizeof(uint32_t)) + (ndims_64 * sizeof(T)); + nnodes_per_sector = SECTOR_LEN / max_node_len; + + diskann::cout << "medoid: " << medoid << "B" << std::endl; + diskann::cout << "max_node_len: " << max_node_len << "B" << std::endl; + diskann::cout << "nnodes_per_sector: " << nnodes_per_sector << "B" << 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); + uint32_t &nnbrs = *(uint32_t *)(node_buf.get() + ndims_64 * sizeof(T)); + uint32_t *nhood_buf = (uint32_t *)(node_buf.get() + (ndims_64 * sizeof(T)) + sizeof(uint32_t)); + + // number of sectors (1 for meta data) + uint64_t n_sectors = ROUND_UP(npts_64, nnodes_per_sector) / nnodes_per_sector; + uint64_t n_reorder_sectors = 0; + uint64_t n_data_nodes_per_sector = 0; + + if (append_reorder_data) + { + n_data_nodes_per_sector = SECTOR_LEN / (ndims_reorder_file * sizeof(float)); + n_reorder_sectors = ROUND_UP(npts_64, n_data_nodes_per_sector) / n_data_nodes_per_sector; } - // flush sector to disk - diskann_writer.write(sector_buf.get(), SECTOR_LEN); - } - if (append_reorder_data) { - diskann::cout << "Index written. Appending reorder data..." << std::endl; + uint64_t disk_index_file_size = (n_sectors + n_reorder_sectors + 1) * SECTOR_LEN; + + std::vector 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((uint64_t)append_reorder_data); + if (append_reorder_data) + { + output_file_meta.push_back(n_sectors + 1); + output_file_meta.push_back(ndims_reorder_file); + output_file_meta.push_back(n_data_nodes_per_sector); + } + output_file_meta.push_back(disk_index_file_size); - auto vec_len = ndims_reorder_file * sizeof(float); - std::unique_ptr vec_buf = std::make_unique(vec_len); + diskann_writer.write(sector_buf.get(), SECTOR_LEN); - for (uint64_t sector = 0; sector < n_reorder_sectors; sector++) { - if (sector % 100000 == 0) { - diskann::cout << "Reorder data Sector #" << sector << "written" - << std::endl; - } - - memset(sector_buf.get(), 0, SECTOR_LEN); - - for (uint64_t sector_node_id = 0; - sector_node_id < n_data_nodes_per_sector && sector_node_id < npts_64; - sector_node_id++) { - memset(vec_buf.get(), 0, vec_len); - reorder_data_reader.read(vec_buf.get(), vec_len); - - // copy node buf into sector_node_buf - memcpy(sector_buf.get() + (sector_node_id * vec_len), vec_buf.get(), - vec_len); - } - // flush sector to disk - diskann_writer.write(sector_buf.get(), SECTOR_LEN); + std::unique_ptr cur_node_coords = std::make_unique(ndims_64); + diskann::cout << "# sectors: " << n_sectors << std::endl; + uint64_t cur_node_id = 0; + for (uint64_t sector = 0; sector < n_sectors; sector++) + { + if (sector % 100000 == 0) + { + diskann::cout << "Sector #" << sector << "written" << std::endl; + } + memset(sector_buf.get(), 0, SECTOR_LEN); + for (uint64_t sector_node_id = 0; sector_node_id < nnodes_per_sector && cur_node_id < npts_64; sector_node_id++) + { + memset(node_buf.get(), 0, max_node_len); + // read cur node's nnbrs + vamana_reader.read((char *)&nnbrs, sizeof(uint32_t)); + + // sanity checks on nnbrs + assert(nnbrs > 0); + assert(nnbrs <= width_u32); + + // read node's nhood + vamana_reader.read((char *)nhood_buf, (std::min)(nnbrs, width_u32) * sizeof(uint32_t)); + if (nnbrs > width_u32) + { + vamana_reader.seekg((nnbrs - width_u32) * sizeof(uint32_t), vamana_reader.cur); + } + + // write coords of node first + // T *node_coords = data + ((uint64_t) ndims_64 * cur_node_id); + 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 + *(uint32_t *)(node_buf.get() + ndims_64 * sizeof(T)) = (std::min)(nnbrs, width_u32); + + // write nhood next + memcpy(node_buf.get() + ndims_64 * sizeof(T) + sizeof(uint32_t), nhood_buf, + (std::min)(nnbrs, width_u32) * sizeof(uint32_t)); + + // get offset into sector_buf + char *sector_node_buf = sector_buf.get() + (sector_node_id * max_node_len); + + // copy node buf into sector_node_buf + memcpy(sector_node_buf, node_buf.get(), max_node_len); + cur_node_id++; + } + // flush sector to disk + diskann_writer.write(sector_buf.get(), SECTOR_LEN); + } + if (append_reorder_data) + { + diskann::cout << "Index written. Appending reorder data..." << std::endl; + + auto vec_len = ndims_reorder_file * sizeof(float); + std::unique_ptr vec_buf = std::make_unique(vec_len); + + for (uint64_t sector = 0; sector < n_reorder_sectors; sector++) + { + if (sector % 100000 == 0) + { + diskann::cout << "Reorder data Sector #" << sector << "written" << std::endl; + } + + memset(sector_buf.get(), 0, SECTOR_LEN); + + for (uint64_t sector_node_id = 0; sector_node_id < n_data_nodes_per_sector && sector_node_id < npts_64; + sector_node_id++) + { + memset(vec_buf.get(), 0, vec_len); + reorder_data_reader.read(vec_buf.get(), vec_len); + + // copy node buf into sector_node_buf + memcpy(sector_buf.get() + (sector_node_id * vec_len), vec_buf.get(), vec_len); + } + // flush sector to disk + diskann_writer.write(sector_buf.get(), SECTOR_LEN); + } } - } - diskann_writer.close(); - diskann::save_bin(output_file, output_file_meta.data(), - output_file_meta.size(), 1, 0); - diskann::cout << "Output disk index file written to " << output_file - << std::endl; + diskann_writer.close(); + diskann::save_bin(output_file, output_file_meta.data(), output_file_meta.size(), 1, 0); + diskann::cout << "Output disk index file written to " << output_file << std::endl; } template -int build_disk_index(const char *dataFilePath, const char *indexFilePath, - const char *indexBuildParameters, - diskann::Metric compareMetric, bool use_opq, - const std::string &codebook_prefix, bool use_filters, - const std::string &label_file, - const std::string &universal_label, - const uint32_t filter_threshold, const uint32_t Lf) { - std::stringstream parser; - parser << std::string(indexBuildParameters); - std::string cur_param; - std::vector param_list; - while (parser >> cur_param) { - param_list.push_back(cur_param); - } - if (param_list.size() < 5 || param_list.size() > 9) { - diskann::cout - << "Correct usage of parameters is R (max degree)\n" - "L (indexing list size, better if >= R)\n" - "B (RAM limit of final index in GB)\n" - "M (memory limit while indexing)\n" - "T (number of threads for indexing)\n" - "B' (PQ bytes for disk index: optional parameter for " - "very large dimensional data)\n" - "reorder (set true to include full precision in data file" - ": optional paramter, use only when using disk PQ\n" - "build_PQ_byte (number of PQ bytes for inde build; set 0 to use " - "full precision vectors)\n" - "QD Quantized Dimension to overwrite the derived dim from B " - << std::endl; - return -1; - } - - if (!std::is_same::value && - compareMetric == diskann::Metric::INNER_PRODUCT) { - std::stringstream stream; - stream << "DiskANN currently only supports floating point data for Max " - "Inner Product Search. " - << std::endl; - throw diskann::ANNException(stream.str(), -1); - } - - size_t disk_pq_dims = 0; - bool use_disk_pq = false; - size_t build_pq_bytes = 0; - - // if there is a 6th parameter, it means we compress the disk index - // vectors also using PQ data (for very large dimensionality data). If the - // provided parameter is 0, it means we store full vectors. - if (param_list.size() > 5) { - disk_pq_dims = atoi(param_list[5].c_str()); - use_disk_pq = true; - if (disk_pq_dims == 0) use_disk_pq = false; - } - - bool reorder_data = false; - if (param_list.size() >= 7) { - if (1 == atoi(param_list[6].c_str())) { - reorder_data = true; +int build_disk_index(const char *dataFilePath, const char *indexFilePath, const char *indexBuildParameters, + diskann::Metric compareMetric, bool use_opq, const std::string &codebook_prefix, bool use_filters, + const std::string &label_file, const std::string &universal_label, const uint32_t filter_threshold, + const uint32_t Lf) +{ + std::stringstream parser; + parser << std::string(indexBuildParameters); + std::string cur_param; + std::vector param_list; + while (parser >> cur_param) + { + param_list.push_back(cur_param); + } + if (param_list.size() < 5 || param_list.size() > 9) + { + diskann::cout << "Correct usage of parameters is R (max degree)\n" + "L (indexing list size, better if >= R)\n" + "B (RAM limit of final index in GB)\n" + "M (memory limit while indexing)\n" + "T (number of threads for indexing)\n" + "B' (PQ bytes for disk index: optional parameter for " + "very large dimensional data)\n" + "reorder (set true to include full precision in data file" + ": optional paramter, use only when using disk PQ\n" + "build_PQ_byte (number of PQ bytes for inde build; set 0 to use " + "full precision vectors)\n" + "QD Quantized Dimension to overwrite the derived dim from B " + << std::endl; + return -1; + } + + if (!std::is_same::value && compareMetric == diskann::Metric::INNER_PRODUCT) + { + std::stringstream stream; + stream << "DiskANN currently only supports floating point data for Max " + "Inner Product Search. " + << std::endl; + throw diskann::ANNException(stream.str(), -1); + } + + size_t disk_pq_dims = 0; + bool use_disk_pq = false; + size_t build_pq_bytes = 0; + + // if there is a 6th parameter, it means we compress the disk index + // vectors also using PQ data (for very large dimensionality data). If the + // provided parameter is 0, it means we store full vectors. + if (param_list.size() > 5) + { + disk_pq_dims = atoi(param_list[5].c_str()); + use_disk_pq = true; + if (disk_pq_dims == 0) + use_disk_pq = false; + } + + bool reorder_data = false; + if (param_list.size() >= 7) + { + if (1 == atoi(param_list[6].c_str())) + { + reorder_data = true; + } + } + + if (param_list.size() >= 8) + { + build_pq_bytes = atoi(param_list[7].c_str()); + } + + std::string base_file(dataFilePath); + std::string data_file_to_use = base_file; + std::string labels_file_original = label_file; + std::string index_prefix_path(indexFilePath); + std::string labels_file_to_use = index_prefix_path + "_label_formatted.txt"; + std::string pq_pivots_path_base = codebook_prefix; + std::string pq_pivots_path = file_exists(pq_pivots_path_base) ? pq_pivots_path_base + "_pq_pivots.bin" + : index_prefix_path + "_pq_pivots.bin"; + std::string pq_compressed_vectors_path = index_prefix_path + "_pq_compressed.bin"; + std::string mem_index_path = index_prefix_path + "_mem.index"; + std::string disk_index_path = index_prefix_path + "_disk.index"; + std::string medoids_path = disk_index_path + "_medoids.bin"; + std::string centroids_path = disk_index_path + "_centroids.bin"; + + std::string labels_to_medoids_path = disk_index_path + "_labels_to_medoids.txt"; + std::string mem_labels_file = mem_index_path + "_labels.txt"; + std::string disk_labels_file = disk_index_path + "_labels.txt"; + std::string mem_univ_label_file = mem_index_path + "_universal_label.txt"; + std::string disk_univ_label_file = disk_index_path + "_universal_label.txt"; + std::string disk_labels_int_map_file = disk_index_path + "_labels_map.txt"; + std::string dummy_remap_file = disk_index_path + "_dummy_remap.txt"; // remap will be used if we break-up points of + // high label-density to create copies + + std::string sample_base_prefix = index_prefix_path + "_sample"; + // optional, used if disk index file must store pq data + std::string disk_pq_pivots_path = index_prefix_path + "_disk.index_pq_pivots.bin"; + // optional, used if disk index must store pq data + std::string disk_pq_compressed_vectors_path = index_prefix_path + "_disk.index_pq_compressed.bin"; + + // output a new base file which contains extra dimension with sqrt(1 - + // ||x||^2/M^2) for every x, M is max norm of all points. Extra space on + // disk needed! + if (compareMetric == diskann::Metric::INNER_PRODUCT) + { + Timer timer; + std::cout << "Using Inner Product search, so need to pre-process base " + "data into temp file. Please ensure there is additional " + "(n*(d+1)*4) bytes for storing pre-processed base vectors, " + "apart from the intermin indices and final index." + << std::endl; + std::string prepped_base = index_prefix_path + "_prepped_base.bin"; + data_file_to_use = prepped_base; + float max_norm_of_base = diskann::prepare_base_for_inner_products(base_file, prepped_base); + std::string norm_file = disk_index_path + "_max_base_norm.bin"; + diskann::save_bin(norm_file, &max_norm_of_base, 1, 1); + diskann::cout << timer.elapsed_seconds_for_step("preprocessing data for inner product") << std::endl; + } + + uint32_t R = (uint32_t)atoi(param_list[0].c_str()); + uint32_t L = (uint32_t)atoi(param_list[1].c_str()); + + 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; + return -1; + } + 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; + return -1; + } + uint32_t num_threads = (uint32_t)atoi(param_list[4].c_str()); + + if (num_threads != 0) + { + omp_set_num_threads(num_threads); + mkl_set_num_threads(num_threads); + } + + 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; + + auto s = std::chrono::high_resolution_clock::now(); + + // If there is filter support, we break-up points which have too many labels + // into replica dummy points which evenly distribute the filters. The rest + // of index build happens on the augmented base and labels + std::string augmented_data_file, augmented_labels_file; + if (use_filters) + { + convert_labels_string_to_int(labels_file_original, labels_file_to_use, disk_labels_int_map_file, + universal_label); + augmented_data_file = index_prefix_path + "_augmented_data.bin"; + augmented_labels_file = index_prefix_path + "_augmented_labels.txt"; + if (filter_threshold != 0) + { + dummy_remap_file = index_prefix_path + "_dummy_remap.txt"; + breakup_dense_points(data_file_to_use, labels_file_to_use, filter_threshold, augmented_data_file, + augmented_labels_file, + dummy_remap_file); // RKNOTE: This has large memory footprint, + // need to make this streaming + data_file_to_use = augmented_data_file; + labels_file_to_use = augmented_labels_file; + } } - } - - if (param_list.size() >= 8) { - build_pq_bytes = atoi(param_list[7].c_str()); - } - - std::string base_file(dataFilePath); - std::string data_file_to_use = base_file; - std::string labels_file_original = label_file; - std::string index_prefix_path(indexFilePath); - std::string labels_file_to_use = index_prefix_path + "_label_formatted.txt"; - std::string pq_pivots_path_base = codebook_prefix; - std::string pq_pivots_path = file_exists(pq_pivots_path_base) - ? pq_pivots_path_base + "_pq_pivots.bin" - : index_prefix_path + "_pq_pivots.bin"; - std::string pq_compressed_vectors_path = - index_prefix_path + "_pq_compressed.bin"; - std::string mem_index_path = index_prefix_path + "_mem.index"; - std::string disk_index_path = index_prefix_path + "_disk.index"; - std::string medoids_path = disk_index_path + "_medoids.bin"; - std::string centroids_path = disk_index_path + "_centroids.bin"; - - std::string labels_to_medoids_path = - disk_index_path + "_labels_to_medoids.txt"; - std::string mem_labels_file = mem_index_path + "_labels.txt"; - std::string disk_labels_file = disk_index_path + "_labels.txt"; - std::string mem_univ_label_file = mem_index_path + "_universal_label.txt"; - std::string disk_univ_label_file = disk_index_path + "_universal_label.txt"; - std::string disk_labels_int_map_file = disk_index_path + "_labels_map.txt"; - std::string dummy_remap_file = - disk_index_path + - "_dummy_remap.txt"; // remap will be used if we break-up points of - // high label-density to create copies - - std::string sample_base_prefix = index_prefix_path + "_sample"; - // optional, used if disk index file must store pq data - std::string disk_pq_pivots_path = - index_prefix_path + "_disk.index_pq_pivots.bin"; - // optional, used if disk index must store pq data - std::string disk_pq_compressed_vectors_path = - index_prefix_path + "_disk.index_pq_compressed.bin"; - - // output a new base file which contains extra dimension with sqrt(1 - - // ||x||^2/M^2) for every x, M is max norm of all points. Extra space on - // disk needed! - if (compareMetric == diskann::Metric::INNER_PRODUCT) { + + size_t points_num, dim; + Timer timer; - std::cout << "Using Inner Product search, so need to pre-process base " - "data into temp file. Please ensure there is additional " - "(n*(d+1)*4) bytes for storing pre-processed base vectors, " - "apart from the intermin indices and final index." - << std::endl; - std::string prepped_base = index_prefix_path + "_prepped_base.bin"; - data_file_to_use = prepped_base; - float max_norm_of_base = - diskann::prepare_base_for_inner_products(base_file, prepped_base); - std::string norm_file = disk_index_path + "_max_base_norm.bin"; - diskann::save_bin(norm_file, &max_norm_of_base, 1, 1); - diskann::cout << timer.elapsed_seconds_for_step( - "preprocessing data for inner product") + diskann::get_bin_metadata(data_file_to_use.c_str(), points_num, dim); + const double p_val = ((double)MAX_PQ_TRAINING_SET_SIZE / (double)points_num); + + if (use_disk_pq) + { + generate_disk_quantized_data(data_file_to_use, disk_pq_pivots_path, disk_pq_compressed_vectors_path, + compareMetric, p_val, disk_pq_dims); + } + size_t num_pq_chunks = (size_t)(std::floor)(uint64_t(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; + + if (param_list.size() >= 9 && atoi(param_list[8].c_str()) <= MAX_PQ_CHUNKS && atoi(param_list[8].c_str()) > 0) + { + std::cout << "Use quantized dimension (QD) to overwrite derived quantized " + "dimension from search_DRAM_budget (B)" << std::endl; - } - - uint32_t R = (uint32_t)atoi(param_list[0].c_str()); - uint32_t L = (uint32_t)atoi(param_list[1].c_str()); - - 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; - return -1; - } - 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; - return -1; - } - uint32_t num_threads = (uint32_t)atoi(param_list[4].c_str()); - - if (num_threads != 0) { - omp_set_num_threads(num_threads); - mkl_set_num_threads(num_threads); - } - - 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; - - auto s = std::chrono::high_resolution_clock::now(); - - // If there is filter support, we break-up points which have too many labels - // into replica dummy points which evenly distribute the filters. The rest - // of index build happens on the augmented base and labels - std::string augmented_data_file, augmented_labels_file; - if (use_filters) { - convert_labels_string_to_int(labels_file_original, labels_file_to_use, - disk_labels_int_map_file, universal_label); - augmented_data_file = index_prefix_path + "_augmented_data.bin"; - augmented_labels_file = index_prefix_path + "_augmented_labels.txt"; - if (filter_threshold != 0) { - dummy_remap_file = index_prefix_path + "_dummy_remap.txt"; - breakup_dense_points( - data_file_to_use, labels_file_to_use, filter_threshold, - augmented_data_file, augmented_labels_file, - dummy_remap_file); // RKNOTE: This has large memory footprint, - // need to make this streaming - data_file_to_use = augmented_data_file; - labels_file_to_use = augmented_labels_file; + num_pq_chunks = atoi(param_list[8].c_str()); } - } - - size_t points_num, dim; - - Timer timer; - diskann::get_bin_metadata(data_file_to_use.c_str(), points_num, dim); - const double p_val = ((double)MAX_PQ_TRAINING_SET_SIZE / (double)points_num); - - if (use_disk_pq) { - generate_disk_quantized_data(data_file_to_use, disk_pq_pivots_path, - disk_pq_compressed_vectors_path, - compareMetric, p_val, disk_pq_dims); - } - size_t num_pq_chunks = - (size_t)(std::floor)(uint64_t(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; - - if (param_list.size() >= 9 && atoi(param_list[8].c_str()) <= MAX_PQ_CHUNKS && - atoi(param_list[8].c_str()) > 0) { - std::cout << "Use quantized dimension (QD) to overwrite derived quantized " - "dimension from search_DRAM_budget (B)" - << std::endl; - num_pq_chunks = atoi(param_list[8].c_str()); - } - - diskann::cout << "Compressing " << dim << "-dimensional data into " - << num_pq_chunks << " bytes per vector." << std::endl; - - generate_quantized_data(data_file_to_use, pq_pivots_path, - pq_compressed_vectors_path, compareMetric, p_val, - num_pq_chunks, use_opq, codebook_prefix); - diskann::cout << timer.elapsed_seconds_for_step("generating quantized data") - << std::endl; + + diskann::cout << "Compressing " << dim << "-dimensional data into " << num_pq_chunks << " bytes per vector." + << std::endl; + + generate_quantized_data(data_file_to_use, pq_pivots_path, pq_compressed_vectors_path, compareMetric, p_val, + num_pq_chunks, use_opq, codebook_prefix); + diskann::cout << timer.elapsed_seconds_for_step("generating quantized data") << std::endl; // Gopal. Splitting diskann_dll into separate DLLs for search and build. // This code should only be available in the "build" DLL. -#if defined(RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS) && \ - defined(DISKANN_BUILD) - MallocExtension::instance()->ReleaseFreeMemory(); +#if defined(RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS) && defined(DISKANN_BUILD) + MallocExtension::instance()->ReleaseFreeMemory(); #endif - timer.reset(); - diskann::build_merged_vamana_index( - data_file_to_use.c_str(), diskann::Metric::L2, L, R, p_val, - indexing_ram_budget, mem_index_path, medoids_path, centroids_path, - build_pq_bytes, use_opq, num_threads, use_filters, labels_file_to_use, - labels_to_medoids_path, universal_label, Lf); - diskann::cout << timer.elapsed_seconds_for_step( - "building merged vamana index") - << std::endl; - - timer.reset(); - if (!use_disk_pq) { - diskann::create_disk_layout(data_file_to_use.c_str(), mem_index_path, - disk_index_path); - } else { - if (!reorder_data) - diskann::create_disk_layout(disk_pq_compressed_vectors_path, - mem_index_path, disk_index_path); + timer.reset(); + diskann::build_merged_vamana_index(data_file_to_use.c_str(), diskann::Metric::L2, L, R, p_val, + indexing_ram_budget, mem_index_path, medoids_path, centroids_path, + build_pq_bytes, use_opq, num_threads, use_filters, labels_file_to_use, + labels_to_medoids_path, universal_label, Lf); + diskann::cout << timer.elapsed_seconds_for_step("building merged vamana index") << std::endl; + + timer.reset(); + if (!use_disk_pq) + { + diskann::create_disk_layout(data_file_to_use.c_str(), mem_index_path, disk_index_path); + } else - diskann::create_disk_layout(disk_pq_compressed_vectors_path, - mem_index_path, disk_index_path, - data_file_to_use.c_str()); - } - diskann::cout << timer.elapsed_seconds_for_step("generating disk layout") - << std::endl; - - 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; - gen_random_slice(data_file_to_use.c_str(), sample_base_prefix, - sample_sampling_rate); - if (use_filters) { - copy_file(labels_file_to_use, disk_labels_file); - std::remove(mem_labels_file.c_str()); - if (universal_label != "") { - copy_file(mem_univ_label_file, disk_univ_label_file); - std::remove(mem_univ_label_file.c_str()); + { + if (!reorder_data) + diskann::create_disk_layout(disk_pq_compressed_vectors_path, mem_index_path, disk_index_path); + else + diskann::create_disk_layout(disk_pq_compressed_vectors_path, mem_index_path, disk_index_path, + data_file_to_use.c_str()); + } + diskann::cout << timer.elapsed_seconds_for_step("generating disk layout") << std::endl; + + 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; + gen_random_slice(data_file_to_use.c_str(), sample_base_prefix, sample_sampling_rate); + if (use_filters) + { + copy_file(labels_file_to_use, disk_labels_file); + std::remove(mem_labels_file.c_str()); + if (universal_label != "") + { + copy_file(mem_univ_label_file, disk_univ_label_file); + std::remove(mem_univ_label_file.c_str()); + } + std::remove(augmented_data_file.c_str()); + std::remove(augmented_labels_file.c_str()); + std::remove(labels_file_to_use.c_str()); } - std::remove(augmented_data_file.c_str()); - std::remove(augmented_labels_file.c_str()); - std::remove(labels_file_to_use.c_str()); - } - std::remove(mem_index_path.c_str()); - if (use_disk_pq) std::remove(disk_pq_compressed_vectors_path.c_str()); + std::remove(mem_index_path.c_str()); + if (use_disk_pq) + std::remove(disk_pq_compressed_vectors_path.c_str()); - auto e = std::chrono::high_resolution_clock::now(); - std::chrono::duration diff = e - s; - diskann::cout << "Indexing time: " << diff.count() << std::endl; + auto e = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = e - s; + diskann::cout << "Indexing time: " << diff.count() << std::endl; - return 0; + return 0; } -template DISKANN_DLLEXPORT void create_disk_layout( - const std::string base_file, const std::string mem_index_file, - const std::string output_file, const std::string reorder_data_file); -template DISKANN_DLLEXPORT void create_disk_layout( - const std::string base_file, const std::string mem_index_file, - const std::string output_file, const std::string reorder_data_file); -template DISKANN_DLLEXPORT void create_disk_layout( - const std::string base_file, const std::string mem_index_file, - const std::string output_file, const std::string reorder_data_file); - -template DISKANN_DLLEXPORT int8_t *load_warmup( - const std::string &cache_warmup_file, uint64_t &warmup_num, - uint64_t warmup_dim, uint64_t warmup_aligned_dim); -template DISKANN_DLLEXPORT uint8_t *load_warmup( - const std::string &cache_warmup_file, uint64_t &warmup_num, - uint64_t warmup_dim, uint64_t warmup_aligned_dim); -template DISKANN_DLLEXPORT float *load_warmup( - const std::string &cache_warmup_file, uint64_t &warmup_num, - uint64_t warmup_dim, uint64_t warmup_aligned_dim); +template DISKANN_DLLEXPORT void create_disk_layout(const std::string base_file, + const std::string mem_index_file, + const std::string output_file, + const std::string reorder_data_file); +template DISKANN_DLLEXPORT void create_disk_layout(const std::string base_file, + const std::string mem_index_file, + const std::string output_file, + const std::string reorder_data_file); +template DISKANN_DLLEXPORT void create_disk_layout(const std::string base_file, const std::string mem_index_file, + const std::string output_file, + const std::string reorder_data_file); + +template DISKANN_DLLEXPORT int8_t *load_warmup(const std::string &cache_warmup_file, uint64_t &warmup_num, + uint64_t warmup_dim, uint64_t warmup_aligned_dim); +template DISKANN_DLLEXPORT uint8_t *load_warmup(const std::string &cache_warmup_file, uint64_t &warmup_num, + uint64_t warmup_dim, uint64_t warmup_aligned_dim); +template DISKANN_DLLEXPORT float *load_warmup(const std::string &cache_warmup_file, uint64_t &warmup_num, + uint64_t warmup_dim, uint64_t warmup_aligned_dim); #ifdef EXEC_ENV_OLS -template DISKANN_DLLEXPORT int8_t *load_warmup( - MemoryMappedFiles &files, const std::string &cache_warmup_file, - uint64_t &warmup_num, uint64_t warmup_dim, uint64_t warmup_aligned_dim); -template DISKANN_DLLEXPORT uint8_t *load_warmup( - MemoryMappedFiles &files, const std::string &cache_warmup_file, - uint64_t &warmup_num, uint64_t warmup_dim, uint64_t warmup_aligned_dim); -template DISKANN_DLLEXPORT float *load_warmup( - MemoryMappedFiles &files, const std::string &cache_warmup_file, - uint64_t &warmup_num, uint64_t warmup_dim, uint64_t warmup_aligned_dim); +template DISKANN_DLLEXPORT int8_t *load_warmup(MemoryMappedFiles &files, const std::string &cache_warmup_file, + uint64_t &warmup_num, uint64_t warmup_dim, + uint64_t warmup_aligned_dim); +template DISKANN_DLLEXPORT uint8_t *load_warmup(MemoryMappedFiles &files, const std::string &cache_warmup_file, + uint64_t &warmup_num, uint64_t warmup_dim, + uint64_t warmup_aligned_dim); +template DISKANN_DLLEXPORT float *load_warmup(MemoryMappedFiles &files, const std::string &cache_warmup_file, + 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, - int8_t *tuning_sample, uint64_t tuning_sample_num, - uint64_t tuning_sample_aligned_dim, uint32_t L, uint32_t nthreads, - uint32_t start_bw); + std::unique_ptr> &pFlashIndex, int8_t *tuning_sample, + uint64_t tuning_sample_num, uint64_t 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, uint64_t tuning_sample_num, - uint64_t tuning_sample_aligned_dim, uint32_t L, uint32_t nthreads, - uint32_t start_bw); + std::unique_ptr> &pFlashIndex, uint8_t *tuning_sample, + uint64_t tuning_sample_num, uint64_t 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, uint64_t tuning_sample_num, - uint64_t tuning_sample_aligned_dim, uint32_t L, uint32_t nthreads, - uint32_t start_bw); + std::unique_ptr> &pFlashIndex, float *tuning_sample, + uint64_t tuning_sample_num, uint64_t 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, - int8_t *tuning_sample, uint64_t tuning_sample_num, - uint64_t tuning_sample_aligned_dim, uint32_t L, uint32_t nthreads, - uint32_t start_bw); + std::unique_ptr> &pFlashIndex, int8_t *tuning_sample, + uint64_t tuning_sample_num, uint64_t 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, uint64_t tuning_sample_num, - uint64_t tuning_sample_aligned_dim, uint32_t L, uint32_t nthreads, - uint32_t start_bw); + std::unique_ptr> &pFlashIndex, uint8_t *tuning_sample, + uint64_t tuning_sample_num, uint64_t 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, uint64_t tuning_sample_num, - uint64_t tuning_sample_aligned_dim, uint32_t L, uint32_t nthreads, - uint32_t start_bw); - -template DISKANN_DLLEXPORT int build_disk_index( - const char *dataFilePath, const char *indexFilePath, - const char *indexBuildParameters, diskann::Metric compareMetric, - bool use_opq, const std::string &codebook_prefix, bool use_filters, - const std::string &label_file, const std::string &universal_label, - const uint32_t filter_threshold, const uint32_t Lf); -template DISKANN_DLLEXPORT int build_disk_index( - const char *dataFilePath, const char *indexFilePath, - const char *indexBuildParameters, diskann::Metric compareMetric, - bool use_opq, const std::string &codebook_prefix, bool use_filters, - const std::string &label_file, const std::string &universal_label, - const uint32_t filter_threshold, const uint32_t Lf); -template DISKANN_DLLEXPORT int build_disk_index( - const char *dataFilePath, const char *indexFilePath, - const char *indexBuildParameters, diskann::Metric compareMetric, - bool use_opq, const std::string &codebook_prefix, bool use_filters, - const std::string &label_file, const std::string &universal_label, - const uint32_t filter_threshold, const uint32_t Lf); + std::unique_ptr> &pFlashIndex, float *tuning_sample, + uint64_t tuning_sample_num, uint64_t tuning_sample_aligned_dim, uint32_t L, uint32_t nthreads, uint32_t start_bw); + +template DISKANN_DLLEXPORT int build_disk_index(const char *dataFilePath, const char *indexFilePath, + const char *indexBuildParameters, + diskann::Metric compareMetric, bool use_opq, + const std::string &codebook_prefix, bool use_filters, + const std::string &label_file, + const std::string &universal_label, + const uint32_t filter_threshold, const uint32_t Lf); +template DISKANN_DLLEXPORT int build_disk_index(const char *dataFilePath, const char *indexFilePath, + const char *indexBuildParameters, + diskann::Metric compareMetric, bool use_opq, + const std::string &codebook_prefix, bool use_filters, + const std::string &label_file, + const std::string &universal_label, + const uint32_t filter_threshold, const uint32_t Lf); +template DISKANN_DLLEXPORT int build_disk_index(const char *dataFilePath, const char *indexFilePath, + const char *indexBuildParameters, + diskann::Metric compareMetric, bool use_opq, + const std::string &codebook_prefix, bool use_filters, + const std::string &label_file, + const std::string &universal_label, + const uint32_t filter_threshold, const uint32_t Lf); // LabelT = uint16 -template DISKANN_DLLEXPORT int build_disk_index( - const char *dataFilePath, const char *indexFilePath, - const char *indexBuildParameters, diskann::Metric compareMetric, - bool use_opq, const std::string &codebook_prefix, bool use_filters, - const std::string &label_file, const std::string &universal_label, - const uint32_t filter_threshold, const uint32_t Lf); -template DISKANN_DLLEXPORT int build_disk_index( - const char *dataFilePath, const char *indexFilePath, - const char *indexBuildParameters, diskann::Metric compareMetric, - bool use_opq, const std::string &codebook_prefix, bool use_filters, - const std::string &label_file, const std::string &universal_label, - const uint32_t filter_threshold, const uint32_t Lf); -template DISKANN_DLLEXPORT int build_disk_index( - const char *dataFilePath, const char *indexFilePath, - const char *indexBuildParameters, diskann::Metric compareMetric, - bool use_opq, const std::string &codebook_prefix, bool use_filters, - const std::string &label_file, const std::string &universal_label, - const uint32_t filter_threshold, const uint32_t Lf); +template DISKANN_DLLEXPORT int build_disk_index(const char *dataFilePath, const char *indexFilePath, + const char *indexBuildParameters, + diskann::Metric compareMetric, bool use_opq, + const std::string &codebook_prefix, bool use_filters, + const std::string &label_file, + const std::string &universal_label, + const uint32_t filter_threshold, const uint32_t Lf); +template DISKANN_DLLEXPORT int build_disk_index(const char *dataFilePath, const char *indexFilePath, + const char *indexBuildParameters, + diskann::Metric compareMetric, bool use_opq, + const std::string &codebook_prefix, bool use_filters, + const std::string &label_file, + const std::string &universal_label, + const uint32_t filter_threshold, const uint32_t Lf); +template DISKANN_DLLEXPORT int build_disk_index(const char *dataFilePath, const char *indexFilePath, + const char *indexBuildParameters, + diskann::Metric compareMetric, bool use_opq, + const std::string &codebook_prefix, bool use_filters, + const std::string &label_file, + const std::string &universal_label, + const uint32_t filter_threshold, const uint32_t Lf); template DISKANN_DLLEXPORT int build_merged_vamana_index( - std::string base_file, diskann::Metric compareMetric, uint32_t L, - uint32_t R, double sampling_rate, double ram_budget, - std::string mem_index_path, std::string medoids_path, - std::string centroids_file, size_t build_pq_bytes, bool use_opq, - uint32_t num_threads, bool use_filters, const std::string &label_file, - const std::string &labels_to_medoids_file, - const std::string &universal_label, const uint32_t Lf); + std::string base_file, diskann::Metric compareMetric, uint32_t L, uint32_t R, double sampling_rate, + double ram_budget, std::string mem_index_path, std::string medoids_path, std::string centroids_file, + size_t build_pq_bytes, bool use_opq, uint32_t num_threads, bool use_filters, const std::string &label_file, + const std::string &labels_to_medoids_file, const std::string &universal_label, const uint32_t Lf); template DISKANN_DLLEXPORT int build_merged_vamana_index( - std::string base_file, diskann::Metric compareMetric, uint32_t L, - uint32_t R, double sampling_rate, double ram_budget, - std::string mem_index_path, std::string medoids_path, - std::string centroids_file, size_t build_pq_bytes, bool use_opq, - uint32_t num_threads, bool use_filters, const std::string &label_file, - const std::string &labels_to_medoids_file, - const std::string &universal_label, const uint32_t Lf); + std::string base_file, diskann::Metric compareMetric, uint32_t L, uint32_t R, double sampling_rate, + double ram_budget, std::string mem_index_path, std::string medoids_path, std::string centroids_file, + size_t build_pq_bytes, bool use_opq, uint32_t num_threads, bool use_filters, const std::string &label_file, + const std::string &labels_to_medoids_file, const std::string &universal_label, const uint32_t Lf); template DISKANN_DLLEXPORT int build_merged_vamana_index( - std::string base_file, diskann::Metric compareMetric, uint32_t L, - uint32_t R, double sampling_rate, double ram_budget, - std::string mem_index_path, std::string medoids_path, - std::string centroids_file, size_t build_pq_bytes, bool use_opq, - uint32_t num_threads, bool use_filters, const std::string &label_file, - const std::string &labels_to_medoids_file, - const std::string &universal_label, const uint32_t Lf); + std::string base_file, diskann::Metric compareMetric, uint32_t L, uint32_t R, double sampling_rate, + double ram_budget, std::string mem_index_path, std::string medoids_path, std::string centroids_file, + size_t build_pq_bytes, bool use_opq, uint32_t num_threads, bool use_filters, const std::string &label_file, + const std::string &labels_to_medoids_file, const std::string &universal_label, const uint32_t Lf); // Label=16_t template DISKANN_DLLEXPORT int build_merged_vamana_index( - std::string base_file, diskann::Metric compareMetric, uint32_t L, - uint32_t R, double sampling_rate, double ram_budget, - std::string mem_index_path, std::string medoids_path, - std::string centroids_file, size_t build_pq_bytes, bool use_opq, - uint32_t num_threads, bool use_filters, const std::string &label_file, - const std::string &labels_to_medoids_file, - const std::string &universal_label, const uint32_t Lf); + std::string base_file, diskann::Metric compareMetric, uint32_t L, uint32_t R, double sampling_rate, + double ram_budget, std::string mem_index_path, std::string medoids_path, std::string centroids_file, + size_t build_pq_bytes, bool use_opq, uint32_t num_threads, bool use_filters, const std::string &label_file, + const std::string &labels_to_medoids_file, const std::string &universal_label, const uint32_t Lf); template DISKANN_DLLEXPORT int build_merged_vamana_index( - std::string base_file, diskann::Metric compareMetric, uint32_t L, - uint32_t R, double sampling_rate, double ram_budget, - std::string mem_index_path, std::string medoids_path, - std::string centroids_file, size_t build_pq_bytes, bool use_opq, - uint32_t num_threads, bool use_filters, const std::string &label_file, - const std::string &labels_to_medoids_file, - const std::string &universal_label, const uint32_t Lf); + std::string base_file, diskann::Metric compareMetric, uint32_t L, uint32_t R, double sampling_rate, + double ram_budget, std::string mem_index_path, std::string medoids_path, std::string centroids_file, + size_t build_pq_bytes, bool use_opq, uint32_t num_threads, bool use_filters, const std::string &label_file, + const std::string &labels_to_medoids_file, const std::string &universal_label, const uint32_t Lf); template DISKANN_DLLEXPORT int build_merged_vamana_index( - std::string base_file, diskann::Metric compareMetric, uint32_t L, - uint32_t R, double sampling_rate, double ram_budget, - std::string mem_index_path, std::string medoids_path, - std::string centroids_file, size_t build_pq_bytes, bool use_opq, - uint32_t num_threads, bool use_filters, const std::string &label_file, - const std::string &labels_to_medoids_file, - const std::string &universal_label, const uint32_t Lf); -}; // namespace diskann + std::string base_file, diskann::Metric compareMetric, uint32_t L, uint32_t R, double sampling_rate, + double ram_budget, std::string mem_index_path, std::string medoids_path, std::string centroids_file, + size_t build_pq_bytes, bool use_opq, uint32_t num_threads, bool use_filters, const std::string &label_file, + const std::string &labels_to_medoids_file, const std::string &universal_label, const uint32_t Lf); +}; // namespace diskann diff --git a/src/distance.cpp b/src/distance.cpp index e6be334e6..d8ae7a0cc 100644 --- a/src/distance.cpp +++ b/src/distance.cpp @@ -19,669 +19,703 @@ #include "logger.h" #include "ann_exception.h" -namespace diskann { +namespace diskann +{ // // Base Class Implementatons // template -float Distance::compare(const T *a, const T *b, const float normA, - const float normB, uint32_t length) const { - throw std::logic_error("This function is not implemented."); +float Distance::compare(const T *a, const T *b, const float normA, const float normB, uint32_t length) const +{ + throw std::logic_error("This function is not implemented."); } -template -uint32_t Distance::post_normalization_dimension( - uint32_t orig_dimension) const { - return orig_dimension; +template uint32_t Distance::post_normalization_dimension(uint32_t orig_dimension) const +{ + return orig_dimension; } -template -diskann::Metric Distance::get_metric() const { - return _distance_metric; +template diskann::Metric Distance::get_metric() const +{ + return _distance_metric; } -template -bool Distance::preprocessing_required() const { - return false; +template bool Distance::preprocessing_required() const +{ + return false; } template -void Distance::preprocess_base_points(T *original_data, - const size_t orig_dim, - const size_t num_points) {} +void Distance::preprocess_base_points(T *original_data, const size_t orig_dim, const size_t num_points) +{ +} -template -void Distance::preprocess_query(const T *query_vec, const size_t query_dim, - T *scratch_query) { - std::memcpy(scratch_query, query_vec, query_dim * sizeof(T)); +template void Distance::preprocess_query(const T *query_vec, const size_t query_dim, T *scratch_query) +{ + std::memcpy(scratch_query, query_vec, query_dim * sizeof(T)); } -template -size_t Distance::get_required_alignment() const { - return _alignment_factor; +template size_t Distance::get_required_alignment() const +{ + return _alignment_factor; } -template -Distance::~Distance() {} +template Distance::~Distance() +{ +} // // Cosine distance functions. // -float DistanceCosineInt8::compare(const int8_t *a, const int8_t *b, - uint32_t length) const { +float DistanceCosineInt8::compare(const int8_t *a, const int8_t *b, uint32_t length) const +{ #ifdef _WINDOWS - return diskann::CosineSimilarity2(a, b, length); + return diskann::CosineSimilarity2(a, b, length); #else - int magA = 0, magB = 0, scalarProduct = 0; - for (uint32_t i = 0; i < length; i++) { - magA += ((int32_t)a[i]) * ((int32_t)a[i]); - magB += ((int32_t)b[i]) * ((int32_t)b[i]); - scalarProduct += ((int32_t)a[i]) * ((int32_t)b[i]); - } - // similarity == 1-cosine distance - return 1.0f - (float)(scalarProduct / (sqrt(magA) * sqrt(magB))); + int magA = 0, magB = 0, scalarProduct = 0; + for (uint32_t i = 0; i < length; i++) + { + magA += ((int32_t)a[i]) * ((int32_t)a[i]); + magB += ((int32_t)b[i]) * ((int32_t)b[i]); + scalarProduct += ((int32_t)a[i]) * ((int32_t)b[i]); + } + // similarity == 1-cosine distance + return 1.0f - (float)(scalarProduct / (sqrt(magA) * sqrt(magB))); #endif } -float DistanceCosineFloat::compare(const float *a, const float *b, - uint32_t length) const { +float DistanceCosineFloat::compare(const float *a, const float *b, uint32_t length) const +{ #ifdef _WINDOWS - return diskann::CosineSimilarity2(a, b, length); + return diskann::CosineSimilarity2(a, b, length); #else - float magA = 0, magB = 0, scalarProduct = 0; - for (uint32_t i = 0; i < length; i++) { - magA += (a[i]) * (a[i]); - magB += (b[i]) * (b[i]); - scalarProduct += (a[i]) * (b[i]); - } - // similarity == 1-cosine distance - return 1.0f - (scalarProduct / (sqrt(magA) * sqrt(magB))); + float magA = 0, magB = 0, scalarProduct = 0; + for (uint32_t i = 0; i < length; i++) + { + magA += (a[i]) * (a[i]); + magB += (b[i]) * (b[i]); + scalarProduct += (a[i]) * (b[i]); + } + // similarity == 1-cosine distance + return 1.0f - (scalarProduct / (sqrt(magA) * sqrt(magB))); #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))); +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 { +float DistanceL2Int8::compare(const int8_t *a, const int8_t *b, uint32_t size) 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]; + __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 - int32_t result = 0; + int32_t result = 0; #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; + 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 #else - int32_t result = 0; + int32_t result = 0; #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; + 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; +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; + 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 DistanceL2Float::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); +float DistanceL2Float::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 DistanceL2Float::compare(const float *a, const float *b, - uint32_t size) const { +float DistanceL2Float::compare(const float *a, const float *b, uint32_t size) const +{ #endif - float result = 0; + 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); - - sum = _mm256_fmadd_ps(tmp_vec, tmp_vec, sum); - } - - // horizontal add sum - result = _mm256_reduce_add_ps(sum); + // 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); + + 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 (int32_t i = 0; i < (int32_t)size; i++) { - result += (a[i] - b[i]) * (a[i] - b[i]); - } + for (int32_t i = 0; i < (int32_t)size; i++) + { + result += (a[i] - b[i]) * (a[i] - b[i]); + } #endif - return result; + return result; } -template -float SlowDistanceL2::compare(const T *a, const T *b, - uint32_t length) const { - float result = 0.0f; - for (uint32_t i = 0; i < length; i++) { - result += ((float)(a[i] - b[i])) * (a[i] - b[i]); - } - return result; +template float SlowDistanceL2::compare(const T *a, const T *b, uint32_t length) const +{ + float result = 0.0f; + for (uint32_t i = 0; i < length; i++) + { + result += ((float)(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 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]; +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 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; +float AVXDistanceL2Float::compare(const float *, const float *, uint32_t) const +{ + return 0; } #endif -template -float DistanceInnerProduct::inner_product(const T *a, const T *b, - uint32_t size) const { - if (!std::is_floating_point::value) { - diskann::cerr << "ERROR: Inner Product only defined for float currently." - << std::endl; - throw diskann::ANNException( - "ERROR: Inner Product only defined for float currently.", -1, - __FUNCSIG__, __FILE__, __LINE__); - } - - float result = 0; +template float DistanceInnerProduct::inner_product(const T *a, const T *b, uint32_t size) const +{ + if (!std::is_floating_point::value) + { + diskann::cerr << "ERROR: Inner Product only defined for float currently." << std::endl; + throw diskann::ANNException("ERROR: Inner Product only defined for float currently.", -1, __FUNCSIG__, __FILE__, + __LINE__); + } + + float result = 0; #ifdef __GNUC__ #ifdef USE_AVX2 -#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; - uint32_t D = (size + 7) & ~7U; - uint32_t DR = D % 16; - uint32_t 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 (uint32_t 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]; +#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; + uint32_t D = (size + 7) & ~7U; + uint32_t DR = D % 16; + uint32_t 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 (uint32_t 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; - uint32_t D = (size + 3) & ~3U; - uint32_t DR = D % 16; - uint32_t 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) { +#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; + uint32_t D = (size + 3) & ~3U; + uint32_t DR = D % 16; + uint32_t 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); + SSE_DOT(e_l + 8, e_r + 8, sum, l2, r2); case 8: - SSE_DOT(e_l + 4, e_r + 4, sum, l1, r1); + SSE_DOT(e_l + 4, e_r + 4, sum, l1, r1); case 4: - SSE_DOT(e_l, e_r, sum, l0, r0); + SSE_DOT(e_l, e_r, sum, l0, r0); default: - break; - } - for (uint32_t 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]; + break; + } + for (uint32_t 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++; - } + 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; + return result; } -template -float DistanceFastL2::compare(const T *a, const T *b, float norm, - uint32_t size) const { - float result = -2 * DistanceInnerProduct::inner_product(a, b, size); - result += norm; - return result; +template float DistanceFastL2::compare(const T *a, const T *b, float norm, uint32_t size) const +{ + float result = -2 * DistanceInnerProduct::inner_product(a, b, size); + result += norm; + return result; } -template -float DistanceFastL2::norm(const T *a, uint32_t size) const { - if (!std::is_floating_point::value) { - diskann::cerr << "ERROR: FastL2 only defined for float currently." - << std::endl; - throw diskann::ANNException( - "ERROR: FastL2 only defined for float currently.", -1, __FUNCSIG__, - __FILE__, __LINE__); - } - float result = 0; +template float DistanceFastL2::norm(const T *a, uint32_t size) const +{ + if (!std::is_floating_point::value) + { + diskann::cerr << "ERROR: FastL2 only defined for float currently." << std::endl; + throw diskann::ANNException("ERROR: FastL2 only defined for float currently.", -1, __FUNCSIG__, __FILE__, + __LINE__); + } + 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; - uint32_t D = (size + 7) & ~7U; - uint32_t DR = D % 16; - uint32_t 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 (uint32_t 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]; +#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; + uint32_t D = (size + 7) & ~7U; + uint32_t DR = D % 16; + uint32_t 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 (uint32_t 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; - uint32_t D = (size + 3) & ~3U; - uint32_t DR = D % 16; - uint32_t 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) { +#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; + uint32_t D = (size + 3) & ~3U; + uint32_t DR = D % 16; + uint32_t 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); + SSE_L2NORM(e_l + 8, sum, l2); case 8: - SSE_L2NORM(e_l + 4, sum, l1); + SSE_L2NORM(e_l + 4, sum, l1); case 4: - SSE_L2NORM(e_l, sum, l0); + SSE_L2NORM(e_l, sum, l0); default: - break; - } - for (uint32_t 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]; + break; + } + for (uint32_t 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++; - } + 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; + return result; } -float AVXDistanceInnerProductFloat::compare(const float *a, const float *b, - uint32_t size) const { - float result = 0.0f; -#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; - uint32_t D = (size + 7) & ~7U; - uint32_t DR = D % 16; - uint32_t 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 AVXDistanceInnerProductFloat::compare(const float *a, const float *b, uint32_t size) const +{ + float result = 0.0f; +#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; + uint32_t D = (size + 7) & ~7U; + uint32_t DR = D % 16; + uint32_t 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; #ifndef _WINDOWS - float unpack[8] __attribute__((aligned(32))) = {0, 0, 0, 0, 0, 0, 0, 0}; + float unpack[8] __attribute__((aligned(32))) = {0, 0, 0, 0, 0, 0, 0, 0}; #else - __declspec(align(32)) float unpack[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + __declspec(align(32)) float unpack[8] = {0, 0, 0, 0, 0, 0, 0, 0}; #endif - sum = _mm256_loadu_ps(unpack); - if (DR) { - AVX_DOT(e_l, e_r, sum, l0, r0); - } + sum = _mm256_loadu_ps(unpack); + if (DR) + { + AVX_DOT(e_l, e_r, sum, l0, r0); + } - for (uint32_t 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]; + for (uint32_t 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]; - return -result; + return -result; } -uint32_t AVXNormalizedCosineDistanceFloat::post_normalization_dimension( - uint32_t orig_dimension) const { - return orig_dimension; +uint32_t AVXNormalizedCosineDistanceFloat::post_normalization_dimension(uint32_t orig_dimension) const +{ + return orig_dimension; } -bool AVXNormalizedCosineDistanceFloat::preprocessing_required() const { - return true; +bool AVXNormalizedCosineDistanceFloat::preprocessing_required() const +{ + return true; } -void AVXNormalizedCosineDistanceFloat::preprocess_base_points( - float *original_data, const size_t orig_dim, const size_t num_points) { - for (uint32_t i = 0; i < num_points; i++) { - normalize((float *)(original_data + i * orig_dim), orig_dim); - } +void AVXNormalizedCosineDistanceFloat::preprocess_base_points(float *original_data, const size_t orig_dim, + const size_t num_points) +{ + for (uint32_t i = 0; i < num_points; i++) + { + normalize((float *)(original_data + i * orig_dim), orig_dim); + } } -void AVXNormalizedCosineDistanceFloat::preprocess_query(const float *query_vec, - const size_t query_dim, - float *query_scratch) { - normalize_and_copy(query_vec, query_dim, query_scratch); +void AVXNormalizedCosineDistanceFloat::preprocess_query(const float *query_vec, const size_t query_dim, + float *query_scratch) +{ + normalize_and_copy(query_vec, query_dim, query_scratch); } -void AVXNormalizedCosineDistanceFloat::normalize_and_copy( - const float *query_vec, const uint32_t query_dim, - float *query_target) const { - float norm = get_norm(query_vec, query_dim); +void AVXNormalizedCosineDistanceFloat::normalize_and_copy(const float *query_vec, const uint32_t query_dim, + float *query_target) const +{ + float norm = get_norm(query_vec, query_dim); - for (uint32_t i = 0; i < query_dim; i++) { - query_target[i] = query_vec[i] / norm; - } + for (uint32_t i = 0; i < query_dim; i++) + { + query_target[i] = query_vec[i] / norm; + } } // 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 DistanceL2Float" - << std::endl; - return new diskann::DistanceL2Float(); - } 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::SlowDistanceL2(); - } - } else if (m == diskann::Metric::COSINE) { - diskann::cout << "Cosine: Using either AVX or AVX2 implementation" - << std::endl; - return new diskann::DistanceCosineFloat(); - } else if (m == diskann::Metric::INNER_PRODUCT) { - diskann::cout << "Inner product: Using AVX2 implementation " - "AVXDistanceInnerProductFloat" - << std::endl; - return new diskann::AVXDistanceInnerProductFloat(); - } else if (m == diskann::Metric::FAST_L2) { - diskann::cout << "Fast_L2: Using AVX2 implementation with norm " - "memoization DistanceFastL2" - << std::endl; - return new diskann::DistanceFastL2(); - } else { - std::stringstream stream; - stream << "Only L2, cosine, and inner product supported for floating " - "point vectors as of now." - << std::endl; - diskann::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 << "L2: Using AVX2 distance computation DistanceL2Float" << std::endl; + return new diskann::DistanceL2Float(); + } + 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::SlowDistanceL2(); + } + } + else if (m == diskann::Metric::COSINE) + { + diskann::cout << "Cosine: Using either AVX or AVX2 implementation" << std::endl; + return new diskann::DistanceCosineFloat(); + } + else if (m == diskann::Metric::INNER_PRODUCT) + { + diskann::cout << "Inner product: Using AVX2 implementation " + "AVXDistanceInnerProductFloat" + << std::endl; + return new diskann::AVXDistanceInnerProductFloat(); + } + else if (m == diskann::Metric::FAST_L2) + { + diskann::cout << "Fast_L2: Using AVX2 implementation with norm " + "memoization DistanceFastL2" + << std::endl; + return new diskann::DistanceFastL2(); + } + else + { + std::stringstream stream; + stream << "Only L2, cosine, and inner product supported for floating " + "point vectors as of now." + << std::endl; + diskann::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 DistanceL2Int8." - << 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 " - "SlowDistanceL2Int." - << std::endl; - return new diskann::SlowDistanceL2(); - } - } else if (m == diskann::Metric::COSINE) { - diskann::cout << "Using either AVX or AVX2 for Cosine similarity " - "DistanceCosineInt8." - << std::endl; - return new diskann::DistanceCosineInt8(); - } else { - std::stringstream stream; - stream << "Only L2 and cosine supported for signed byte vectors." - << std::endl; - diskann::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 DistanceL2Int8." << 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 " + "SlowDistanceL2Int." + << std::endl; + return new diskann::SlowDistanceL2(); + } + } + else if (m == diskann::Metric::COSINE) + { + diskann::cout << "Using either AVX or AVX2 for Cosine similarity " + "DistanceCosineInt8." + << std::endl; + return new diskann::DistanceCosineInt8(); + } + else + { + std::stringstream stream; + stream << "Only L2 and cosine supported for signed byte vectors." << std::endl; + diskann::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) { +template <> diskann::Distance *get_distance_function(diskann::Metric m) +{ + if (m == diskann::Metric::L2) + { #ifdef _WINDOWS - diskann::cout - << "WARNING: AVX/AVX2 distance function not defined for Uint8. " - "Using " - "slow version. " - "Contact gopalsr@microsoft.com if you need AVX/AVX2 support." - << std::endl; + diskann::cout << "WARNING: AVX/AVX2 distance function not defined for Uint8. " + "Using " + "slow version. " + "Contact gopalsr@microsoft.com if you need AVX/AVX2 support." + << std::endl; #endif - return new diskann::DistanceL2UInt8(); - } else if (m == diskann::Metric::COSINE) { - diskann::cout - << "AVX/AVX2 distance function not defined for Uint8. Using " - "slow version SlowDistanceCosineUint8() " - "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 supported for uint32_t byte vectors." - << std::endl; - diskann::cerr << stream.str() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); - } + return new diskann::DistanceL2UInt8(); + } + else if (m == diskann::Metric::COSINE) + { + diskann::cout << "AVX/AVX2 distance function not defined for Uint8. Using " + "slow version SlowDistanceCosineUint8() " + "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 supported for uint32_t byte vectors." << std::endl; + diskann::cerr << stream.str() << std::endl; + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } } template DISKANN_DLLEXPORT class DistanceInnerProduct; @@ -696,4 +730,4 @@ template DISKANN_DLLEXPORT class SlowDistanceL2; template DISKANN_DLLEXPORT class SlowDistanceL2; template DISKANN_DLLEXPORT class SlowDistanceL2; -} // namespace diskann +} // namespace diskann diff --git a/src/dll/dllmain.cpp b/src/dll/dllmain.cpp index 826323fd0..9f5ce4420 100644 --- a/src/dll/dllmain.cpp +++ b/src/dll/dllmain.cpp @@ -1,14 +1,15 @@ // dllmain.cpp : Defines the entry point for the DLL application. #include -BOOL APIENTRY DllMain(HMODULE hModule, DWORD ul_reason_for_call, - LPVOID lpReserved) { - switch (ul_reason_for_call) { +BOOL APIENTRY DllMain(HMODULE hModule, DWORD ul_reason_for_call, LPVOID lpReserved) +{ + switch (ul_reason_for_call) + { case DLL_PROCESS_ATTACH: case DLL_THREAD_ATTACH: case DLL_THREAD_DETACH: case DLL_PROCESS_DETACH: - break; - } - return TRUE; + break; + } + return TRUE; } diff --git a/src/filter_utils.cpp b/src/filter_utils.cpp index c0406c352..965762d1f 100644 --- a/src/filter_utils.cpp +++ b/src/filter_utils.cpp @@ -14,7 +14,8 @@ #include "parameters.h" #include "utils.h" -namespace diskann { +namespace diskann +{ /* * Using passed in parameters and files generated from step 3, * builds a vanilla diskANN index for each label. @@ -23,124 +24,119 @@ namespace diskann { * final_index_path_prefix + "_" + label */ template -void generate_label_indices(path input_data_path, path final_index_path_prefix, - label_set all_labels, uint32_t R, uint32_t L, - float alpha, uint32_t num_threads) { - diskann::IndexWriteParameters label_index_build_parameters = - diskann::IndexWriteParametersBuilder(L, R) - .with_saturate_graph(false) - .with_alpha(alpha) - .with_num_threads(num_threads) - .build(); - - std::cout << "Generating indices per label..." << std::endl; - // for each label, build an index on resp. points - double total_indexing_time = 0.0, indexing_percentage = 0.0; - std::cout.setstate(std::ios_base::failbit); - diskann::cout.setstate(std::ios_base::failbit); - for (const auto &lbl : all_labels) { - path curr_label_input_data_path(input_data_path + "_" + lbl); - path curr_label_index_path(final_index_path_prefix + "_" + lbl); - - size_t number_of_label_points, dimension; - diskann::get_bin_metadata(curr_label_input_data_path, - number_of_label_points, dimension); - diskann::Index index(diskann::Metric::L2, dimension, - number_of_label_points, false, false); - - auto index_build_timer = std::chrono::high_resolution_clock::now(); - index.build(curr_label_input_data_path.c_str(), number_of_label_points, - label_index_build_parameters); - std::chrono::duration current_indexing_time = - std::chrono::high_resolution_clock::now() - index_build_timer; - - total_indexing_time += current_indexing_time.count(); - indexing_percentage += (1 / (double)all_labels.size()); - print_progress(indexing_percentage); - - index.save(curr_label_index_path.c_str()); - } - std::cout.clear(); - diskann::cout.clear(); - - std::cout << "\nDone. Generated per-label indices in " << total_indexing_time - << " seconds\n" - << std::endl; +void generate_label_indices(path input_data_path, path final_index_path_prefix, label_set all_labels, uint32_t R, + uint32_t L, float alpha, uint32_t num_threads) +{ + diskann::IndexWriteParameters label_index_build_parameters = diskann::IndexWriteParametersBuilder(L, R) + .with_saturate_graph(false) + .with_alpha(alpha) + .with_num_threads(num_threads) + .build(); + + std::cout << "Generating indices per label..." << std::endl; + // for each label, build an index on resp. points + double total_indexing_time = 0.0, indexing_percentage = 0.0; + std::cout.setstate(std::ios_base::failbit); + diskann::cout.setstate(std::ios_base::failbit); + for (const auto &lbl : all_labels) + { + path curr_label_input_data_path(input_data_path + "_" + lbl); + path curr_label_index_path(final_index_path_prefix + "_" + lbl); + + size_t number_of_label_points, dimension; + diskann::get_bin_metadata(curr_label_input_data_path, number_of_label_points, dimension); + diskann::Index index(diskann::Metric::L2, dimension, number_of_label_points, false, false); + + auto index_build_timer = std::chrono::high_resolution_clock::now(); + index.build(curr_label_input_data_path.c_str(), number_of_label_points, label_index_build_parameters); + std::chrono::duration current_indexing_time = + std::chrono::high_resolution_clock::now() - index_build_timer; + + total_indexing_time += current_indexing_time.count(); + indexing_percentage += (1 / (double)all_labels.size()); + print_progress(indexing_percentage); + + index.save(curr_label_index_path.c_str()); + } + std::cout.clear(); + diskann::cout.clear(); + + std::cout << "\nDone. Generated per-label indices in " << total_indexing_time << " seconds\n" << std::endl; } // for use on systems without writev (i.e. Windows) template -tsl::robin_map> -generate_label_specific_vector_files_compat( - path input_data_path, - tsl::robin_map labels_to_number_of_points, - std::vector point_ids_to_labels, label_set all_labels) { - auto file_writing_timer = std::chrono::high_resolution_clock::now(); - std::ifstream input_data_stream(input_data_path); - - uint32_t number_of_points, dimension; - input_data_stream.read((char *)&number_of_points, sizeof(uint32_t)); - input_data_stream.read((char *)&dimension, sizeof(uint32_t)); - const uint32_t VECTOR_SIZE = dimension * sizeof(T); - if (number_of_points != point_ids_to_labels.size()) { - std::cerr << "Error: number of points in labels file and data file differ." - << std::endl; - throw; - } - - tsl::robin_map labels_to_vectors; - tsl::robin_map labels_to_curr_vector; - tsl::robin_map> label_id_to_orig_id; - - for (const auto &lbl : all_labels) { - uint32_t number_of_label_pts = labels_to_number_of_points[lbl]; - char *vectors = (char *)malloc(number_of_label_pts * VECTOR_SIZE); - if (vectors == nullptr) { - throw; +tsl::robin_map> generate_label_specific_vector_files_compat( + path input_data_path, tsl::robin_map labels_to_number_of_points, + std::vector point_ids_to_labels, label_set all_labels) +{ + auto file_writing_timer = std::chrono::high_resolution_clock::now(); + std::ifstream input_data_stream(input_data_path); + + uint32_t number_of_points, dimension; + input_data_stream.read((char *)&number_of_points, sizeof(uint32_t)); + input_data_stream.read((char *)&dimension, sizeof(uint32_t)); + const uint32_t VECTOR_SIZE = dimension * sizeof(T); + if (number_of_points != point_ids_to_labels.size()) + { + std::cerr << "Error: number of points in labels file and data file differ." << std::endl; + throw; + } + + tsl::robin_map labels_to_vectors; + tsl::robin_map labels_to_curr_vector; + tsl::robin_map> label_id_to_orig_id; + + for (const auto &lbl : all_labels) + { + uint32_t number_of_label_pts = labels_to_number_of_points[lbl]; + char *vectors = (char *)malloc(number_of_label_pts * VECTOR_SIZE); + if (vectors == nullptr) + { + throw; + } + labels_to_vectors[lbl] = vectors; + labels_to_curr_vector[lbl] = 0; + label_id_to_orig_id[lbl].reserve(number_of_label_pts); } - labels_to_vectors[lbl] = vectors; - labels_to_curr_vector[lbl] = 0; - label_id_to_orig_id[lbl].reserve(number_of_label_pts); - } - - for (uint32_t point_id = 0; point_id < number_of_points; point_id++) { - char *curr_vector = (char *)malloc(VECTOR_SIZE); - input_data_stream.read(curr_vector, VECTOR_SIZE); - for (const auto &lbl : point_ids_to_labels[point_id]) { - char *curr_label_vector_ptr = - labels_to_vectors[lbl] + (labels_to_curr_vector[lbl] * VECTOR_SIZE); - memcpy(curr_label_vector_ptr, curr_vector, VECTOR_SIZE); - labels_to_curr_vector[lbl]++; - label_id_to_orig_id[lbl].push_back(point_id); + + for (uint32_t point_id = 0; point_id < number_of_points; point_id++) + { + char *curr_vector = (char *)malloc(VECTOR_SIZE); + input_data_stream.read(curr_vector, VECTOR_SIZE); + for (const auto &lbl : point_ids_to_labels[point_id]) + { + char *curr_label_vector_ptr = labels_to_vectors[lbl] + (labels_to_curr_vector[lbl] * VECTOR_SIZE); + memcpy(curr_label_vector_ptr, curr_vector, VECTOR_SIZE); + labels_to_curr_vector[lbl]++; + label_id_to_orig_id[lbl].push_back(point_id); + } + free(curr_vector); + } + + for (const auto &lbl : all_labels) + { + path curr_label_input_data_path(input_data_path + "_" + lbl); + uint32_t number_of_label_pts = labels_to_number_of_points[lbl]; + + std::ofstream label_file_stream; + label_file_stream.exceptions(std::ios::badbit | std::ios::failbit); + label_file_stream.open(curr_label_input_data_path, std::ios_base::binary); + label_file_stream.write((char *)&number_of_label_pts, sizeof(uint32_t)); + label_file_stream.write((char *)&dimension, sizeof(uint32_t)); + label_file_stream.write((char *)labels_to_vectors[lbl], number_of_label_pts * VECTOR_SIZE); + + label_file_stream.close(); + free(labels_to_vectors[lbl]); } - free(curr_vector); - } - - for (const auto &lbl : all_labels) { - path curr_label_input_data_path(input_data_path + "_" + lbl); - uint32_t number_of_label_pts = labels_to_number_of_points[lbl]; - - std::ofstream label_file_stream; - label_file_stream.exceptions(std::ios::badbit | std::ios::failbit); - label_file_stream.open(curr_label_input_data_path, std::ios_base::binary); - label_file_stream.write((char *)&number_of_label_pts, sizeof(uint32_t)); - label_file_stream.write((char *)&dimension, sizeof(uint32_t)); - label_file_stream.write((char *)labels_to_vectors[lbl], - number_of_label_pts * VECTOR_SIZE); - - label_file_stream.close(); - free(labels_to_vectors[lbl]); - } - input_data_stream.close(); - - std::chrono::duration file_writing_time = - std::chrono::high_resolution_clock::now() - file_writing_timer; - std::cout << "generated " << all_labels.size() - << " label-specific vector files for index building in time " - << file_writing_time.count() << "\n" - << std::endl; - - return label_id_to_orig_id; + input_data_stream.close(); + + std::chrono::duration file_writing_time = std::chrono::high_resolution_clock::now() - file_writing_timer; + std::cout << "generated " << all_labels.size() << " label-specific vector files for index building in time " + << file_writing_time.count() << "\n" + << std::endl; + + return label_id_to_orig_id; } /* @@ -148,37 +144,36 @@ generate_label_specific_vector_files_compat( * * Returns both the graph index and the size of the file in bytes. */ -load_label_index_return_values load_label_index( - path label_index_path, uint32_t label_number_of_points) { - std::ifstream label_index_stream; - label_index_stream.exceptions(std::ios::badbit | std::ios::failbit); - label_index_stream.open(label_index_path, std::ios::binary); - - uint64_t index_file_size, index_num_frozen_points; - uint32_t index_max_observed_degree, index_entry_point; - const size_t INDEX_METADATA = 2 * sizeof(uint64_t) + 2 * sizeof(uint32_t); - label_index_stream.read((char *)&index_file_size, sizeof(uint64_t)); - label_index_stream.read((char *)&index_max_observed_degree, sizeof(uint32_t)); - label_index_stream.read((char *)&index_entry_point, sizeof(uint32_t)); - label_index_stream.read((char *)&index_num_frozen_points, sizeof(uint64_t)); - size_t bytes_read = INDEX_METADATA; - - std::vector> label_index(label_number_of_points); - uint32_t nodes_read = 0; - while (bytes_read != index_file_size) { - uint32_t current_node_num_neighbors; - label_index_stream.read((char *)¤t_node_num_neighbors, - sizeof(uint32_t)); - nodes_read++; - - std::vector current_node_neighbors(current_node_num_neighbors); - label_index_stream.read((char *)current_node_neighbors.data(), - current_node_num_neighbors * sizeof(uint32_t)); - label_index[nodes_read - 1].swap(current_node_neighbors); - bytes_read += sizeof(uint32_t) * (current_node_num_neighbors + 1); - } - - return std::make_tuple(label_index, index_file_size); +load_label_index_return_values load_label_index(path label_index_path, uint32_t label_number_of_points) +{ + std::ifstream label_index_stream; + label_index_stream.exceptions(std::ios::badbit | std::ios::failbit); + label_index_stream.open(label_index_path, std::ios::binary); + + uint64_t index_file_size, index_num_frozen_points; + uint32_t index_max_observed_degree, index_entry_point; + const size_t INDEX_METADATA = 2 * sizeof(uint64_t) + 2 * sizeof(uint32_t); + label_index_stream.read((char *)&index_file_size, sizeof(uint64_t)); + label_index_stream.read((char *)&index_max_observed_degree, sizeof(uint32_t)); + label_index_stream.read((char *)&index_entry_point, sizeof(uint32_t)); + label_index_stream.read((char *)&index_num_frozen_points, sizeof(uint64_t)); + size_t bytes_read = INDEX_METADATA; + + std::vector> label_index(label_number_of_points); + uint32_t nodes_read = 0; + while (bytes_read != index_file_size) + { + uint32_t current_node_num_neighbors; + label_index_stream.read((char *)¤t_node_num_neighbors, sizeof(uint32_t)); + nodes_read++; + + std::vector current_node_neighbors(current_node_num_neighbors); + label_index_stream.read((char *)current_node_neighbors.data(), current_node_num_neighbors * sizeof(uint32_t)); + label_index[nodes_read - 1].swap(current_node_neighbors); + bytes_read += sizeof(uint32_t) * (current_node_num_neighbors + 1); + } + + return std::make_tuple(label_index, index_file_size); } /* @@ -190,95 +185,100 @@ load_label_index_return_values load_label_index( * 2. map: key is label, value is number of points with the label * 3. the label universe as a set */ -parse_label_file_return_values parse_label_file(path label_data_path, - std::string universal_label) { - std::ifstream label_data_stream(label_data_path); - std::string line, token; - uint32_t line_cnt = 0; - - // allows us to reserve space for the points_to_labels vector - while (std::getline(label_data_stream, line)) line_cnt++; - label_data_stream.clear(); - label_data_stream.seekg(0, std::ios::beg); - - // values to return - std::vector point_ids_to_labels(line_cnt); - tsl::robin_map labels_to_number_of_points; - label_set all_labels; - - std::vector points_with_universal_label; - line_cnt = 0; - while (std::getline(label_data_stream, line)) { - std::istringstream current_labels_comma_separated(line); - label_set current_labels; - - // get point id - uint32_t point_id = line_cnt; - - // parse comma separated labels - bool current_universal_label_check = false; - while (getline(current_labels_comma_separated, token, ',')) { - token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); - token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); - - // if token is empty, there's no labels for the point - if (token == universal_label) { - points_with_universal_label.push_back(point_id); - current_universal_label_check = true; - } else { - all_labels.insert(token); - current_labels.insert(token); - labels_to_number_of_points[token]++; - } +parse_label_file_return_values parse_label_file(path label_data_path, std::string universal_label) +{ + std::ifstream label_data_stream(label_data_path); + std::string line, token; + uint32_t line_cnt = 0; + + // allows us to reserve space for the points_to_labels vector + while (std::getline(label_data_stream, line)) + line_cnt++; + label_data_stream.clear(); + label_data_stream.seekg(0, std::ios::beg); + + // values to return + std::vector point_ids_to_labels(line_cnt); + tsl::robin_map labels_to_number_of_points; + label_set all_labels; + + std::vector points_with_universal_label; + line_cnt = 0; + while (std::getline(label_data_stream, line)) + { + std::istringstream current_labels_comma_separated(line); + label_set current_labels; + + // get point id + uint32_t point_id = line_cnt; + + // parse comma separated labels + bool current_universal_label_check = false; + while (getline(current_labels_comma_separated, token, ',')) + { + token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); + token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); + + // if token is empty, there's no labels for the point + if (token == universal_label) + { + points_with_universal_label.push_back(point_id); + current_universal_label_check = true; + } + else + { + all_labels.insert(token); + current_labels.insert(token); + labels_to_number_of_points[token]++; + } + } + + if (current_labels.size() <= 0 && !current_universal_label_check) + { + std::cerr << "Error: " << point_id << " has no labels." << std::endl; + exit(-1); + } + point_ids_to_labels[point_id] = current_labels; + line_cnt++; } - if (current_labels.size() <= 0 && !current_universal_label_check) { - std::cerr << "Error: " << point_id << " has no labels." << std::endl; - exit(-1); + // for every point with universal label, set its label set to all labels + // also, increment the count for number of points a label has + for (const auto &point_id : points_with_universal_label) + { + point_ids_to_labels[point_id] = all_labels; + for (const auto &lbl : all_labels) + labels_to_number_of_points[lbl]++; } - point_ids_to_labels[point_id] = current_labels; - line_cnt++; - } - - // for every point with universal label, set its label set to all labels - // also, increment the count for number of points a label has - for (const auto &point_id : points_with_universal_label) { - point_ids_to_labels[point_id] = all_labels; - for (const auto &lbl : all_labels) labels_to_number_of_points[lbl]++; - } - - std::cout << "Identified " << all_labels.size() << " distinct label(s) for " - << point_ids_to_labels.size() << " points\n" - << std::endl; - - return std::make_tuple(point_ids_to_labels, labels_to_number_of_points, - all_labels); + + std::cout << "Identified " << all_labels.size() << " distinct label(s) for " << point_ids_to_labels.size() + << " points\n" + << std::endl; + + return std::make_tuple(point_ids_to_labels, labels_to_number_of_points, all_labels); } -template DISKANN_DLLEXPORT void generate_label_indices( - path input_data_path, path final_index_path_prefix, label_set all_labels, - uint32_t R, uint32_t L, float alpha, uint32_t num_threads); -template DISKANN_DLLEXPORT void generate_label_indices( - path input_data_path, path final_index_path_prefix, label_set all_labels, - uint32_t R, uint32_t L, float alpha, uint32_t num_threads); -template DISKANN_DLLEXPORT void generate_label_indices( - path input_data_path, path final_index_path_prefix, label_set all_labels, - uint32_t R, uint32_t L, float alpha, uint32_t num_threads); +template DISKANN_DLLEXPORT void generate_label_indices(path input_data_path, path final_index_path_prefix, + label_set all_labels, uint32_t R, uint32_t L, float alpha, + uint32_t num_threads); +template DISKANN_DLLEXPORT void generate_label_indices(path input_data_path, path final_index_path_prefix, + label_set all_labels, uint32_t R, uint32_t L, + float alpha, uint32_t num_threads); +template DISKANN_DLLEXPORT void generate_label_indices(path input_data_path, path final_index_path_prefix, + label_set all_labels, uint32_t R, uint32_t L, + float alpha, uint32_t num_threads); template DISKANN_DLLEXPORT tsl::robin_map> -generate_label_specific_vector_files_compat( - path input_data_path, - tsl::robin_map labels_to_number_of_points, - std::vector point_ids_to_labels, label_set all_labels); +generate_label_specific_vector_files_compat(path input_data_path, + tsl::robin_map labels_to_number_of_points, + std::vector point_ids_to_labels, label_set all_labels); template DISKANN_DLLEXPORT tsl::robin_map> -generate_label_specific_vector_files_compat( - path input_data_path, - tsl::robin_map labels_to_number_of_points, - std::vector point_ids_to_labels, label_set all_labels); +generate_label_specific_vector_files_compat(path input_data_path, + tsl::robin_map labels_to_number_of_points, + std::vector point_ids_to_labels, label_set all_labels); template DISKANN_DLLEXPORT tsl::robin_map> -generate_label_specific_vector_files_compat( - path input_data_path, - tsl::robin_map labels_to_number_of_points, - std::vector point_ids_to_labels, label_set all_labels); +generate_label_specific_vector_files_compat(path input_data_path, + tsl::robin_map labels_to_number_of_points, + std::vector point_ids_to_labels, label_set all_labels); -} // namespace diskann \ No newline at end of file +} // namespace diskann \ No newline at end of file diff --git a/src/in_mem_data_store.cpp b/src/in_mem_data_store.cpp index 54757b17a..3cc386c4d 100644 --- a/src/in_mem_data_store.cpp +++ b/src/in_mem_data_store.cpp @@ -6,364 +6,363 @@ #include "utils.h" -namespace diskann { +namespace diskann +{ template -InMemDataStore::InMemDataStore( - const location_t num_points, const size_t dim, - std::shared_ptr> distance_metric) - : AbstractDataStore(num_points, dim), - _distance_fn(distance_metric) { - _aligned_dim = ROUND_UP(dim, _distance_fn->get_required_alignment()); - alloc_aligned(((void **)&_data), - this->_capacity * _aligned_dim * sizeof(data_t), - 8 * sizeof(data_t)); - std::memset(_data, 0, this->_capacity * _aligned_dim * sizeof(data_t)); +InMemDataStore::InMemDataStore(const location_t num_points, const size_t dim, + std::shared_ptr> distance_metric) + : AbstractDataStore(num_points, dim), _distance_fn(distance_metric) +{ + _aligned_dim = ROUND_UP(dim, _distance_fn->get_required_alignment()); + alloc_aligned(((void **)&_data), this->_capacity * _aligned_dim * sizeof(data_t), 8 * sizeof(data_t)); + std::memset(_data, 0, this->_capacity * _aligned_dim * sizeof(data_t)); } -template -InMemDataStore::~InMemDataStore() { - if (_data != nullptr) { - aligned_free(this->_data); - } +template InMemDataStore::~InMemDataStore() +{ + if (_data != nullptr) + { + aligned_free(this->_data); + } } -template -size_t InMemDataStore::get_aligned_dim() const { - return _aligned_dim; +template size_t InMemDataStore::get_aligned_dim() const +{ + return _aligned_dim; } -template -size_t InMemDataStore::get_alignment_factor() const { - return _distance_fn->get_required_alignment(); +template size_t InMemDataStore::get_alignment_factor() const +{ + return _distance_fn->get_required_alignment(); } -template -location_t InMemDataStore::load(const std::string &filename) { - return load_impl(filename); +template location_t InMemDataStore::load(const std::string &filename) +{ + return load_impl(filename); } #ifdef EXEC_ENV_OLS -template -location_t Index::load_impl(AlignedFileReader &reader) { - size_t file_dim, file_num_points; - - diskann::get_bin_metadata(reader, file_num_points, file_dim); - - if (file_dim != _dim) { - std::stringstream stream; - stream << "ERROR: Driver requests loading " << _dim << " dimension," - << "but file has " << file_dim << " dimension." << std::endl; - diskann::cerr << stream.str() << std::endl; - aligned_free(_data); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); - } +template location_t Index::load_impl(AlignedFileReader &reader) +{ + size_t file_dim, file_num_points; + + diskann::get_bin_metadata(reader, file_num_points, file_dim); + + if (file_dim != _dim) + { + std::stringstream stream; + stream << "ERROR: Driver requests loading " << _dim << " dimension," + << "but file has " << file_dim << " dimension." << std::endl; + diskann::cerr << stream.str() << std::endl; + aligned_free(_data); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } - if (file_num_points > _max_points + _num_frozen_pts) { - resize(file_num_points - _num_frozen_pts); - } + if (file_num_points > _max_points + _num_frozen_pts) + { + resize(file_num_points - _num_frozen_pts); + } - return file_num_points; + return file_num_points; } #endif -template -location_t InMemDataStore::load_impl(const std::string &filename) { - size_t file_dim, file_num_points; - if (!file_exists(filename)) { - std::stringstream stream; - stream << "ERROR: data file " << filename << " does not exist." - << std::endl; - diskann::cerr << stream.str() << std::endl; - aligned_free(_data); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); - } - diskann::get_bin_metadata(filename, file_num_points, file_dim); - - if (file_dim != this->_dim) { - std::stringstream stream; - stream << "ERROR: Driver requests loading " << this->_dim << " dimension," - << "but file has " << file_dim << " dimension." << std::endl; - diskann::cerr << stream.str() << std::endl; - aligned_free(_data); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); - } +template location_t InMemDataStore::load_impl(const std::string &filename) +{ + size_t file_dim, file_num_points; + if (!file_exists(filename)) + { + std::stringstream stream; + stream << "ERROR: data file " << filename << " does not exist." << std::endl; + diskann::cerr << stream.str() << std::endl; + aligned_free(_data); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } + diskann::get_bin_metadata(filename, file_num_points, file_dim); + + if (file_dim != this->_dim) + { + std::stringstream stream; + stream << "ERROR: Driver requests loading " << this->_dim << " dimension," + << "but file has " << file_dim << " dimension." << std::endl; + diskann::cerr << stream.str() << std::endl; + aligned_free(_data); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } - if (file_num_points > this->capacity()) { - this->resize(file_num_points); - } + if (file_num_points > this->capacity()) + { + this->resize(file_num_points); + } - copy_aligned_data_from_file(filename.c_str(), _data, file_num_points, - file_dim, _aligned_dim); + copy_aligned_data_from_file(filename.c_str(), _data, file_num_points, file_dim, _aligned_dim); - return file_num_points; + return file_num_points; } -template -size_t InMemDataStore::save(const std::string &filename, - const location_t num_points) { - return save_data_in_base_dimensions(filename, _data, num_points, - this->get_dims(), this->get_aligned_dim(), - 0U); +template size_t InMemDataStore::save(const std::string &filename, const location_t num_points) +{ + return save_data_in_base_dimensions(filename, _data, num_points, this->get_dims(), this->get_aligned_dim(), 0U); } -template -void InMemDataStore::populate_data(const data_t *vectors, - const location_t num_pts) { - for (location_t i = 0; i < num_pts; i++) { - memset(_data + i * _aligned_dim, 0, _aligned_dim * sizeof(data_t)); - std::memmove(_data + i * _aligned_dim, vectors + i * this->_dim, - this->_dim * sizeof(data_t)); - } - - if (_distance_fn->preprocessing_required()) { - _distance_fn->preprocess_base_points(_data, this->_aligned_dim, num_pts); - } +template void InMemDataStore::populate_data(const data_t *vectors, const location_t num_pts) +{ + for (location_t i = 0; i < num_pts; i++) + { + memset(_data + i * _aligned_dim, 0, _aligned_dim * sizeof(data_t)); + std::memmove(_data + i * _aligned_dim, vectors + i * this->_dim, this->_dim * sizeof(data_t)); + } + + if (_distance_fn->preprocessing_required()) + { + _distance_fn->preprocess_base_points(_data, this->_aligned_dim, num_pts); + } } -template -void InMemDataStore::populate_data(const std::string &filename, - const size_t offset) { - size_t npts, ndim; - copy_aligned_data_from_file(filename.c_str(), _data, npts, ndim, _aligned_dim, - offset); - - if ((location_t)npts > this->capacity()) { - std::stringstream ss; - ss << "Number of points in the file: " << filename - << " is greater than the capacity of data store: " << this->capacity() - << ". Must invoke resize before calling populate_data()" << std::endl; - throw diskann::ANNException(ss.str(), -1); - } - - if ((location_t)ndim != this->get_dims()) { - std::stringstream ss; - ss << "Number of dimensions of a point in the file: " << filename - << " is not equal to dimensions of data store: " << this->capacity() - << "." << std::endl; - throw diskann::ANNException(ss.str(), -1); - } - - if (_distance_fn->preprocessing_required()) { - _distance_fn->preprocess_base_points(_data, this->_aligned_dim, - this->capacity()); - } +template void InMemDataStore::populate_data(const std::string &filename, const size_t offset) +{ + size_t npts, ndim; + copy_aligned_data_from_file(filename.c_str(), _data, npts, ndim, _aligned_dim, offset); + + if ((location_t)npts > this->capacity()) + { + std::stringstream ss; + ss << "Number of points in the file: " << filename + << " is greater than the capacity of data store: " << this->capacity() + << ". Must invoke resize before calling populate_data()" << std::endl; + throw diskann::ANNException(ss.str(), -1); + } + + if ((location_t)ndim != this->get_dims()) + { + std::stringstream ss; + ss << "Number of dimensions of a point in the file: " << filename + << " is not equal to dimensions of data store: " << this->capacity() << "." << std::endl; + throw diskann::ANNException(ss.str(), -1); + } + + if (_distance_fn->preprocessing_required()) + { + _distance_fn->preprocess_base_points(_data, this->_aligned_dim, this->capacity()); + } } template -void InMemDataStore::extract_data_to_bin(const std::string &filename, - const location_t num_points) { - save_data_in_base_dimensions(filename, _data, num_points, this->get_dims(), - this->get_aligned_dim(), 0U); +void InMemDataStore::extract_data_to_bin(const std::string &filename, const location_t num_points) +{ + save_data_in_base_dimensions(filename, _data, num_points, this->get_dims(), this->get_aligned_dim(), 0U); } -template -void InMemDataStore::get_vector(const location_t i, - data_t *dest) const { - memcpy(dest, _data + i * _aligned_dim, this->_dim * sizeof(data_t)); +template void InMemDataStore::get_vector(const location_t i, data_t *dest) const +{ + memcpy(dest, _data + i * _aligned_dim, this->_dim * sizeof(data_t)); } -template -void InMemDataStore::set_vector(const location_t loc, - const data_t *const vector) { - size_t offset_in_data = loc * _aligned_dim; - memset(_data + offset_in_data, 0, _aligned_dim * sizeof(data_t)); - memcpy(_data + offset_in_data, vector, this->_dim * sizeof(data_t)); - if (_distance_fn->preprocessing_required()) { - _distance_fn->preprocess_base_points(_data + offset_in_data, _aligned_dim, - 1); - } +template void InMemDataStore::set_vector(const location_t loc, const data_t *const vector) +{ + size_t offset_in_data = loc * _aligned_dim; + memset(_data + offset_in_data, 0, _aligned_dim * sizeof(data_t)); + memcpy(_data + offset_in_data, vector, this->_dim * sizeof(data_t)); + if (_distance_fn->preprocessing_required()) + { + _distance_fn->preprocess_base_points(_data + offset_in_data, _aligned_dim, 1); + } } -template -void InMemDataStore::prefetch_vector(const location_t loc) { - diskann::prefetch_vector((const char *)_data + _aligned_dim * (size_t)loc, - sizeof(data_t) * _aligned_dim); +template void InMemDataStore::prefetch_vector(const location_t loc) +{ + diskann::prefetch_vector((const char *)_data + _aligned_dim * (size_t)loc, sizeof(data_t) * _aligned_dim); } -template -float InMemDataStore::get_distance(const data_t *query, - const location_t loc) const { - return _distance_fn->compare(query, _data + _aligned_dim * loc, _aligned_dim); +template float InMemDataStore::get_distance(const data_t *query, const location_t loc) const +{ + return _distance_fn->compare(query, _data + _aligned_dim * loc, _aligned_dim); } template -void InMemDataStore::get_distance(const data_t *query, - const location_t *locations, - const uint32_t location_count, - float *distances) const { - for (location_t i = 0; i < location_count; i++) { - distances[i] = _distance_fn->compare( - query, _data + locations[i] * _aligned_dim, this->_aligned_dim); - } +void InMemDataStore::get_distance(const data_t *query, const location_t *locations, + const uint32_t location_count, float *distances) const +{ + for (location_t i = 0; i < location_count; i++) + { + distances[i] = _distance_fn->compare(query, _data + locations[i] * _aligned_dim, this->_aligned_dim); + } } template -float InMemDataStore::get_distance(const location_t loc1, - const location_t loc2) const { - return _distance_fn->compare(_data + loc1 * _aligned_dim, - _data + loc2 * _aligned_dim, this->_aligned_dim); +float InMemDataStore::get_distance(const location_t loc1, const location_t loc2) const +{ + return _distance_fn->compare(_data + loc1 * _aligned_dim, _data + loc2 * _aligned_dim, this->_aligned_dim); } -template -location_t InMemDataStore::expand(const location_t new_size) { - 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); - } +template location_t InMemDataStore::expand(const location_t new_size) +{ + 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, _data, this->capacity() * _aligned_dim * sizeof(data_t)); - aligned_free(_data); - _data = new_data; + data_t *new_data; + alloc_aligned((void **)&new_data, new_size * _aligned_dim * sizeof(data_t), 8 * sizeof(data_t)); + memcpy(new_data, _data, this->capacity() * _aligned_dim * sizeof(data_t)); + aligned_free(_data); + _data = new_data; #else - realloc_aligned((void **)&_data, new_size * _aligned_dim * sizeof(data_t), - 8 * sizeof(data_t)); + realloc_aligned((void **)&_data, new_size * _aligned_dim * sizeof(data_t), 8 * sizeof(data_t)); #endif - this->_capacity = new_size; - return this->_capacity; + this->_capacity = new_size; + return this->_capacity; } -template -location_t InMemDataStore::shrink(const location_t new_size) { - 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); - } +template location_t InMemDataStore::shrink(const location_t new_size) +{ + 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, _data, new_size * _aligned_dim * sizeof(data_t)); - aligned_free(_data); - _data = new_data; + data_t *new_data; + alloc_aligned((void **)&new_data, new_size * _aligned_dim * sizeof(data_t), 8 * sizeof(data_t)); + memcpy(new_data, _data, new_size * _aligned_dim * sizeof(data_t)); + aligned_free(_data); + _data = new_data; #else - realloc_aligned((void **)&_data, new_size * _aligned_dim * sizeof(data_t), - 8 * sizeof(data_t)); + realloc_aligned((void **)&_data, new_size * _aligned_dim * sizeof(data_t), 8 * sizeof(data_t)); #endif - this->_capacity = new_size; - return this->_capacity; + this->_capacity = new_size; + return this->_capacity; } template -void InMemDataStore::move_vectors(const location_t old_location_start, - const location_t new_location_start, - const location_t num_locations) { - if (num_locations == 0 || old_location_start == new_location_start) { - return; - } - - /* // Update pointers to the moved nodes. Note: the computation is correct - even - // when new_location_start < old_location_start given the C++ uint32_t - // integer arithmetic rules. - const uint32_t location_delta = new_location_start - old_location_start; - */ - // 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; +void InMemDataStore::move_vectors(const location_t old_location_start, const location_t new_location_start, + const location_t num_locations) +{ + if (num_locations == 0 || old_location_start == new_location_start) + { + return; } - } 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; + + /* // Update pointers to the moved nodes. Note: the computation is correct + even + // when new_location_start < old_location_start given the C++ uint32_t + // integer arithmetic rules. + const uint32_t location_delta = new_location_start - old_location_start; + */ + // 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(_data + _aligned_dim * mem_clear_loc_start, 0, - sizeof(data_t) * _aligned_dim * - (mem_clear_loc_end_limit - mem_clear_loc_start)); + // Use memmove to handle overlapping ranges. + copy_vectors(old_location_start, new_location_start, num_locations); + memset(_data + _aligned_dim * mem_clear_loc_start, 0, + sizeof(data_t) * _aligned_dim * (mem_clear_loc_end_limit - mem_clear_loc_start)); } template -void InMemDataStore::copy_vectors(const location_t from_loc, - const location_t to_loc, - const location_t num_points) { - assert(from_loc < this->_capacity); - assert(to_loc < this->_capacity); - assert(num_points < this->_capacity); - memmove(_data + _aligned_dim * to_loc, _data + _aligned_dim * from_loc, - num_points * _aligned_dim * sizeof(data_t)); +void InMemDataStore::copy_vectors(const location_t from_loc, const location_t to_loc, + const location_t num_points) +{ + assert(from_loc < this->_capacity); + assert(to_loc < this->_capacity); + assert(num_points < this->_capacity); + memmove(_data + _aligned_dim * to_loc, _data + _aligned_dim * from_loc, num_points * _aligned_dim * sizeof(data_t)); } -template -location_t InMemDataStore::calculate_medoid() const { - // allocate and init centroid - float *center = new float[_aligned_dim]; - for (size_t j = 0; j < _aligned_dim; j++) center[j] = 0; +template location_t InMemDataStore::calculate_medoid() const +{ + // allocate and init centroid + float *center = new float[_aligned_dim]; + for (size_t j = 0; j < _aligned_dim; j++) + center[j] = 0; + + for (size_t i = 0; i < this->capacity(); i++) + for (size_t j = 0; j < _aligned_dim; j++) + center[j] += (float)_data[i * _aligned_dim + j]; - for (size_t i = 0; i < this->capacity(); i++) for (size_t j = 0; j < _aligned_dim; j++) - center[j] += (float)_data[i * _aligned_dim + j]; - - for (size_t j = 0; j < _aligned_dim; j++) - center[j] /= (float)this->capacity(); - - // compute all to one distance - float *distances = new float[this->capacity()]; - - // TODO: REFACTOR. Removing pragma might make this slow. Must revisit. - // Problem is that we need to pass num_threads here, it is not clear - // if data store must be aware of threads! - // #pragma omp parallel for schedule(static, 65536) - for (int64_t i = 0; i < (int64_t)this->capacity(); i++) { - // extract point and distance reference - float &dist = distances[i]; - const data_t *cur_vec = _data + (i * (size_t)_aligned_dim); - dist = 0; - float diff = 0; - for (size_t j = 0; j < _aligned_dim; j++) { - diff = (center[j] - (float)cur_vec[j]) * (center[j] - (float)cur_vec[j]); - dist += diff; + center[j] /= (float)this->capacity(); + + // compute all to one distance + float *distances = new float[this->capacity()]; + + // TODO: REFACTOR. Removing pragma might make this slow. Must revisit. + // Problem is that we need to pass num_threads here, it is not clear + // if data store must be aware of threads! + // #pragma omp parallel for schedule(static, 65536) + for (int64_t i = 0; i < (int64_t)this->capacity(); i++) + { + // extract point and distance reference + float &dist = distances[i]; + const data_t *cur_vec = _data + (i * (size_t)_aligned_dim); + dist = 0; + float diff = 0; + for (size_t j = 0; j < _aligned_dim; j++) + { + diff = (center[j] - (float)cur_vec[j]) * (center[j] - (float)cur_vec[j]); + dist += diff; + } } - } - // find imin - uint32_t min_idx = 0; - float min_dist = distances[0]; - for (uint32_t i = 1; i < this->capacity(); i++) { - if (distances[i] < min_dist) { - min_idx = i; - min_dist = distances[i]; + // find imin + uint32_t min_idx = 0; + float min_dist = distances[0]; + for (uint32_t i = 1; i < this->capacity(); i++) + { + if (distances[i] < min_dist) + { + min_idx = i; + min_dist = distances[i]; + } } - } - delete[] distances; - delete[] center; - return min_idx; + delete[] distances; + delete[] center; + return min_idx; } -template -Distance *InMemDataStore::get_dist_fn() { - return this->_distance_fn.get(); +template Distance *InMemDataStore::get_dist_fn() +{ + return this->_distance_fn.get(); } template DISKANN_DLLEXPORT class InMemDataStore; template DISKANN_DLLEXPORT class InMemDataStore; template DISKANN_DLLEXPORT class InMemDataStore; -} // namespace diskann \ No newline at end of file +} // namespace diskann \ No newline at end of file diff --git a/src/in_mem_graph_store.cpp b/src/in_mem_graph_store.cpp index ad9efbbf4..cb83481d3 100644 --- a/src/in_mem_graph_store.cpp +++ b/src/in_mem_graph_store.cpp @@ -4,18 +4,26 @@ #include "in_mem_graph_store.h" #include "utils.h" -namespace diskann { +namespace diskann +{ -InMemGraphStore::InMemGraphStore(const size_t max_pts) - : AbstractGraphStore(max_pts) {} +InMemGraphStore::InMemGraphStore(const size_t max_pts) : AbstractGraphStore(max_pts) +{ +} -int InMemGraphStore::load(const std::string &index_path_prefix) {} -int InMemGraphStore::store(const std::string &index_path_prefix) {} +int InMemGraphStore::load(const std::string &index_path_prefix) +{ +} +int InMemGraphStore::store(const std::string &index_path_prefix) +{ +} -void InMemGraphStore::get_adj_list(const location_t i, - std::vector &neighbors) {} +void InMemGraphStore::get_adj_list(const location_t i, std::vector &neighbors) +{ +} -void InMemGraphStore::set_adj_list(const location_t i, - std::vector &neighbors) {} +void InMemGraphStore::set_adj_list(const location_t i, std::vector &neighbors) +{ +} -} // namespace diskann \ No newline at end of file +} // namespace diskann \ No newline at end of file diff --git a/src/index.cpp b/src/index.cpp index da3b3939d..58d82c57d 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -11,8 +11,7 @@ #include "memory_mapper.h" #include "timer.h" #include "windows_customizations.h" -#if defined(RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS) && \ - defined(DISKANN_BUILD) +#if defined(RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS) && defined(DISKANN_BUILD) #include "gperftools/malloc_extension.h" #endif @@ -23,2790 +22,2926 @@ #define MAX_POINTS_FOR_USING_BITSET 10000000 -namespace diskann { +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 IndexWriteParameters &indexParams, - const uint32_t initial_search_list_size, const uint32_t search_threads, - const bool enable_tags, const bool concurrent_consolidate, - const bool pq_dist_build, const size_t num_pq_chunks, const bool use_opq) - : Index(m, dim, max_points, dynamic_index, enable_tags, - concurrent_consolidate, pq_dist_build, num_pq_chunks, use_opq, - indexParams.num_frozen_points) { - _indexingQueueSize = indexParams.search_list_size; - _indexingRange = indexParams.max_degree; - _indexingMaxC = indexParams.max_occlusion_size; - _indexingAlpha = indexParams.alpha; - _filterIndexingQueueSize = indexParams.filter_list_size; - - uint32_t num_threads_indx = indexParams.num_threads; - uint32_t num_scratch_spaces = search_threads + num_threads_indx; - - initialize_query_scratch(num_scratch_spaces, initial_search_list_size, - _indexingQueueSize, _indexingRange, _indexingMaxC, - dim); -} +Index::Index(Metric m, const size_t dim, const size_t max_points, const bool dynamic_index, + const IndexWriteParameters &indexParams, const uint32_t initial_search_list_size, + const uint32_t search_threads, const bool enable_tags, const bool concurrent_consolidate, + const bool pq_dist_build, const size_t num_pq_chunks, const bool use_opq) + : Index(m, dim, max_points, dynamic_index, enable_tags, concurrent_consolidate, pq_dist_build, num_pq_chunks, + use_opq, indexParams.num_frozen_points) +{ + _indexingQueueSize = indexParams.search_list_size; + _indexingRange = indexParams.max_degree; + _indexingMaxC = indexParams.max_occlusion_size; + _indexingAlpha = indexParams.alpha; + _filterIndexingQueueSize = indexParams.filter_list_size; + + uint32_t num_threads_indx = indexParams.num_threads; + uint32_t num_scratch_spaces = search_threads + num_threads_indx; + + initialize_query_scratch(num_scratch_spaces, initial_search_list_size, _indexingQueueSize, _indexingRange, + _indexingMaxC, dim); +} + +template +Index::Index(Metric m, const size_t dim, const size_t max_points, 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 size_t num_frozen_pts) + : _dist_metric(m), _dim(dim), _max_points(max_points), _num_frozen_pts(num_frozen_pts), + _dynamic_index(dynamic_index), _enable_tags(enable_tags), _indexingMaxC(DEFAULT_MAXC), _query_scratch(nullptr), + _pq_dist(pq_dist_build), _use_opq(use_opq), _num_pq_chunks(num_pq_chunks), + _delete_set(new tsl::robin_set), _conc_consolidate(concurrent_consolidate) +{ + if (dynamic_index && !enable_tags) + { + throw ANNException("ERROR: Dynamic Indexing must have tags enabled.", -1, __FUNCSIG__, __FILE__, __LINE__); + } -template -Index::Index(Metric m, const size_t dim, - const size_t max_points, 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 size_t num_frozen_pts) - : _dist_metric(m), - _dim(dim), - _max_points(max_points), - _num_frozen_pts(num_frozen_pts), - _dynamic_index(dynamic_index), - _enable_tags(enable_tags), - _indexingMaxC(DEFAULT_MAXC), - _query_scratch(nullptr), - _pq_dist(pq_dist_build), - _use_opq(use_opq), - _num_pq_chunks(num_pq_chunks), - _delete_set(new tsl::robin_set), - _conc_consolidate(concurrent_consolidate) { - if (dynamic_index && !enable_tags) { - throw ANNException("ERROR: Dynamic Indexing must have tags enabled.", -1, - __FUNCSIG__, __FILE__, __LINE__); - } - - if (_pq_dist) { - if (dynamic_index) - throw ANNException( - "ERROR: Dynamic Indexing not supported with PQ distance based " - "index construction", - -1, __FUNCSIG__, __FILE__, __LINE__); - if (m == diskann::Metric::INNER_PRODUCT) - throw ANNException( - "ERROR: Inner product metrics not yet supported " - "with PQ distance " - "base index", - -1, __FUNCSIG__, __FILE__, __LINE__); - } - - if (dynamic_index && _num_frozen_pts == 0) { - _num_frozen_pts = 1; - } - // Sanity check. While logically it is correct, max_points = 0 causes - // downstream problems. - if (_max_points == 0) { - _max_points = 1; - } - const size_t total_internal_points = _max_points + _num_frozen_pts; - - if (_pq_dist) { - if (_num_pq_chunks > _dim) - throw diskann::ANNException("ERROR: num_pq_chunks > dim", -1, __FUNCSIG__, - __FILE__, __LINE__); - alloc_aligned(((void **)&_pq_data), - total_internal_points * _num_pq_chunks * sizeof(char), - 8 * sizeof(char)); - std::memset(_pq_data, 0, - total_internal_points * _num_pq_chunks * sizeof(char)); - } - - _start = (uint32_t)_max_points; - - _final_graph.resize(total_internal_points); - - // This should come from a factory. - if (m == diskann::Metric::COSINE && std::is_floating_point::value) { - // This is safe because T is float inside the if block. - this->_distance.reset( - (Distance *)new AVXNormalizedCosineDistanceFloat()); - this->_normalize_vecs = true; - diskann::cout << "Normalizing vectors and using L2 for cosine " - "AVXNormalizedCosineDistanceFloat()." - << std::endl; - } else { - this->_distance.reset((Distance *)get_distance_function(m)); - } - // REFACTOR: TODO This should move to a factory method. + if (_pq_dist) + { + if (dynamic_index) + throw ANNException("ERROR: Dynamic Indexing not supported with PQ distance based " + "index construction", + -1, __FUNCSIG__, __FILE__, __LINE__); + if (m == diskann::Metric::INNER_PRODUCT) + throw ANNException("ERROR: Inner product metrics not yet supported " + "with PQ distance " + "base index", + -1, __FUNCSIG__, __FILE__, __LINE__); + } + + if (dynamic_index && _num_frozen_pts == 0) + { + _num_frozen_pts = 1; + } + // Sanity check. While logically it is correct, max_points = 0 causes + // downstream problems. + if (_max_points == 0) + { + _max_points = 1; + } + const size_t total_internal_points = _max_points + _num_frozen_pts; + + if (_pq_dist) + { + if (_num_pq_chunks > _dim) + throw diskann::ANNException("ERROR: num_pq_chunks > dim", -1, __FUNCSIG__, __FILE__, __LINE__); + alloc_aligned(((void **)&_pq_data), total_internal_points * _num_pq_chunks * sizeof(char), 8 * sizeof(char)); + std::memset(_pq_data, 0, total_internal_points * _num_pq_chunks * sizeof(char)); + } + + _start = (uint32_t)_max_points; + + _final_graph.resize(total_internal_points); - _data_store = std::make_unique>( - (location_t)total_internal_points, _dim, this->_distance); + // This should come from a factory. + if (m == diskann::Metric::COSINE && std::is_floating_point::value) + { + // This is safe because T is float inside the if block. + this->_distance.reset((Distance *)new AVXNormalizedCosineDistanceFloat()); + this->_normalize_vecs = true; + diskann::cout << "Normalizing vectors and using L2 for cosine " + "AVXNormalizedCosineDistanceFloat()." + << std::endl; + } + else + { + this->_distance.reset((Distance *)get_distance_function(m)); + } + // REFACTOR: TODO This should move to a factory method. + + _data_store = + std::make_unique>((location_t)total_internal_points, _dim, this->_distance); - _locks = std::vector(total_internal_points); + _locks = std::vector(total_internal_points); - if (enable_tags) { - _location_to_tag.reserve(total_internal_points); - _tag_to_location.reserve(total_internal_points); - } + if (enable_tags) + { + _location_to_tag.reserve(total_internal_points); + _tag_to_location.reserve(total_internal_points); + } } -template -Index::~Index() { - // Ensure that no other activity is happening before dtor() - std::unique_lock ul(_update_lock); - std::unique_lock cl(_consolidate_lock); - std::unique_lock tl(_tag_lock); - std::unique_lock dl(_delete_lock); - - for (auto &lock : _locks) { - LockGuard lg(lock); - } - - // if (this->_distance != nullptr) - //{ - // delete this->_distance; - // this->_distance = nullptr; - // } - // REFACTOR - - if (_opt_graph != nullptr) { - delete[] _opt_graph; - } - - if (!_query_scratch.empty()) { - ScratchStoreManager> manager(_query_scratch); - manager.destroy(); - } +template Index::~Index() +{ + // Ensure that no other activity is happening before dtor() + std::unique_lock ul(_update_lock); + std::unique_lock cl(_consolidate_lock); + std::unique_lock tl(_tag_lock); + std::unique_lock dl(_delete_lock); + + for (auto &lock : _locks) + { + LockGuard lg(lock); + } + + // if (this->_distance != nullptr) + //{ + // delete this->_distance; + // this->_distance = nullptr; + // } + // REFACTOR + + if (_opt_graph != nullptr) + { + delete[] _opt_graph; + } + + if (!_query_scratch.empty()) + { + ScratchStoreManager> manager(_query_scratch); + manager.destroy(); + } } template -void Index::initialize_query_scratch(uint32_t num_threads, - uint32_t search_l, - uint32_t indexing_l, - uint32_t r, uint32_t maxc, - size_t dim) { - for (uint32_t i = 0; i < num_threads; i++) { - auto scratch = new InMemQueryScratch( - search_l, indexing_l, r, maxc, dim, _data_store->get_aligned_dim(), - _data_store->get_alignment_factor(), _pq_dist); - _query_scratch.push(scratch); - } +void Index::initialize_query_scratch(uint32_t num_threads, uint32_t search_l, uint32_t indexing_l, + uint32_t r, uint32_t maxc, size_t dim) +{ + for (uint32_t i = 0; i < num_threads; i++) + { + auto scratch = new InMemQueryScratch(search_l, indexing_l, r, maxc, dim, _data_store->get_aligned_dim(), + _data_store->get_alignment_factor(), _pq_dist); + _query_scratch.push(scratch); + } } -template -size_t Index::save_tags(std::string tags_file) { - 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 (uint32_t i = 0; i < _nd; i++) { - TagT tag; - if (_location_to_tag.try_get(i, tag)) { - tag_data[i] = tag; - } else { - // catering to future when tagT can be any type. - std::memset((char *)&tag_data[i], 0, sizeof(TagT)); - } - } - if (_num_frozen_pts > 0) { - std::memset((char *)&tag_data[_start], 0, sizeof(TagT) * _num_frozen_pts); - } - try { - tag_bytes_written = - save_bin(tags_file, tag_data, _nd + _num_frozen_pts, 1); - } catch (std::system_error &e) { - throw FileException(tags_file, e, __FUNCSIG__, __FILE__, __LINE__); - } - delete[] tag_data; - return tag_bytes_written; +template size_t Index::save_tags(std::string tags_file) +{ + 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 (uint32_t i = 0; i < _nd; i++) + { + TagT tag; + if (_location_to_tag.try_get(i, tag)) + { + tag_data[i] = tag; + } + else + { + // catering to future when tagT can be any type. + std::memset((char *)&tag_data[i], 0, sizeof(TagT)); + } + } + if (_num_frozen_pts > 0) + { + std::memset((char *)&tag_data[_start], 0, sizeof(TagT) * _num_frozen_pts); + } + try + { + tag_bytes_written = save_bin(tags_file, tag_data, _nd + _num_frozen_pts, 1); + } + catch (std::system_error &e) + { + throw FileException(tags_file, e, __FUNCSIG__, __FILE__, __LINE__); + } + delete[] tag_data; + return tag_bytes_written; } -template -size_t Index::save_data(std::string data_file) { - // Note: at this point, either _nd == _max_points or any frozen points have - // been temporarily moved to _nd, so _nd + _num_frozen_points is the valid - // location limit. - return _data_store->save(data_file, _nd + _num_frozen_pts); +template size_t Index::save_data(std::string data_file) +{ + // Note: at this point, either _nd == _max_points or any frozen points have + // been temporarily moved to _nd, so _nd + _num_frozen_points is the valid + // location limit. + return _data_store->save(data_file, _nd + _num_frozen_pts); } // 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 uint32_t) -template -size_t Index::save_graph(std::string graph_file) { - std::ofstream out; - open_file_to_write(out, graph_file); - - size_t file_offset = 0; // we will use this if we want - out.seekp(file_offset, out.beg); - size_t index_size = 24; - uint32_t max_degree = 0; - out.write((char *)&index_size, sizeof(uint64_t)); - out.write((char *)&_max_observed_degree, sizeof(uint32_t)); - uint32_t ep_u32 = _start; - out.write((char *)&ep_u32, sizeof(uint32_t)); - out.write((char *)&_num_frozen_pts, sizeof(size_t)); - // Note: at this point, either _nd == _max_points or any frozen points have - // been temporarily moved to _nd, so _nd + _num_frozen_points is the valid - // location limit. - for (uint32_t i = 0; i < _nd + _num_frozen_pts; i++) { - uint32_t GK = (uint32_t)_final_graph[i].size(); - out.write((char *)&GK, sizeof(uint32_t)); - out.write((char *)_final_graph[i].data(), GK * sizeof(uint32_t)); - max_degree = _final_graph[i].size() > max_degree - ? (uint32_t)_final_graph[i].size() - : max_degree; - index_size += (size_t)(sizeof(uint32_t) * (GK + 1)); - } - out.seekp(file_offset, out.beg); - out.write((char *)&index_size, sizeof(uint64_t)); - out.write((char *)&max_degree, sizeof(uint32_t)); - out.close(); - return index_size; // number of bytes written +template size_t Index::save_graph(std::string graph_file) +{ + std::ofstream out; + open_file_to_write(out, graph_file); + + size_t file_offset = 0; // we will use this if we want + out.seekp(file_offset, out.beg); + size_t index_size = 24; + uint32_t max_degree = 0; + out.write((char *)&index_size, sizeof(uint64_t)); + out.write((char *)&_max_observed_degree, sizeof(uint32_t)); + uint32_t ep_u32 = _start; + out.write((char *)&ep_u32, sizeof(uint32_t)); + out.write((char *)&_num_frozen_pts, sizeof(size_t)); + // Note: at this point, either _nd == _max_points or any frozen points have + // been temporarily moved to _nd, so _nd + _num_frozen_points is the valid + // location limit. + for (uint32_t i = 0; i < _nd + _num_frozen_pts; i++) + { + uint32_t GK = (uint32_t)_final_graph[i].size(); + out.write((char *)&GK, sizeof(uint32_t)); + out.write((char *)_final_graph[i].data(), GK * sizeof(uint32_t)); + max_degree = _final_graph[i].size() > max_degree ? (uint32_t)_final_graph[i].size() : max_degree; + index_size += (size_t)(sizeof(uint32_t) * (GK + 1)); + } + out.seekp(file_offset, out.beg); + out.write((char *)&index_size, sizeof(uint64_t)); + out.write((char *)&max_degree, sizeof(uint32_t)); + out.close(); + return index_size; // number of bytes written } template -size_t Index::save_delete_list(const std::string &filename) { - if (_delete_set->size() == 0) { - return 0; - } - std::unique_ptr delete_list = - std::make_unique(_delete_set->size()); - uint32_t i = 0; - for (auto &del : *_delete_set) { - delete_list[i++] = del; - } - return save_bin(filename, delete_list.get(), _delete_set->size(), - 1); +size_t Index::save_delete_list(const std::string &filename) +{ + if (_delete_set->size() == 0) + { + return 0; + } + std::unique_ptr delete_list = std::make_unique(_delete_set->size()); + uint32_t i = 0; + for (auto &del : *_delete_set) + { + delete_list[i++] = del; + } + return save_bin(filename, delete_list.get(), _delete_set->size(), 1); } template -void Index::save(const char *filename, - bool compact_before_save) { - diskann::Timer timer; - - std::unique_lock ul(_update_lock); - std::unique_lock cl(_consolidate_lock); - std::unique_lock tl(_tag_lock); - std::unique_lock dl(_delete_lock); - - if (compact_before_save) { - compact_data(); - compact_frozen_point(); - } else { - if (!_data_compacted) { - throw ANNException( - "Index save for non-compacted index is not yet implemented", -1, - __FUNCSIG__, __FILE__, __LINE__); - } - } - - if (!_save_as_one_file) { - if (_filtered_index) { - if (_label_to_medoid_id.size() > 0) { - std::ofstream medoid_writer(std::string(filename) + - "_labels_to_medoids.txt"); - if (medoid_writer.fail()) { - throw diskann::ANNException( - std::string("Failed to open file ") + filename, -1); - } - for (auto iter : _label_to_medoid_id) { - medoid_writer << iter.first << ", " << iter.second << std::endl; - } - medoid_writer.close(); - } - - if (_use_universal_label) { - std::ofstream universal_label_writer(std::string(filename) + - "_universal_label.txt"); - assert(universal_label_writer.is_open()); - universal_label_writer << _universal_label << std::endl; - universal_label_writer.close(); - } - - if (_pts_to_labels.size() > 0) { - std::ofstream label_writer(std::string(filename) + "_labels.txt"); - assert(label_writer.is_open()); - for (uint32_t i = 0; i < _pts_to_labels.size(); i++) { - for (uint32_t j = 0; j < (_pts_to_labels[i].size() - 1); j++) { - label_writer << _pts_to_labels[i][j] << ","; - } - if (_pts_to_labels[i].size() != 0) - label_writer << _pts_to_labels[i][_pts_to_labels[i].size() - 1]; - label_writer << std::endl; - } - label_writer.close(); - } - } - - 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 { - diskann::cout << "Save index in a single file currently not supported. " - "Not saving the index." - << std::endl; - } +void Index::save(const char *filename, bool compact_before_save) +{ + diskann::Timer timer; + + std::unique_lock ul(_update_lock); + std::unique_lock cl(_consolidate_lock); + std::unique_lock tl(_tag_lock); + std::unique_lock dl(_delete_lock); + + if (compact_before_save) + { + compact_data(); + compact_frozen_point(); + } + else + { + if (!_data_compacted) + { + throw ANNException("Index save for non-compacted index is not yet implemented", -1, __FUNCSIG__, __FILE__, + __LINE__); + } + } + + if (!_save_as_one_file) + { + if (_filtered_index) + { + if (_label_to_medoid_id.size() > 0) + { + std::ofstream medoid_writer(std::string(filename) + "_labels_to_medoids.txt"); + if (medoid_writer.fail()) + { + throw diskann::ANNException(std::string("Failed to open file ") + filename, -1); + } + for (auto iter : _label_to_medoid_id) + { + medoid_writer << iter.first << ", " << iter.second << std::endl; + } + medoid_writer.close(); + } + + if (_use_universal_label) + { + std::ofstream universal_label_writer(std::string(filename) + "_universal_label.txt"); + assert(universal_label_writer.is_open()); + universal_label_writer << _universal_label << std::endl; + universal_label_writer.close(); + } + + if (_pts_to_labels.size() > 0) + { + std::ofstream label_writer(std::string(filename) + "_labels.txt"); + assert(label_writer.is_open()); + for (uint32_t i = 0; i < _pts_to_labels.size(); i++) + { + for (uint32_t j = 0; j < (_pts_to_labels[i].size() - 1); j++) + { + label_writer << _pts_to_labels[i][j] << ","; + } + if (_pts_to_labels[i].size() != 0) + label_writer << _pts_to_labels[i][_pts_to_labels[i].size() - 1]; + label_writer << std::endl; + } + label_writer.close(); + } + } - // If frozen points were temporarily compacted to _nd, move back to - // _max_points. - reposition_frozen_point_to_end(); + 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 + { + diskann::cout << "Save index in a single file currently not supported. " + "Not saving the index." + << std::endl; + } - diskann::cout << "Time taken for save: " << timer.elapsed() / 1000000.0 - << "s." << std::endl; + // If frozen points were temporarily compacted to _nd, move back to + // _max_points. + reposition_frozen_point_to_end(); + + diskann::cout << "Time taken for save: " << timer.elapsed() / 1000000.0 << "s." << std::endl; } #ifdef EXEC_ENV_OLS template -size_t Index::load_tags(AlignedFileReader &reader) { +size_t Index::load_tags(AlignedFileReader &reader) +{ #else template -size_t Index::load_tags(const std::string tag_filename) { - 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__); - } +size_t Index::load_tags(const std::string tag_filename) +{ + 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__); + } #endif - if (!_enable_tags) { - diskann::cout << "Tags not loaded as tags not enabled." << std::endl; - return 0; - } + 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; + size_t file_dim, file_num_points; + TagT *tag_data; #ifdef EXEC_ENV_OLS - load_bin(reader, tag_data, file_num_points, file_dim); + load_bin(reader, tag_data, file_num_points, file_dim); #else - load_bin(std::string(tag_filename), tag_data, file_num_points, - file_dim); + load_bin(std::string(tag_filename), tag_data, file_num_points, file_dim); #endif - if (file_dim != 1) { - std::stringstream stream; - stream << "ERROR: Found " << file_dim << " dimensions for tags," - << "but tag file must have 1 dimension." << std::endl; - diskann::cerr << stream.str() << std::endl; - delete[] tag_data; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); - } + if (file_dim != 1) + { + std::stringstream stream; + stream << "ERROR: Found " << file_dim << " dimensions for tags," + << "but tag file must have 1 dimension." << std::endl; + diskann::cerr << stream.str() << std::endl; + delete[] tag_data; + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } - const size_t num_data_points = file_num_points - _num_frozen_pts; - _location_to_tag.reserve(num_data_points); - _tag_to_location.reserve(num_data_points); - for (uint32_t i = 0; i < (uint32_t)num_data_points; i++) { - TagT tag = *(tag_data + i); - if (_delete_set->find(i) == _delete_set->end()) { - _location_to_tag.set(i, tag); - _tag_to_location[tag] = i; + const size_t num_data_points = file_num_points - _num_frozen_pts; + _location_to_tag.reserve(num_data_points); + _tag_to_location.reserve(num_data_points); + for (uint32_t i = 0; i < (uint32_t)num_data_points; i++) + { + TagT tag = *(tag_data + i); + if (_delete_set->find(i) == _delete_set->end()) + { + _location_to_tag.set(i, tag); + _tag_to_location[tag] = i; + } } - } - diskann::cout << "Tags loaded." << std::endl; - delete[] tag_data; - return file_num_points; + diskann::cout << "Tags loaded." << std::endl; + delete[] tag_data; + return file_num_points; } template #ifdef EXEC_ENV_OLS -size_t Index::load_data(AlignedFileReader &reader) { +size_t Index::load_data(AlignedFileReader &reader) +{ #else -size_t Index::load_data(std::string filename) { +size_t Index::load_data(std::string filename) +{ #endif - size_t file_dim, file_num_points; + size_t file_dim, file_num_points; #ifdef EXEC_ENV_OLS - diskann::get_bin_metadata(reader, file_num_points, file_dim); + diskann::get_bin_metadata(reader, file_num_points, file_dim); #else - if (!file_exists(filename)) { - std::stringstream stream; - stream << "ERROR: data file " << filename << " does not exist." - << std::endl; - diskann::cerr << stream.str() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); - } - diskann::get_bin_metadata(filename, file_num_points, file_dim); + if (!file_exists(filename)) + { + std::stringstream stream; + stream << "ERROR: data file " << filename << " does not exist." << std::endl; + diskann::cerr << stream.str() << std::endl; + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } + diskann::get_bin_metadata(filename, file_num_points, file_dim); #endif - // since we are loading a new dataset, _empty_slots must be cleared - _empty_slots.clear(); + // since we are loading a new dataset, _empty_slots must be cleared + _empty_slots.clear(); - if (file_dim != _dim) { - std::stringstream stream; - stream << "ERROR: Driver requests loading " << _dim << " dimension," - << "but file has " << file_dim << " dimension." << std::endl; - diskann::cerr << stream.str() << std::endl; - 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; + diskann::cerr << stream.str() << std::endl; + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } - if (file_num_points > _max_points + _num_frozen_pts) { - // update and tag lock acquired in load() before calling load_data - resize(file_num_points - _num_frozen_pts); - } + if (file_num_points > _max_points + _num_frozen_pts) + { + // update and tag lock acquired in load() before calling load_data + resize(file_num_points - _num_frozen_pts); + } #ifdef EXEC_ENV_OLS - // REFACTOR TODO: Must figure out how to support aligned reader in a clean - // manner. - copy_aligned_data_from_file(reader, _data, file_num_points, file_dim, - _aligned_dim); + // REFACTOR TODO: Must figure out how to support aligned reader in a clean + // manner. + copy_aligned_data_from_file(reader, _data, file_num_points, file_dim, _aligned_dim); #else - _data_store->load(filename); // offset == 0. + _data_store->load(filename); // offset == 0. #endif - return file_num_points; + return file_num_points; } #ifdef EXEC_ENV_OLS template -size_t Index::load_delete_set(AlignedFileReader &reader) { +size_t Index::load_delete_set(AlignedFileReader &reader) +{ #else template -size_t Index::load_delete_set(const std::string &filename) { +size_t Index::load_delete_set(const std::string &filename) +{ #endif - std::unique_ptr delete_list; - size_t npts, ndim; + std::unique_ptr delete_list; + size_t npts, ndim; #ifdef EXEC_ENV_OLS - diskann::load_bin(reader, delete_list, npts, ndim); + diskann::load_bin(reader, delete_list, npts, ndim); #else - diskann::load_bin(filename, delete_list, npts, ndim); + diskann::load_bin(filename, delete_list, npts, ndim); #endif - assert(ndim == 1); - for (uint32_t i = 0; i < npts; i++) { - _delete_set->insert(delete_list[i]); - } - return npts; + assert(ndim == 1); + for (uint32_t i = 0; i < npts; i++) + { + _delete_set->insert(delete_list[i]); + } + return npts; } // load the index from file and update the max_degree, cur (navigating // node loc), and _final_graph (adjacency list) template #ifdef EXEC_ENV_OLS -void Index::load(AlignedFileReader &reader, - uint32_t num_threads, uint32_t search_l) { +void Index::load(AlignedFileReader &reader, uint32_t num_threads, uint32_t search_l) +{ #else -void Index::load(const char *filename, uint32_t num_threads, - uint32_t search_l) { +void Index::load(const char *filename, uint32_t num_threads, uint32_t search_l) +{ #endif - std::unique_lock ul(_update_lock); - std::unique_lock cl(_consolidate_lock); - std::unique_lock tl(_tag_lock); - std::unique_lock dl(_delete_lock); + std::unique_lock ul(_update_lock); + std::unique_lock cl(_consolidate_lock); + std::unique_lock tl(_tag_lock); + std::unique_lock dl(_delete_lock); - _has_built = true; + _has_built = true; - size_t tags_file_num_pts = 0, graph_num_pts = 0, data_file_num_pts = 0, - label_num_pts = 0; + size_t tags_file_num_pts = 0, graph_num_pts = 0, data_file_num_pts = 0, label_num_pts = 0; - std::string mem_index_file(filename); - std::string labels_file = mem_index_file + "_labels.txt"; - std::string labels_to_medoids = mem_index_file + "_labels_to_medoids.txt"; - std::string labels_map_file = mem_index_file + "_labels_map.txt"; + std::string mem_index_file(filename); + std::string labels_file = mem_index_file + "_labels.txt"; + std::string labels_to_medoids = mem_index_file + "_labels_to_medoids.txt"; + std::string labels_map_file = mem_index_file + "_labels_map.txt"; - if (!_save_as_one_file) { - // For DLVS Store, we will not support saving the index in multiple - // files. + if (!_save_as_one_file) + { + // For DLVS Store, we will not support saving the index in multiple + // files. #ifndef EXEC_ENV_OLS - 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); + 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); #endif - } else { - diskann::cout << "Single index file saving/loading support not yet " - "enabled. Not loading the index." + } + else + { + diskann::cout << "Single index file saving/loading support not yet " + "enabled. Not loading the index." + << std::endl; + return; + } + + 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; + diskann::cerr << stream.str() << std::endl; + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } + + if (file_exists(labels_file)) + { + _label_map = load_label_map(labels_map_file); + parse_label_file(labels_file, label_num_pts); + assert(label_num_pts == data_file_num_pts); + if (file_exists(labels_to_medoids)) + { + std::ifstream medoid_stream(labels_to_medoids); + std::string line, token; + uint32_t line_cnt = 0; + + _label_to_medoid_id.clear(); + + while (std::getline(medoid_stream, line)) + { + std::istringstream iss(line); + uint32_t cnt = 0; + uint32_t medoid = 0; + LabelT label; + while (std::getline(iss, token, ',')) + { + token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); + token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); + LabelT token_as_num = std::stoul(token); + if (cnt == 0) + label = token_as_num; + else + medoid = token_as_num; + cnt++; + } + _label_to_medoid_id[label] = medoid; + line_cnt++; + } + } + + std::string universal_label_file(filename); + universal_label_file += "_universal_label.txt"; + if (file_exists(universal_label_file)) + { + std::ifstream universal_label_reader(universal_label_file); + universal_label_reader >> _universal_label; + _use_universal_label = true; + universal_label_reader.close(); + } + } + + _nd = data_file_num_pts - _num_frozen_pts; + _empty_slots.clear(); + _empty_slots.reserve(_max_points); + for (auto i = _nd; i < _max_points; i++) + { + _empty_slots.insert((uint32_t)i); + } + + reposition_frozen_point_to_end(); + diskann::cout << "Num frozen points:" << _num_frozen_pts << " _nd: " << _nd << " _start: " << _start + << " size(_location_to_tag): " << _location_to_tag.size() + << " size(_tag_to_location):" << _tag_to_location.size() << " Max points: " << _max_points << std::endl; - return; - } - 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; - diskann::cerr << stream.str() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); - } - - if (file_exists(labels_file)) { - _label_map = load_label_map(labels_map_file); - parse_label_file(labels_file, label_num_pts); - assert(label_num_pts == data_file_num_pts); - if (file_exists(labels_to_medoids)) { - std::ifstream medoid_stream(labels_to_medoids); - std::string line, token; - uint32_t line_cnt = 0; - - _label_to_medoid_id.clear(); - - while (std::getline(medoid_stream, line)) { - std::istringstream iss(line); - uint32_t cnt = 0; - uint32_t medoid = 0; - LabelT label; - while (std::getline(iss, token, ',')) { - token.erase(std::remove(token.begin(), token.end(), '\n'), - token.end()); - token.erase(std::remove(token.begin(), token.end(), '\r'), - token.end()); - LabelT token_as_num = std::stoul(token); - if (cnt == 0) - label = token_as_num; - else - medoid = token_as_num; - cnt++; - } - _label_to_medoid_id[label] = medoid; - line_cnt++; - } - } - - std::string universal_label_file(filename); - universal_label_file += "_universal_label.txt"; - if (file_exists(universal_label_file)) { - std::ifstream universal_label_reader(universal_label_file); - universal_label_reader >> _universal_label; - _use_universal_label = true; - universal_label_reader.close(); - } - } - - _nd = data_file_num_pts - _num_frozen_pts; - _empty_slots.clear(); - _empty_slots.reserve(_max_points); - for (auto i = _nd; i < _max_points; i++) { - _empty_slots.insert((uint32_t)i); - } - - reposition_frozen_point_to_end(); - diskann::cout << "Num frozen points:" << _num_frozen_pts << " _nd: " << _nd - << " _start: " << _start - << " size(_location_to_tag): " << _location_to_tag.size() - << " size(_tag_to_location):" << _tag_to_location.size() - << " Max points: " << _max_points << std::endl; - - // For incremental index, _query_scratch is initialized in the constructor. - // For the bulk index, the params required to initialize _query_scratch - // are known only at load time, hence this check and the call to - // initialize_q_s(). - if (_query_scratch.size() == 0) { - initialize_query_scratch(num_threads, search_l, search_l, - (uint32_t)_max_range_of_loaded_graph, - _indexingMaxC, _dim); - } + // For incremental index, _query_scratch is initialized in the constructor. + // For the bulk index, the params required to initialize _query_scratch + // are known only at load time, hence this check and the call to + // initialize_q_s(). + if (_query_scratch.size() == 0) + { + initialize_query_scratch(num_threads, search_l, search_l, (uint32_t)_max_range_of_loaded_graph, _indexingMaxC, + _dim); + } } #ifndef EXEC_ENV_OLS template -size_t Index::get_graph_num_frozen_points( - const std::string &graph_file) { - size_t expected_file_size; - uint32_t max_observed_degree, start; - size_t file_frozen_pts; +size_t Index::get_graph_num_frozen_points(const std::string &graph_file) +{ + size_t expected_file_size; + uint32_t max_observed_degree, start; + size_t file_frozen_pts; - std::ifstream in; - in.exceptions(std::ios::badbit | std::ios::failbit); + std::ifstream in; + in.exceptions(std::ios::badbit | std::ios::failbit); - in.open(graph_file, std::ios::binary); - in.read((char *)&expected_file_size, sizeof(size_t)); - in.read((char *)&max_observed_degree, sizeof(uint32_t)); - in.read((char *)&start, sizeof(uint32_t)); - in.read((char *)&file_frozen_pts, sizeof(size_t)); + in.open(graph_file, std::ios::binary); + in.read((char *)&expected_file_size, sizeof(size_t)); + in.read((char *)&max_observed_degree, sizeof(uint32_t)); + in.read((char *)&start, sizeof(uint32_t)); + in.read((char *)&file_frozen_pts, sizeof(size_t)); - return file_frozen_pts; + return file_frozen_pts; } #endif #ifdef EXEC_ENV_OLS template -size_t Index::load_graph(AlignedFileReader &reader, - size_t expected_num_points) { +size_t Index::load_graph(AlignedFileReader &reader, size_t expected_num_points) +{ #else template -size_t Index::load_graph(std::string filename, - size_t expected_num_points) { +size_t Index::load_graph(std::string filename, size_t expected_num_points) +{ #endif - size_t expected_file_size; - size_t file_frozen_pts; + size_t expected_file_size; + size_t file_frozen_pts; #ifdef EXEC_ENV_OLS - int header_size = 2 * sizeof(size_t) + 2 * sizeof(uint32_t); - std::unique_ptr header = std::make_unique(header_size); - read_array(reader, header.get(), header_size); - - expected_file_size = *((size_t *)header.get()); - _max_observed_degree = *((uint32_t *)(header.get() + sizeof(size_t))); - _start = *((uint32_t *)(header.get() + sizeof(size_t) + sizeof(uint32_t))); - file_frozen_pts = *((size_t *)(header.get() + sizeof(size_t) + - sizeof(uint32_t) + sizeof(uint32_t))); + int header_size = 2 * sizeof(size_t) + 2 * sizeof(uint32_t); + std::unique_ptr header = std::make_unique(header_size); + read_array(reader, header.get(), header_size); + + expected_file_size = *((size_t *)header.get()); + _max_observed_degree = *((uint32_t *)(header.get() + sizeof(size_t))); + _start = *((uint32_t *)(header.get() + sizeof(size_t) + sizeof(uint32_t))); + file_frozen_pts = *((size_t *)(header.get() + sizeof(size_t) + sizeof(uint32_t) + sizeof(uint32_t))); #else - size_t file_offset = 0; // will need this for single file format support - std::ifstream in; - in.exceptions(std::ios::badbit | std::ios::failbit); - in.open(filename, std::ios::binary); - in.seekg(file_offset, in.beg); - in.read((char *)&expected_file_size, sizeof(size_t)); - in.read((char *)&_max_observed_degree, sizeof(uint32_t)); - in.read((char *)&_start, sizeof(uint32_t)); - in.read((char *)&file_frozen_pts, sizeof(size_t)); - size_t vamana_metadata_size = - sizeof(size_t) + sizeof(uint32_t) + sizeof(uint32_t) + sizeof(size_t); + size_t file_offset = 0; // will need this for single file format support + std::ifstream in; + in.exceptions(std::ios::badbit | std::ios::failbit); + in.open(filename, std::ios::binary); + in.seekg(file_offset, in.beg); + in.read((char *)&expected_file_size, sizeof(size_t)); + in.read((char *)&_max_observed_degree, sizeof(uint32_t)); + in.read((char *)&_start, sizeof(uint32_t)); + in.read((char *)&file_frozen_pts, sizeof(size_t)); + size_t vamana_metadata_size = sizeof(size_t) + sizeof(uint32_t) + sizeof(uint32_t) + sizeof(size_t); #endif - diskann::cout << "From graph header, expected_file_size: " - << expected_file_size - << ", _max_observed_degree: " << _max_observed_degree - << ", _start: " << _start - << ", file_frozen_pts: " << file_frozen_pts << std::endl; + diskann::cout << "From graph header, expected_file_size: " << expected_file_size + << ", _max_observed_degree: " << _max_observed_degree << ", _start: " << _start + << ", file_frozen_pts: " << file_frozen_pts << std::endl; - 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; + 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; + } + diskann::cerr << stream.str() << std::endl; + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } - diskann::cerr << stream.str() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); - } #ifdef EXEC_ENV_OLS - diskann::cout << "Loading vamana graph from reader..." << std::flush; + diskann::cout << "Loading vamana graph from reader..." << std::flush; #else - diskann::cout << "Loading vamana graph " << filename << "..." << std::flush; + diskann::cout << "Loading vamana graph " << filename << "..." << std::flush; #endif - const size_t expected_max_points = expected_num_points - file_frozen_pts; + const size_t expected_max_points = expected_num_points - file_frozen_pts; - // If user provides more points than max_points - // resize the _final_graph to the larger size. - if (_max_points < expected_max_points) { - diskann::cout << "Number of points in data: " << expected_max_points - << " is greater than max_points: " << _max_points - << " Setting max points to: " << expected_max_points - << std::endl; - _final_graph.resize(expected_max_points + _num_frozen_pts); - _max_points = expected_max_points; - } + // If user provides more points than max_points + // resize the _final_graph to the larger size. + if (_max_points < expected_max_points) + { + diskann::cout << "Number of points in data: " << expected_max_points + << " is greater than max_points: " << _max_points + << " Setting max points to: " << expected_max_points << std::endl; + _final_graph.resize(expected_max_points + _num_frozen_pts); + _max_points = expected_max_points; + } #ifdef EXEC_ENV_OLS - uint32_t nodes_read = 0; - size_t cc = 0; - size_t graph_offset = header_size; - while (nodes_read < expected_num_points) { - uint32_t k; - read_value(reader, k, graph_offset); - graph_offset += sizeof(uint32_t); - std::vector tmp(k); - tmp.reserve(k); - read_array(reader, tmp.data(), k, graph_offset); - graph_offset += k * sizeof(uint32_t); - cc += k; - _final_graph[nodes_read].swap(tmp); - nodes_read++; - if (nodes_read % 1000000 == 0) { - diskann::cout << "." << std::flush; - } - if (k > _max_range_of_loaded_graph) { - _max_range_of_loaded_graph = k; - } - } + uint32_t nodes_read = 0; + size_t cc = 0; + size_t graph_offset = header_size; + while (nodes_read < expected_num_points) + { + uint32_t k; + read_value(reader, k, graph_offset); + graph_offset += sizeof(uint32_t); + std::vector tmp(k); + tmp.reserve(k); + read_array(reader, tmp.data(), k, graph_offset); + graph_offset += k * sizeof(uint32_t); + cc += k; + _final_graph[nodes_read].swap(tmp); + nodes_read++; + if (nodes_read % 1000000 == 0) + { + diskann::cout << "." << std::flush; + } + if (k > _max_range_of_loaded_graph) + { + _max_range_of_loaded_graph = k; + } + } #else - size_t bytes_read = vamana_metadata_size; - size_t cc = 0; - uint32_t nodes_read = 0; - while (bytes_read != expected_file_size) { - uint32_t k; - in.read((char *)&k, sizeof(uint32_t)); - - if (k == 0) { - diskann::cerr << "ERROR: Point found with no out-neighbors, point#" - << nodes_read << std::endl; - } - - cc += k; - ++nodes_read; - std::vector tmp(k); - tmp.reserve(k); - in.read((char *)tmp.data(), k * sizeof(uint32_t)); - _final_graph[nodes_read - 1].swap(tmp); - bytes_read += sizeof(uint32_t) * ((size_t)k + 1); - if (nodes_read % 10000000 == 0) diskann::cout << "." << std::flush; - if (k > _max_range_of_loaded_graph) { - _max_range_of_loaded_graph = k; - } - } + size_t bytes_read = vamana_metadata_size; + size_t cc = 0; + uint32_t nodes_read = 0; + while (bytes_read != expected_file_size) + { + uint32_t k; + in.read((char *)&k, sizeof(uint32_t)); + + if (k == 0) + { + diskann::cerr << "ERROR: Point found with no out-neighbors, point#" << nodes_read << std::endl; + } + + cc += k; + ++nodes_read; + std::vector tmp(k); + tmp.reserve(k); + in.read((char *)tmp.data(), k * sizeof(uint32_t)); + _final_graph[nodes_read - 1].swap(tmp); + bytes_read += sizeof(uint32_t) * ((size_t)k + 1); + if (nodes_read % 10000000 == 0) + diskann::cout << "." << std::flush; + if (k > _max_range_of_loaded_graph) + { + _max_range_of_loaded_graph = k; + } + } #endif - diskann::cout << "done. Index has " << nodes_read << " nodes and " << cc - << " out-edges, _start is set to " << _start << std::endl; - return nodes_read; + diskann::cout << "done. Index has " << nodes_read << " nodes and " << cc << " out-edges, _start is set to " + << _start << std::endl; + return nodes_read; } -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; - } +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; + } - location_t location = _tag_to_location[tag]; - _data_store->get_vector(location, vec); + location_t location = _tag_to_location[tag]; + _data_store->get_vector(location, vec); - return 0; + return 0; } -template -uint32_t Index::calculate_entry_point() { - // TODO: need to compute medoid with PQ data too, for now sample at random - if (_pq_dist) { - size_t r = (size_t)rand() * (size_t)RAND_MAX + (size_t)rand(); - return (uint32_t)(r % (size_t)_nd); - } +template uint32_t Index::calculate_entry_point() +{ + // TODO: need to compute medoid with PQ data too, for now sample at random + if (_pq_dist) + { + size_t r = (size_t)rand() * (size_t)RAND_MAX + (size_t)rand(); + return (uint32_t)(r % (size_t)_nd); + } - // TODO: This function does not support multi-threaded calculation of medoid. - // Must revisit if perf is a concern. - return _data_store->calculate_medoid(); + // TODO: This function does not support multi-threaded calculation of medoid. + // Must revisit if perf is a concern. + return _data_store->calculate_medoid(); } -template -std::vector Index::get_init_ids() { - std::vector init_ids; - init_ids.reserve(1 + _num_frozen_pts); +template std::vector Index::get_init_ids() +{ + std::vector init_ids; + init_ids.reserve(1 + _num_frozen_pts); - init_ids.emplace_back(_start); + init_ids.emplace_back(_start); - for (uint32_t frozen = _max_points; frozen < _max_points + _num_frozen_pts; - frozen++) { - if (frozen != _start) { - init_ids.emplace_back(frozen); + for (uint32_t frozen = _max_points; frozen < _max_points + _num_frozen_pts; frozen++) + { + if (frozen != _start) + { + init_ids.emplace_back(frozen); + } } - } - return init_ids; + return init_ids; } template std::pair Index::iterate_to_fixed_point( - const T *query, const uint32_t Lsize, const std::vector &init_ids, - InMemQueryScratch *scratch, bool use_filter, - const std::vector &filter_label, bool search_invocation) { - std::vector &expanded_nodes = scratch->pool(); - NeighborPriorityQueue &best_L_nodes = scratch->best_l_nodes(); - best_L_nodes.reserve(Lsize); - tsl::robin_set &inserted_into_pool_rs = - scratch->inserted_into_pool_rs(); - boost::dynamic_bitset<> &inserted_into_pool_bs = - scratch->inserted_into_pool_bs(); - std::vector &id_scratch = scratch->id_scratch(); - std::vector &dist_scratch = scratch->dist_scratch(); - assert(id_scratch.size() == 0); - - // REFACTOR - // T *aligned_query = scratch->aligned_query(); - // memcpy(aligned_query, query, _dim * sizeof(T)); - // if (_normalize_vecs) - //{ - // normalize((float *)aligned_query, _dim); - // } - - T *aligned_query = scratch->aligned_query(); - - float *query_float = nullptr; - float *query_rotated = nullptr; - float *pq_dists = nullptr; - uint8_t *pq_coord_scratch = nullptr; - // Intialize PQ related scratch to use PQ based distances - if (_pq_dist) { - // Get scratch spaces - PQScratch *pq_query_scratch = scratch->pq_scratch(); - query_float = pq_query_scratch->aligned_query_float; - query_rotated = pq_query_scratch->rotated_query; - pq_dists = pq_query_scratch->aligned_pqtable_dist_scratch; - - // Copy query vector to float and then to "rotated" query - for (size_t d = 0; d < _dim; d++) { - query_float[d] = (float)aligned_query[d]; - } - pq_query_scratch->set(_dim, aligned_query); - - // center the query and rotate if we have a rotation matrix - _pq_table.preprocess_query(query_rotated); - _pq_table.populate_chunk_distances(query_rotated, pq_dists); - - pq_coord_scratch = pq_query_scratch->aligned_pq_coord_scratch; - } - - if (expanded_nodes.size() > 0 || id_scratch.size() > 0) { - throw ANNException("ERROR: Clear scratch space before passing.", -1, - __FUNCSIG__, __FILE__, __LINE__); - } - - // Decide whether to use bitset or robin set to mark visited nodes - auto total_num_points = _max_points + _num_frozen_pts; - bool fast_iterate = total_num_points <= MAX_POINTS_FOR_USING_BITSET; - - if (fast_iterate) { - if (inserted_into_pool_bs.size() < total_num_points) { - // hopefully using 2X will reduce the number of allocations. - auto resize_size = 2 * total_num_points > MAX_POINTS_FOR_USING_BITSET - ? MAX_POINTS_FOR_USING_BITSET - : 2 * total_num_points; - inserted_into_pool_bs.resize(resize_size); - } - } - - // Lambda to determine if a node has been visited - auto is_not_visited = [this, fast_iterate, &inserted_into_pool_bs, - &inserted_into_pool_rs](const uint32_t id) { - return fast_iterate - ? inserted_into_pool_bs[id] == 0 - : inserted_into_pool_rs.find(id) == inserted_into_pool_rs.end(); - }; - - // Lambda to batch compute query<-> node distances in PQ space - auto compute_dists = [this, pq_coord_scratch, pq_dists]( - const std::vector &ids, - std::vector &dists_out) { - diskann::aggregate_coords(ids, this->_pq_data, this->_num_pq_chunks, - pq_coord_scratch); - diskann::pq_dist_lookup(pq_coord_scratch, ids.size(), this->_num_pq_chunks, - pq_dists, dists_out); - }; - - // Initialize the candidate pool with starting points - for (auto id : init_ids) { - if (id >= _max_points + _num_frozen_pts) { - diskann::cerr << "Out of range loc found as an edge : " << id - << std::endl; - throw diskann::ANNException(std::string("Wrong loc") + std::to_string(id), - -1, __FUNCSIG__, __FILE__, __LINE__); - } - - if (use_filter) { - std::vector common_filters; - auto &x = _pts_to_labels[id]; - std::set_intersection(filter_label.begin(), filter_label.end(), x.begin(), - x.end(), std::back_inserter(common_filters)); - if (_use_universal_label) { - if (std::find(filter_label.begin(), filter_label.end(), - _universal_label) != filter_label.end() || - std::find(x.begin(), x.end(), _universal_label) != x.end()) - common_filters.emplace_back(_universal_label); - } - - if (common_filters.size() == 0) continue; - } - - if (is_not_visited(id)) { - if (fast_iterate) { - inserted_into_pool_bs[id] = 1; - } else { - inserted_into_pool_rs.insert(id); - } - - float distance; - if (_pq_dist) { - pq_dist_lookup(pq_coord_scratch, 1, this->_num_pq_chunks, pq_dists, - &distance); - } else { - distance = _data_store->get_distance(aligned_query, id); - } - Neighbor nn = Neighbor(id, distance); - best_L_nodes.insert(nn); - } - } - - uint32_t hops = 0; - uint32_t cmps = 0; - - while (best_L_nodes.has_unexpanded_node()) { - auto nbr = best_L_nodes.closest_unexpanded(); - auto n = nbr.id; - - // Add node to expanded nodes to create pool for prune later - if (!search_invocation) { - if (!use_filter) { - expanded_nodes.emplace_back(nbr); - } else { // in filter based indexing, the same point might invoke - // multiple iterate_to_fixed_points, so need to be careful - // not to add the same item to pool multiple times. - if (std::find(expanded_nodes.begin(), expanded_nodes.end(), nbr) == - expanded_nodes.end()) { - expanded_nodes.emplace_back(nbr); - } - } - } - - // Find which of the nodes in des have not been visited before - id_scratch.clear(); - dist_scratch.clear(); - { - if (_dynamic_index) _locks[n].lock(); - for (auto id : _final_graph[n]) { - assert(id < _max_points + _num_frozen_pts); - - if (use_filter) { - // NOTE: NEED TO CHECK IF THIS CORRECT WITH NEW LOCKS. - std::vector common_filters; - auto &x = _pts_to_labels[id]; - std::set_intersection(filter_label.begin(), filter_label.end(), - x.begin(), x.end(), - std::back_inserter(common_filters)); - if (_use_universal_label) { - if (std::find(filter_label.begin(), filter_label.end(), - _universal_label) != filter_label.end() || - std::find(x.begin(), x.end(), _universal_label) != x.end()) - common_filters.emplace_back(_universal_label); - } - - if (common_filters.size() == 0) continue; - } - - if (is_not_visited(id)) { - id_scratch.push_back(id); - } - } - - if (_dynamic_index) _locks[n].unlock(); - } - - // Mark nodes visited - for (auto id : id_scratch) { - if (fast_iterate) { - inserted_into_pool_bs[id] = 1; - } else { - inserted_into_pool_rs.insert(id); - } - } + const T *query, const uint32_t Lsize, const std::vector &init_ids, InMemQueryScratch *scratch, + bool use_filter, const std::vector &filter_label, bool search_invocation) +{ + std::vector &expanded_nodes = scratch->pool(); + NeighborPriorityQueue &best_L_nodes = scratch->best_l_nodes(); + best_L_nodes.reserve(Lsize); + tsl::robin_set &inserted_into_pool_rs = scratch->inserted_into_pool_rs(); + boost::dynamic_bitset<> &inserted_into_pool_bs = scratch->inserted_into_pool_bs(); + std::vector &id_scratch = scratch->id_scratch(); + std::vector &dist_scratch = scratch->dist_scratch(); + assert(id_scratch.size() == 0); + + // REFACTOR + // T *aligned_query = scratch->aligned_query(); + // memcpy(aligned_query, query, _dim * sizeof(T)); + // if (_normalize_vecs) + //{ + // normalize((float *)aligned_query, _dim); + // } - // Compute distances to unvisited nodes in the expansion - if (_pq_dist) { - assert(dist_scratch.capacity() >= id_scratch.size()); - compute_dists(id_scratch, dist_scratch); - } else { - assert(dist_scratch.size() == 0); - for (size_t m = 0; m < id_scratch.size(); ++m) { - uint32_t id = id_scratch[m]; + T *aligned_query = scratch->aligned_query(); - if (m + 1 < id_scratch.size()) { - auto nextn = id_scratch[m + 1]; - _data_store->prefetch_vector(nextn); + float *query_float = nullptr; + float *query_rotated = nullptr; + float *pq_dists = nullptr; + uint8_t *pq_coord_scratch = nullptr; + // Intialize PQ related scratch to use PQ based distances + if (_pq_dist) + { + // Get scratch spaces + PQScratch *pq_query_scratch = scratch->pq_scratch(); + query_float = pq_query_scratch->aligned_query_float; + query_rotated = pq_query_scratch->rotated_query; + pq_dists = pq_query_scratch->aligned_pqtable_dist_scratch; + + // Copy query vector to float and then to "rotated" query + for (size_t d = 0; d < _dim; d++) + { + query_float[d] = (float)aligned_query[d]; } + pq_query_scratch->set(_dim, aligned_query); - dist_scratch.push_back(_data_store->get_distance(aligned_query, id)); - } + // center the query and rotate if we have a rotation matrix + _pq_table.preprocess_query(query_rotated); + _pq_table.populate_chunk_distances(query_rotated, pq_dists); + + pq_coord_scratch = pq_query_scratch->aligned_pq_coord_scratch; } - cmps += id_scratch.size(); - // Insert pairs into the pool of candidates - for (size_t m = 0; m < id_scratch.size(); ++m) { - best_L_nodes.insert(Neighbor(id_scratch[m], dist_scratch[m])); + if (expanded_nodes.size() > 0 || id_scratch.size() > 0) + { + throw ANNException("ERROR: Clear scratch space before passing.", -1, __FUNCSIG__, __FILE__, __LINE__); } - } - return std::make_pair(hops, cmps); -} -template -void Index::search_for_point_and_prune( - int location, uint32_t Lindex, std::vector &pruned_list, - InMemQueryScratch *scratch, bool use_filter, uint32_t filteredLindex) { - const std::vector init_ids = get_init_ids(); - const std::vector unused_filter_label; + // Decide whether to use bitset or robin set to mark visited nodes + auto total_num_points = _max_points + _num_frozen_pts; + bool fast_iterate = total_num_points <= MAX_POINTS_FOR_USING_BITSET; - if (!use_filter) { - _data_store->get_vector(location, scratch->aligned_query()); - iterate_to_fixed_point(scratch->aligned_query(), Lindex, init_ids, scratch, - false, unused_filter_label, false); - } else { - std::vector filter_specific_start_nodes; - for (auto &x : _pts_to_labels[location]) - filter_specific_start_nodes.emplace_back(_label_to_medoid_id[x]); + if (fast_iterate) + { + if (inserted_into_pool_bs.size() < total_num_points) + { + // hopefully using 2X will reduce the number of allocations. + auto resize_size = + 2 * total_num_points > MAX_POINTS_FOR_USING_BITSET ? MAX_POINTS_FOR_USING_BITSET : 2 * total_num_points; + inserted_into_pool_bs.resize(resize_size); + } + } - _data_store->get_vector(location, scratch->aligned_query()); - iterate_to_fixed_point(scratch->aligned_query(), filteredLindex, - filter_specific_start_nodes, scratch, true, - _pts_to_labels[location], false); - } + // Lambda to determine if a node has been visited + auto is_not_visited = [this, fast_iterate, &inserted_into_pool_bs, &inserted_into_pool_rs](const uint32_t id) { + return fast_iterate ? inserted_into_pool_bs[id] == 0 + : inserted_into_pool_rs.find(id) == inserted_into_pool_rs.end(); + }; - auto &pool = scratch->pool(); + // Lambda to batch compute query<-> node distances in PQ space + auto compute_dists = [this, pq_coord_scratch, pq_dists](const std::vector &ids, + std::vector &dists_out) { + diskann::aggregate_coords(ids, this->_pq_data, this->_num_pq_chunks, pq_coord_scratch); + diskann::pq_dist_lookup(pq_coord_scratch, ids.size(), this->_num_pq_chunks, pq_dists, dists_out); + }; - for (uint32_t i = 0; i < pool.size(); i++) { - if (pool[i].id == (uint32_t)location) { - pool.erase(pool.begin() + i); - i--; - } - } + // Initialize the candidate pool with starting points + for (auto id : init_ids) + { + if (id >= _max_points + _num_frozen_pts) + { + diskann::cerr << "Out of range loc found as an edge : " << id << std::endl; + throw diskann::ANNException(std::string("Wrong loc") + std::to_string(id), -1, __FUNCSIG__, __FILE__, + __LINE__); + } - if (pruned_list.size() > 0) { - throw diskann::ANNException("ERROR: non-empty pruned_list passed", -1, - __FUNCSIG__, __FILE__, __LINE__); - } + if (use_filter) + { + std::vector common_filters; + auto &x = _pts_to_labels[id]; + std::set_intersection(filter_label.begin(), filter_label.end(), x.begin(), x.end(), + std::back_inserter(common_filters)); + if (_use_universal_label) + { + if (std::find(filter_label.begin(), filter_label.end(), _universal_label) != filter_label.end() || + std::find(x.begin(), x.end(), _universal_label) != x.end()) + common_filters.emplace_back(_universal_label); + } - prune_neighbors(location, pool, pruned_list, scratch); + if (common_filters.size() == 0) + continue; + } - assert(!pruned_list.empty()); - assert(_final_graph.size() == _max_points + _num_frozen_pts); -} + if (is_not_visited(id)) + { + if (fast_iterate) + { + inserted_into_pool_bs[id] = 1; + } + else + { + inserted_into_pool_rs.insert(id); + } -template -void Index::occlude_list( - const uint32_t location, std::vector &pool, const float alpha, - const uint32_t degree, const uint32_t maxc, std::vector &result, - InMemQueryScratch *scratch, - const tsl::robin_set *const delete_set_ptr) { - if (pool.size() == 0) return; - - // Truncate pool at maxc and initialize scratch spaces - assert(std::is_sorted(pool.begin(), pool.end())); - assert(result.size() == 0); - if (pool.size() > maxc) pool.resize(maxc); - std::vector &occlude_factor = scratch->occlude_factor(); - // occlude_list can be called with the same scratch more than once by - // search_for_point_and_add_link through inter_insert. - occlude_factor.clear(); - // Initialize occlude_factor to pool.size() many 0.0f values for correctness - occlude_factor.insert(occlude_factor.end(), pool.size(), 0.0f); - - float cur_alpha = 1; - while (cur_alpha <= alpha && result.size() < degree) { - // used for MIPS, where we store a value of eps in cur_alpha to - // denote pruned out entries which we can skip in later rounds. - float eps = cur_alpha + 0.01f; - - for (auto iter = pool.begin(); result.size() < degree && iter != pool.end(); - ++iter) { - if (occlude_factor[iter - pool.begin()] > cur_alpha) { - continue; - } - // Set the entry to float::max so that is not considered again - occlude_factor[iter - pool.begin()] = std::numeric_limits::max(); - // Add the entry to the result if its not been deleted, and doesn't - // add a self loop - if (delete_set_ptr == nullptr || - delete_set_ptr->find(iter->id) == delete_set_ptr->end()) { - if (iter->id != location) { - result.push_back(iter->id); - } - } - - // Update occlude factor for points from iter+1 to pool.end() - for (auto iter2 = iter + 1; iter2 != pool.end(); iter2++) { - auto t = iter2 - pool.begin(); - if (occlude_factor[t] > alpha) continue; - - bool prune_allowed = true; - if (_filtered_index) { - uint32_t a = iter->id; - uint32_t b = iter2->id; - for (auto &x : _pts_to_labels[b]) { - if (std::find(_pts_to_labels[a].begin(), _pts_to_labels[a].end(), - x) == _pts_to_labels[a].end()) { - prune_allowed = false; + float distance; + if (_pq_dist) + { + pq_dist_lookup(pq_coord_scratch, 1, this->_num_pq_chunks, pq_dists, &distance); } - if (!prune_allowed) break; - } + else + { + distance = _data_store->get_distance(aligned_query, id); + } + Neighbor nn = Neighbor(id, distance); + best_L_nodes.insert(nn); } - if (!prune_allowed) continue; + } - float djk = _data_store->get_distance(iter2->id, iter->id); - if (_dist_metric == diskann::Metric::L2 || - _dist_metric == diskann::Metric::COSINE) { - occlude_factor[t] = - (djk == 0) ? std::numeric_limits::max() - : std::max(occlude_factor[t], iter2->distance / djk); - } else if (_dist_metric == diskann::Metric::INNER_PRODUCT) { - // Improvization for flipping max and min dist for MIPS - float x = -iter2->distance; - float y = -djk; - if (y > cur_alpha * x) { - occlude_factor[t] = std::max(occlude_factor[t], eps); - } + uint32_t hops = 0; + uint32_t cmps = 0; + + while (best_L_nodes.has_unexpanded_node()) + { + auto nbr = best_L_nodes.closest_unexpanded(); + auto n = nbr.id; + + // Add node to expanded nodes to create pool for prune later + if (!search_invocation) + { + if (!use_filter) + { + expanded_nodes.emplace_back(nbr); + } + else + { // in filter based indexing, the same point might invoke + // multiple iterate_to_fixed_points, so need to be careful + // not to add the same item to pool multiple times. + if (std::find(expanded_nodes.begin(), expanded_nodes.end(), nbr) == expanded_nodes.end()) + { + expanded_nodes.emplace_back(nbr); + } + } } - } - } - cur_alpha *= 1.2f; - } -} -template -void Index::prune_neighbors(const uint32_t location, - std::vector &pool, - std::vector &pruned_list, - InMemQueryScratch *scratch) { - prune_neighbors(location, pool, _indexingRange, _indexingMaxC, _indexingAlpha, - pruned_list, scratch); -} + // Find which of the nodes in des have not been visited before + id_scratch.clear(); + dist_scratch.clear(); + { + if (_dynamic_index) + _locks[n].lock(); + for (auto id : _final_graph[n]) + { + assert(id < _max_points + _num_frozen_pts); + + if (use_filter) + { + // NOTE: NEED TO CHECK IF THIS CORRECT WITH NEW LOCKS. + std::vector common_filters; + auto &x = _pts_to_labels[id]; + std::set_intersection(filter_label.begin(), filter_label.end(), x.begin(), x.end(), + std::back_inserter(common_filters)); + if (_use_universal_label) + { + if (std::find(filter_label.begin(), filter_label.end(), _universal_label) != + filter_label.end() || + std::find(x.begin(), x.end(), _universal_label) != x.end()) + common_filters.emplace_back(_universal_label); + } + + if (common_filters.size() == 0) + continue; + } + + if (is_not_visited(id)) + { + id_scratch.push_back(id); + } + } -template -void Index::prune_neighbors( - const uint32_t location, std::vector &pool, const uint32_t range, - const uint32_t max_candidate_size, const float alpha, - std::vector &pruned_list, InMemQueryScratch *scratch) { - if (pool.size() == 0) { - // if the pool is empty, behave like a noop - pruned_list.clear(); - return; - } - - _max_observed_degree = (std::max)(_max_observed_degree, range); - - // If using _pq_build, over-write the PQ distances with actual distances - if (_pq_dist) { - for (auto &ngh : pool) - ngh.distance = _data_store->get_distance(ngh.id, location); - } - - // sort the pool based on distance to query and prune it with occlude_list - std::sort(pool.begin(), pool.end()); - pruned_list.clear(); - pruned_list.reserve(range); - - if (pool.begin()->distance == 0) { - diskann::cerr - << "Warning: a candidate with distance 0 found in prune_neighbors" - << std::endl; - } - occlude_list(location, pool, alpha, range, max_candidate_size, pruned_list, - scratch); - assert(pruned_list.size() <= range); - - if (_saturate_graph && alpha > 1) { - for (const auto &node : pool) { - if (pruned_list.size() >= range) break; - if ((std::find(pruned_list.begin(), pruned_list.end(), node.id) == - pruned_list.end()) && - node.id != location) - pruned_list.push_back(node.id); - } - } -} + if (_dynamic_index) + _locks[n].unlock(); + } -template -void Index::inter_insert(uint32_t n, - std::vector &pruned_list, - const uint32_t range, - InMemQueryScratch *scratch) { - const auto &src_pool = pruned_list; - - assert(!src_pool.empty()); - - for (auto des : src_pool) { - // des.loc is the loc of the neighbors of n - assert(des < _max_points + _num_frozen_pts); - // des_pool contains the neighbors of the neighbors of n - std::vector copy_of_neighbors; - bool prune_needed = false; - { - LockGuard guard(_locks[des]); - auto &des_pool = _final_graph[des]; - if (std::find(des_pool.begin(), des_pool.end(), n) == des_pool.end()) { - if (des_pool.size() < (uint64_t)(GRAPH_SLACK_FACTOR * range)) { - des_pool.emplace_back(n); - prune_needed = false; - } else { - copy_of_neighbors.reserve(des_pool.size() + 1); - copy_of_neighbors = des_pool; - copy_of_neighbors.push_back(n); - prune_needed = true; - } - } - } // des lock is released by this point - - if (prune_needed) { - tsl::robin_set dummy_visited(0); - std::vector dummy_pool(0); - - size_t reserveSize = - (size_t)(std::ceil(1.05 * GRAPH_SLACK_FACTOR * range)); - dummy_visited.reserve(reserveSize); - dummy_pool.reserve(reserveSize); - - for (auto cur_nbr : copy_of_neighbors) { - if (dummy_visited.find(cur_nbr) == dummy_visited.end() && - cur_nbr != des) { - float dist = _data_store->get_distance(des, cur_nbr); - dummy_pool.emplace_back(Neighbor(cur_nbr, dist)); - dummy_visited.insert(cur_nbr); - } - } - std::vector new_out_neighbors; - prune_neighbors(des, dummy_pool, new_out_neighbors, scratch); - { - LockGuard guard(_locks[des]); - - _final_graph[des] = new_out_neighbors; - } - } - } -} + // Mark nodes visited + for (auto id : id_scratch) + { + if (fast_iterate) + { + inserted_into_pool_bs[id] = 1; + } + else + { + inserted_into_pool_rs.insert(id); + } + } -template -void Index::inter_insert(uint32_t n, - std::vector &pruned_list, - InMemQueryScratch *scratch) { - inter_insert(n, pruned_list, _indexingRange, scratch); + // Compute distances to unvisited nodes in the expansion + if (_pq_dist) + { + assert(dist_scratch.capacity() >= id_scratch.size()); + compute_dists(id_scratch, dist_scratch); + } + else + { + assert(dist_scratch.size() == 0); + for (size_t m = 0; m < id_scratch.size(); ++m) + { + uint32_t id = id_scratch[m]; + + if (m + 1 < id_scratch.size()) + { + auto nextn = id_scratch[m + 1]; + _data_store->prefetch_vector(nextn); + } + + dist_scratch.push_back(_data_store->get_distance(aligned_query, id)); + } + } + cmps += id_scratch.size(); + + // Insert pairs into the pool of candidates + for (size_t m = 0; m < id_scratch.size(); ++m) + { + best_L_nodes.insert(Neighbor(id_scratch[m], dist_scratch[m])); + } + } + return std::make_pair(hops, cmps); } template -void Index::link(IndexWriteParameters ¶meters) { - uint32_t num_threads = parameters.num_threads; - if (num_threads != 0) omp_set_num_threads(num_threads); - - _saturate_graph = parameters.saturate_graph; - - _indexingQueueSize = parameters.search_list_size; - _filterIndexingQueueSize = parameters.filter_list_size; - _indexingRange = parameters.max_degree; - _indexingMaxC = parameters.max_occlusion_size; - _indexingAlpha = parameters.alpha; - - /* visit_order is a vector that is initialized to the entire graph */ - std::vector visit_order; - std::vector pool, tmp; - tsl::robin_set visited; - visit_order.reserve(_nd + _num_frozen_pts); - for (uint32_t i = 0; i < (uint32_t)_nd; i++) { - visit_order.emplace_back(i); - } - - // If there are any frozen points, add them all. - for (uint32_t frozen = _max_points; frozen < _max_points + _num_frozen_pts; - frozen++) { - visit_order.emplace_back(frozen); - } - - // if there are frozen points, the first such one is set to be the _start - if (_num_frozen_pts > 0) - _start = (uint32_t)_max_points; - else - _start = calculate_entry_point(); +void Index::search_for_point_and_prune(int location, uint32_t Lindex, + std::vector &pruned_list, + InMemQueryScratch *scratch, bool use_filter, + uint32_t filteredLindex) +{ + const std::vector init_ids = get_init_ids(); + const std::vector unused_filter_label; - for (size_t p = 0; p < _nd; p++) { - _final_graph[p].reserve( - (size_t)(std::ceil(_indexingRange * GRAPH_SLACK_FACTOR * 1.05))); - } - - diskann::Timer link_timer; + if (!use_filter) + { + _data_store->get_vector(location, scratch->aligned_query()); + iterate_to_fixed_point(scratch->aligned_query(), Lindex, init_ids, scratch, false, unused_filter_label, false); + } + else + { + std::vector filter_specific_start_nodes; + for (auto &x : _pts_to_labels[location]) + filter_specific_start_nodes.emplace_back(_label_to_medoid_id[x]); -#pragma omp parallel for schedule(dynamic, 2048) - for (int64_t node_ctr = 0; node_ctr < (int64_t)(visit_order.size()); - node_ctr++) { - auto node = visit_order[node_ctr]; + _data_store->get_vector(location, scratch->aligned_query()); + iterate_to_fixed_point(scratch->aligned_query(), filteredLindex, filter_specific_start_nodes, scratch, true, + _pts_to_labels[location], false); + } - ScratchStoreManager> manager(_query_scratch); - auto scratch = manager.scratch_space(); + auto &pool = scratch->pool(); - std::vector pruned_list; - if (_filtered_index) { - search_for_point_and_prune(node, _indexingQueueSize, pruned_list, scratch, - _filtered_index, _filterIndexingQueueSize); - } else { - search_for_point_and_prune(node, _indexingQueueSize, pruned_list, - scratch); - } + for (uint32_t i = 0; i < pool.size(); i++) { - LockGuard guard(_locks[node]); - _final_graph[node].reserve( - (size_t)(_indexingRange * GRAPH_SLACK_FACTOR * 1.05)); - _final_graph[node] = pruned_list; - assert(_final_graph[node].size() <= _indexingRange); + if (pool[i].id == (uint32_t)location) + { + pool.erase(pool.begin() + i); + i--; + } } - inter_insert(node, pruned_list, scratch); - - if (node_ctr % 100000 == 0) { - diskann::cout << "\r" << (100.0 * node_ctr) / (visit_order.size()) - << "% of index build completed." << std::flush; + if (pruned_list.size() > 0) + { + throw diskann::ANNException("ERROR: non-empty pruned_list passed", -1, __FUNCSIG__, __FILE__, __LINE__); } - } - if (_nd > 0) { - diskann::cout << "Starting final cleanup.." << std::flush; - } -#pragma omp parallel for schedule(dynamic, 2048) - for (int64_t node_ctr = 0; node_ctr < (int64_t)(visit_order.size()); - node_ctr++) { - auto node = visit_order[node_ctr]; - if (_final_graph[node].size() > _indexingRange) { - ScratchStoreManager> manager(_query_scratch); - auto scratch = manager.scratch_space(); - - 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 = _data_store->get_distance(node, cur_nbr); - dummy_pool.emplace_back(Neighbor(cur_nbr, dist)); - dummy_visited.insert(cur_nbr); - } - } - prune_neighbors(node, dummy_pool, new_out_neighbors, scratch); - - _final_graph[node].clear(); - for (auto id : new_out_neighbors) _final_graph[node].emplace_back(id); - } - } - if (_nd > 0) { - diskann::cout << "done. Link time: " - << ((double)link_timer.elapsed() / (double)1000000) << "s" - << std::endl; - } + prune_neighbors(location, pool, pruned_list, scratch); + + assert(!pruned_list.empty()); + assert(_final_graph.size() == _max_points + _num_frozen_pts); } template -void Index::prune_all_neighbors( - const uint32_t max_degree, const uint32_t max_occlusion_size, - const float alpha) { - const uint32_t range = max_degree; - const uint32_t maxc = max_occlusion_size; +void Index::occlude_list(const uint32_t location, std::vector &pool, const float alpha, + const uint32_t degree, const uint32_t maxc, std::vector &result, + InMemQueryScratch *scratch, + const tsl::robin_set *const delete_set_ptr) +{ + if (pool.size() == 0) + return; - _filtered_index = true; + // Truncate pool at maxc and initialize scratch spaces + assert(std::is_sorted(pool.begin(), pool.end())); + assert(result.size() == 0); + if (pool.size() > maxc) + pool.resize(maxc); + std::vector &occlude_factor = scratch->occlude_factor(); + // occlude_list can be called with the same scratch more than once by + // search_for_point_and_add_link through inter_insert. + occlude_factor.clear(); + // Initialize occlude_factor to pool.size() many 0.0f values for correctness + occlude_factor.insert(occlude_factor.end(), pool.size(), 0.0f); - diskann::Timer timer; -#pragma omp parallel for - for (int64_t node = 0; node < (int64_t)(_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; - - ScratchStoreManager> manager(_query_scratch); - auto scratch = manager.scratch_space(); + float cur_alpha = 1; + while (cur_alpha <= alpha && result.size() < degree) + { + // used for MIPS, where we store a value of eps in cur_alpha to + // denote pruned out entries which we can skip in later rounds. + float eps = cur_alpha + 0.01f; + + for (auto iter = pool.begin(); result.size() < degree && iter != pool.end(); ++iter) + { + if (occlude_factor[iter - pool.begin()] > cur_alpha) + { + continue; + } + // Set the entry to float::max so that is not considered again + occlude_factor[iter - pool.begin()] = std::numeric_limits::max(); + // Add the entry to the result if its not been deleted, and doesn't + // add a self loop + if (delete_set_ptr == nullptr || delete_set_ptr->find(iter->id) == delete_set_ptr->end()) + { + if (iter->id != location) + { + result.push_back(iter->id); + } + } - for (auto cur_nbr : _final_graph[node]) { - if (dummy_visited.find(cur_nbr) == dummy_visited.end() && - cur_nbr != node) { - float dist = _data_store->get_distance(node, cur_nbr); - dummy_pool.emplace_back(Neighbor(cur_nbr, dist)); - dummy_visited.insert(cur_nbr); - } - } - - prune_neighbors((uint32_t)node, dummy_pool, range, maxc, alpha, - new_out_neighbors, scratch); - _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 < _max_points + _num_frozen_pts; i++) { - if (i < _nd || i >= _max_points) { - const std::vector &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; - } + // Update occlude factor for points from iter+1 to pool.end() + for (auto iter2 = iter + 1; iter2 != pool.end(); iter2++) + { + auto t = iter2 - pool.begin(); + if (occlude_factor[t] > alpha) + continue; + + bool prune_allowed = true; + if (_filtered_index) + { + uint32_t a = iter->id; + uint32_t b = iter2->id; + for (auto &x : _pts_to_labels[b]) + { + if (std::find(_pts_to_labels[a].begin(), _pts_to_labels[a].end(), x) == _pts_to_labels[a].end()) + { + prune_allowed = false; + } + if (!prune_allowed) + break; + } + } + if (!prune_allowed) + continue; + + float djk = _data_store->get_distance(iter2->id, iter->id); + if (_dist_metric == diskann::Metric::L2 || _dist_metric == diskann::Metric::COSINE) + { + occlude_factor[t] = (djk == 0) ? std::numeric_limits::max() + : std::max(occlude_factor[t], iter2->distance / djk); + } + else if (_dist_metric == diskann::Metric::INNER_PRODUCT) + { + // Improvization for flipping max and min dist for MIPS + float x = -iter2->distance; + float y = -djk; + if (y > cur_alpha * x) + { + occlude_factor[t] = std::max(occlude_factor[t], eps); + } + } + } + } + cur_alpha *= 1.2f; + } } -// REFACTOR template -void Index::set_start_points(const T *data, - size_t data_count) { - std::unique_lock ul(_update_lock); - std::unique_lock tl(_tag_lock); - if (_nd > 0) - throw ANNException("Can not set starting point for a non-empty index", -1, - __FUNCSIG__, __FILE__, __LINE__); - - if (data_count != _num_frozen_pts * _dim) - throw ANNException("Invalid number of points", -1, __FUNCSIG__, __FILE__, - __LINE__); - - // memcpy(_data + _aligned_dim * _max_points, data, _aligned_dim * - // sizeof(T) * _num_frozen_pts); - for (location_t i = _max_points; i < _max_points + _num_frozen_pts; i++) { - _data_store->set_vector(i, data + i * _dim); - } - _has_built = true; - diskann::cout << "Index start points set: #" << _num_frozen_pts << std::endl; +void Index::prune_neighbors(const uint32_t location, std::vector &pool, + std::vector &pruned_list, InMemQueryScratch *scratch) +{ + prune_neighbors(location, pool, _indexingRange, _indexingMaxC, _indexingAlpha, pruned_list, scratch); } template -void Index::set_start_points_at_random(T radius, - uint32_t random_seed) { - std::mt19937 gen{random_seed}; - std::normal_distribution<> d{0.0, 1.0}; +void Index::prune_neighbors(const uint32_t location, std::vector &pool, const uint32_t range, + const uint32_t max_candidate_size, const float alpha, + std::vector &pruned_list, InMemQueryScratch *scratch) +{ + if (pool.size() == 0) + { + // if the pool is empty, behave like a noop + pruned_list.clear(); + return; + } - std::vector points_data; - points_data.reserve(_dim * _num_frozen_pts); - std::vector real_vec(_dim); + _max_observed_degree = (std::max)(_max_observed_degree, range); - for (size_t frozen_point = 0; frozen_point < _num_frozen_pts; - frozen_point++) { - double norm_sq = 0.0; - for (size_t i = 0; i < _dim; ++i) { - auto r = d(gen); - real_vec[i] = r; - norm_sq += r * r; + // If using _pq_build, over-write the PQ distances with actual distances + if (_pq_dist) + { + for (auto &ngh : pool) + ngh.distance = _data_store->get_distance(ngh.id, location); } - const double norm = std::sqrt(norm_sq); - for (auto iter : real_vec) - points_data.push_back(static_cast(iter * radius / norm)); - } + // sort the pool based on distance to query and prune it with occlude_list + std::sort(pool.begin(), pool.end()); + pruned_list.clear(); + pruned_list.reserve(range); - set_start_points(points_data.data(), points_data.size()); -} + if (pool.begin()->distance == 0) + { + diskann::cerr << "Warning: a candidate with distance 0 found in prune_neighbors" << std::endl; + } + occlude_list(location, pool, alpha, range, max_candidate_size, pruned_list, scratch); + assert(pruned_list.size() <= range); -template -void Index::build_with_data_populated( - IndexWriteParameters ¶meters, const std::vector &tags) { - diskann::cout << "Starting index build with " << _nd << " points... " - << std::endl; + if (_saturate_graph && alpha > 1) + { + for (const auto &node : pool) + { + if (pruned_list.size() >= range) + break; + if ((std::find(pruned_list.begin(), pruned_list.end(), node.id) == pruned_list.end()) && + node.id != location) + pruned_list.push_back(node.id); + } + } +} - if (_nd < 1) - throw ANNException("Error: Trying to build an index with 0 points", -1, - __FUNCSIG__, __FILE__, __LINE__); +template +void Index::inter_insert(uint32_t n, std::vector &pruned_list, const uint32_t range, + InMemQueryScratch *scratch) +{ + const auto &src_pool = pruned_list; - if (_enable_tags && tags.size() != _nd) { - std::stringstream stream; - stream << "ERROR: Driver requests loading " << _nd << " points from file," - << "but tags vector is of size " << tags.size() << "." << std::endl; - diskann::cerr << stream.str() << std::endl; - 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]] = (uint32_t)i; - _location_to_tag.set(static_cast(i), tags[i]); - } - } - - uint32_t index_R = parameters.max_degree; - uint32_t num_threads_index = parameters.num_threads; - uint32_t index_L = parameters.search_list_size; - uint32_t maxc = parameters.max_occlusion_size; - - if (_query_scratch.size() == 0) { - initialize_query_scratch(5 + num_threads_index, index_L, index_L, index_R, - maxc, _data_store->get_aligned_dim()); - } - - generate_frozen_point(); - link(parameters); - - size_t max = 0, min = SIZE_MAX, 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++; - } - diskann::cout << "Index built with degree: max:" << max - << " avg:" << (float)total / (float)(_nd + _num_frozen_pts) - << " min:" << min << " count(deg<2):" << cnt << std::endl; - - _max_observed_degree = std::max((uint32_t)max, _max_observed_degree); - _has_built = true; + assert(!src_pool.empty()); + + for (auto des : src_pool) + { + // des.loc is the loc of the neighbors of n + assert(des < _max_points + _num_frozen_pts); + // des_pool contains the neighbors of the neighbors of n + std::vector copy_of_neighbors; + bool prune_needed = false; + { + LockGuard guard(_locks[des]); + auto &des_pool = _final_graph[des]; + if (std::find(des_pool.begin(), des_pool.end(), n) == des_pool.end()) + { + if (des_pool.size() < (uint64_t)(GRAPH_SLACK_FACTOR * range)) + { + des_pool.emplace_back(n); + prune_needed = false; + } + else + { + copy_of_neighbors.reserve(des_pool.size() + 1); + copy_of_neighbors = des_pool; + copy_of_neighbors.push_back(n); + prune_needed = true; + } + } + } // des lock is released by this point + + if (prune_needed) + { + tsl::robin_set dummy_visited(0); + std::vector dummy_pool(0); + + size_t reserveSize = (size_t)(std::ceil(1.05 * GRAPH_SLACK_FACTOR * range)); + dummy_visited.reserve(reserveSize); + dummy_pool.reserve(reserveSize); + + for (auto cur_nbr : copy_of_neighbors) + { + if (dummy_visited.find(cur_nbr) == dummy_visited.end() && cur_nbr != des) + { + float dist = _data_store->get_distance(des, cur_nbr); + dummy_pool.emplace_back(Neighbor(cur_nbr, dist)); + dummy_visited.insert(cur_nbr); + } + } + std::vector new_out_neighbors; + prune_neighbors(des, dummy_pool, new_out_neighbors, scratch); + { + LockGuard guard(_locks[des]); + + _final_graph[des] = new_out_neighbors; + } + } + } +} + +template +void Index::inter_insert(uint32_t n, std::vector &pruned_list, InMemQueryScratch *scratch) +{ + inter_insert(n, pruned_list, _indexingRange, scratch); } template -void Index::build(const T *data, - const size_t num_points_to_load, - IndexWriteParameters ¶meters, - const std::vector &tags) { - if (num_points_to_load == 0) { - 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); - - { +void Index::link(IndexWriteParameters ¶meters) +{ + uint32_t num_threads = parameters.num_threads; + if (num_threads != 0) + omp_set_num_threads(num_threads); + + _saturate_graph = parameters.saturate_graph; + + _indexingQueueSize = parameters.search_list_size; + _filterIndexingQueueSize = parameters.filter_list_size; + _indexingRange = parameters.max_degree; + _indexingMaxC = parameters.max_occlusion_size; + _indexingAlpha = parameters.alpha; + + /* visit_order is a vector that is initialized to the entire graph */ + std::vector visit_order; + std::vector pool, tmp; + tsl::robin_set visited; + visit_order.reserve(_nd + _num_frozen_pts); + for (uint32_t i = 0; i < (uint32_t)_nd; i++) + { + visit_order.emplace_back(i); + } + + // If there are any frozen points, add them all. + for (uint32_t frozen = _max_points; frozen < _max_points + _num_frozen_pts; frozen++) + { + visit_order.emplace_back(frozen); + } + + // if there are frozen points, the first such one is set to be the _start + if (_num_frozen_pts > 0) + _start = (uint32_t)_max_points; + else + _start = calculate_entry_point(); + + for (size_t p = 0; p < _nd; p++) + { + _final_graph[p].reserve((size_t)(std::ceil(_indexingRange * GRAPH_SLACK_FACTOR * 1.05))); + } + + diskann::Timer link_timer; + +#pragma omp parallel for schedule(dynamic, 2048) + for (int64_t node_ctr = 0; node_ctr < (int64_t)(visit_order.size()); node_ctr++) + { + auto node = visit_order[node_ctr]; + + ScratchStoreManager> manager(_query_scratch); + auto scratch = manager.scratch_space(); + + std::vector pruned_list; + if (_filtered_index) + { + search_for_point_and_prune(node, _indexingQueueSize, pruned_list, scratch, _filtered_index, + _filterIndexingQueueSize); + } + else + { + search_for_point_and_prune(node, _indexingQueueSize, pruned_list, scratch); + } + { + LockGuard guard(_locks[node]); + _final_graph[node].reserve((size_t)(_indexingRange * GRAPH_SLACK_FACTOR * 1.05)); + _final_graph[node] = pruned_list; + assert(_final_graph[node].size() <= _indexingRange); + } + + inter_insert(node, pruned_list, scratch); + + if (node_ctr % 100000 == 0) + { + diskann::cout << "\r" << (100.0 * node_ctr) / (visit_order.size()) << "% of index build completed." + << std::flush; + } + } + + if (_nd > 0) + { + diskann::cout << "Starting final cleanup.." << std::flush; + } +#pragma omp parallel for schedule(dynamic, 2048) + for (int64_t node_ctr = 0; node_ctr < (int64_t)(visit_order.size()); node_ctr++) + { + auto node = visit_order[node_ctr]; + if (_final_graph[node].size() > _indexingRange) + { + ScratchStoreManager> manager(_query_scratch); + auto scratch = manager.scratch_space(); + + 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 = _data_store->get_distance(node, cur_nbr); + dummy_pool.emplace_back(Neighbor(cur_nbr, dist)); + dummy_visited.insert(cur_nbr); + } + } + prune_neighbors(node, dummy_pool, new_out_neighbors, scratch); + + _final_graph[node].clear(); + for (auto id : new_out_neighbors) + _final_graph[node].emplace_back(id); + } + } + if (_nd > 0) + { + diskann::cout << "done. Link time: " << ((double)link_timer.elapsed() / (double)1000000) << "s" << std::endl; + } +} + +template +void Index::prune_all_neighbors(const uint32_t max_degree, const uint32_t max_occlusion_size, + const float alpha) +{ + const uint32_t range = max_degree; + const uint32_t maxc = max_occlusion_size; + + _filtered_index = true; + + diskann::Timer timer; +#pragma omp parallel for + for (int64_t node = 0; node < (int64_t)(_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; + + ScratchStoreManager> manager(_query_scratch); + auto scratch = manager.scratch_space(); + + for (auto cur_nbr : _final_graph[node]) + { + if (dummy_visited.find(cur_nbr) == dummy_visited.end() && cur_nbr != node) + { + float dist = _data_store->get_distance(node, cur_nbr); + dummy_pool.emplace_back(Neighbor(cur_nbr, dist)); + dummy_visited.insert(cur_nbr); + } + } + + prune_neighbors((uint32_t)node, dummy_pool, range, maxc, alpha, new_out_neighbors, scratch); + _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 < _max_points + _num_frozen_pts; i++) + { + if (i < _nd || i >= _max_points) + { + const std::vector &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; + } +} + +// REFACTOR +template +void Index::set_start_points(const T *data, size_t data_count) +{ + std::unique_lock ul(_update_lock); std::unique_lock tl(_tag_lock); - _nd = num_points_to_load; + if (_nd > 0) + throw ANNException("Can not set starting point for a non-empty index", -1, __FUNCSIG__, __FILE__, __LINE__); - _data_store->populate_data(data, num_points_to_load); + if (data_count != _num_frozen_pts * _dim) + throw ANNException("Invalid number of points", -1, __FUNCSIG__, __FILE__, __LINE__); - // REFACTOR - // memcpy((char *)_data, (char *)data, _aligned_dim * _nd * sizeof(T)); - // if (_normalize_vecs) - //{ - // for (size_t i = 0; i < num_points_to_load; i++) - // { - // normalize(_data + _aligned_dim * i, _aligned_dim); - // } - // } - } + // memcpy(_data + _aligned_dim * _max_points, data, _aligned_dim * + // sizeof(T) * _num_frozen_pts); + for (location_t i = _max_points; i < _max_points + _num_frozen_pts; i++) + { + _data_store->set_vector(i, data + i * _dim); + } + _has_built = true; + diskann::cout << "Index start points set: #" << _num_frozen_pts << std::endl; +} - build_with_data_populated(parameters, tags); +template +void Index::set_start_points_at_random(T radius, uint32_t random_seed) +{ + std::mt19937 gen{random_seed}; + std::normal_distribution<> d{0.0, 1.0}; + + std::vector points_data; + points_data.reserve(_dim * _num_frozen_pts); + std::vector real_vec(_dim); + + for (size_t frozen_point = 0; frozen_point < _num_frozen_pts; frozen_point++) + { + double norm_sq = 0.0; + for (size_t i = 0; i < _dim; ++i) + { + auto r = d(gen); + real_vec[i] = r; + norm_sq += r * r; + } + + const double norm = std::sqrt(norm_sq); + for (auto iter : real_vec) + points_data.push_back(static_cast(iter * radius / norm)); + } + + set_start_points(points_data.data(), points_data.size()); } template -void Index::build(const char *filename, - const size_t num_points_to_load, - IndexWriteParameters ¶meters, - const std::vector &tags) { - std::unique_lock ul(_update_lock); - if (num_points_to_load == 0) - throw ANNException("Do not call build with 0 points", -1, __FUNCSIG__, - __FILE__, __LINE__); - - if (!file_exists(filename)) { - std::stringstream stream; - stream << "ERROR: Data file " << filename << " does not exist." - << std::endl; - diskann::cerr << stream.str() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); - } - - size_t file_num_points, file_dim; - if (filename == nullptr) { - throw diskann::ANNException("Can not build with an empty file", -1, - __FUNCSIG__, __FILE__, __LINE__); - } - - diskann::get_bin_metadata(filename, file_num_points, file_dim); - if (file_num_points > _max_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; +void Index::build_with_data_populated(IndexWriteParameters ¶meters, const std::vector &tags) +{ + diskann::cout << "Starting index build with " << _nd << " points... " << std::endl; - if (_pq_dist) aligned_free(_pq_data); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); - } + if (_nd < 1) + throw ANNException("Error: Trying to build an index with 0 points", -1, __FUNCSIG__, __FILE__, __LINE__); - if (num_points_to_load > file_num_points) { - std::stringstream stream; - stream << "ERROR: Driver requests loading " << num_points_to_load - << " points and file has only " << file_num_points << " points." - << std::endl; + if (_enable_tags && tags.size() != _nd) + { + std::stringstream stream; + stream << "ERROR: Driver requests loading " << _nd << " points from file," + << "but tags vector is of size " << tags.size() << "." << std::endl; + diskann::cerr << stream.str() << std::endl; + 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]] = (uint32_t)i; + _location_to_tag.set(static_cast(i), tags[i]); + } + } - if (_pq_dist) aligned_free(_pq_data); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); - } + uint32_t index_R = parameters.max_degree; + uint32_t num_threads_index = parameters.num_threads; + uint32_t index_L = parameters.search_list_size; + uint32_t maxc = parameters.max_occlusion_size; - if (file_dim != _dim) { - std::stringstream stream; - stream << "ERROR: Driver requests loading " << _dim << " dimension," - << "but file has " << file_dim << " dimension." << std::endl; - diskann::cerr << stream.str() << std::endl; + if (_query_scratch.size() == 0) + { + initialize_query_scratch(5 + num_threads_index, index_L, index_L, index_R, maxc, + _data_store->get_aligned_dim()); + } + + generate_frozen_point(); + link(parameters); + + size_t max = 0, min = SIZE_MAX, 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++; + } + diskann::cout << "Index built with degree: max:" << max << " avg:" << (float)total / (float)(_nd + _num_frozen_pts) + << " min:" << min << " count(deg<2):" << cnt << std::endl; + + _max_observed_degree = std::max((uint32_t)max, _max_observed_degree); + _has_built = true; +} - if (_pq_dist) aligned_free(_pq_data); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); - } - - if (_pq_dist) { - double p_val = std::min( - 1.0, ((double)MAX_PQ_TRAINING_SET_SIZE / (double)file_num_points)); - - std::string suffix = _use_opq ? "_opq" : "_pq"; - suffix += std::to_string(_num_pq_chunks); - auto pq_pivots_file = std::string(filename) + suffix + "_pivots.bin"; - auto pq_compressed_file = - std::string(filename) + suffix + "_compressed.bin"; - generate_quantized_data(std::string(filename), pq_pivots_file, - pq_compressed_file, _dist_metric, p_val, - _num_pq_chunks, _use_opq); - - copy_aligned_data_from_file(pq_compressed_file.c_str(), _pq_data, - file_num_points, _num_pq_chunks, - _num_pq_chunks); +template +void Index::build(const T *data, const size_t num_points_to_load, IndexWriteParameters ¶meters, + const std::vector &tags) +{ + if (num_points_to_load == 0) + { + 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); + + { + std::unique_lock tl(_tag_lock); + _nd = num_points_to_load; + + _data_store->populate_data(data, num_points_to_load); + + // REFACTOR + // memcpy((char *)_data, (char *)data, _aligned_dim * _nd * sizeof(T)); + // if (_normalize_vecs) + //{ + // for (size_t i = 0; i < num_points_to_load; i++) + // { + // normalize(_data + _aligned_dim * i, _aligned_dim); + // } + // } + } + + build_with_data_populated(parameters, tags); +} + +template +void Index::build(const char *filename, const size_t num_points_to_load, + IndexWriteParameters ¶meters, const std::vector &tags) +{ + std::unique_lock ul(_update_lock); + if (num_points_to_load == 0) + throw ANNException("Do not call build with 0 points", -1, __FUNCSIG__, __FILE__, __LINE__); + + if (!file_exists(filename)) + { + std::stringstream stream; + stream << "ERROR: Data file " << filename << " does not exist." << std::endl; + diskann::cerr << stream.str() << std::endl; + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } + + size_t file_num_points, file_dim; + if (filename == nullptr) + { + throw diskann::ANNException("Can not build with an empty file", -1, __FUNCSIG__, __FILE__, __LINE__); + } + + diskann::get_bin_metadata(filename, file_num_points, file_dim); + if (file_num_points > _max_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; + + if (_pq_dist) + aligned_free(_pq_data); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } + + if (num_points_to_load > file_num_points) + { + std::stringstream stream; + stream << "ERROR: Driver requests loading " << num_points_to_load << " points and file has only " + << file_num_points << " points." << std::endl; + + if (_pq_dist) + aligned_free(_pq_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; + diskann::cerr << stream.str() << std::endl; + + if (_pq_dist) + aligned_free(_pq_data); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } + + if (_pq_dist) + { + double p_val = std::min(1.0, ((double)MAX_PQ_TRAINING_SET_SIZE / (double)file_num_points)); + + std::string suffix = _use_opq ? "_opq" : "_pq"; + suffix += std::to_string(_num_pq_chunks); + auto pq_pivots_file = std::string(filename) + suffix + "_pivots.bin"; + auto pq_compressed_file = std::string(filename) + suffix + "_compressed.bin"; + generate_quantized_data(std::string(filename), pq_pivots_file, pq_compressed_file, _dist_metric, p_val, + _num_pq_chunks, _use_opq); + + copy_aligned_data_from_file(pq_compressed_file.c_str(), _pq_data, file_num_points, _num_pq_chunks, + _num_pq_chunks); #ifdef EXEC_ENV_OLS - throw ANNException( - "load_pq_centroid_bin should not be called when " - "EXEC_ENV_OLS is defined.", - -1, __FUNCSIG__, __FILE__, __LINE__); + throw ANNException("load_pq_centroid_bin should not be called when " + "EXEC_ENV_OLS is defined.", + -1, __FUNCSIG__, __FILE__, __LINE__); #else - _pq_table.load_pq_centroid_bin(pq_pivots_file.c_str(), _num_pq_chunks); + _pq_table.load_pq_centroid_bin(pq_pivots_file.c_str(), _num_pq_chunks); #endif - } + } - _data_store->populate_data(filename, 0U); - diskann::cout << "Using only first " << num_points_to_load << " from file.. " - << std::endl; + _data_store->populate_data(filename, 0U); + diskann::cout << "Using only first " << num_points_to_load << " from file.. " << std::endl; - { - std::unique_lock tl(_tag_lock); - _nd = num_points_to_load; - } - build_with_data_populated(parameters, tags); + { + std::unique_lock tl(_tag_lock); + _nd = num_points_to_load; + } + build_with_data_populated(parameters, tags); } template -void Index::build(const char *filename, - const size_t num_points_to_load, - IndexWriteParameters ¶meters, - const char *tag_filename) { - std::vector tags; +void Index::build(const char *filename, const size_t num_points_to_load, + IndexWriteParameters ¶meters, const char *tag_filename) +{ + std::vector tags; - if (_enable_tags) { - std::unique_lock tl(_tag_lock); - if (tag_filename == nullptr) { - throw ANNException("Tag filename is null, while _enable_tags is set", -1, - __FUNCSIG__, __FILE__, __LINE__); - } 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, insufficient to populate tags for " - << num_points_to_load << " points to load"; - throw diskann::ANNException(sstream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); - } - for (size_t i = 0; i < num_points_to_load; i++) { - tags.push_back(tag_data[i]); + if (_enable_tags) + { + std::unique_lock tl(_tag_lock); + if (tag_filename == nullptr) + { + throw ANNException("Tag filename is null, while _enable_tags is set", -1, __FUNCSIG__, __FILE__, __LINE__); + } + 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, insufficient to populate tags for " << num_points_to_load + << " points to load"; + throw diskann::ANNException(sstream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } + for (size_t i = 0; i < num_points_to_load; i++) + { + tags.push_back(tag_data[i]); + } + delete[] tag_data; + } + else + { + throw diskann::ANNException(std::string("Tag file") + tag_filename + " does not exist", -1, __FUNCSIG__, + __FILE__, __LINE__); + } } - delete[] tag_data; - } else { - throw diskann::ANNException( - std::string("Tag file") + tag_filename + " does not exist", -1, - __FUNCSIG__, __FILE__, __LINE__); - } } - } - build(filename, num_points_to_load, parameters, tags); + build(filename, num_points_to_load, parameters, tags); } template -std::unordered_map Index::load_label_map( - const std::string &labels_map_file) { - std::unordered_map string_to_int_mp; - std::ifstream map_reader(labels_map_file); - std::string line, token; - LabelT token_as_num; - std::string label_str; - while (std::getline(map_reader, line)) { - std::istringstream iss(line); - getline(iss, token, '\t'); - label_str = token; - getline(iss, token, '\t'); - token_as_num = std::stoul(token); - string_to_int_mp[label_str] = token_as_num; - } - return string_to_int_mp; +std::unordered_map Index::load_label_map(const std::string &labels_map_file) +{ + std::unordered_map string_to_int_mp; + std::ifstream map_reader(labels_map_file); + std::string line, token; + LabelT token_as_num; + std::string label_str; + while (std::getline(map_reader, line)) + { + std::istringstream iss(line); + getline(iss, token, '\t'); + label_str = token; + getline(iss, token, '\t'); + token_as_num = std::stoul(token); + string_to_int_mp[label_str] = token_as_num; + } + return string_to_int_mp; } template -LabelT Index::get_converted_label( - const std::string &raw_label) { - if (_label_map.find(raw_label) != _label_map.end()) { - return _label_map[raw_label]; - } - std::stringstream stream; - stream << "Unable to find label in the Label Map"; - diskann::cerr << stream.str() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); +LabelT Index::get_converted_label(const std::string &raw_label) +{ + if (_label_map.find(raw_label) != _label_map.end()) + { + return _label_map[raw_label]; + } + std::stringstream stream; + stream << "Unable to find label in the Label Map"; + diskann::cerr << stream.str() << std::endl; + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } template -void Index::parse_label_file(const std::string &label_file, - size_t &num_points) { - // Format of Label txt file: filters with comma separators - - std::ifstream infile(label_file); - if (infile.fail()) { - throw diskann::ANNException( - std::string("Failed to open file ") + label_file, -1); - } - - std::string line, token; - uint32_t line_cnt = 0; - - while (std::getline(infile, line)) { - line_cnt++; - } - _pts_to_labels.resize(line_cnt, std::vector()); - - infile.clear(); - infile.seekg(0, std::ios::beg); - line_cnt = 0; - - while (std::getline(infile, line)) { - std::istringstream iss(line); - std::vector lbls(0); - getline(iss, token, '\t'); - std::istringstream new_iss(token); - while (getline(new_iss, token, ',')) { - token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); - token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); - LabelT token_as_num = std::stoul(token); - lbls.push_back(token_as_num); - _labels.insert(token_as_num); - } - if (lbls.size() <= 0) { - diskann::cout << "No label found"; - exit(-1); - } - std::sort(lbls.begin(), lbls.end()); - _pts_to_labels[line_cnt] = lbls; - line_cnt++; - } - num_points = (size_t)line_cnt; - diskann::cout << "Identified " << _labels.size() << " distinct label(s)" - << std::endl; +void Index::parse_label_file(const std::string &label_file, size_t &num_points) +{ + // Format of Label txt file: filters with comma separators + + std::ifstream infile(label_file); + if (infile.fail()) + { + throw diskann::ANNException(std::string("Failed to open file ") + label_file, -1); + } + + std::string line, token; + uint32_t line_cnt = 0; + + while (std::getline(infile, line)) + { + line_cnt++; + } + _pts_to_labels.resize(line_cnt, std::vector()); + + infile.clear(); + infile.seekg(0, std::ios::beg); + line_cnt = 0; + + while (std::getline(infile, line)) + { + std::istringstream iss(line); + std::vector lbls(0); + getline(iss, token, '\t'); + std::istringstream new_iss(token); + while (getline(new_iss, token, ',')) + { + token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); + token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); + LabelT token_as_num = std::stoul(token); + lbls.push_back(token_as_num); + _labels.insert(token_as_num); + } + if (lbls.size() <= 0) + { + diskann::cout << "No label found"; + exit(-1); + } + std::sort(lbls.begin(), lbls.end()); + _pts_to_labels[line_cnt] = lbls; + line_cnt++; + } + num_points = (size_t)line_cnt; + diskann::cout << "Identified " << _labels.size() << " distinct label(s)" << std::endl; } template -void Index::set_universal_label(const LabelT &label) { - _use_universal_label = true; - _universal_label = label; +void Index::set_universal_label(const LabelT &label) +{ + _use_universal_label = true; + _universal_label = label; } template -void Index::build_filtered_index( - const char *filename, const std::string &label_file, - const size_t num_points_to_load, IndexWriteParameters ¶meters, - const std::vector &tags) { - _labels_file = label_file; - _filtered_index = true; - _label_to_medoid_id.clear(); - size_t num_points_labels = 0; - parse_label_file( - label_file, - num_points_labels); // determines medoid for each label and identifies - // the points to label mapping - - std::unordered_map> label_to_points; - - for (typename tsl::robin_set::size_type lbl = 0; lbl < _labels.size(); - lbl++) { - auto itr = _labels.begin(); - std::advance(itr, lbl); - auto &x = *itr; - - std::vector labeled_points; - for (uint32_t point_id = 0; point_id < num_points_to_load; point_id++) { - bool pt_has_lbl = std::find(_pts_to_labels[point_id].begin(), - _pts_to_labels[point_id].end(), - x) != _pts_to_labels[point_id].end(); - - bool pt_has_univ_lbl = - (_use_universal_label && - (std::find(_pts_to_labels[point_id].begin(), - _pts_to_labels[point_id].end(), - _universal_label) != _pts_to_labels[point_id].end())); - - if (pt_has_lbl || pt_has_univ_lbl) { - labeled_points.emplace_back(point_id); - } - } - label_to_points[x] = labeled_points; - } - - uint32_t num_cands = 25; - for (auto itr = _labels.begin(); itr != _labels.end(); itr++) { - uint32_t best_medoid_count = std::numeric_limits::max(); - auto &curr_label = *itr; - uint32_t best_medoid; - auto labeled_points = label_to_points[curr_label]; - for (uint32_t cnd = 0; cnd < num_cands; cnd++) { - uint32_t cur_cnd = labeled_points[rand() % labeled_points.size()]; - uint32_t cur_cnt = std::numeric_limits::max(); - if (_medoid_counts.find(cur_cnd) == _medoid_counts.end()) { - _medoid_counts[cur_cnd] = 0; - cur_cnt = 0; - } else { - cur_cnt = _medoid_counts[cur_cnd]; - } - if (cur_cnt < best_medoid_count) { - best_medoid_count = cur_cnt; - best_medoid = cur_cnd; - } - } - _label_to_medoid_id[curr_label] = best_medoid; - _medoid_counts[best_medoid]++; - } - - this->build(filename, num_points_to_load, parameters, tags); +void Index::build_filtered_index(const char *filename, const std::string &label_file, + const size_t num_points_to_load, IndexWriteParameters ¶meters, + const std::vector &tags) +{ + _labels_file = label_file; + _filtered_index = true; + _label_to_medoid_id.clear(); + size_t num_points_labels = 0; + parse_label_file(label_file, + num_points_labels); // determines medoid for each label and identifies + // the points to label mapping + + std::unordered_map> label_to_points; + + for (typename tsl::robin_set::size_type lbl = 0; lbl < _labels.size(); lbl++) + { + auto itr = _labels.begin(); + std::advance(itr, lbl); + auto &x = *itr; + + std::vector labeled_points; + for (uint32_t point_id = 0; point_id < num_points_to_load; point_id++) + { + bool pt_has_lbl = std::find(_pts_to_labels[point_id].begin(), _pts_to_labels[point_id].end(), x) != + _pts_to_labels[point_id].end(); + + bool pt_has_univ_lbl = + (_use_universal_label && (std::find(_pts_to_labels[point_id].begin(), _pts_to_labels[point_id].end(), + _universal_label) != _pts_to_labels[point_id].end())); + + if (pt_has_lbl || pt_has_univ_lbl) + { + labeled_points.emplace_back(point_id); + } + } + label_to_points[x] = labeled_points; + } + + uint32_t num_cands = 25; + for (auto itr = _labels.begin(); itr != _labels.end(); itr++) + { + uint32_t best_medoid_count = std::numeric_limits::max(); + auto &curr_label = *itr; + uint32_t best_medoid; + auto labeled_points = label_to_points[curr_label]; + for (uint32_t cnd = 0; cnd < num_cands; cnd++) + { + uint32_t cur_cnd = labeled_points[rand() % labeled_points.size()]; + uint32_t cur_cnt = std::numeric_limits::max(); + if (_medoid_counts.find(cur_cnd) == _medoid_counts.end()) + { + _medoid_counts[cur_cnd] = 0; + cur_cnt = 0; + } + else + { + cur_cnt = _medoid_counts[cur_cnd]; + } + if (cur_cnt < best_medoid_count) + { + best_medoid_count = cur_cnt; + best_medoid = cur_cnd; + } + } + _label_to_medoid_id[curr_label] = best_medoid; + _medoid_counts[best_medoid]++; + } + + this->build(filename, num_points_to_load, parameters, tags); } template template -std::pair Index::search(const T *query, - const size_t K, - const uint32_t L, - IdType *indices, - float *distances) { - if (K > (uint64_t)L) { - throw ANNException("Set L to a value of at least K", -1, __FUNCSIG__, - __FILE__, __LINE__); - } - - ScratchStoreManager> manager(_query_scratch); - auto scratch = manager.scratch_space(); - - if (L > scratch->get_L()) { - diskann::cout << "Attempting to expand query scratch_space. Was created " - << "with Lsize: " << scratch->get_L() - << " but search L is: " << L << std::endl; - scratch->resize_for_new_L(L); - diskann::cout << "Resize completed. New scratch->L is " << scratch->get_L() - << std::endl; - } +std::pair Index::search(const T *query, const size_t K, const uint32_t L, + IdType *indices, float *distances) +{ + if (K > (uint64_t)L) + { + throw ANNException("Set L to a value of at least K", -1, __FUNCSIG__, __FILE__, __LINE__); + } + + ScratchStoreManager> manager(_query_scratch); + auto scratch = manager.scratch_space(); + + if (L > scratch->get_L()) + { + diskann::cout << "Attempting to expand query scratch_space. Was created " + << "with Lsize: " << scratch->get_L() << " but search L is: " << L << std::endl; + scratch->resize_for_new_L(L); + diskann::cout << "Resize completed. New scratch->L is " << scratch->get_L() << std::endl; + } - const std::vector unused_filter_label; - const std::vector init_ids = get_init_ids(); + const std::vector unused_filter_label; + const std::vector init_ids = get_init_ids(); - std::shared_lock lock(_update_lock); + std::shared_lock lock(_update_lock); - _distance->preprocess_query(query, _data_store->get_dims(), - scratch->aligned_query()); - auto retval = - iterate_to_fixed_point(scratch->aligned_query(), L, init_ids, scratch, - false, unused_filter_label, true); + _distance->preprocess_query(query, _data_store->get_dims(), scratch->aligned_query()); + auto retval = + iterate_to_fixed_point(scratch->aligned_query(), L, init_ids, scratch, false, unused_filter_label, true); - NeighborPriorityQueue &best_L_nodes = scratch->best_l_nodes(); + NeighborPriorityQueue &best_L_nodes = scratch->best_l_nodes(); - size_t pos = 0; - for (size_t i = 0; i < best_L_nodes.size(); ++i) { - if (best_L_nodes[i].id < _max_points) { - // safe because Index uses uint32_t ids internally - // and IDType will be uint32_t or uint64_t - indices[pos] = (IdType)best_L_nodes[i].id; - if (distances != nullptr) { + size_t pos = 0; + for (size_t i = 0; i < best_L_nodes.size(); ++i) + { + if (best_L_nodes[i].id < _max_points) + { + // safe because Index uses uint32_t ids internally + // and IDType will be uint32_t or uint64_t + indices[pos] = (IdType)best_L_nodes[i].id; + if (distances != nullptr) + { #ifdef EXEC_ENV_OLS - // DLVS expects negative distances - distances[pos] = best_L_nodes[i].distance; + // DLVS expects negative distances + distances[pos] = best_L_nodes[i].distance; #else - distances[pos] = _dist_metric == diskann::Metric::INNER_PRODUCT - ? -1 * best_L_nodes[i].distance - : best_L_nodes[i].distance; + distances[pos] = _dist_metric == diskann::Metric::INNER_PRODUCT ? -1 * best_L_nodes[i].distance + : best_L_nodes[i].distance; #endif - } - pos++; + } + pos++; + } + if (pos == K) + break; + } + if (pos < K) + { + diskann::cerr << "Found fewer than K elements for query" << std::endl; } - if (pos == K) break; - } - if (pos < K) { - diskann::cerr << "Found fewer than K elements for query" << std::endl; - } - return retval; + return retval; } template template -std::pair Index::search_with_filters( - const T *query, const LabelT &filter_label, const size_t K, - const uint32_t L, IdType *indices, float *distances) { - if (K > (uint64_t)L) { - throw ANNException("Set L to a value of at least K", -1, __FUNCSIG__, - __FILE__, __LINE__); - } - - ScratchStoreManager> manager(_query_scratch); - auto scratch = manager.scratch_space(); - - if (L > scratch->get_L()) { - diskann::cout << "Attempting to expand query scratch_space. Was created " - << "with Lsize: " << scratch->get_L() - << " but search L is: " << L << std::endl; - scratch->resize_for_new_L(L); - diskann::cout << "Resize completed. New scratch->L is " << scratch->get_L() - << std::endl; - } - - std::vector filter_vec; - std::vector init_ids = get_init_ids(); - - std::shared_lock lock(_update_lock); - - if (_label_to_medoid_id.find(filter_label) != _label_to_medoid_id.end()) { - init_ids.emplace_back(_label_to_medoid_id[filter_label]); - } else { - diskann::cout << "No filtered medoid found. exitting " - << std::endl; // RKNOTE: If universal label found start there - throw diskann::ANNException("No filtered medoid found. exitting ", -1); - } - filter_vec.emplace_back(filter_label); - - // REFACTOR - // T *aligned_query = scratch->aligned_query(); - // memcpy(aligned_query, query, _dim * sizeof(T)); - _distance->preprocess_query(query, _data_store->get_dims(), - scratch->aligned_query()); - auto retval = iterate_to_fixed_point(scratch->aligned_query(), L, init_ids, - scratch, true, filter_vec, true); - - auto best_L_nodes = scratch->best_l_nodes(); - - size_t pos = 0; - for (size_t i = 0; i < best_L_nodes.size(); ++i) { - if (best_L_nodes[i].id < _max_points) { - // safe because Index uses uint32_t ids internally - // and IDType will be uint32_t or uint64_t - indices[pos] = (IdType)best_L_nodes[i].id; - if (distances != nullptr) { +std::pair Index::search_with_filters(const T *query, const LabelT &filter_label, + const size_t K, const uint32_t L, + IdType *indices, float *distances) +{ + if (K > (uint64_t)L) + { + throw ANNException("Set L to a value of at least K", -1, __FUNCSIG__, __FILE__, __LINE__); + } + + ScratchStoreManager> manager(_query_scratch); + auto scratch = manager.scratch_space(); + + if (L > scratch->get_L()) + { + diskann::cout << "Attempting to expand query scratch_space. Was created " + << "with Lsize: " << scratch->get_L() << " but search L is: " << L << std::endl; + scratch->resize_for_new_L(L); + diskann::cout << "Resize completed. New scratch->L is " << scratch->get_L() << std::endl; + } + + std::vector filter_vec; + std::vector init_ids = get_init_ids(); + + std::shared_lock lock(_update_lock); + + if (_label_to_medoid_id.find(filter_label) != _label_to_medoid_id.end()) + { + init_ids.emplace_back(_label_to_medoid_id[filter_label]); + } + else + { + diskann::cout << "No filtered medoid found. exitting " + << std::endl; // RKNOTE: If universal label found start there + throw diskann::ANNException("No filtered medoid found. exitting ", -1); + } + filter_vec.emplace_back(filter_label); + + // REFACTOR + // T *aligned_query = scratch->aligned_query(); + // memcpy(aligned_query, query, _dim * sizeof(T)); + _distance->preprocess_query(query, _data_store->get_dims(), scratch->aligned_query()); + auto retval = iterate_to_fixed_point(scratch->aligned_query(), L, init_ids, scratch, true, filter_vec, true); + + auto best_L_nodes = scratch->best_l_nodes(); + + size_t pos = 0; + for (size_t i = 0; i < best_L_nodes.size(); ++i) + { + if (best_L_nodes[i].id < _max_points) + { + // safe because Index uses uint32_t ids internally + // and IDType will be uint32_t or uint64_t + indices[pos] = (IdType)best_L_nodes[i].id; + if (distances != nullptr) + { #ifdef EXEC_ENV_OLS - // DLVS expects negative distances - distances[pos] = best_L_nodes[i].distance; + // DLVS expects negative distances + distances[pos] = best_L_nodes[i].distance; #else - distances[pos] = _dist_metric == diskann::Metric::INNER_PRODUCT - ? -1 * best_L_nodes[i].distance - : best_L_nodes[i].distance; + distances[pos] = _dist_metric == diskann::Metric::INNER_PRODUCT ? -1 * best_L_nodes[i].distance + : best_L_nodes[i].distance; #endif - } - pos++; + } + pos++; + } + if (pos == K) + break; + } + if (pos < K) + { + diskann::cerr << "Found fewer than K elements for query" << std::endl; } - if (pos == K) break; - } - if (pos < K) { - diskann::cerr << "Found fewer than K elements for query" << std::endl; - } - return retval; + return retval; } template -size_t Index::search_with_tags(const T *query, - const uint64_t K, - const uint32_t L, TagT *tags, - float *distances, - std::vector &res_vectors) { - if (K > (uint64_t)L) { - throw ANNException("Set L to a value of at least K", -1, __FUNCSIG__, - __FILE__, __LINE__); - } - ScratchStoreManager> manager(_query_scratch); - auto scratch = manager.scratch_space(); - - if (L > scratch->get_L()) { - diskann::cout << "Attempting to expand query scratch_space. Was created " - << "with Lsize: " << scratch->get_L() - << " but search L is: " << L << std::endl; - scratch->resize_for_new_L(L); - diskann::cout << "Resize completed. New scratch->L is " << scratch->get_L() - << std::endl; - } +size_t Index::search_with_tags(const T *query, const uint64_t K, const uint32_t L, TagT *tags, + float *distances, std::vector &res_vectors) +{ + if (K > (uint64_t)L) + { + throw ANNException("Set L to a value of at least K", -1, __FUNCSIG__, __FILE__, __LINE__); + } + ScratchStoreManager> manager(_query_scratch); + auto scratch = manager.scratch_space(); - std::shared_lock ul(_update_lock); + if (L > scratch->get_L()) + { + diskann::cout << "Attempting to expand query scratch_space. Was created " + << "with Lsize: " << scratch->get_L() << " but search L is: " << L << std::endl; + scratch->resize_for_new_L(L); + diskann::cout << "Resize completed. New scratch->L is " << scratch->get_L() << std::endl; + } - const std::vector init_ids = get_init_ids(); - const std::vector unused_filter_label; + std::shared_lock ul(_update_lock); - _distance->preprocess_query(query, _data_store->get_dims(), - scratch->aligned_query()); - iterate_to_fixed_point(scratch->aligned_query(), L, init_ids, scratch, false, - unused_filter_label, true); + const std::vector init_ids = get_init_ids(); + const std::vector unused_filter_label; - NeighborPriorityQueue &best_L_nodes = scratch->best_l_nodes(); - assert(best_L_nodes.size() <= L); + _distance->preprocess_query(query, _data_store->get_dims(), scratch->aligned_query()); + iterate_to_fixed_point(scratch->aligned_query(), L, init_ids, scratch, false, unused_filter_label, true); - std::shared_lock tl(_tag_lock); + NeighborPriorityQueue &best_L_nodes = scratch->best_l_nodes(); + assert(best_L_nodes.size() <= L); - size_t pos = 0; - for (size_t i = 0; i < best_L_nodes.size(); ++i) { - auto node = best_L_nodes[i]; + std::shared_lock tl(_tag_lock); - TagT tag; - if (_location_to_tag.try_get(node.id, tag)) { - tags[pos] = tag; + size_t pos = 0; + for (size_t i = 0; i < best_L_nodes.size(); ++i) + { + auto node = best_L_nodes[i]; + + TagT tag; + if (_location_to_tag.try_get(node.id, tag)) + { + tags[pos] = tag; - if (res_vectors.size() > 0) { - _data_store->get_vector(node.id, res_vectors[pos]); - } + if (res_vectors.size() > 0) + { + _data_store->get_vector(node.id, res_vectors[pos]); + } - if (distances != nullptr) { + if (distances != nullptr) + { #ifdef EXEC_ENV_OLS - distances[pos] = node.distance; // DLVS expects negative distances + distances[pos] = node.distance; // DLVS expects negative distances #else - distances[pos] = - _dist_metric == INNER_PRODUCT ? -1 * node.distance : node.distance; + distances[pos] = _dist_metric == INNER_PRODUCT ? -1 * node.distance : node.distance; #endif - } - pos++; - // If res_vectors.size() < k, clip at the value. - if (pos == K || pos == res_vectors.size()) break; + } + pos++; + // If res_vectors.size() < k, clip at the value. + if (pos == K || pos == res_vectors.size()) + break; + } } - } - return pos; + return pos; } -template -size_t Index::get_num_points() { - std::shared_lock tl(_tag_lock); - return _nd; +template size_t Index::get_num_points() +{ + std::shared_lock tl(_tag_lock); + return _nd; } -template -size_t Index::get_max_points() { - std::shared_lock tl(_tag_lock); - return _max_points; +template size_t Index::get_max_points() +{ + std::shared_lock tl(_tag_lock); + return _max_points; } -template -void Index::generate_frozen_point() { - if (_num_frozen_pts == 0) return; - - if (_num_frozen_pts > 1) { - throw ANNException( - "More than one frozen point not supported in generate_frozen_point", -1, - __FUNCSIG__, __FILE__, __LINE__); - } - - if (_nd == 0) { - throw ANNException("ERROR: Can not pick a frozen point since nd=0", -1, - __FUNCSIG__, __FILE__, __LINE__); - } - size_t res = calculate_entry_point(); - - if (_pq_dist) { - // copy the PQ data corresponding to the point returned by - // calculate_entry_point - memcpy(_pq_data + _max_points * _num_pq_chunks, - _pq_data + res * _num_pq_chunks, - _num_pq_chunks * DIV_ROUND_UP(NUM_PQ_BITS, 8)); - } else { - _data_store->copy_vectors(res, _max_points, 1); - } +template void Index::generate_frozen_point() +{ + if (_num_frozen_pts == 0) + return; + + if (_num_frozen_pts > 1) + { + throw ANNException("More than one frozen point not supported in generate_frozen_point", -1, __FUNCSIG__, + __FILE__, __LINE__); + } + + if (_nd == 0) + { + throw ANNException("ERROR: Can not pick a frozen point since nd=0", -1, __FUNCSIG__, __FILE__, __LINE__); + } + size_t res = calculate_entry_point(); + + if (_pq_dist) + { + // copy the PQ data corresponding to the point returned by + // calculate_entry_point + memcpy(_pq_data + _max_points * _num_pq_chunks, _pq_data + res * _num_pq_chunks, + _num_pq_chunks * DIV_ROUND_UP(NUM_PQ_BITS, 8)); + } + else + { + _data_store->copy_vectors(res, _max_points, 1); + } } -template -int Index::enable_delete() { - assert(_enable_tags); +template int Index::enable_delete() +{ + assert(_enable_tags); - if (!_enable_tags) { - diskann::cerr << "Tags must be instantiated for deletions" << std::endl; - return -2; - } + if (!_enable_tags) + { + diskann::cerr << "Tags must be instantiated for deletions" << std::endl; + return -2; + } - std::unique_lock ul(_update_lock); - std::unique_lock tl(_tag_lock); - std::unique_lock dl(_delete_lock); + std::unique_lock ul(_update_lock); + std::unique_lock tl(_tag_lock); + std::unique_lock dl(_delete_lock); - if (_data_compacted) { - for (uint32_t slot = (uint32_t)_nd; slot < _max_points; ++slot) { - _empty_slots.insert(slot); + if (_data_compacted) + { + for (uint32_t slot = (uint32_t)_nd; slot < _max_points; ++slot) + { + _empty_slots.insert(slot); + } } - } - return 0; + return 0; } template -inline void Index::process_delete( - const tsl::robin_set &old_delete_set, size_t loc, - const uint32_t range, const uint32_t maxc, const float alpha, - InMemQueryScratch *scratch) { - tsl::robin_set &expanded_nodes_set = scratch->expanded_nodes_set(); - std::vector &expanded_nghrs_vec = scratch->expanded_nodes_vec(); - - // If this condition were not true, deadlock could result - assert(old_delete_set.find(loc) == old_delete_set.end()); - - std::vector adj_list; - { - // Acquire and release lock[loc] before acquiring locks for neighbors - std::unique_lock adj_list_lock; - if (_conc_consolidate) - adj_list_lock = std::unique_lock(_locks[loc]); - adj_list = _final_graph[loc]; - } - - bool modify = false; - for (auto ngh : adj_list) { - if (old_delete_set.find(ngh) == old_delete_set.end()) { - expanded_nodes_set.insert(ngh); - } else { - modify = true; - - std::unique_lock ngh_lock; - if (_conc_consolidate) - ngh_lock = std::unique_lock(_locks[ngh]); - for (auto j : _final_graph[ngh]) - if (j != loc && old_delete_set.find(j) == old_delete_set.end()) - expanded_nodes_set.insert(j); - } - } - - if (modify) { - if (expanded_nodes_set.size() <= range) { - std::unique_lock adj_list_lock(_locks[loc]); - _final_graph[loc].clear(); - for (auto &ngh : expanded_nodes_set) _final_graph[loc].push_back(ngh); - } else { - // Create a pool of Neighbor candidates from the expanded_nodes_set - expanded_nghrs_vec.reserve(expanded_nodes_set.size()); - for (auto &ngh : expanded_nodes_set) { - expanded_nghrs_vec.emplace_back(ngh, - _data_store->get_distance(loc, ngh)); - } - std::sort(expanded_nghrs_vec.begin(), expanded_nghrs_vec.end()); - std::vector &occlude_list_output = - scratch->occlude_list_output(); - occlude_list(loc, expanded_nghrs_vec, alpha, range, maxc, - occlude_list_output, scratch, &old_delete_set); - std::unique_lock adj_list_lock(_locks[loc]); - _final_graph[loc] = occlude_list_output; - } - } +inline void Index::process_delete(const tsl::robin_set &old_delete_set, size_t loc, + const uint32_t range, const uint32_t maxc, const float alpha, + InMemQueryScratch *scratch) +{ + tsl::robin_set &expanded_nodes_set = scratch->expanded_nodes_set(); + std::vector &expanded_nghrs_vec = scratch->expanded_nodes_vec(); + + // If this condition were not true, deadlock could result + assert(old_delete_set.find(loc) == old_delete_set.end()); + + std::vector adj_list; + { + // Acquire and release lock[loc] before acquiring locks for neighbors + std::unique_lock adj_list_lock; + if (_conc_consolidate) + adj_list_lock = std::unique_lock(_locks[loc]); + adj_list = _final_graph[loc]; + } + + bool modify = false; + for (auto ngh : adj_list) + { + if (old_delete_set.find(ngh) == old_delete_set.end()) + { + expanded_nodes_set.insert(ngh); + } + else + { + modify = true; + + std::unique_lock ngh_lock; + if (_conc_consolidate) + ngh_lock = std::unique_lock(_locks[ngh]); + for (auto j : _final_graph[ngh]) + if (j != loc && old_delete_set.find(j) == old_delete_set.end()) + expanded_nodes_set.insert(j); + } + } + + if (modify) + { + if (expanded_nodes_set.size() <= range) + { + std::unique_lock adj_list_lock(_locks[loc]); + _final_graph[loc].clear(); + for (auto &ngh : expanded_nodes_set) + _final_graph[loc].push_back(ngh); + } + else + { + // Create a pool of Neighbor candidates from the expanded_nodes_set + expanded_nghrs_vec.reserve(expanded_nodes_set.size()); + for (auto &ngh : expanded_nodes_set) + { + expanded_nghrs_vec.emplace_back(ngh, _data_store->get_distance(loc, ngh)); + } + std::sort(expanded_nghrs_vec.begin(), expanded_nghrs_vec.end()); + std::vector &occlude_list_output = scratch->occlude_list_output(); + occlude_list(loc, expanded_nghrs_vec, alpha, range, maxc, occlude_list_output, scratch, &old_delete_set); + std::unique_lock adj_list_lock(_locks[loc]); + _final_graph[loc] = occlude_list_output; + } + } } // Returns number of live points left after consolidation template -consolidation_report Index::consolidate_deletes( - const IndexWriteParameters ¶ms) { - if (!_enable_tags) - throw diskann::ANNException("Point tag array not instantiated", -1, - __FUNCSIG__, __FILE__, __LINE__); +consolidation_report Index::consolidate_deletes(const IndexWriteParameters ¶ms) +{ + if (!_enable_tags) + throw diskann::ANNException("Point tag array not instantiated", -1, __FUNCSIG__, __FILE__, __LINE__); - { - std::shared_lock ul(_update_lock); - std::shared_lock tl(_tag_lock); - std::shared_lock dl(_delete_lock); - if (_empty_slots.size() + _nd != _max_points) { - std::string err = "#empty slots + nd != max points"; - diskann::cerr << err << std::endl; - throw ANNException(err, -1, __FUNCSIG__, __FILE__, __LINE__); - } - - if (_location_to_tag.size() + _delete_set->size() != _nd) { - diskann::cerr << "Error: _location_to_tag.size (" - << _location_to_tag.size() << ") + _delete_set->size (" - << _delete_set->size() << ") != _nd(" << _nd << ") "; - return consolidation_report( - diskann::consolidation_report::status_code::INCONSISTENT_COUNT_ERROR, - 0, 0, 0, 0, 0, 0, 0); - } - - if (_location_to_tag.size() != _tag_to_location.size()) { - throw diskann::ANNException( - "_location_to_tag and _tag_to_location not of same size", -1, - __FUNCSIG__, __FILE__, __LINE__); - } - } - - std::unique_lock update_lock(_update_lock, - std::defer_lock); - if (!_conc_consolidate) update_lock.lock(); - - std::unique_lock cl(_consolidate_lock, - std::defer_lock); - if (!cl.try_lock()) { - diskann::cerr - << "Consildate delete function failed to acquire consolidate lock" - << std::endl; - return consolidation_report( - diskann::consolidation_report::status_code::LOCK_FAIL, 0, 0, 0, 0, 0, 0, - 0); - } - - diskann::cout << "Starting consolidate_deletes... "; - - std::unique_ptr> old_delete_set( - new tsl::robin_set); - { - std::unique_lock dl(_delete_lock); - std::swap(_delete_set, old_delete_set); - } - - if (old_delete_set->find(_start) != old_delete_set->end()) { - throw diskann::ANNException("ERROR: start node has been deleted", -1, - __FUNCSIG__, __FILE__, __LINE__); - } - - const uint32_t range = params.max_degree; - const uint32_t maxc = params.max_occlusion_size; - const float alpha = params.alpha; - const uint32_t num_threads = - params.num_threads == 0 ? omp_get_num_threads() : params.num_threads; - - uint32_t num_calls_to_process_delete = 0; - diskann::Timer timer; + { + std::shared_lock ul(_update_lock); + std::shared_lock tl(_tag_lock); + std::shared_lock dl(_delete_lock); + if (_empty_slots.size() + _nd != _max_points) + { + std::string err = "#empty slots + nd != max points"; + diskann::cerr << err << std::endl; + throw ANNException(err, -1, __FUNCSIG__, __FILE__, __LINE__); + } + + if (_location_to_tag.size() + _delete_set->size() != _nd) + { + diskann::cerr << "Error: _location_to_tag.size (" << _location_to_tag.size() << ") + _delete_set->size (" + << _delete_set->size() << ") != _nd(" << _nd << ") "; + return consolidation_report(diskann::consolidation_report::status_code::INCONSISTENT_COUNT_ERROR, 0, 0, 0, + 0, 0, 0, 0); + } + + if (_location_to_tag.size() != _tag_to_location.size()) + { + throw diskann::ANNException("_location_to_tag and _tag_to_location not of same size", -1, __FUNCSIG__, + __FILE__, __LINE__); + } + } + + std::unique_lock update_lock(_update_lock, std::defer_lock); + if (!_conc_consolidate) + update_lock.lock(); + + std::unique_lock cl(_consolidate_lock, std::defer_lock); + if (!cl.try_lock()) + { + diskann::cerr << "Consildate delete function failed to acquire consolidate lock" << std::endl; + return consolidation_report(diskann::consolidation_report::status_code::LOCK_FAIL, 0, 0, 0, 0, 0, 0, 0); + } + + diskann::cout << "Starting consolidate_deletes... "; + + std::unique_ptr> old_delete_set(new tsl::robin_set); + { + std::unique_lock dl(_delete_lock); + std::swap(_delete_set, old_delete_set); + } + + if (old_delete_set->find(_start) != old_delete_set->end()) + { + throw diskann::ANNException("ERROR: start node has been deleted", -1, __FUNCSIG__, __FILE__, __LINE__); + } + + const uint32_t range = params.max_degree; + const uint32_t maxc = params.max_occlusion_size; + const float alpha = params.alpha; + const uint32_t num_threads = params.num_threads == 0 ? omp_get_num_threads() : params.num_threads; + + uint32_t num_calls_to_process_delete = 0; + diskann::Timer timer; #pragma omp parallel for num_threads(num_threads) schedule(dynamic, 8192) reduction(+ : num_calls_to_process_delete) - for (int64_t loc = 0; loc < (int64_t)_max_points; loc++) { - if (old_delete_set->find((uint32_t)loc) == old_delete_set->end() && - !_empty_slots.is_in_set((uint32_t)loc)) { - ScratchStoreManager> manager(_query_scratch); - auto scratch = manager.scratch_space(); - process_delete(*old_delete_set, loc, range, maxc, alpha, scratch); - num_calls_to_process_delete += 1; - } - } - for (int64_t loc = _max_points; - loc < (int64_t)(_max_points + _num_frozen_pts); loc++) { - ScratchStoreManager> manager(_query_scratch); - auto scratch = manager.scratch_space(); - process_delete(*old_delete_set, loc, range, maxc, alpha, scratch); - num_calls_to_process_delete += 1; - } + for (int64_t loc = 0; loc < (int64_t)_max_points; loc++) + { + if (old_delete_set->find((uint32_t)loc) == old_delete_set->end() && !_empty_slots.is_in_set((uint32_t)loc)) + { + ScratchStoreManager> manager(_query_scratch); + auto scratch = manager.scratch_space(); + process_delete(*old_delete_set, loc, range, maxc, alpha, scratch); + num_calls_to_process_delete += 1; + } + } + for (int64_t loc = _max_points; loc < (int64_t)(_max_points + _num_frozen_pts); loc++) + { + ScratchStoreManager> manager(_query_scratch); + auto scratch = manager.scratch_space(); + process_delete(*old_delete_set, loc, range, maxc, alpha, scratch); + num_calls_to_process_delete += 1; + } - std::unique_lock tl(_tag_lock); - size_t ret_nd = release_locations(*old_delete_set); - size_t max_points = _max_points; - size_t empty_slots_size = _empty_slots.size(); + std::unique_lock tl(_tag_lock); + size_t ret_nd = release_locations(*old_delete_set); + size_t max_points = _max_points; + size_t empty_slots_size = _empty_slots.size(); - std::shared_lock dl(_delete_lock); - size_t delete_set_size = _delete_set->size(); - size_t old_delete_set_size = old_delete_set->size(); + std::shared_lock dl(_delete_lock); + size_t delete_set_size = _delete_set->size(); + size_t old_delete_set_size = old_delete_set->size(); - if (!_conc_consolidate) { - update_lock.unlock(); - } + if (!_conc_consolidate) + { + update_lock.unlock(); + } - double duration = timer.elapsed() / 1000000.0; - diskann::cout << " done in " << duration << " seconds." << std::endl; - return consolidation_report( - diskann::consolidation_report::status_code::SUCCESS, ret_nd, max_points, - empty_slots_size, old_delete_set_size, delete_set_size, - num_calls_to_process_delete, duration); + double duration = timer.elapsed() / 1000000.0; + diskann::cout << " done in " << duration << " seconds." << std::endl; + return consolidation_report(diskann::consolidation_report::status_code::SUCCESS, ret_nd, max_points, + empty_slots_size, old_delete_set_size, delete_set_size, num_calls_to_process_delete, + duration); } -template -void Index::compact_frozen_point() { - if (_nd < _max_points && _num_frozen_pts > 0) { - reposition_points(_max_points, _nd, _num_frozen_pts); - _start = (uint32_t)_nd; - } +template void Index::compact_frozen_point() +{ + if (_nd < _max_points && _num_frozen_pts > 0) + { + reposition_points(_max_points, _nd, _num_frozen_pts); + _start = (uint32_t)_nd; + } } // Should be called after acquiring _update_lock -template -void Index::compact_data() { - if (!_dynamic_index) - throw ANNException("Can not compact a non-dynamic index", -1, __FUNCSIG__, - __FILE__, __LINE__); - - if (_data_compacted) { - diskann::cerr - << "Warning! Calling compact_data() when _data_compacted is true!" - << std::endl; - return; - } - - if (_delete_set->size() > 0) { - throw ANNException( - "Can not compact data when index has non-empty _delete_set of " - "size: " + - std::to_string(_delete_set->size()), - -1, __FUNCSIG__, __FILE__, __LINE__); - } - - diskann::Timer timer; - - std::vector new_location = - std::vector(_max_points + _num_frozen_pts, UINT32_MAX); - - uint32_t new_counter = 0; - std::set empty_locations; - for (uint32_t old_location = 0; old_location < _max_points; old_location++) { - if (_location_to_tag.contains(old_location)) { - new_location[old_location] = new_counter; - new_counter++; - } else { - empty_locations.insert(old_location); - } - } - for (uint32_t old_location = _max_points; - old_location < _max_points + _num_frozen_pts; old_location++) { - new_location[old_location] = old_location; - } - - // If start node is removed, throw an exception - if (_start < _max_points && !_location_to_tag.contains(_start)) { - throw diskann::ANNException("ERROR: Start node deleted.", -1, __FUNCSIG__, - __FILE__, __LINE__); - } - - size_t num_dangling = 0; - for (uint32_t old = 0; old < _max_points + _num_frozen_pts; ++old) { - std::vector new_adj_list; - - if ((new_location[old] < _max_points) // If point continues to exist - || (old >= _max_points && old < _max_points + _num_frozen_pts)) { - new_adj_list.reserve(_final_graph[old].size()); - for (auto ngh_iter : _final_graph[old]) { - if (empty_locations.find(ngh_iter) != empty_locations.end()) { - ++num_dangling; - diskann::cerr << "Error in compact_data(). _final_graph[" << old - << "] has neighbor " << ngh_iter - << " which is a location not associated with any tag." - << std::endl; - } else { - new_adj_list.push_back(new_location[ngh_iter]); - } - } - _final_graph[old].swap(new_adj_list); - - // Move the data and adj list to the correct position - if (new_location[old] != old) { - assert(new_location[old] < old); - _final_graph[new_location[old]].swap(_final_graph[old]); - - _data_store->copy_vectors(old, new_location[old], 1); - } - } else { - _final_graph[old].clear(); - } - } - diskann::cerr << "#dangling references after data compaction: " - << num_dangling << std::endl; - - _tag_to_location.clear(); - for (auto pos = _location_to_tag.find_first(); pos.is_valid(); - pos = _location_to_tag.find_next(pos)) { - const auto tag = _location_to_tag.get(pos); - _tag_to_location[tag] = new_location[pos._key]; - } - _location_to_tag.clear(); - for (const auto &iter : _tag_to_location) { - _location_to_tag.set(iter.second, iter.first); - } - - for (size_t old = _nd; old < _max_points; ++old) { - _final_graph[old].clear(); - } - _empty_slots.clear(); - for (auto i = _nd; i < _max_points; i++) { - _empty_slots.insert((uint32_t)i); - } - _data_compacted = true; - diskann::cout << "Time taken for compact_data: " << timer.elapsed() / 1000000. - << "s." << std::endl; +template void Index::compact_data() +{ + if (!_dynamic_index) + throw ANNException("Can not compact a non-dynamic index", -1, __FUNCSIG__, __FILE__, __LINE__); + + if (_data_compacted) + { + diskann::cerr << "Warning! Calling compact_data() when _data_compacted is true!" << std::endl; + return; + } + + if (_delete_set->size() > 0) + { + throw ANNException("Can not compact data when index has non-empty _delete_set of " + "size: " + + std::to_string(_delete_set->size()), + -1, __FUNCSIG__, __FILE__, __LINE__); + } + + diskann::Timer timer; + + std::vector new_location = std::vector(_max_points + _num_frozen_pts, UINT32_MAX); + + uint32_t new_counter = 0; + std::set empty_locations; + for (uint32_t old_location = 0; old_location < _max_points; old_location++) + { + if (_location_to_tag.contains(old_location)) + { + new_location[old_location] = new_counter; + new_counter++; + } + else + { + empty_locations.insert(old_location); + } + } + for (uint32_t old_location = _max_points; old_location < _max_points + _num_frozen_pts; old_location++) + { + new_location[old_location] = old_location; + } + + // If start node is removed, throw an exception + if (_start < _max_points && !_location_to_tag.contains(_start)) + { + throw diskann::ANNException("ERROR: Start node deleted.", -1, __FUNCSIG__, __FILE__, __LINE__); + } + + size_t num_dangling = 0; + for (uint32_t old = 0; old < _max_points + _num_frozen_pts; ++old) + { + std::vector new_adj_list; + + if ((new_location[old] < _max_points) // If point continues to exist + || (old >= _max_points && old < _max_points + _num_frozen_pts)) + { + new_adj_list.reserve(_final_graph[old].size()); + for (auto ngh_iter : _final_graph[old]) + { + if (empty_locations.find(ngh_iter) != empty_locations.end()) + { + ++num_dangling; + diskann::cerr << "Error in compact_data(). _final_graph[" << old << "] has neighbor " << ngh_iter + << " which is a location not associated with any tag." << std::endl; + } + else + { + new_adj_list.push_back(new_location[ngh_iter]); + } + } + _final_graph[old].swap(new_adj_list); + + // Move the data and adj list to the correct position + if (new_location[old] != old) + { + assert(new_location[old] < old); + _final_graph[new_location[old]].swap(_final_graph[old]); + + _data_store->copy_vectors(old, new_location[old], 1); + } + } + else + { + _final_graph[old].clear(); + } + } + diskann::cerr << "#dangling references after data compaction: " << num_dangling << std::endl; + + _tag_to_location.clear(); + for (auto pos = _location_to_tag.find_first(); pos.is_valid(); pos = _location_to_tag.find_next(pos)) + { + const auto tag = _location_to_tag.get(pos); + _tag_to_location[tag] = new_location[pos._key]; + } + _location_to_tag.clear(); + for (const auto &iter : _tag_to_location) + { + _location_to_tag.set(iter.second, iter.first); + } + + for (size_t old = _nd; old < _max_points; ++old) + { + _final_graph[old].clear(); + } + _empty_slots.clear(); + for (auto i = _nd; i < _max_points; i++) + { + _empty_slots.insert((uint32_t)i); + } + _data_compacted = true; + diskann::cout << "Time taken for compact_data: " << timer.elapsed() / 1000000. << "s." << std::endl; } // // Caller must hold unique _tag_lock and _delete_lock before calling this // -template -int Index::reserve_location() { - if (_nd >= _max_points) { - return -1; - } - uint32_t location; - if (_data_compacted && _empty_slots.is_empty()) { - // This code path is encountered when enable_delete hasn't been - // called yet, so no points have been deleted and _empty_slots - // hasn't been filled in. In that case, just keep assigning - // consecutive locations. - location = (uint32_t)_nd; - } else { - assert(_empty_slots.size() != 0); - assert(_empty_slots.size() + _nd == _max_points); - - location = _empty_slots.pop_any(); - _delete_set->erase(location); - } - - ++_nd; - return location; -} +template int Index::reserve_location() +{ + if (_nd >= _max_points) + { + return -1; + } + uint32_t location; + if (_data_compacted && _empty_slots.is_empty()) + { + // This code path is encountered when enable_delete hasn't been + // called yet, so no points have been deleted and _empty_slots + // hasn't been filled in. In that case, just keep assigning + // consecutive locations. + location = (uint32_t)_nd; + } + else + { + assert(_empty_slots.size() != 0); + assert(_empty_slots.size() + _nd == _max_points); -template -size_t Index::release_location(int location) { - if (_empty_slots.is_in_set(location)) - throw ANNException( - "Trying to release location, but location already in empty slots", -1, - __FUNCSIG__, __FILE__, __LINE__); - _empty_slots.insert(location); + location = _empty_slots.pop_any(); + _delete_set->erase(location); + } - _nd--; - return _nd; + ++_nd; + return location; } -template -size_t Index::release_locations( - const tsl::robin_set &locations) { - for (auto location : locations) { +template size_t Index::release_location(int location) +{ if (_empty_slots.is_in_set(location)) - throw ANNException( - "Trying to release location, but location " - "already in empty slots", - -1, __FUNCSIG__, __FILE__, __LINE__); + throw ANNException("Trying to release location, but location already in empty slots", -1, __FUNCSIG__, __FILE__, + __LINE__); _empty_slots.insert(location); _nd--; - } - - if (_empty_slots.size() + _nd != _max_points) - throw ANNException("#empty slots + nd != max points", -1, __FUNCSIG__, - __FILE__, __LINE__); - - return _nd; + return _nd; } template -void Index::reposition_points(uint32_t old_location_start, - uint32_t new_location_start, - uint32_t num_locations) { - if (num_locations == 0 || old_location_start == new_location_start) { - return; - } - - // Update pointers to the moved nodes. Note: the computation is correct even - // when new_location_start < old_location_start given the C++ uint32_t - // integer arithmetic rules. - const uint32_t location_delta = new_location_start - old_location_start; - - for (uint32_t i = 0; i < _max_points + _num_frozen_pts; i++) - for (auto &loc : _final_graph[i]) - if (loc >= old_location_start && loc < old_location_start + num_locations) - loc += location_delta; - - // 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; - - // Move the adjacency lists. Make sure that overlapping ranges are handled - // correctly. - if (new_location_start < old_location_start) { - // New location before the old location: copy the entries in order - // to avoid modifying locations that are yet to be copied. - for (uint32_t loc_offset = 0; loc_offset < num_locations; loc_offset++) { - assert(_final_graph[new_location_start + loc_offset].empty()); - _final_graph[new_location_start + loc_offset].swap( - _final_graph[old_location_start + loc_offset]); - } - - // 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 { - // Old location after the new location: copy from the end of the range - // to avoid modifying locations that are yet to be copied. - for (uint32_t loc_offset = num_locations; loc_offset > 0; loc_offset--) { - assert(_final_graph[new_location_start + loc_offset - 1u].empty()); - _final_graph[new_location_start + loc_offset - 1u].swap( - _final_graph[old_location_start + loc_offset - 1u]); - } - - // 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; - } - } - _data_store->move_vectors(old_location_start, new_location_start, - num_locations); -} +size_t Index::release_locations(const tsl::robin_set &locations) +{ + for (auto location : locations) + { + if (_empty_slots.is_in_set(location)) + throw ANNException("Trying to release location, but location " + "already in empty slots", + -1, __FUNCSIG__, __FILE__, __LINE__); + _empty_slots.insert(location); -template -void Index::reposition_frozen_point_to_end() { - if (_num_frozen_pts == 0) return; + _nd--; + } - if (_nd == _max_points) { - diskann::cout - << "Not repositioning frozen point as it is already at the end." - << std::endl; - return; - } + if (_empty_slots.size() + _nd != _max_points) + throw ANNException("#empty slots + nd != max points", -1, __FUNCSIG__, __FILE__, __LINE__); - reposition_points((uint32_t)_nd, (uint32_t)_max_points, - (uint32_t)_num_frozen_pts); - _start = (uint32_t)_max_points; + return _nd; } template -void Index::resize(size_t new_max_points) { - const size_t new_internal_points = new_max_points + _num_frozen_pts; - auto start = std::chrono::high_resolution_clock::now(); - assert(_empty_slots.size() == - 0); // should not resize if there are empty slots. - - _data_store->resize(new_internal_points); - _final_graph.resize(new_internal_points); - _locks = std::vector(new_internal_points); - - if (_num_frozen_pts != 0) { - reposition_points((uint32_t)_max_points, (uint32_t)new_max_points, - (uint32_t)_num_frozen_pts); - _start = (uint32_t)new_max_points; - } - - _max_points = new_max_points; - _empty_slots.reserve(_max_points); - for (auto i = _nd; i < _max_points; i++) { - _empty_slots.insert((uint32_t)i); - } - - auto stop = std::chrono::high_resolution_clock::now(); - diskann::cout << "Resizing took: " - << std::chrono::duration(stop - start).count() << "s" - << std::endl; +void Index::reposition_points(uint32_t old_location_start, uint32_t new_location_start, + uint32_t num_locations) +{ + if (num_locations == 0 || old_location_start == new_location_start) + { + return; + } + + // Update pointers to the moved nodes. Note: the computation is correct even + // when new_location_start < old_location_start given the C++ uint32_t + // integer arithmetic rules. + const uint32_t location_delta = new_location_start - old_location_start; + + for (uint32_t i = 0; i < _max_points + _num_frozen_pts; i++) + for (auto &loc : _final_graph[i]) + if (loc >= old_location_start && loc < old_location_start + num_locations) + loc += location_delta; + + // 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; + + // Move the adjacency lists. Make sure that overlapping ranges are handled + // correctly. + if (new_location_start < old_location_start) + { + // New location before the old location: copy the entries in order + // to avoid modifying locations that are yet to be copied. + for (uint32_t loc_offset = 0; loc_offset < num_locations; loc_offset++) + { + assert(_final_graph[new_location_start + loc_offset].empty()); + _final_graph[new_location_start + loc_offset].swap(_final_graph[old_location_start + loc_offset]); + } + + // 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 + { + // Old location after the new location: copy from the end of the range + // to avoid modifying locations that are yet to be copied. + for (uint32_t loc_offset = num_locations; loc_offset > 0; loc_offset--) + { + assert(_final_graph[new_location_start + loc_offset - 1u].empty()); + _final_graph[new_location_start + loc_offset - 1u].swap(_final_graph[old_location_start + loc_offset - 1u]); + } + + // 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; + } + } + _data_store->move_vectors(old_location_start, new_location_start, num_locations); } -template -int Index::insert_point(const T *point, const TagT tag) { - assert(_has_built); - if (tag == static_cast(0)) { - throw diskann::ANNException( - "Do not insert point with tag 0. That is " - "reserved for points hidden " - "from the user.", - -1, __FUNCSIG__, __FILE__, __LINE__); - } - - std::shared_lock shared_ul(_update_lock); - std::unique_lock tl(_tag_lock); - std::unique_lock dl(_delete_lock); - - // Find a vacant location in the data array to insert the new point - auto location = reserve_location(); - if (location == -1) { -#if EXPAND_IF_FULL - dl.unlock(); - tl.unlock(); - shared_ul.unlock(); +template void Index::reposition_frozen_point_to_end() +{ + if (_num_frozen_pts == 0) + return; + if (_nd == _max_points) { - std::unique_lock ul(_update_lock); - tl.lock(); - dl.lock(); + diskann::cout << "Not repositioning frozen point as it is already at the end." << std::endl; + return; + } + + reposition_points((uint32_t)_nd, (uint32_t)_max_points, (uint32_t)_num_frozen_pts); + _start = (uint32_t)_max_points; +} - if (_nd >= _max_points) { - auto new_max_points = (size_t)(_max_points * INDEX_GROWTH_FACTOR); - resize(new_max_points); - } +template void Index::resize(size_t new_max_points) +{ + const size_t new_internal_points = new_max_points + _num_frozen_pts; + auto start = std::chrono::high_resolution_clock::now(); + assert(_empty_slots.size() == 0); // should not resize if there are empty slots. - dl.unlock(); - tl.unlock(); - ul.unlock(); + _data_store->resize(new_internal_points); + _final_graph.resize(new_internal_points); + _locks = std::vector(new_internal_points); + + if (_num_frozen_pts != 0) + { + reposition_points((uint32_t)_max_points, (uint32_t)new_max_points, (uint32_t)_num_frozen_pts); + _start = (uint32_t)new_max_points; + } + + _max_points = new_max_points; + _empty_slots.reserve(_max_points); + for (auto i = _nd; i < _max_points; i++) + { + _empty_slots.insert((uint32_t)i); } - shared_ul.lock(); - tl.lock(); - dl.lock(); + auto stop = std::chrono::high_resolution_clock::now(); + diskann::cout << "Resizing took: " << std::chrono::duration(stop - start).count() << "s" << std::endl; +} - location = reserve_location(); - if (location == -1) { - throw diskann::ANNException( - "Cannot reserve location even after " - "expanding graph. Terminating.", - -1, __FUNCSIG__, __FILE__, __LINE__); +template +int Index::insert_point(const T *point, const TagT tag) +{ + assert(_has_built); + if (tag == static_cast(0)) + { + throw diskann::ANNException("Do not insert point with tag 0. That is " + "reserved for points hidden " + "from the user.", + -1, __FUNCSIG__, __FILE__, __LINE__); } + + std::shared_lock shared_ul(_update_lock); + std::unique_lock tl(_tag_lock); + std::unique_lock dl(_delete_lock); + + // Find a vacant location in the data array to insert the new point + auto location = reserve_location(); + if (location == -1) + { +#if EXPAND_IF_FULL + dl.unlock(); + tl.unlock(); + shared_ul.unlock(); + + { + std::unique_lock ul(_update_lock); + tl.lock(); + dl.lock(); + + if (_nd >= _max_points) + { + auto new_max_points = (size_t)(_max_points * INDEX_GROWTH_FACTOR); + resize(new_max_points); + } + + dl.unlock(); + tl.unlock(); + ul.unlock(); + } + + shared_ul.lock(); + tl.lock(); + dl.lock(); + + location = reserve_location(); + if (location == -1) + { + throw diskann::ANNException("Cannot reserve location even after " + "expanding graph. Terminating.", + -1, __FUNCSIG__, __FILE__, __LINE__); + } #else - return -1; + return -1; #endif - } - dl.unlock(); - - // Insert tag and mapping to location - if (_enable_tags) { - if (_tag_to_location.find(tag) != _tag_to_location.end()) { - release_location(location); - return -1; } + dl.unlock(); - _tag_to_location[tag] = location; - _location_to_tag.set(location, tag); - } - tl.unlock(); - - _data_store->set_vector(location, point); + // Insert tag and mapping to location + if (_enable_tags) + { + if (_tag_to_location.find(tag) != _tag_to_location.end()) + { + release_location(location); + return -1; + } - // Find and add appropriate graph edges - ScratchStoreManager> manager(_query_scratch); - auto scratch = manager.scratch_space(); - std::vector pruned_list; - if (_filtered_index) { - search_for_point_and_prune(location, _indexingQueueSize, pruned_list, - scratch, true, _filterIndexingQueueSize); - } else { - search_for_point_and_prune(location, _indexingQueueSize, pruned_list, - scratch); - } - { - std::shared_lock tlock(_tag_lock, std::defer_lock); - if (_conc_consolidate) tlock.lock(); + _tag_to_location[tag] = location; + _location_to_tag.set(location, tag); + } + tl.unlock(); - LockGuard guard(_locks[location]); - _final_graph[location].clear(); - _final_graph[location].reserve( - (size_t)(_indexingRange * GRAPH_SLACK_FACTOR * 1.05)); + _data_store->set_vector(location, point); - for (auto link : pruned_list) { - if (_conc_consolidate) - if (!_location_to_tag.contains(link)) continue; - _final_graph[location].emplace_back(link); + // Find and add appropriate graph edges + ScratchStoreManager> manager(_query_scratch); + auto scratch = manager.scratch_space(); + std::vector pruned_list; + if (_filtered_index) + { + search_for_point_and_prune(location, _indexingQueueSize, pruned_list, scratch, true, _filterIndexingQueueSize); } - assert(_final_graph[location].size() <= _indexingRange); + else + { + search_for_point_and_prune(location, _indexingQueueSize, pruned_list, scratch); + } + { + std::shared_lock tlock(_tag_lock, std::defer_lock); + if (_conc_consolidate) + tlock.lock(); + + LockGuard guard(_locks[location]); + _final_graph[location].clear(); + _final_graph[location].reserve((size_t)(_indexingRange * GRAPH_SLACK_FACTOR * 1.05)); + + for (auto link : pruned_list) + { + if (_conc_consolidate) + if (!_location_to_tag.contains(link)) + continue; + _final_graph[location].emplace_back(link); + } + assert(_final_graph[location].size() <= _indexingRange); - if (_conc_consolidate) tlock.unlock(); - } + if (_conc_consolidate) + tlock.unlock(); + } - inter_insert(location, pruned_list, scratch); + inter_insert(location, pruned_list, scratch); - return 0; + return 0; } -template -int Index::lazy_delete(const TagT &tag) { - std::shared_lock ul(_update_lock); - std::unique_lock tl(_tag_lock); - std::unique_lock dl(_delete_lock); - _data_compacted = false; +template int Index::lazy_delete(const TagT &tag) +{ + std::shared_lock ul(_update_lock); + std::unique_lock tl(_tag_lock); + std::unique_lock dl(_delete_lock); + _data_compacted = false; - if (_tag_to_location.find(tag) == _tag_to_location.end()) { - diskann::cerr << "Delete tag not found " << tag << std::endl; - return -1; - } - assert(_tag_to_location[tag] < _max_points); + if (_tag_to_location.find(tag) == _tag_to_location.end()) + { + diskann::cerr << "Delete tag not found " << tag << std::endl; + return -1; + } + assert(_tag_to_location[tag] < _max_points); - const auto location = _tag_to_location[tag]; - _delete_set->insert(location); - _location_to_tag.erase(location); - _tag_to_location.erase(tag); + const auto location = _tag_to_location[tag]; + _delete_set->insert(location); + _location_to_tag.erase(location); + _tag_to_location.erase(tag); - return 0; + return 0; } template -void Index::lazy_delete(const std::vector &tags, - std::vector &failed_tags) { - if (failed_tags.size() > 0) { - throw ANNException("failed_tags should be passed as an empty list", -1, - __FUNCSIG__, __FILE__, __LINE__); - } - std::shared_lock ul(_update_lock); - std::unique_lock tl(_tag_lock); - std::unique_lock dl(_delete_lock); - _data_compacted = false; - - for (auto tag : tags) { - if (_tag_to_location.find(tag) == _tag_to_location.end()) { - failed_tags.push_back(tag); - } else { - const auto location = _tag_to_location[tag]; - _delete_set->insert(location); - _location_to_tag.erase(location); - _tag_to_location.erase(tag); - } - } +void Index::lazy_delete(const std::vector &tags, std::vector &failed_tags) +{ + if (failed_tags.size() > 0) + { + throw ANNException("failed_tags should be passed as an empty list", -1, __FUNCSIG__, __FILE__, __LINE__); + } + std::shared_lock ul(_update_lock); + std::unique_lock tl(_tag_lock); + std::unique_lock dl(_delete_lock); + _data_compacted = false; + + for (auto tag : tags) + { + if (_tag_to_location.find(tag) == _tag_to_location.end()) + { + failed_tags.push_back(tag); + } + else + { + const auto location = _tag_to_location[tag]; + _delete_set->insert(location); + _location_to_tag.erase(location); + _tag_to_location.erase(tag); + } + } } -template -bool Index::is_index_saved() { - return _is_saved; +template bool Index::is_index_saved() +{ + return _is_saved; } template -void Index::get_active_tags( - tsl::robin_set &active_tags) { - active_tags.clear(); - std::shared_lock tl(_tag_lock); - for (auto iter : _tag_to_location) { - active_tags.insert(iter.first); - } +void Index::get_active_tags(tsl::robin_set &active_tags) +{ + active_tags.clear(); + std::shared_lock tl(_tag_lock); + for (auto iter : _tag_to_location) + { + active_tags.insert(iter.first); + } } -template -void Index::print_status() { - std::shared_lock ul(_update_lock); - std::shared_lock cl(_consolidate_lock); - std::shared_lock tl(_tag_lock); - std::shared_lock dl(_delete_lock); - - 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 - << std::endl; - diskann::cout << "---------------------------------------------------------" - "------------" - << std::endl; +template void Index::print_status() +{ + std::shared_lock ul(_update_lock); + std::shared_lock cl(_consolidate_lock); + std::shared_lock tl(_tag_lock); + std::shared_lock dl(_delete_lock); + + 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 << std::endl; + diskann::cout << "---------------------------------------------------------" + "------------" + << std::endl; } -template -void Index::count_nodes_at_bfs_levels() { - std::unique_lock ul(_update_lock); +template void Index::count_nodes_at_bfs_levels() +{ + std::unique_lock ul(_update_lock); - boost::dynamic_bitset<> visited(_max_points + _num_frozen_pts); + boost::dynamic_bitset<> visited(_max_points + _num_frozen_pts); - size_t MAX_BFS_LEVELS = 32; - auto bfs_sets = new tsl::robin_set[MAX_BFS_LEVELS]; + size_t MAX_BFS_LEVELS = 32; + auto bfs_sets = new tsl::robin_set[MAX_BFS_LEVELS]; - bfs_sets[0].insert(_start); - visited.set(_start); + bfs_sets[0].insert(_start); + visited.set(_start); - for (uint32_t i = _max_points; i < _max_points + _num_frozen_pts; ++i) { - if (i != _start) { - bfs_sets[0].insert(i); - visited.set(i); + for (uint32_t i = _max_points; i < _max_points + _num_frozen_pts; ++i) + { + if (i != _start) + { + bfs_sets[0].insert(i); + visited.set(i); + } } - } - for (size_t l = 0; l < MAX_BFS_LEVELS - 1; ++l) { - diskann::cout << "Number of nodes at BFS level " << l << " is " - << bfs_sets[l].size() << std::endl; - if (bfs_sets[l].size() == 0) break; - for (auto node : bfs_sets[l]) { - for (auto nghbr : _final_graph[node]) { - if (!visited.test(nghbr)) { - visited.set(nghbr); - bfs_sets[l + 1].insert(nghbr); + for (size_t l = 0; l < MAX_BFS_LEVELS - 1; ++l) + { + diskann::cout << "Number of nodes at BFS level " << l << " is " << bfs_sets[l].size() << std::endl; + if (bfs_sets[l].size() == 0) + break; + for (auto node : bfs_sets[l]) + { + for (auto nghbr : _final_graph[node]) + { + if (!visited.test(nghbr)) + { + visited.set(nghbr); + bfs_sets[l + 1].insert(nghbr); + } + } } - } } - } - delete[] bfs_sets; + delete[] bfs_sets; } // REFACTOR: This should be an OptimizedDataStore class, dummy impl here for @@ -2816,42 +2951,38 @@ void Index::count_nodes_at_bfs_levels() { //} // REFACTOR: This should be an OptimizedDataStore class -template -void Index::optimize_index_layout() { // use after build or load - if (_dynamic_index) { - throw diskann::ANNException( - "Optimize_index_layout not implemented for dyanmic indices", -1, - __FUNCSIG__, __FILE__, __LINE__); - } - - float *cur_vec = new float[_data_store->get_aligned_dim()]; - std::memset(cur_vec, 0, _data_store->get_aligned_dim() * sizeof(float)); - _data_len = (_data_store->get_aligned_dim() + 1) * sizeof(float); - _neighbor_len = (_max_observed_degree + 1) * sizeof(uint32_t); - _node_size = _data_len + _neighbor_len; - _opt_graph = new char[_node_size * _nd]; - DistanceFastL2 *dist_fast = - (DistanceFastL2 *)_data_store->get_dist_fn(); - for (uint32_t i = 0; i < _nd; i++) { - char *cur_node_offset = _opt_graph + i * _node_size; - _data_store->get_vector(i, (T *)cur_vec); - float cur_norm = - dist_fast->norm((T *)cur_vec, _data_store->get_aligned_dim()); - std::memcpy(cur_node_offset, &cur_norm, sizeof(float)); - std::memcpy(cur_node_offset + sizeof(float), cur_vec, - _data_len - sizeof(float)); - - cur_node_offset += _data_len; - uint32_t k = _final_graph[i].size(); - std::memcpy(cur_node_offset, &k, sizeof(uint32_t)); - std::memcpy(cur_node_offset + sizeof(uint32_t), _final_graph[i].data(), - k * sizeof(uint32_t)); - std::vector().swap(_final_graph[i]); - } - _final_graph.clear(); - _final_graph.shrink_to_fit(); - delete[] cur_vec; +template void Index::optimize_index_layout() +{ // use after build or load + if (_dynamic_index) + { + throw diskann::ANNException("Optimize_index_layout not implemented for dyanmic indices", -1, __FUNCSIG__, + __FILE__, __LINE__); + } + + float *cur_vec = new float[_data_store->get_aligned_dim()]; + std::memset(cur_vec, 0, _data_store->get_aligned_dim() * sizeof(float)); + _data_len = (_data_store->get_aligned_dim() + 1) * sizeof(float); + _neighbor_len = (_max_observed_degree + 1) * sizeof(uint32_t); + _node_size = _data_len + _neighbor_len; + _opt_graph = new char[_node_size * _nd]; + DistanceFastL2 *dist_fast = (DistanceFastL2 *)_data_store->get_dist_fn(); + for (uint32_t i = 0; i < _nd; i++) + { + char *cur_node_offset = _opt_graph + i * _node_size; + _data_store->get_vector(i, (T *)cur_vec); + float cur_norm = dist_fast->norm((T *)cur_vec, _data_store->get_aligned_dim()); + std::memcpy(cur_node_offset, &cur_norm, sizeof(float)); + std::memcpy(cur_node_offset + sizeof(float), cur_vec, _data_len - sizeof(float)); + + cur_node_offset += _data_len; + uint32_t k = _final_graph[i].size(); + std::memcpy(cur_node_offset, &k, sizeof(uint32_t)); + std::memcpy(cur_node_offset + sizeof(uint32_t), _final_graph[i].data(), k * sizeof(uint32_t)); + std::vector().swap(_final_graph[i]); + } + _final_graph.clear(); + _final_graph.shrink_to_fit(); + delete[] cur_vec; } // REFACTOR: once optimized layout becomes its own Data+Graph store, we should @@ -2863,85 +2994,90 @@ void Index -void Index::search_with_optimized_layout(const T *query, - size_t K, size_t L, - uint32_t *indices) { - DistanceFastL2 *dist_fast = - (DistanceFastL2 *)_data_store->get_dist_fn(); - - NeighborPriorityQueue retset(L); - std::vector init_ids(L); - - boost::dynamic_bitset<> flags{_nd, 0}; - uint32_t tmp_l = 0; - uint32_t *neighbors = - (uint32_t *)(_opt_graph + _node_size * _start + _data_len); - uint32_t 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) { - uint32_t id = rand() % _nd; - if (flags[id]) continue; - flags[id] = true; - init_ids[tmp_l] = id; - tmp_l++; - } - - for (uint32_t i = 0; i < init_ids.size(); i++) { - uint32_t id = init_ids[i]; - if (id >= _nd) continue; - _mm_prefetch(_opt_graph + _node_size * id, _MM_HINT_T0); - } - L = 0; - for (uint32_t i = 0; i < init_ids.size(); i++) { - uint32_t 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, - (uint32_t)_data_store->get_aligned_dim()); - retset.insert(Neighbor(id, dist)); - flags[id] = true; - L++; - } - - while (retset.has_unexpanded_node()) { - auto nbr = retset.closest_unexpanded(); - auto n = nbr.id; - _mm_prefetch(_opt_graph + _node_size * n + _data_len, _MM_HINT_T0); - neighbors = (uint32_t *)(_opt_graph + _node_size * n + _data_len); - uint32_t MaxM = *neighbors; +void Index::search_with_optimized_layout(const T *query, size_t K, size_t L, uint32_t *indices) +{ + DistanceFastL2 *dist_fast = (DistanceFastL2 *)_data_store->get_dist_fn(); + + NeighborPriorityQueue retset(L); + std::vector init_ids(L); + + boost::dynamic_bitset<> flags{_nd, 0}; + uint32_t tmp_l = 0; + uint32_t *neighbors = (uint32_t *)(_opt_graph + _node_size * _start + _data_len); + uint32_t MaxM_ep = *neighbors; neighbors++; - for (uint32_t m = 0; m < MaxM; ++m) - _mm_prefetch(_opt_graph + _node_size * neighbors[m], _MM_HINT_T0); - for (uint32_t m = 0; m < MaxM; ++m) { - uint32_t id = neighbors[m]; - if (flags[id]) continue; - flags[id] = 1; - T *data = (T *)(_opt_graph + _node_size * id); - float norm = *data; - data++; - float dist = dist_fast->compare(query, data, norm, - (uint32_t)_data_store->get_aligned_dim()); - Neighbor nn(id, dist); - retset.insert(nn); - } - } - - for (size_t i = 0; i < K; i++) { - indices[i] = retset[i].id; - } + + 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) + { + uint32_t id = rand() % _nd; + if (flags[id]) + continue; + flags[id] = true; + init_ids[tmp_l] = id; + tmp_l++; + } + + for (uint32_t i = 0; i < init_ids.size(); i++) + { + uint32_t id = init_ids[i]; + if (id >= _nd) + continue; + _mm_prefetch(_opt_graph + _node_size * id, _MM_HINT_T0); + } + L = 0; + for (uint32_t i = 0; i < init_ids.size(); i++) + { + uint32_t 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, (uint32_t)_data_store->get_aligned_dim()); + retset.insert(Neighbor(id, dist)); + flags[id] = true; + L++; + } + + while (retset.has_unexpanded_node()) + { + auto nbr = retset.closest_unexpanded(); + auto n = nbr.id; + _mm_prefetch(_opt_graph + _node_size * n + _data_len, _MM_HINT_T0); + neighbors = (uint32_t *)(_opt_graph + _node_size * n + _data_len); + uint32_t MaxM = *neighbors; + neighbors++; + for (uint32_t m = 0; m < MaxM; ++m) + _mm_prefetch(_opt_graph + _node_size * neighbors[m], _MM_HINT_T0); + for (uint32_t m = 0; m < MaxM; ++m) + { + uint32_t id = neighbors[m]; + if (flags[id]) + continue; + flags[id] = 1; + T *data = (T *)(_opt_graph + _node_size * id); + float norm = *data; + data++; + float dist = dist_fast->compare(query, data, norm, (uint32_t)_data_store->get_aligned_dim()); + Neighbor nn(id, dist); + retset.insert(nn); + } + } + + for (size_t i = 0; i < K; i++) + { + indices[i] = retset[i].id; + } } /* Internals of the library */ -template -const float Index::INDEX_GROWTH_FACTOR = 1.5f; +template const float Index::INDEX_GROWTH_FACTOR = 1.5f; // EXPORTS template DISKANN_DLLEXPORT class Index; @@ -2970,252 +3106,132 @@ template DISKANN_DLLEXPORT class Index; template DISKANN_DLLEXPORT class Index; template DISKANN_DLLEXPORT class Index; -template DISKANN_DLLEXPORT std::pair -Index::search(const float *query, - const size_t K, - const uint32_t L, - uint64_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search(const float *query, - const size_t K, - const uint32_t L, - uint32_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search(const uint8_t *query, - const size_t K, - const uint32_t L, - uint64_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search(const uint8_t *query, - const size_t K, - const uint32_t L, - uint32_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search(const int8_t *query, - const size_t K, - const uint32_t L, - uint64_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search(const int8_t *query, - const size_t K, - const uint32_t L, - uint32_t *indices, - float *distances); +template DISKANN_DLLEXPORT std::pair Index::search( + const float *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair Index::search( + const float *query, const size_t K, const uint32_t L, uint32_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair Index::search( + const uint8_t *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair Index::search( + const uint8_t *query, const size_t K, const uint32_t L, uint32_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair Index::search( + const int8_t *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair Index::search( + const int8_t *query, const size_t K, const uint32_t L, uint32_t *indices, float *distances); // TagT==uint32_t -template DISKANN_DLLEXPORT std::pair -Index::search(const float *query, - const size_t K, - const uint32_t L, - uint64_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search(const float *query, - const size_t K, - const uint32_t L, - uint32_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search(const uint8_t *query, - const size_t K, - const uint32_t L, - uint64_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search(const uint8_t *query, - const size_t K, - const uint32_t L, - uint32_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search(const int8_t *query, - const size_t K, - const uint32_t L, - uint64_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search(const int8_t *query, - const size_t K, - const uint32_t L, - uint32_t *indices, - float *distances); - -template DISKANN_DLLEXPORT std::pair -Index::search_with_filters( - const float *query, const uint32_t &filter_label, const size_t K, - const uint32_t L, uint64_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search_with_filters( - const float *query, const uint32_t &filter_label, const size_t K, - const uint32_t L, uint32_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search_with_filters( - const uint8_t *query, const uint32_t &filter_label, const size_t K, - const uint32_t L, uint64_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search_with_filters( - const uint8_t *query, const uint32_t &filter_label, const size_t K, - const uint32_t L, uint32_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search_with_filters( - const int8_t *query, const uint32_t &filter_label, const size_t K, - const uint32_t L, uint64_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search_with_filters( - const int8_t *query, const uint32_t &filter_label, const size_t K, - const uint32_t L, uint32_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair Index::search( + const float *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair Index::search( + const float *query, const size_t K, const uint32_t L, uint32_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair Index::search( + const uint8_t *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair Index::search( + const uint8_t *query, const size_t K, const uint32_t L, uint32_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair Index::search( + const int8_t *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair Index::search( + const int8_t *query, const size_t K, const uint32_t L, uint32_t *indices, float *distances); + +template DISKANN_DLLEXPORT std::pair Index::search_with_filters< + uint64_t>(const float *query, const uint32_t &filter_label, const size_t K, const uint32_t L, uint64_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair Index::search_with_filters< + uint32_t>(const float *query, const uint32_t &filter_label, const size_t K, const uint32_t L, uint32_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair Index::search_with_filters< + uint64_t>(const uint8_t *query, const uint32_t &filter_label, const size_t K, const uint32_t L, uint64_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair Index::search_with_filters< + uint32_t>(const uint8_t *query, const uint32_t &filter_label, const size_t K, const uint32_t L, uint32_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair Index::search_with_filters< + uint64_t>(const int8_t *query, const uint32_t &filter_label, const size_t K, const uint32_t L, uint64_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair Index::search_with_filters< + uint32_t>(const int8_t *query, const uint32_t &filter_label, const size_t K, const uint32_t L, uint32_t *indices, + float *distances); // TagT==uint32_t -template DISKANN_DLLEXPORT std::pair -Index::search_with_filters( - const float *query, const uint32_t &filter_label, const size_t K, - const uint32_t L, uint64_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search_with_filters( - const float *query, const uint32_t &filter_label, const size_t K, - const uint32_t L, uint32_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search_with_filters( - const uint8_t *query, const uint32_t &filter_label, const size_t K, - const uint32_t L, uint64_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search_with_filters( - const uint8_t *query, const uint32_t &filter_label, const size_t K, - const uint32_t L, uint32_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search_with_filters( - const int8_t *query, const uint32_t &filter_label, const size_t K, - const uint32_t L, uint64_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search_with_filters( - const int8_t *query, const uint32_t &filter_label, const size_t K, - const uint32_t L, uint32_t *indices, float *distances); - -template DISKANN_DLLEXPORT std::pair -Index::search(const float *query, - const size_t K, - const uint32_t L, - uint64_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search(const float *query, - const size_t K, - const uint32_t L, - uint32_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search(const uint8_t *query, - const size_t K, - const uint32_t L, - uint64_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search(const uint8_t *query, - const size_t K, - const uint32_t L, - uint32_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search(const int8_t *query, - const size_t K, - const uint32_t L, - uint64_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search(const int8_t *query, - const size_t K, - const uint32_t L, - uint32_t *indices, - float *distances); +template DISKANN_DLLEXPORT std::pair Index::search_with_filters< + uint64_t>(const float *query, const uint32_t &filter_label, const size_t K, const uint32_t L, uint64_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair Index::search_with_filters< + uint32_t>(const float *query, const uint32_t &filter_label, const size_t K, const uint32_t L, uint32_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair Index::search_with_filters< + uint64_t>(const uint8_t *query, const uint32_t &filter_label, const size_t K, const uint32_t L, uint64_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair Index::search_with_filters< + uint32_t>(const uint8_t *query, const uint32_t &filter_label, const size_t K, const uint32_t L, uint32_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair Index::search_with_filters< + uint64_t>(const int8_t *query, const uint32_t &filter_label, const size_t K, const uint32_t L, uint64_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair Index::search_with_filters< + uint32_t>(const int8_t *query, const uint32_t &filter_label, const size_t K, const uint32_t L, uint32_t *indices, + float *distances); + +template DISKANN_DLLEXPORT std::pair Index::search( + const float *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair Index::search( + const float *query, const size_t K, const uint32_t L, uint32_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair Index::search( + const uint8_t *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair Index::search( + const uint8_t *query, const size_t K, const uint32_t L, uint32_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair Index::search( + const int8_t *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair Index::search( + const int8_t *query, const size_t K, const uint32_t L, uint32_t *indices, float *distances); // TagT==uint32_t -template DISKANN_DLLEXPORT std::pair -Index::search(const float *query, - const size_t K, - const uint32_t L, - uint64_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search(const float *query, - const size_t K, - const uint32_t L, - uint32_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search(const uint8_t *query, - const size_t K, - const uint32_t L, - uint64_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search(const uint8_t *query, - const size_t K, - const uint32_t L, - uint32_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search(const int8_t *query, - const size_t K, - const uint32_t L, - uint64_t *indices, - float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search(const int8_t *query, - const size_t K, - const uint32_t L, - uint32_t *indices, - float *distances); - -template DISKANN_DLLEXPORT std::pair -Index::search_with_filters( - const float *query, const uint16_t &filter_label, const size_t K, - const uint32_t L, uint64_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search_with_filters( - const float *query, const uint16_t &filter_label, const size_t K, - const uint32_t L, uint32_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search_with_filters( - const uint8_t *query, const uint16_t &filter_label, const size_t K, - const uint32_t L, uint64_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search_with_filters( - const uint8_t *query, const uint16_t &filter_label, const size_t K, - const uint32_t L, uint32_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search_with_filters( - const int8_t *query, const uint16_t &filter_label, const size_t K, - const uint32_t L, uint64_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search_with_filters( - const int8_t *query, const uint16_t &filter_label, const size_t K, - const uint32_t L, uint32_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair Index::search( + const float *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair Index::search( + const float *query, const size_t K, const uint32_t L, uint32_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair Index::search( + const uint8_t *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair Index::search( + const uint8_t *query, const size_t K, const uint32_t L, uint32_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair Index::search( + const int8_t *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances); +template DISKANN_DLLEXPORT std::pair Index::search( + const int8_t *query, const size_t K, const uint32_t L, uint32_t *indices, float *distances); + +template DISKANN_DLLEXPORT std::pair Index::search_with_filters< + uint64_t>(const float *query, const uint16_t &filter_label, const size_t K, const uint32_t L, uint64_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair Index::search_with_filters< + uint32_t>(const float *query, const uint16_t &filter_label, const size_t K, const uint32_t L, uint32_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair Index::search_with_filters< + uint64_t>(const uint8_t *query, const uint16_t &filter_label, const size_t K, const uint32_t L, uint64_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair Index::search_with_filters< + uint32_t>(const uint8_t *query, const uint16_t &filter_label, const size_t K, const uint32_t L, uint32_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair Index::search_with_filters< + uint64_t>(const int8_t *query, const uint16_t &filter_label, const size_t K, const uint32_t L, uint64_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair Index::search_with_filters< + uint32_t>(const int8_t *query, const uint16_t &filter_label, const size_t K, const uint32_t L, uint32_t *indices, + float *distances); // TagT==uint32_t -template DISKANN_DLLEXPORT std::pair -Index::search_with_filters( - const float *query, const uint16_t &filter_label, const size_t K, - const uint32_t L, uint64_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search_with_filters( - const float *query, const uint16_t &filter_label, const size_t K, - const uint32_t L, uint32_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search_with_filters( - const uint8_t *query, const uint16_t &filter_label, const size_t K, - const uint32_t L, uint64_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search_with_filters( - const uint8_t *query, const uint16_t &filter_label, const size_t K, - const uint32_t L, uint32_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search_with_filters( - const int8_t *query, const uint16_t &filter_label, const size_t K, - const uint32_t L, uint64_t *indices, float *distances); -template DISKANN_DLLEXPORT std::pair -Index::search_with_filters( - const int8_t *query, const uint16_t &filter_label, const size_t K, - const uint32_t L, uint32_t *indices, float *distances); - -} // namespace diskann +template DISKANN_DLLEXPORT std::pair Index::search_with_filters< + uint64_t>(const float *query, const uint16_t &filter_label, const size_t K, const uint32_t L, uint64_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair Index::search_with_filters< + uint32_t>(const float *query, const uint16_t &filter_label, const size_t K, const uint32_t L, uint32_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair Index::search_with_filters< + uint64_t>(const uint8_t *query, const uint16_t &filter_label, const size_t K, const uint32_t L, uint64_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair Index::search_with_filters< + uint32_t>(const uint8_t *query, const uint16_t &filter_label, const size_t K, const uint32_t L, uint32_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair Index::search_with_filters< + uint64_t>(const int8_t *query, const uint16_t &filter_label, const size_t K, const uint32_t L, uint64_t *indices, + float *distances); +template DISKANN_DLLEXPORT std::pair Index::search_with_filters< + uint32_t>(const int8_t *query, const uint16_t &filter_label, const size_t K, const uint32_t L, uint32_t *indices, + float *distances); + +} // namespace diskann diff --git a/src/linux_aligned_file_reader.cpp b/src/linux_aligned_file_reader.cpp index 5fef8384f..47c7cb1fb 100644 --- a/src/linux_aligned_file_reader.cpp +++ b/src/linux_aligned_file_reader.cpp @@ -10,188 +10,215 @@ #include "utils.h" #define MAX_EVENTS 1024 -namespace { +namespace +{ typedef struct io_event io_event_t; typedef struct iocb iocb_t; -void execute_io(io_context_t ctx, int fd, std::vector &read_reqs, - uint64_t n_retries = 0) { +void execute_io(io_context_t ctx, int fd, std::vector &read_reqs, uint64_t n_retries = 0) +{ #ifdef DEBUG - for (auto &req : read_reqs) { - assert(IS_ALIGNED(req.len, 512)); - // std::cout << "request:"<= req.len); - } + for (auto &req : read_reqs) + { + assert(IS_ALIGNED(req.len, 512)); + // std::cout << "request:"<= req.len); + } #endif - // break-up requests into chunks of size MAX_EVENTS each - uint64_t n_iters = ROUND_UP(read_reqs.size(), MAX_EVENTS) / MAX_EVENTS; - for (uint64_t iter = 0; iter < n_iters; iter++) { - uint64_t n_ops = std::min((uint64_t)read_reqs.size() - (iter * MAX_EVENTS), - (uint64_t)MAX_EVENTS); - std::vector cbs(n_ops, nullptr); - std::vector evts(n_ops); - std::vector cb(n_ops); - for (uint64_t j = 0; j < n_ops; j++) { - io_prep_pread(cb.data() + j, fd, read_reqs[j + iter * MAX_EVENTS].buf, - read_reqs[j + iter * MAX_EVENTS].len, - read_reqs[j + iter * MAX_EVENTS].offset); - } + // break-up requests into chunks of size MAX_EVENTS each + uint64_t n_iters = ROUND_UP(read_reqs.size(), MAX_EVENTS) / MAX_EVENTS; + for (uint64_t iter = 0; iter < n_iters; iter++) + { + uint64_t n_ops = std::min((uint64_t)read_reqs.size() - (iter * MAX_EVENTS), (uint64_t)MAX_EVENTS); + std::vector cbs(n_ops, nullptr); + std::vector evts(n_ops); + std::vector cb(n_ops); + for (uint64_t j = 0; j < n_ops; j++) + { + io_prep_pread(cb.data() + j, fd, read_reqs[j + iter * MAX_EVENTS].buf, read_reqs[j + iter * MAX_EVENTS].len, + read_reqs[j + iter * MAX_EVENTS].offset); + } - // initialize `cbs` using `cb` array - // + // initialize `cbs` using `cb` array + // - for (uint64_t i = 0; i < n_ops; i++) { - cbs[i] = cb.data() + i; - } + for (uint64_t i = 0; i < n_ops; i++) + { + cbs[i] = cb.data() + i; + } - uint64_t n_tries = 0; - while (n_tries <= n_retries) { - // issue reads - int64_t ret = io_submit(ctx, (int64_t)n_ops, cbs.data()); - // if requests didn't get accepted - 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"; - exit(-1); - } else { - // wait on io_getevents - ret = io_getevents(ctx, (int64_t)n_ops, (int64_t)n_ops, evts.data(), - nullptr); - // if requests didn't complete - if (ret != (int64_t)n_ops) { - std::cerr << "io_getevents() failed; returned " << ret - << ", expected=" << n_ops << ", ernno=" << errno << "=" - << ::strerror(-ret) << ", try #" << n_tries + 1; - exit(-1); - } else { - break; + uint64_t n_tries = 0; + while (n_tries <= n_retries) + { + // issue reads + int64_t ret = io_submit(ctx, (int64_t)n_ops, cbs.data()); + // if requests didn't get accepted + 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"; + exit(-1); + } + else + { + // wait on io_getevents + ret = io_getevents(ctx, (int64_t)n_ops, (int64_t)n_ops, evts.data(), nullptr); + // if requests didn't complete + if (ret != (int64_t)n_ops) + { + std::cerr << "io_getevents() failed; returned " << ret << ", expected=" << n_ops + << ", ernno=" << errno << "=" << ::strerror(-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); + // 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); + } + */ } - */ - } } -} // namespace - -LinuxAlignedFileReader::LinuxAlignedFileReader() { this->file_desc = -1; } - -LinuxAlignedFileReader::~LinuxAlignedFileReader() { - int64_t ret; - // check to make sure file_desc is closed - ret = ::fcntl(this->file_desc, F_GETFD); - if (ret == -1) { - if (errno != EBADF) { - std::cerr << "close() not called" << std::endl; - // close file desc - ret = ::close(this->file_desc); - // error checks - if (ret == -1) { - std::cerr << "close() failed; returned " << ret << ", errno=" << errno - << ":" << ::strerror(errno) << std::endl; - } +} // namespace + +LinuxAlignedFileReader::LinuxAlignedFileReader() +{ + this->file_desc = -1; +} + +LinuxAlignedFileReader::~LinuxAlignedFileReader() +{ + int64_t ret; + // check to make sure file_desc is closed + ret = ::fcntl(this->file_desc, F_GETFD); + if (ret == -1) + { + if (errno != EBADF) + { + std::cerr << "close() not called" << std::endl; + // close file desc + ret = ::close(this->file_desc); + // error checks + if (ret == -1) + { + std::cerr << "close() failed; returned " << ret << ", errno=" << errno << ":" << ::strerror(errno) + << std::endl; + } + } } - } } -io_context_t &LinuxAlignedFileReader::get_ctx() { - std::unique_lock lk(ctx_mut); - // perform checks only in DEBUG mode - if (ctx_map.find(std::this_thread::get_id()) == ctx_map.end()) { - std::cerr << "bad thread access; returning -1 as io_context_t" << std::endl; - return this->bad_ctx; - } else { - return ctx_map[std::this_thread::get_id()]; - } +io_context_t &LinuxAlignedFileReader::get_ctx() +{ + std::unique_lock lk(ctx_mut); + // perform checks only in DEBUG mode + if (ctx_map.find(std::this_thread::get_id()) == ctx_map.end()) + { + std::cerr << "bad thread access; returning -1 as io_context_t" << std::endl; + return this->bad_ctx; + } + else + { + return ctx_map[std::this_thread::get_id()]; + } } -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; - int ret = io_setup(MAX_EVENTS, &ctx); - if (ret != 0) { - lk.unlock(); - assert(errno != EAGAIN); - assert(errno != ENOMEM); - std::cerr << "io_setup() failed; returned " << ret << ", errno=" << errno - << ":" << ::strerror(errno) << std::endl; - } else { - diskann::cout << "allocating ctx: " << ctx << " to thread-id:" << my_id +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; + int ret = io_setup(MAX_EVENTS, &ctx); + if (ret != 0) + { + lk.unlock(); + assert(errno != EAGAIN); + assert(errno != ENOMEM); + std::cerr << "io_setup() failed; returned " << ret << ", errno=" << errno << ":" << ::strerror(errno) << std::endl; - ctx_map[my_id] = ctx; - } - lk.unlock(); + } + else + { + diskann::cout << "allocating ctx: " << ctx << " to thread-id:" << my_id << std::endl; + ctx_map[my_id] = ctx; + } + lk.unlock(); } -void LinuxAlignedFileReader::deregister_thread() { - auto my_id = std::this_thread::get_id(); - std::unique_lock lk(ctx_mut); - assert(ctx_map.find(my_id) != ctx_map.end()); - - lk.unlock(); - io_context_t ctx = this->get_ctx(); - io_destroy(ctx); - // assert(ret == 0); - lk.lock(); - ctx_map.erase(my_id); - std::cerr << "returned ctx from thread-id:" << my_id << std::endl; - lk.unlock(); -} +void LinuxAlignedFileReader::deregister_thread() +{ + auto my_id = std::this_thread::get_id(); + std::unique_lock lk(ctx_mut); + assert(ctx_map.find(my_id) != ctx_map.end()); -void LinuxAlignedFileReader::deregister_all_threads() { - std::unique_lock lk(ctx_mut); - for (auto x = ctx_map.begin(); x != ctx_map.end(); x++) { - io_context_t ctx = x.value(); + lk.unlock(); + io_context_t ctx = this->get_ctx(); io_destroy(ctx); // assert(ret == 0); - // lk.lock(); - // ctx_map.erase(my_id); - // std::cerr << "returned ctx from thread-id:" << my_id << std::endl; - } - ctx_map.clear(); - // lk.unlock(); + lk.lock(); + ctx_map.erase(my_id); + 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 x = ctx_map.begin(); x != ctx_map.end(); x++) + { + io_context_t ctx = x.value(); + io_destroy(ctx); + // assert(ret == 0); + // lk.lock(); + // ctx_map.erase(my_id); + // std::cerr << "returned ctx from thread-id:" << my_id << std::endl; + } + ctx_map.clear(); + // lk.unlock(); } -void LinuxAlignedFileReader::open(const std::string &fname) { - int flags = O_DIRECT | O_RDONLY | O_LARGEFILE; - this->file_desc = ::open(fname.c_str(), flags); - // error checks - assert(this->file_desc != -1); - std::cerr << "Opened file : " << fname << std::endl; +void LinuxAlignedFileReader::open(const std::string &fname) +{ + int flags = O_DIRECT | O_RDONLY | O_LARGEFILE; + this->file_desc = ::open(fname.c_str(), flags); + // error checks + assert(this->file_desc != -1); + std::cerr << "Opened file : " << fname << std::endl; } -void LinuxAlignedFileReader::close() { - // int64_t ret; +void LinuxAlignedFileReader::close() +{ + // int64_t ret; - // check to make sure file_desc is closed - ::fcntl(this->file_desc, F_GETFD); - // assert(ret != -1); + // check to make sure file_desc is closed + ::fcntl(this->file_desc, F_GETFD); + // assert(ret != -1); - ::close(this->file_desc); - // assert(ret != -1); + ::close(this->file_desc); + // assert(ret != -1); } -void LinuxAlignedFileReader::read(std::vector &read_reqs, - io_context_t &ctx, bool async) { - if (async == true) { - diskann::cout << "Async currently not supported in linux." << std::endl; - } - assert(this->file_desc != -1); - execute_io(ctx, this->file_desc, read_reqs); +void LinuxAlignedFileReader::read(std::vector &read_reqs, io_context_t &ctx, bool async) +{ + if (async == true) + { + diskann::cout << "Async currently not supported in linux." << std::endl; + } + assert(this->file_desc != -1); + execute_io(ctx, this->file_desc, read_reqs); } diff --git a/src/logger.cpp b/src/logger.cpp index 5ab15b84e..1444487f7 100644 --- a/src/logger.cpp +++ b/src/logger.cpp @@ -7,7 +7,8 @@ #include "logger_impl.h" #include "windows_customizations.h" -namespace diskann { +namespace diskann +{ DISKANN_DLLEXPORT ANNStreamBuf coutBuff(stdout); DISKANN_DLLEXPORT ANNStreamBuf cerrBuff(stderr); @@ -18,76 +19,85 @@ DISKANN_DLLEXPORT std::basic_ostream cerr(&cerrBuff); #ifdef EXEC_ENV_OLS std::function g_logger; -void SetCustomLogger(std::function logger) { - g_logger = logger; +void SetCustomLogger(std::function logger) +{ + g_logger = logger; } #endif -ANNStreamBuf::ANNStreamBuf(FILE *fp) { - if (fp == nullptr) { - throw diskann::ANNException( - "File pointer passed to ANNStreamBuf() cannot be null", -1); - } - if (fp != stdout && fp != stderr) { - throw diskann::ANNException( - "The custom logger only supports stdout and stderr.", -1); - } - _fp = fp; - _logLevel = (_fp == stdout) ? LogLevel::LL_Info : LogLevel::LL_Error; +ANNStreamBuf::ANNStreamBuf(FILE *fp) +{ + if (fp == nullptr) + { + throw diskann::ANNException("File pointer passed to ANNStreamBuf() cannot be null", -1); + } + if (fp != stdout && fp != stderr) + { + throw diskann::ANNException("The custom logger only supports stdout and stderr.", -1); + } + _fp = fp; + _logLevel = (_fp == stdout) ? LogLevel::LL_Info : LogLevel::LL_Error; #ifdef EXEC_ENV_OLS - _buf = new char[BUFFER_SIZE + 1]; // See comment in the header + _buf = new char[BUFFER_SIZE + 1]; // See comment in the header #else - _buf = new char[BUFFER_SIZE]; // See comment in the header + _buf = new char[BUFFER_SIZE]; // See comment in the header #endif - std::memset(_buf, 0, (BUFFER_SIZE) * sizeof(char)); - setp(_buf, _buf + BUFFER_SIZE - 1); + std::memset(_buf, 0, (BUFFER_SIZE) * sizeof(char)); + setp(_buf, _buf + BUFFER_SIZE - 1); } -ANNStreamBuf::~ANNStreamBuf() { - sync(); - _fp = nullptr; // we'll not close because we can't. - delete[] _buf; +ANNStreamBuf::~ANNStreamBuf() +{ + sync(); + _fp = nullptr; // we'll not close because we can't. + delete[] _buf; } -int ANNStreamBuf::overflow(int c) { - std::lock_guard lock(_mutex); - if (c != EOF) { - *pptr() = (char)c; - pbump(1); - } - flush(); - return c; +int ANNStreamBuf::overflow(int c) +{ + std::lock_guard lock(_mutex); + if (c != EOF) + { + *pptr() = (char)c; + pbump(1); + } + flush(); + return c; } -int ANNStreamBuf::sync() { - std::lock_guard lock(_mutex); - flush(); - return 0; +int ANNStreamBuf::sync() +{ + std::lock_guard lock(_mutex); + flush(); + return 0; } -int ANNStreamBuf::underflow() { - throw diskann::ANNException( - "Attempt to read on streambuf meant only for writing.", -1); +int ANNStreamBuf::underflow() +{ + throw diskann::ANNException("Attempt to read on streambuf meant only for writing.", -1); } -int ANNStreamBuf::flush() { - const int num = (int)(pptr() - pbase()); - logImpl(pbase(), num); - pbump(-num); - return num; +int ANNStreamBuf::flush() +{ + const int num = (int)(pptr() - pbase()); + logImpl(pbase(), num); + pbump(-num); + return num; } -void ANNStreamBuf::logImpl(char *str, int num) { +void ANNStreamBuf::logImpl(char *str, int num) +{ #ifdef EXEC_ENV_OLS - str[num] = '\0'; // Safe. See the c'tor. - // Invoke the OLS custom logging function. - if (g_logger) { - g_logger(_logLevel, str); - } + str[num] = '\0'; // Safe. See the c'tor. + // Invoke the OLS custom logging function. + if (g_logger) + { + g_logger(_logLevel, str); + } #else - fwrite(str, sizeof(char), num, _fp); - fflush(_fp); + fwrite(str, sizeof(char), num, _fp); + fflush(_fp); #endif } -} // namespace diskann +} // namespace diskann diff --git a/src/math_utils.cpp b/src/math_utils.cpp index c08fd66b1..7481da848 100644 --- a/src/math_utils.cpp +++ b/src/math_utils.cpp @@ -8,42 +8,47 @@ #include "logger.h" #include "utils.h" -namespace math_utils { - -float calc_distance(float *vec_1, float *vec_2, size_t dim) { - float dist = 0; - for (size_t j = 0; j < dim; j++) { - dist += (vec_1[j] - vec_2[j]) * (vec_1[j] - vec_2[j]); - } - return dist; +namespace math_utils +{ + +float calc_distance(float *vec_1, float *vec_2, size_t dim) +{ + float dist = 0; + for (size_t j = 0; j < dim; j++) + { + dist += (vec_1[j] - vec_2[j]) * (vec_1[j] - vec_2[j]); + } + return dist; } // compute l2-squared norms of data stored in row major num_points * dim, // needs // to be pre-allocated -void compute_vecs_l2sq(float *vecs_l2sq, float *data, const size_t num_points, - const size_t dim) { +void compute_vecs_l2sq(float *vecs_l2sq, float *data, const size_t num_points, const size_t dim) +{ #pragma omp parallel for schedule(static, 8192) - for (int64_t n_iter = 0; n_iter < (int64_t)num_points; n_iter++) { - vecs_l2sq[n_iter] = cblas_snrm2((MKL_INT)dim, (data + (n_iter * dim)), 1); - vecs_l2sq[n_iter] *= vecs_l2sq[n_iter]; - } + for (int64_t n_iter = 0; n_iter < (int64_t)num_points; n_iter++) + { + vecs_l2sq[n_iter] = cblas_snrm2((MKL_INT)dim, (data + (n_iter * dim)), 1); + vecs_l2sq[n_iter] *= vecs_l2sq[n_iter]; + } } -void rotate_data_randomly(float *data, size_t num_points, size_t dim, - float *rot_mat, float *&new_mat, bool transpose_rot) { - CBLAS_TRANSPOSE transpose = CblasNoTrans; - if (transpose_rot) { - diskann::cout << "Transposing rotation matrix.." << std::flush; - transpose = CblasTrans; - } - diskann::cout << "done Rotating data with random matrix.." << std::flush; +void rotate_data_randomly(float *data, size_t num_points, size_t dim, float *rot_mat, float *&new_mat, + bool transpose_rot) +{ + CBLAS_TRANSPOSE transpose = CblasNoTrans; + if (transpose_rot) + { + diskann::cout << "Transposing rotation matrix.." << std::flush; + transpose = CblasTrans; + } + diskann::cout << "done Rotating data with random matrix.." << std::flush; - cblas_sgemm(CblasRowMajor, CblasNoTrans, transpose, (MKL_INT)num_points, - (MKL_INT)dim, (MKL_INT)dim, 1.0, data, (MKL_INT)dim, rot_mat, - (MKL_INT)dim, 0, new_mat, (MKL_INT)dim); + cblas_sgemm(CblasRowMajor, CblasNoTrans, transpose, (MKL_INT)num_points, (MKL_INT)dim, (MKL_INT)dim, 1.0, data, + (MKL_INT)dim, rot_mat, (MKL_INT)dim, 0, new_mat, (MKL_INT)dim); - diskann::cout << "done." << std::endl; + diskann::cout << "done." << std::endl; } // calculate k closest centers to data of num_points * dim (row major) @@ -56,70 +61,77 @@ void rotate_data_randomly(float *data, size_t num_points, size_t dim, // Default value of k is 1 // Ideally used only by compute_closest_centers -void compute_closest_centers_in_block( - const float *const data, const size_t num_points, const size_t dim, - const float *const centers, const size_t num_centers, - const float *const docs_l2sq, const float *const centers_l2sq, - uint32_t *center_index, float *const dist_matrix, size_t k) { - if (k > num_centers) { - diskann::cout << "ERROR: k (" << k << ") > num_center(" << num_centers - << ")" << std::endl; - return; - } - - float *ones_a = new float[num_centers]; - float *ones_b = new float[num_points]; - - for (size_t i = 0; i < num_centers; i++) { - ones_a[i] = 1.0; - } - for (size_t i = 0; i < num_points; i++) { - ones_b[i] = 1.0; - } - - cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, (MKL_INT)num_points, - (MKL_INT)num_centers, (MKL_INT)1, 1.0f, docs_l2sq, (MKL_INT)1, - ones_a, (MKL_INT)1, 0.0f, dist_matrix, (MKL_INT)num_centers); - - cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, (MKL_INT)num_points, - (MKL_INT)num_centers, (MKL_INT)1, 1.0f, ones_b, (MKL_INT)1, - centers_l2sq, (MKL_INT)1, 1.0f, dist_matrix, - (MKL_INT)num_centers); - - cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, (MKL_INT)num_points, - (MKL_INT)num_centers, (MKL_INT)dim, -2.0f, data, (MKL_INT)dim, - centers, (MKL_INT)dim, 1.0f, dist_matrix, (MKL_INT)num_centers); - - if (k == 1) { +void compute_closest_centers_in_block(const float *const data, const size_t num_points, const size_t dim, + const float *const centers, const size_t num_centers, + const float *const docs_l2sq, const float *const centers_l2sq, + uint32_t *center_index, float *const dist_matrix, size_t k) +{ + if (k > num_centers) + { + diskann::cout << "ERROR: k (" << k << ") > num_center(" << num_centers << ")" << std::endl; + return; + } + + float *ones_a = new float[num_centers]; + float *ones_b = new float[num_points]; + + for (size_t i = 0; i < num_centers; i++) + { + ones_a[i] = 1.0; + } + for (size_t i = 0; i < num_points; i++) + { + ones_b[i] = 1.0; + } + + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, (MKL_INT)num_points, (MKL_INT)num_centers, (MKL_INT)1, 1.0f, + docs_l2sq, (MKL_INT)1, ones_a, (MKL_INT)1, 0.0f, dist_matrix, (MKL_INT)num_centers); + + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, (MKL_INT)num_points, (MKL_INT)num_centers, (MKL_INT)1, 1.0f, + ones_b, (MKL_INT)1, centers_l2sq, (MKL_INT)1, 1.0f, dist_matrix, (MKL_INT)num_centers); + + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, (MKL_INT)num_points, (MKL_INT)num_centers, (MKL_INT)dim, -2.0f, + data, (MKL_INT)dim, centers, (MKL_INT)dim, 1.0f, dist_matrix, (MKL_INT)num_centers); + + if (k == 1) + { #pragma omp parallel for schedule(static, 8192) - for (int64_t i = 0; i < (int64_t)num_points; i++) { - float min = std::numeric_limits::max(); - float *current = dist_matrix + (i * num_centers); - for (size_t j = 0; j < num_centers; j++) { - if (current[j] < min) { - center_index[i] = (uint32_t)j; - min = current[j]; + for (int64_t i = 0; i < (int64_t)num_points; i++) + { + float min = std::numeric_limits::max(); + float *current = dist_matrix + (i * num_centers); + for (size_t j = 0; j < num_centers; j++) + { + if (current[j] < min) + { + center_index[i] = (uint32_t)j; + min = current[j]; + } + } } - } } - } else { + else + { #pragma omp parallel for schedule(static, 8192) - for (int64_t i = 0; i < (int64_t)num_points; i++) { - std::priority_queue top_k_queue; - float *current = dist_matrix + (i * num_centers); - for (size_t j = 0; j < num_centers; j++) { - PivotContainer this_piv(j, current[j]); - top_k_queue.push(this_piv); - } - for (size_t j = 0; j < k; j++) { - PivotContainer this_piv = top_k_queue.top(); - center_index[i * k + j] = (uint32_t)this_piv.piv_id; - top_k_queue.pop(); - } + for (int64_t i = 0; i < (int64_t)num_points; i++) + { + std::priority_queue top_k_queue; + float *current = dist_matrix + (i * num_centers); + for (size_t j = 0; j < num_centers; j++) + { + PivotContainer this_piv(j, current[j]); + top_k_queue.push(this_piv); + } + for (size_t j = 0; j < k; j++) + { + PivotContainer this_piv = top_k_queue.top(); + center_index[i * k + j] = (uint32_t)this_piv.piv_id; + top_k_queue.pop(); + } + } } - } - delete[] ones_a; - delete[] ones_b; + delete[] ones_a; + delete[] ones_b; } // Given data in num_points * new_dim row major @@ -132,93 +144,92 @@ void compute_closest_centers_in_block( // indices is an empty vector. Additionally, if pts_norms_squared is not null, // then it will assume that point norms are pre-computed and use those values -void compute_closest_centers(float *data, size_t num_points, size_t dim, - float *pivot_data, size_t num_centers, size_t k, - uint32_t *closest_centers_ivf, - std::vector *inverted_index, - float *pts_norms_squared) { - if (k > num_centers) { - diskann::cout << "ERROR: k (" << k << ") > num_center(" << num_centers - << ")" << std::endl; - return; - } - - bool is_norm_given_for_pts = (pts_norms_squared != NULL); - - float *pivs_norms_squared = new float[num_centers]; - if (!is_norm_given_for_pts) pts_norms_squared = new float[num_points]; - - 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; - - if (!is_norm_given_for_pts) - math_utils::compute_vecs_l2sq(pts_norms_squared, data, num_points, dim); - math_utils::compute_vecs_l2sq(pivs_norms_squared, pivot_data, num_centers, - dim); - uint32_t *closest_centers = new uint32_t[PAR_BLOCK_SIZE * k]; - float *distance_matrix = new float[num_centers * PAR_BLOCK_SIZE]; - - for (size_t cur_blk = 0; cur_blk < N_BLOCKS; cur_blk++) { - float *data_cur_blk = data + cur_blk * PAR_BLOCK_SIZE * dim; - size_t num_pts_blk = - std::min(PAR_BLOCK_SIZE, num_points - cur_blk * PAR_BLOCK_SIZE); - float *pts_norms_blk = pts_norms_squared + cur_blk * PAR_BLOCK_SIZE; - - math_utils::compute_closest_centers_in_block( - data_cur_blk, num_pts_blk, dim, pivot_data, num_centers, pts_norms_blk, - pivs_norms_squared, closest_centers, distance_matrix, k); +void compute_closest_centers(float *data, size_t num_points, size_t dim, float *pivot_data, size_t num_centers, + size_t k, uint32_t *closest_centers_ivf, std::vector *inverted_index, + float *pts_norms_squared) +{ + if (k > num_centers) + { + diskann::cout << "ERROR: k (" << k << ") > num_center(" << num_centers << ")" << std::endl; + return; + } + + bool is_norm_given_for_pts = (pts_norms_squared != NULL); + + float *pivs_norms_squared = new float[num_centers]; + if (!is_norm_given_for_pts) + pts_norms_squared = new float[num_points]; + + 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; + + if (!is_norm_given_for_pts) + math_utils::compute_vecs_l2sq(pts_norms_squared, data, num_points, dim); + math_utils::compute_vecs_l2sq(pivs_norms_squared, pivot_data, num_centers, dim); + uint32_t *closest_centers = new uint32_t[PAR_BLOCK_SIZE * k]; + float *distance_matrix = new float[num_centers * PAR_BLOCK_SIZE]; + + for (size_t cur_blk = 0; cur_blk < N_BLOCKS; cur_blk++) + { + float *data_cur_blk = data + cur_blk * PAR_BLOCK_SIZE * dim; + size_t num_pts_blk = std::min(PAR_BLOCK_SIZE, num_points - cur_blk * PAR_BLOCK_SIZE); + float *pts_norms_blk = pts_norms_squared + cur_blk * PAR_BLOCK_SIZE; + + math_utils::compute_closest_centers_in_block(data_cur_blk, num_pts_blk, dim, pivot_data, num_centers, + pts_norms_blk, pivs_norms_squared, closest_centers, + distance_matrix, k); #pragma omp parallel for schedule(static, 1) - for (int64_t j = cur_blk * PAR_BLOCK_SIZE; - j < std::min((int64_t)num_points, - (int64_t)((cur_blk + 1) * PAR_BLOCK_SIZE)); - j++) { - for (size_t l = 0; l < k; l++) { - size_t this_center_id = - closest_centers[(j - cur_blk * PAR_BLOCK_SIZE) * k + l]; - closest_centers_ivf[j * k + l] = (uint32_t)this_center_id; - if (inverted_index != NULL) { + for (int64_t j = cur_blk * PAR_BLOCK_SIZE; + j < std::min((int64_t)num_points, (int64_t)((cur_blk + 1) * PAR_BLOCK_SIZE)); j++) + { + for (size_t l = 0; l < k; l++) + { + size_t this_center_id = closest_centers[(j - cur_blk * PAR_BLOCK_SIZE) * k + l]; + closest_centers_ivf[j * k + l] = (uint32_t)this_center_id; + if (inverted_index != NULL) + { #pragma omp critical - inverted_index[this_center_id].push_back(j); + inverted_index[this_center_id].push_back(j); + } + } } - } } - } - delete[] closest_centers; - delete[] distance_matrix; - delete[] pivs_norms_squared; - if (!is_norm_given_for_pts) delete[] pts_norms_squared; + delete[] closest_centers; + delete[] distance_matrix; + delete[] pivs_norms_squared; + if (!is_norm_given_for_pts) + delete[] pts_norms_squared; } // if to_subtract is 1, will subtract nearest center from each row. Else will // add. Output will be in data_load iself. // Nearest centers need to be provided in closst_centers. -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; +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 < (int64_t)num_points; n_iter++) { - for (size_t d_iter = 0; d_iter < dim; d_iter++) { - if (to_subtract == 1) - data_load[n_iter * dim + d_iter] = - data_load[n_iter * dim + d_iter] - - cur_pivot_data[closest_centers[n_iter] * dim + d_iter]; - else - data_load[n_iter * dim + d_iter] = - data_load[n_iter * dim + d_iter] + - cur_pivot_data[closest_centers[n_iter] * dim + d_iter]; + for (int64_t n_iter = 0; n_iter < (int64_t)num_points; n_iter++) + { + for (size_t d_iter = 0; d_iter < dim; d_iter++) + { + if (to_subtract == 1) + data_load[n_iter * dim + d_iter] = + data_load[n_iter * dim + d_iter] - cur_pivot_data[closest_centers[n_iter] * dim + d_iter]; + else + data_load[n_iter * dim + d_iter] = + data_load[n_iter * dim + d_iter] + cur_pivot_data[closest_centers[n_iter] * dim + d_iter]; + } } - } } -} // namespace math_utils +} // namespace math_utils -namespace kmeans { +namespace kmeans +{ // run Lloyds one iteration // Given data in row major num_points * dim, and centers in row major @@ -228,64 +239,67 @@ namespace kmeans { // closest_centers == NULL, will allocate memory and return. Similarly, if // closest_docs == NULL, will allocate memory and return. -float lloyds_iter(float *data, size_t num_points, size_t dim, float *centers, - size_t num_centers, float *docs_l2sq, - std::vector *closest_docs, - uint32_t *&closest_center) { - bool compute_residual = true; - // Timer timer; +float lloyds_iter(float *data, size_t num_points, size_t dim, float *centers, size_t num_centers, float *docs_l2sq, + std::vector *closest_docs, uint32_t *&closest_center) +{ + bool compute_residual = true; + // Timer timer; - if (closest_center == NULL) closest_center = new uint32_t[num_points]; - if (closest_docs == NULL) - closest_docs = new std::vector[num_centers]; - else - for (size_t c = 0; c < num_centers; ++c) closest_docs[c].clear(); + if (closest_center == NULL) + closest_center = new uint32_t[num_points]; + if (closest_docs == NULL) + closest_docs = new std::vector[num_centers]; + else + for (size_t c = 0; c < num_centers; ++c) + closest_docs[c].clear(); - math_utils::compute_closest_centers(data, num_points, dim, centers, - num_centers, 1, closest_center, - closest_docs, docs_l2sq); + math_utils::compute_closest_centers(data, num_points, dim, centers, num_centers, 1, closest_center, closest_docs, + docs_l2sq); - memset(centers, 0, sizeof(float) * (size_t)num_centers * (size_t)dim); + memset(centers, 0, sizeof(float) * (size_t)num_centers * (size_t)dim); #pragma omp parallel for schedule(static, 1) - for (int64_t c = 0; c < (int64_t)num_centers; ++c) { - float *center = centers + (size_t)c * (size_t)dim; - double *cluster_sum = new double[dim]; - for (size_t i = 0; i < dim; i++) cluster_sum[i] = 0.0; - for (size_t i = 0; i < closest_docs[c].size(); i++) { - float *current = data + ((closest_docs[c][i]) * dim); - for (size_t j = 0; j < dim; j++) { - cluster_sum[j] += (double)current[j]; - } - } - if (closest_docs[c].size() > 0) { - for (size_t i = 0; i < dim; i++) - center[i] = (float)(cluster_sum[i] / ((double)closest_docs[c].size())); + for (int64_t c = 0; c < (int64_t)num_centers; ++c) + { + float *center = centers + (size_t)c * (size_t)dim; + double *cluster_sum = new double[dim]; + for (size_t i = 0; i < dim; i++) + cluster_sum[i] = 0.0; + for (size_t i = 0; i < closest_docs[c].size(); i++) + { + float *current = data + ((closest_docs[c][i]) * dim); + for (size_t j = 0; j < dim; j++) + { + cluster_sum[j] += (double)current[j]; + } + } + if (closest_docs[c].size() > 0) + { + for (size_t i = 0; i < dim; i++) + center[i] = (float)(cluster_sum[i] / ((double)closest_docs[c].size())); + } + delete[] cluster_sum; } - delete[] cluster_sum; - } - float residual = 0.0; - if (compute_residual) { - size_t BUF_PAD = 32; - size_t CHUNK_SIZE = 2 * 8192; - size_t nchunks = - num_points / CHUNK_SIZE + (num_points % CHUNK_SIZE == 0 ? 0 : 1); - std::vector residuals(nchunks * BUF_PAD, 0.0); + float residual = 0.0; + if (compute_residual) + { + size_t BUF_PAD = 32; + size_t CHUNK_SIZE = 2 * 8192; + size_t nchunks = num_points / CHUNK_SIZE + (num_points % CHUNK_SIZE == 0 ? 0 : 1); + std::vector residuals(nchunks * BUF_PAD, 0.0); #pragma omp parallel for schedule(static, 32) - for (int64_t chunk = 0; chunk < (int64_t)nchunks; ++chunk) - for (size_t d = chunk * CHUNK_SIZE; - d < num_points && d < (chunk + 1) * CHUNK_SIZE; ++d) - residuals[chunk * BUF_PAD] += math_utils::calc_distance( - data + (d * dim), centers + (size_t)closest_center[d] * (size_t)dim, - dim); - - for (size_t chunk = 0; chunk < nchunks; ++chunk) - residual += residuals[chunk * BUF_PAD]; - } - - return residual; + for (int64_t chunk = 0; chunk < (int64_t)nchunks; ++chunk) + for (size_t d = chunk * CHUNK_SIZE; d < num_points && d < (chunk + 1) * CHUNK_SIZE; ++d) + residuals[chunk * BUF_PAD] += + math_utils::calc_distance(data + (d * dim), centers + (size_t)closest_center[d] * (size_t)dim, dim); + + for (size_t chunk = 0; chunk < nchunks; ++chunk) + residual += residuals[chunk * BUF_PAD]; + } + + return residual; } // Run Lloyds until max_reps or stopping criterion @@ -295,142 +309,150 @@ float lloyds_iter(float *data, size_t num_points, size_t dim, float *centers, // vector [num_centers], and closest_center = new size_t[num_points] // Final centers are output in centers as row major num_centers * dim // -float run_lloyds(float *data, size_t num_points, size_t dim, float *centers, - const size_t num_centers, const size_t max_reps, - std::vector *closest_docs, uint32_t *closest_center) { - float residual = std::numeric_limits::max(); - bool ret_closest_docs = true; - bool ret_closest_center = true; - if (closest_docs == NULL) { - closest_docs = new std::vector[num_centers]; - ret_closest_docs = false; - } - if (closest_center == NULL) { - closest_center = new uint32_t[num_points]; - ret_closest_center = false; - } - - float *docs_l2sq = new float[num_points]; - math_utils::compute_vecs_l2sq(docs_l2sq, data, num_points, dim); - - float old_residual; - // Timer timer; - for (size_t i = 0; i < max_reps; ++i) { - old_residual = residual; - - residual = lloyds_iter(data, num_points, dim, centers, num_centers, - docs_l2sq, closest_docs, closest_center); - - if (((i != 0) && ((old_residual - residual) / residual) < 0.00001) || - (residual < std::numeric_limits::epsilon())) { - diskann::cout << "Residuals unchanged: " << old_residual << " becomes " - << residual << ". Early termination." << std::endl; - break; +float run_lloyds(float *data, size_t num_points, size_t dim, float *centers, const size_t num_centers, + const size_t max_reps, std::vector *closest_docs, uint32_t *closest_center) +{ + float residual = std::numeric_limits::max(); + bool ret_closest_docs = true; + bool ret_closest_center = true; + if (closest_docs == NULL) + { + closest_docs = new std::vector[num_centers]; + ret_closest_docs = false; + } + if (closest_center == NULL) + { + closest_center = new uint32_t[num_points]; + ret_closest_center = false; + } + + float *docs_l2sq = new float[num_points]; + math_utils::compute_vecs_l2sq(docs_l2sq, data, num_points, dim); + + float old_residual; + // Timer timer; + for (size_t i = 0; i < max_reps; ++i) + { + old_residual = residual; + + residual = lloyds_iter(data, num_points, dim, centers, num_centers, docs_l2sq, closest_docs, closest_center); + + if (((i != 0) && ((old_residual - residual) / residual) < 0.00001) || + (residual < std::numeric_limits::epsilon())) + { + diskann::cout << "Residuals unchanged: " << old_residual << " becomes " << residual + << ". Early termination." << std::endl; + break; + } } - } - delete[] docs_l2sq; - if (!ret_closest_docs) delete[] closest_docs; - if (!ret_closest_center) delete[] closest_center; - return residual; + delete[] docs_l2sq; + if (!ret_closest_docs) + delete[] closest_docs; + if (!ret_closest_center) + delete[] closest_center; + return residual; } // assumes memory allocated for pivot_data as new // float[num_centers*dim] // and select randomly num_centers points as pivots -void selecting_pivots(float *data, size_t num_points, size_t dim, - float *pivot_data, size_t num_centers) { - // pivot_data = new float[num_centers * dim]; - - std::vector picked; - std::random_device rd; - auto x = rd(); - std::mt19937 generator(x); - std::uniform_int_distribution distribution(0, num_points - 1); - - size_t tmp_pivot; - for (size_t j = 0; j < num_centers; j++) { - tmp_pivot = distribution(generator); - if (std::find(picked.begin(), picked.end(), tmp_pivot) != picked.end()) - continue; - picked.push_back(tmp_pivot); - std::memcpy(pivot_data + j * dim, data + tmp_pivot * dim, - dim * sizeof(float)); - } +void selecting_pivots(float *data, size_t num_points, size_t dim, float *pivot_data, size_t num_centers) +{ + // pivot_data = new float[num_centers * dim]; + + std::vector picked; + std::random_device rd; + auto x = rd(); + std::mt19937 generator(x); + std::uniform_int_distribution distribution(0, num_points - 1); + + size_t tmp_pivot; + for (size_t j = 0; j < num_centers; j++) + { + tmp_pivot = distribution(generator); + if (std::find(picked.begin(), picked.end(), tmp_pivot) != picked.end()) + continue; + picked.push_back(tmp_pivot); + std::memcpy(pivot_data + j * dim, data + tmp_pivot * dim, dim * sizeof(float)); + } } -void kmeanspp_selecting_pivots(float *data, size_t num_points, size_t dim, - float *pivot_data, size_t num_centers) { - if (num_points > 1 << 23) { - diskann::cout << "ERROR: n_pts " << num_points - << " currently not supported for k-means++, maximum is " - "8388608. Falling back to random pivot " - "selection." - << std::endl; - selecting_pivots(data, num_points, dim, pivot_data, num_centers); - return; - } - - 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); - size_t init_id = int_dist(generator); - size_t num_picked = 1; - - picked.push_back(init_id); - std::memcpy(pivot_data, data + init_id * dim, dim * sizeof(float)); - - float *dist = new float[num_points]; +void kmeanspp_selecting_pivots(float *data, size_t num_points, size_t dim, float *pivot_data, size_t num_centers) +{ + if (num_points > 1 << 23) + { + diskann::cout << "ERROR: n_pts " << num_points + << " currently not supported for k-means++, maximum is " + "8388608. Falling back to random pivot " + "selection." + << std::endl; + selecting_pivots(data, num_points, dim, pivot_data, num_centers); + return; + } -#pragma omp parallel for schedule(static, 8192) - for (int64_t i = 0; i < (int64_t)num_points; i++) { - dist[i] = - math_utils::calc_distance(data + i * dim, data + init_id * dim, dim); - } + 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); + size_t init_id = int_dist(generator); + size_t num_picked = 1; - double dart_val; - size_t tmp_pivot; - bool sum_flag = false; + picked.push_back(init_id); + std::memcpy(pivot_data, data + init_id * dim, dim * sizeof(float)); - while (num_picked < num_centers) { - dart_val = distribution(generator); + float *dist = new float[num_points]; - double sum = 0; - for (size_t i = 0; i < num_points; i++) { - sum = sum + dist[i]; +#pragma omp parallel for schedule(static, 8192) + for (int64_t i = 0; i < (int64_t)num_points; i++) + { + dist[i] = math_utils::calc_distance(data + i * dim, data + init_id * dim, dim); } - if (sum == 0) sum_flag = true; - dart_val *= sum; + double dart_val; + size_t tmp_pivot; + bool sum_flag = false; - 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]) { - break; - } + while (num_picked < num_centers) + { + dart_val = distribution(generator); - prefix_sum += dist[i]; - } + double sum = 0; + for (size_t i = 0; i < num_points; i++) + { + sum = sum + dist[i]; + } + if (sum == 0) + sum_flag = true; - if (std::find(picked.begin(), picked.end(), tmp_pivot) != picked.end() && - (sum_flag == false)) - continue; - picked.push_back(tmp_pivot); - std::memcpy(pivot_data + num_picked * dim, data + tmp_pivot * dim, - dim * sizeof(float)); + dart_val *= sum; + + 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]) + { + break; + } + + prefix_sum += dist[i]; + } + + if (std::find(picked.begin(), picked.end(), tmp_pivot) != picked.end() && (sum_flag == false)) + continue; + picked.push_back(tmp_pivot); + std::memcpy(pivot_data + num_picked * dim, data + tmp_pivot * dim, dim * sizeof(float)); #pragma omp parallel for schedule(static, 8192) - for (int64_t i = 0; i < (int64_t)num_points; i++) { - dist[i] = - (std::min)(dist[i], math_utils::calc_distance( - data + i * dim, data + tmp_pivot * dim, dim)); + for (int64_t i = 0; i < (int64_t)num_points; i++) + { + dist[i] = (std::min)(dist[i], math_utils::calc_distance(data + i * dim, data + tmp_pivot * dim, dim)); + } + num_picked++; } - num_picked++; - } - delete[] dist; + delete[] dist; } -} // namespace kmeans +} // namespace kmeans diff --git a/src/memory_mapper.cpp b/src/memory_mapper.cpp index 2f6d4fd1f..d1c5ef984 100644 --- a/src/memory_mapper.cpp +++ b/src/memory_mapper.cpp @@ -8,86 +8,100 @@ using namespace diskann; -MemoryMapper::MemoryMapper(const std::string &filename) - : MemoryMapper(filename.c_str()) {} +MemoryMapper::MemoryMapper(const std::string &filename) : MemoryMapper(filename.c_str()) +{ +} -MemoryMapper::MemoryMapper(const char *filename) { +MemoryMapper::MemoryMapper(const char *filename) +{ #ifndef _WINDOWS - _fd = open(filename, O_RDONLY); - if (_fd <= 0) { - std::cerr << "Inner vertices file not found" << std::endl; - return; - } - struct stat sb; - if (fstat(_fd, &sb) != 0) { - std::cerr << "Inner vertices file not dound. " << std::endl; - return; - } - _fileSize = sb.st_size; - diskann::cout << "File Size: " << _fileSize << std::endl; - _buf = (char *)mmap(NULL, _fileSize, PROT_READ, MAP_PRIVATE, _fd, 0); + _fd = open(filename, O_RDONLY); + if (_fd <= 0) + { + std::cerr << "Inner vertices file not found" << std::endl; + return; + } + struct stat sb; + if (fstat(_fd, &sb) != 0) + { + std::cerr << "Inner vertices file not dound. " << std::endl; + return; + } + _fileSize = sb.st_size; + diskann::cout << "File Size: " << _fileSize << std::endl; + _buf = (char *)mmap(NULL, _fileSize, PROT_READ, MAP_PRIVATE, _fd, 0); #else - _bareFile = CreateFileA(filename, GENERIC_READ | GENERIC_EXECUTE, 0, NULL, - OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL); - if (_bareFile == nullptr) { - std::ostringstream message; - message << "CreateFileA(" << filename << ") failed with error " - << GetLastError() << std::endl; - std::cerr << message.str(); - throw std::exception(message.str().c_str()); - } + _bareFile = + CreateFileA(filename, GENERIC_READ | GENERIC_EXECUTE, 0, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL); + if (_bareFile == nullptr) + { + std::ostringstream message; + message << "CreateFileA(" << filename << ") failed with error " << GetLastError() << std::endl; + std::cerr << message.str(); + throw std::exception(message.str().c_str()); + } - _fd = CreateFileMapping(_bareFile, NULL, PAGE_EXECUTE_READ, 0, 0, NULL); - if (_fd == nullptr) { - std::ostringstream message; - message << "CreateFileMapping(" << filename << ") failed with error " - << GetLastError() << std::endl; - std::cerr << message.str() << std::endl; - throw std::exception(message.str().c_str()); - } + _fd = CreateFileMapping(_bareFile, NULL, PAGE_EXECUTE_READ, 0, 0, NULL); + if (_fd == nullptr) + { + std::ostringstream message; + message << "CreateFileMapping(" << filename << ") failed with error " << GetLastError() << std::endl; + std::cerr << message.str() << std::endl; + throw std::exception(message.str().c_str()); + } - _buf = (char *)MapViewOfFile(_fd, FILE_MAP_READ, 0, 0, 0); - if (_buf == nullptr) { - std::ostringstream message; - message << "MapViewOfFile(" << filename - << ") failed with error: " << GetLastError() << std::endl; - std::cerr << message.str() << std::endl; - throw std::exception(message.str().c_str()); - } + _buf = (char *)MapViewOfFile(_fd, FILE_MAP_READ, 0, 0, 0); + if (_buf == nullptr) + { + std::ostringstream message; + message << "MapViewOfFile(" << filename << ") failed with error: " << GetLastError() << std::endl; + std::cerr << message.str() << std::endl; + throw std::exception(message.str().c_str()); + } - LARGE_INTEGER fSize; - if (TRUE == GetFileSizeEx(_bareFile, &fSize)) { - _fileSize = fSize.QuadPart; // take the 64-bit value - diskann::cout << "File Size: " << _fileSize << std::endl; - } else { - std::cerr << "Failed to get size of file " << filename << std::endl; - } + LARGE_INTEGER fSize; + if (TRUE == GetFileSizeEx(_bareFile, &fSize)) + { + _fileSize = fSize.QuadPart; // take the 64-bit value + diskann::cout << "File Size: " << _fileSize << std::endl; + } + else + { + std::cerr << "Failed to get size of file " << filename << std::endl; + } #endif } -char *MemoryMapper::getBuf() { return _buf; } +char *MemoryMapper::getBuf() +{ + return _buf; +} -size_t MemoryMapper::getFileSize() { return _fileSize; } +size_t MemoryMapper::getFileSize() +{ + return _fileSize; +} -MemoryMapper::~MemoryMapper() { +MemoryMapper::~MemoryMapper() +{ #ifndef _WINDOWS - if (munmap(_buf, _fileSize) != 0) - std::cerr << "ERROR unmapping. CHECK!" << std::endl; - close(_fd); + if (munmap(_buf, _fileSize) != 0) + std::cerr << "ERROR unmapping. CHECK!" << std::endl; + close(_fd); #else - if (FALSE == UnmapViewOfFile(_buf)) { - std::cerr << "Unmap view of file failed. Error: " << GetLastError() - << std::endl; - } + if (FALSE == UnmapViewOfFile(_buf)) + { + std::cerr << "Unmap view of file failed. Error: " << GetLastError() << std::endl; + } - if (FALSE == CloseHandle(_fd)) { - std::cerr << "Failed to close memory mapped file. Error: " << GetLastError() - << std::endl; - } + if (FALSE == CloseHandle(_fd)) + { + std::cerr << "Failed to close memory mapped file. Error: " << GetLastError() << std::endl; + } - if (FALSE == CloseHandle(_bareFile)) { - std::cerr << "Failed to close file: " << _fileName - << " Error: " << GetLastError() << std::endl; - } + if (FALSE == CloseHandle(_bareFile)) + { + std::cerr << "Failed to close file: " << _fileName << " Error: " << GetLastError() << std::endl; + } #endif } diff --git a/src/natural_number_map.cpp b/src/natural_number_map.cpp index 59e742d9b..9050831a2 100644 --- a/src/natural_number_map.cpp +++ b/src/natural_number_map.cpp @@ -6,98 +6,104 @@ #include "natural_number_map.h" -namespace diskann { +namespace diskann +{ static constexpr auto invalid_position = boost::dynamic_bitset<>::npos; template natural_number_map::natural_number_map() - : _size(0), _values_bitset(std::make_unique>()) {} + : _size(0), _values_bitset(std::make_unique>()) +{ +} -template -void natural_number_map::reserve(size_t count) { - _values_vector.reserve(count); - _values_bitset->reserve(count); +template void natural_number_map::reserve(size_t count) +{ + _values_vector.reserve(count); + _values_bitset->reserve(count); } -template -size_t natural_number_map::size() const { - return _size; +template size_t natural_number_map::size() const +{ + return _size; } -template -void natural_number_map::set(Key key, Value value) { - if (key >= _values_bitset->size()) { - _values_bitset->resize(static_cast(key) + 1); - _values_vector.resize(_values_bitset->size()); - } - - _values_vector[key] = value; - const bool was_present = _values_bitset->test_set(key, true); - - if (!was_present) { - ++_size; - } +template void natural_number_map::set(Key key, Value value) +{ + if (key >= _values_bitset->size()) + { + _values_bitset->resize(static_cast(key) + 1); + _values_vector.resize(_values_bitset->size()); + } + + _values_vector[key] = value; + const bool was_present = _values_bitset->test_set(key, true); + + if (!was_present) + { + ++_size; + } } -template -void natural_number_map::erase(Key key) { - if (key < _values_bitset->size()) { - const bool was_present = _values_bitset->test_set(key, false); +template void natural_number_map::erase(Key key) +{ + if (key < _values_bitset->size()) + { + const bool was_present = _values_bitset->test_set(key, false); - if (was_present) { - --_size; + if (was_present) + { + --_size; + } } - } } -template -bool natural_number_map::contains(Key key) const { - return key < _values_bitset->size() && _values_bitset->test(key); +template bool natural_number_map::contains(Key key) const +{ + return key < _values_bitset->size() && _values_bitset->test(key); } -template -bool natural_number_map::try_get(Key key, Value &value) const { - if (!contains(key)) { - return false; - } +template bool natural_number_map::try_get(Key key, Value &value) const +{ + if (!contains(key)) + { + return false; + } - value = _values_vector[key]; - return true; + value = _values_vector[key]; + return true; } template -typename natural_number_map::position -natural_number_map::find_first() const { - return position{_size > 0 ? _values_bitset->find_first() : invalid_position, - 0}; +typename natural_number_map::position natural_number_map::find_first() const +{ + return position{_size > 0 ? _values_bitset->find_first() : invalid_position, 0}; } template -typename natural_number_map::position -natural_number_map::find_next( - const position &after_position) const { - return position{after_position._keys_already_enumerated < _size - ? _values_bitset->find_next(after_position._key) - : invalid_position, - after_position._keys_already_enumerated + 1}; +typename natural_number_map::position natural_number_map::find_next( + const position &after_position) const +{ + return position{after_position._keys_already_enumerated < _size ? _values_bitset->find_next(after_position._key) + : invalid_position, + after_position._keys_already_enumerated + 1}; } -template -bool natural_number_map::position::is_valid() const { - return _key != invalid_position; +template bool natural_number_map::position::is_valid() const +{ + return _key != invalid_position; } -template -Value natural_number_map::get(const position &pos) const { - assert(pos.is_valid()); - return _values_vector[pos._key]; +template Value natural_number_map::get(const position &pos) const +{ + assert(pos.is_valid()); + return _values_vector[pos._key]; } -template -void natural_number_map::clear() { - _size = 0; - _values_vector.clear(); - _values_bitset->clear(); +template void natural_number_map::clear() +{ + _size = 0; + _values_vector.clear(); + _values_bitset->clear(); } // Instantiate used templates. @@ -105,4 +111,4 @@ template class natural_number_map; template class natural_number_map; template class natural_number_map; template class natural_number_map; -} // namespace diskann +} // namespace diskann diff --git a/src/natural_number_set.cpp b/src/natural_number_set.cpp index 5cd7d6f71..b36cb5298 100644 --- a/src/natural_number_set.cpp +++ b/src/natural_number_set.cpp @@ -6,63 +6,65 @@ #include "ann_exception.h" #include "natural_number_set.h" -namespace diskann { +namespace diskann +{ template -natural_number_set::natural_number_set() - : _values_bitset(std::make_unique>()) {} +natural_number_set::natural_number_set() : _values_bitset(std::make_unique>()) +{ +} -template -bool natural_number_set::is_empty() const { - return _values_vector.empty(); +template bool natural_number_set::is_empty() const +{ + return _values_vector.empty(); } -template -void natural_number_set::reserve(size_t count) { - _values_vector.reserve(count); - _values_bitset->reserve(count); +template void natural_number_set::reserve(size_t count) +{ + _values_vector.reserve(count); + _values_bitset->reserve(count); } -template -void natural_number_set::insert(T id) { - _values_vector.emplace_back(id); +template void natural_number_set::insert(T id) +{ + _values_vector.emplace_back(id); - if (id >= _values_bitset->size()) - _values_bitset->resize(static_cast(id) + 1); + if (id >= _values_bitset->size()) + _values_bitset->resize(static_cast(id) + 1); - _values_bitset->set(id, true); + _values_bitset->set(id, true); } -template -T natural_number_set::pop_any() { - if (_values_vector.empty()) { - throw diskann::ANNException("No values available", -1, __FUNCSIG__, - __FILE__, __LINE__); - } +template T natural_number_set::pop_any() +{ + if (_values_vector.empty()) + { + throw diskann::ANNException("No values available", -1, __FUNCSIG__, __FILE__, __LINE__); + } - const T id = _values_vector.back(); - _values_vector.pop_back(); + const T id = _values_vector.back(); + _values_vector.pop_back(); - _values_bitset->set(id, false); + _values_bitset->set(id, false); - return id; + return id; } -template -void natural_number_set::clear() { - _values_vector.clear(); - _values_bitset->clear(); +template void natural_number_set::clear() +{ + _values_vector.clear(); + _values_bitset->clear(); } -template -size_t natural_number_set::size() const { - return _values_vector.size(); +template size_t natural_number_set::size() const +{ + return _values_vector.size(); } -template -bool natural_number_set::is_in_set(T id) const { - return _values_bitset->test(id); +template bool natural_number_set::is_in_set(T id) const +{ + return _values_bitset->test(id); } // Instantiate used templates. template class natural_number_set; -} // namespace diskann +} // namespace diskann diff --git a/src/partition.cpp b/src/partition.cpp index 386797e10..f35cadb83 100644 --- a/src/partition.cpp +++ b/src/partition.cpp @@ -11,8 +11,7 @@ #include "tsl/robin_map.h" #include "tsl/robin_set.h" -#if defined(RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS) && \ - defined(DISKANN_BUILD) +#if defined(RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS) && defined(DISKANN_BUILD) #include "gperftools/malloc_extension.h" #endif @@ -32,59 +31,56 @@ // #define SAVE_INFLATED_PQ true template -void gen_random_slice(const std::string base_file, - const std::string output_prefix, double sampling_rate) { - size_t read_blk_size = 64 * 1024 * 1024; - cached_ifstream base_reader(base_file.c_str(), read_blk_size); - std::ofstream sample_writer(std::string(output_prefix + "_data.bin").c_str(), - std::ios::binary); - std::ofstream sample_id_writer( - std::string(output_prefix + "_ids.bin").c_str(), std::ios::binary); - - std::random_device - rd; // Will be used to obtain a seed for the random number engine - auto x = rd(); - std::mt19937 generator( - x); // Standard mersenne_twister_engine seeded with rd() - std::uniform_real_distribution distribution(0, 1); - - size_t npts, nd; - uint32_t npts_u32, nd_u32; - uint32_t num_sampled_pts_u32 = 0; - uint32_t one_const = 1; - - base_reader.read((char *)&npts_u32, sizeof(uint32_t)); - base_reader.read((char *)&nd_u32, sizeof(uint32_t)); - diskann::cout << "Loading base " << base_file << ". #points: " << npts_u32 - << ". #dim: " << nd_u32 << "." << std::endl; - sample_writer.write((char *)&num_sampled_pts_u32, sizeof(uint32_t)); - sample_writer.write((char *)&nd_u32, sizeof(uint32_t)); - sample_id_writer.write((char *)&num_sampled_pts_u32, sizeof(uint32_t)); - sample_id_writer.write((char *)&one_const, sizeof(uint32_t)); - - npts = npts_u32; - nd = nd_u32; - std::unique_ptr cur_row = std::make_unique(nd); - - 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) { - sample_writer.write((char *)cur_row.get(), sizeof(T) * nd); - uint32_t cur_i_u32 = (uint32_t)i; - sample_id_writer.write((char *)&cur_i_u32, sizeof(uint32_t)); - num_sampled_pts_u32++; +void gen_random_slice(const std::string base_file, const std::string output_prefix, double sampling_rate) +{ + size_t read_blk_size = 64 * 1024 * 1024; + cached_ifstream base_reader(base_file.c_str(), read_blk_size); + std::ofstream sample_writer(std::string(output_prefix + "_data.bin").c_str(), std::ios::binary); + std::ofstream sample_id_writer(std::string(output_prefix + "_ids.bin").c_str(), std::ios::binary); + + std::random_device rd; // Will be used to obtain a seed for the random number engine + auto x = rd(); + std::mt19937 generator(x); // Standard mersenne_twister_engine seeded with rd() + std::uniform_real_distribution distribution(0, 1); + + size_t npts, nd; + uint32_t npts_u32, nd_u32; + uint32_t num_sampled_pts_u32 = 0; + uint32_t one_const = 1; + + base_reader.read((char *)&npts_u32, sizeof(uint32_t)); + base_reader.read((char *)&nd_u32, sizeof(uint32_t)); + diskann::cout << "Loading base " << base_file << ". #points: " << npts_u32 << ". #dim: " << nd_u32 << "." + << std::endl; + sample_writer.write((char *)&num_sampled_pts_u32, sizeof(uint32_t)); + sample_writer.write((char *)&nd_u32, sizeof(uint32_t)); + sample_id_writer.write((char *)&num_sampled_pts_u32, sizeof(uint32_t)); + sample_id_writer.write((char *)&one_const, sizeof(uint32_t)); + + npts = npts_u32; + nd = nd_u32; + std::unique_ptr cur_row = std::make_unique(nd); + + 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) + { + sample_writer.write((char *)cur_row.get(), sizeof(T) * nd); + uint32_t cur_i_u32 = (uint32_t)i; + sample_id_writer.write((char *)&cur_i_u32, sizeof(uint32_t)); + num_sampled_pts_u32++; + } } - } - 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); - sample_id_writer.write((char *)&num_sampled_pts_u32, sizeof(uint32_t)); - sample_writer.close(); - sample_id_writer.close(); - diskann::cout << "Wrote " << num_sampled_pts_u32 - << " points to sample file: " << output_prefix + "_data.bin" - << std::endl; + 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); + sample_id_writer.write((char *)&num_sampled_pts_u32, sizeof(uint32_t)); + sample_writer.close(); + sample_id_writer.close(); + diskann::cout << "Wrote " << num_sampled_pts_u32 << " points to sample file: " << output_prefix + "_data.bin" + << std::endl; } // streams data from the file, and samples each vector with probability p_val @@ -96,588 +92,565 @@ void gen_random_slice(const std::string base_file, ************************************/ template -void gen_random_slice(const std::string data_file, double p_val, - float *&sampled_data, size_t &slice_size, size_t &ndims) { - size_t npts; - uint32_t npts32, ndims32; - std::vector> sampled_vectors; - - // amount to read in one shot - size_t read_blk_size = 64 * 1024 * 1024; - // create cached reader + writer - cached_ifstream base_reader(data_file.c_str(), read_blk_size); - - // metadata: npts, ndims - base_reader.read((char *)&npts32, sizeof(uint32_t)); - base_reader.read((char *)&ndims32, sizeof(uint32_t)); - npts = npts32; - ndims = ndims32; - - std::unique_ptr cur_vector_T = std::make_unique(ndims); - p_val = p_val < 1 ? p_val : 1; - - std::random_device rd; // Will be used to obtain a seed for the random number - size_t x = rd(); - std::mt19937 generator((uint32_t)x); - std::uniform_real_distribution distribution(0, 1); - - 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) { - 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); +void gen_random_slice(const std::string data_file, double p_val, float *&sampled_data, size_t &slice_size, + size_t &ndims) +{ + size_t npts; + uint32_t npts32, ndims32; + std::vector> sampled_vectors; + + // amount to read in one shot + size_t read_blk_size = 64 * 1024 * 1024; + // create cached reader + writer + cached_ifstream base_reader(data_file.c_str(), read_blk_size); + + // metadata: npts, ndims + base_reader.read((char *)&npts32, sizeof(uint32_t)); + base_reader.read((char *)&ndims32, sizeof(uint32_t)); + npts = npts32; + ndims = ndims32; + + std::unique_ptr cur_vector_T = std::make_unique(ndims); + p_val = p_val < 1 ? p_val : 1; + + std::random_device rd; // Will be used to obtain a seed for the random number + size_t x = rd(); + std::mt19937 generator((uint32_t)x); + std::uniform_real_distribution distribution(0, 1); + + 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) + { + 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); + } } - } - slice_size = sampled_vectors.size(); - 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]; + slice_size = sampled_vectors.size(); + 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]; + } } - } } // same as above, but samples from the matrix inputdata instead of a file of // npts*ndims to return sampled_data of size slice_size*ndims. template -void gen_random_slice(const T *inputdata, size_t npts, size_t ndims, - double p_val, float *&sampled_data, size_t &slice_size) { - std::vector> sampled_vectors; - const T *cur_vector_T; - - p_val = p_val < 1 ? p_val : 1; - - std::random_device - rd; // Will be used to obtain a seed for the random number engine - size_t x = rd(); - std::mt19937 generator( - (uint32_t)x); // Standard mersenne_twister_engine seeded with rd() - std::uniform_real_distribution distribution(0, 1); - - for (size_t i = 0; i < npts; i++) { - cur_vector_T = inputdata + ndims * i; - float rnd_val = distribution(generator); - if (rnd_val < 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); +void gen_random_slice(const T *inputdata, size_t npts, size_t ndims, double p_val, float *&sampled_data, + size_t &slice_size) +{ + std::vector> sampled_vectors; + const T *cur_vector_T; + + p_val = p_val < 1 ? p_val : 1; + + std::random_device rd; // Will be used to obtain a seed for the random number engine + size_t x = rd(); + std::mt19937 generator((uint32_t)x); // Standard mersenne_twister_engine seeded with rd() + std::uniform_real_distribution distribution(0, 1); + + for (size_t i = 0; i < npts; i++) + { + cur_vector_T = inputdata + ndims * i; + float rnd_val = distribution(generator); + if (rnd_val < 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); + } } - } - slice_size = sampled_vectors.size(); - 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]; + slice_size = sampled_vectors.size(); + 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]; + } } - } } -int estimate_cluster_sizes(float *test_data_float, size_t num_test, - float *pivots, const size_t num_centers, - const size_t test_dim, const size_t k_base, - std::vector &cluster_sizes) { - cluster_sizes.clear(); +int estimate_cluster_sizes(float *test_data_float, size_t num_test, float *pivots, const size_t num_centers, + const size_t test_dim, const size_t k_base, std::vector &cluster_sizes) +{ + cluster_sizes.clear(); - size_t *shard_counts = new size_t[num_centers]; + size_t *shard_counts = new size_t[num_centers]; - for (size_t i = 0; i < num_centers; i++) { - shard_counts[i] = 0; - } + for (size_t i = 0; i < num_centers; i++) + { + shard_counts[i] = 0; + } - size_t block_size = num_test <= BLOCK_SIZE ? num_test : BLOCK_SIZE; - uint32_t *block_closest_centers = new uint32_t[block_size * k_base]; - float *block_data_float; + size_t block_size = num_test <= BLOCK_SIZE ? num_test : BLOCK_SIZE; + uint32_t *block_closest_centers = new uint32_t[block_size * k_base]; + float *block_data_float; - size_t num_blocks = DIV_ROUND_UP(num_test, block_size); + size_t num_blocks = DIV_ROUND_UP(num_test, block_size); - for (size_t block = 0; block < num_blocks; block++) { - size_t start_id = block * block_size; - size_t end_id = (std::min)((block + 1) * block_size, num_test); - size_t cur_blk_size = end_id - start_id; + for (size_t block = 0; block < num_blocks; block++) + { + size_t start_id = block * block_size; + size_t end_id = (std::min)((block + 1) * block_size, num_test); + size_t cur_blk_size = end_id - start_id; - block_data_float = test_data_float + start_id * test_dim; + block_data_float = test_data_float + start_id * test_dim; - math_utils::compute_closest_centers(block_data_float, cur_blk_size, - test_dim, pivots, num_centers, k_base, - block_closest_centers); + math_utils::compute_closest_centers(block_data_float, cur_blk_size, test_dim, pivots, num_centers, k_base, + block_closest_centers); - for (size_t p = 0; p < cur_blk_size; p++) { - for (size_t p1 = 0; p1 < k_base; p1++) { - size_t shard_id = block_closest_centers[p * k_base + p1]; - shard_counts[shard_id]++; - } + for (size_t p = 0; p < cur_blk_size; p++) + { + for (size_t p1 = 0; p1 < k_base; p1++) + { + size_t shard_id = block_closest_centers[p * k_base + p1]; + shard_counts[shard_id]++; + } + } } - } - - diskann::cout << "Estimated cluster sizes: "; - for (size_t i = 0; i < num_centers; i++) { - uint32_t cur_shard_count = (uint32_t)shard_counts[i]; - cluster_sizes.push_back((size_t)cur_shard_count); - diskann::cout << cur_shard_count << " "; - } - diskann::cout << std::endl; - delete[] shard_counts; - delete[] block_closest_centers; - return 0; -} -template -int shard_data_into_clusters(const std::string data_file, float *pivots, - const size_t num_centers, const size_t dim, - const size_t k_base, std::string prefix_path) { - size_t read_blk_size = 64 * 1024 * 1024; - // uint64_t write_blk_size = 64 * 1024 * 1024; - // create cached reader + writer - cached_ifstream base_reader(data_file, read_blk_size); - uint32_t npts32; - uint32_t basedim32; - base_reader.read((char *)&npts32, sizeof(uint32_t)); - base_reader.read((char *)&basedim32, sizeof(uint32_t)); - size_t num_points = npts32; - if (basedim32 != dim) { - diskann::cout << "Error. dimensions dont match for train set and base set" - << std::endl; - return -1; - } - - std::unique_ptr shard_counts = - std::make_unique(num_centers); - std::vector shard_data_writer(num_centers); - std::vector shard_idmap_writer(num_centers); - uint32_t dummy_size = 0; - uint32_t const_one = 1; - - for (size_t i = 0; i < num_centers; i++) { - std::string data_filename = - prefix_path + "_subshard-" + std::to_string(i) + ".bin"; - std::string idmap_filename = - prefix_path + "_subshard-" + std::to_string(i) + "_ids_uint32.bin"; - shard_data_writer[i] = - std::ofstream(data_filename.c_str(), std::ios::binary); - shard_idmap_writer[i] = - std::ofstream(idmap_filename.c_str(), std::ios::binary); - shard_data_writer[i].write((char *)&dummy_size, sizeof(uint32_t)); - shard_data_writer[i].write((char *)&basedim32, sizeof(uint32_t)); - shard_idmap_writer[i].write((char *)&dummy_size, sizeof(uint32_t)); - shard_idmap_writer[i].write((char *)&const_one, sizeof(uint32_t)); - shard_counts[i] = 0; - } - - size_t block_size = num_points <= BLOCK_SIZE ? num_points : BLOCK_SIZE; - std::unique_ptr block_closest_centers = - std::make_unique(block_size * k_base); - std::unique_ptr block_data_T = std::make_unique(block_size * dim); - std::unique_ptr block_data_float = - std::make_unique(block_size * dim); - - size_t num_blocks = DIV_ROUND_UP(num_points, block_size); - - for (size_t block = 0; block < num_blocks; block++) { - size_t start_id = block * block_size; - size_t end_id = (std::min)((block + 1) * block_size, num_points); - size_t cur_blk_size = end_id - start_id; - - base_reader.read((char *)block_data_T.get(), - sizeof(T) * (cur_blk_size * dim)); - diskann::convert_types(block_data_T.get(), block_data_float.get(), - cur_blk_size, dim); - - math_utils::compute_closest_centers(block_data_float.get(), cur_blk_size, - dim, pivots, num_centers, k_base, - block_closest_centers.get()); - - for (size_t p = 0; p < cur_blk_size; p++) { - for (size_t p1 = 0; p1 < k_base; p1++) { - size_t shard_id = block_closest_centers[p * k_base + p1]; - uint32_t original_point_map_id = (uint32_t)(start_id + p); - shard_data_writer[shard_id].write( - (char *)(block_data_T.get() + p * dim), sizeof(T) * dim); - shard_idmap_writer[shard_id].write((char *)&original_point_map_id, - sizeof(uint32_t)); - shard_counts[shard_id]++; - } + diskann::cout << "Estimated cluster sizes: "; + for (size_t i = 0; i < num_centers; i++) + { + uint32_t cur_shard_count = (uint32_t)shard_counts[i]; + cluster_sizes.push_back((size_t)cur_shard_count); + diskann::cout << cur_shard_count << " "; } - } - - size_t total_count = 0; - diskann::cout << "Actual shard sizes: " << std::flush; - for (size_t i = 0; i < num_centers; i++) { - uint32_t cur_shard_count = (uint32_t)shard_counts[i]; - total_count += cur_shard_count; - diskann::cout << cur_shard_count << " "; - shard_data_writer[i].seekp(0); - shard_data_writer[i].write((char *)&cur_shard_count, sizeof(uint32_t)); - shard_data_writer[i].close(); - shard_idmap_writer[i].seekp(0); - shard_idmap_writer[i].write((char *)&cur_shard_count, sizeof(uint32_t)); - shard_idmap_writer[i].close(); - } - - diskann::cout << "\n Partitioned " << num_points - << " with replication factor " << k_base << " to get " - << total_count << " points across " << num_centers << " shards " - << std::endl; - return 0; + diskann::cout << std::endl; + delete[] shard_counts; + delete[] block_closest_centers; + return 0; } -// useful for partitioning large dataset. we first generate only the IDS for -// each shard, and retrieve the actual vectors on demand. template -int shard_data_into_clusters_only_ids(const std::string data_file, - float *pivots, const size_t num_centers, - const size_t dim, const size_t k_base, - std::string prefix_path) { - size_t read_blk_size = 64 * 1024 * 1024; - // uint64_t write_blk_size = 64 * 1024 * 1024; - // create cached reader + writer - cached_ifstream base_reader(data_file, read_blk_size); - uint32_t npts32; - uint32_t basedim32; - base_reader.read((char *)&npts32, sizeof(uint32_t)); - base_reader.read((char *)&basedim32, sizeof(uint32_t)); - size_t num_points = npts32; - if (basedim32 != dim) { - diskann::cout << "Error. dimensions dont match for train set and base set" - << std::endl; - return -1; - } - - std::unique_ptr shard_counts = - std::make_unique(num_centers); - - std::vector shard_idmap_writer(num_centers); - uint32_t dummy_size = 0; - uint32_t const_one = 1; - - for (size_t i = 0; i < num_centers; i++) { - std::string idmap_filename = - prefix_path + "_subshard-" + std::to_string(i) + "_ids_uint32.bin"; - shard_idmap_writer[i] = - std::ofstream(idmap_filename.c_str(), std::ios::binary); - shard_idmap_writer[i].write((char *)&dummy_size, sizeof(uint32_t)); - shard_idmap_writer[i].write((char *)&const_one, sizeof(uint32_t)); - shard_counts[i] = 0; - } - - size_t block_size = num_points <= BLOCK_SIZE ? num_points : BLOCK_SIZE; - std::unique_ptr block_closest_centers = - std::make_unique(block_size * k_base); - std::unique_ptr block_data_T = std::make_unique(block_size * dim); - std::unique_ptr block_data_float = - std::make_unique(block_size * dim); - - size_t num_blocks = DIV_ROUND_UP(num_points, block_size); - - for (size_t block = 0; block < num_blocks; block++) { - size_t start_id = block * block_size; - size_t end_id = (std::min)((block + 1) * block_size, num_points); - size_t cur_blk_size = end_id - start_id; - - base_reader.read((char *)block_data_T.get(), - sizeof(T) * (cur_blk_size * dim)); - diskann::convert_types(block_data_T.get(), block_data_float.get(), - cur_blk_size, dim); - - math_utils::compute_closest_centers(block_data_float.get(), cur_blk_size, - dim, pivots, num_centers, k_base, - block_closest_centers.get()); - - for (size_t p = 0; p < cur_blk_size; p++) { - for (size_t p1 = 0; p1 < k_base; p1++) { - size_t shard_id = block_closest_centers[p * k_base + p1]; - uint32_t original_point_map_id = (uint32_t)(start_id + p); - shard_idmap_writer[shard_id].write((char *)&original_point_map_id, - sizeof(uint32_t)); - shard_counts[shard_id]++; - } +int shard_data_into_clusters(const std::string data_file, float *pivots, const size_t num_centers, const size_t dim, + const size_t k_base, std::string prefix_path) +{ + size_t read_blk_size = 64 * 1024 * 1024; + // uint64_t write_blk_size = 64 * 1024 * 1024; + // create cached reader + writer + cached_ifstream base_reader(data_file, read_blk_size); + uint32_t npts32; + uint32_t basedim32; + base_reader.read((char *)&npts32, sizeof(uint32_t)); + base_reader.read((char *)&basedim32, sizeof(uint32_t)); + size_t num_points = npts32; + if (basedim32 != dim) + { + diskann::cout << "Error. dimensions dont match for train set and base set" << std::endl; + return -1; } - } - - size_t total_count = 0; - diskann::cout << "Actual shard sizes: " << std::flush; - for (size_t i = 0; i < num_centers; i++) { - uint32_t cur_shard_count = (uint32_t)shard_counts[i]; - total_count += cur_shard_count; - diskann::cout << cur_shard_count << " "; - shard_idmap_writer[i].seekp(0); - shard_idmap_writer[i].write((char *)&cur_shard_count, sizeof(uint32_t)); - shard_idmap_writer[i].close(); - } - - diskann::cout << "\n Partitioned " << num_points - << " with replication factor " << k_base << " to get " - << total_count << " points across " << num_centers << " shards " - << std::endl; - return 0; -} -template -int retrieve_shard_data_from_ids(const std::string data_file, - std::string idmap_filename, - std::string data_filename) { - size_t read_blk_size = 64 * 1024 * 1024; - // uint64_t write_blk_size = 64 * 1024 * 1024; - // create cached reader + writer - cached_ifstream base_reader(data_file, read_blk_size); - uint32_t npts32; - uint32_t basedim32; - base_reader.read((char *)&npts32, sizeof(uint32_t)); - base_reader.read((char *)&basedim32, sizeof(uint32_t)); - size_t num_points = npts32; - size_t dim = basedim32; - - uint32_t dummy_size = 0; - - std::ofstream shard_data_writer(data_filename.c_str(), std::ios::binary); - shard_data_writer.write((char *)&dummy_size, sizeof(uint32_t)); - shard_data_writer.write((char *)&basedim32, sizeof(uint32_t)); - - uint32_t *shard_ids; - uint64_t shard_size, tmp; - diskann::load_bin(idmap_filename, shard_ids, shard_size, tmp); - - uint32_t cur_pos = 0; - uint32_t num_written = 0; - std::cout << "Shard has " << shard_size << " points" << std::endl; - - size_t block_size = num_points <= BLOCK_SIZE ? num_points : BLOCK_SIZE; - std::unique_ptr block_data_T = std::make_unique(block_size * dim); - - size_t num_blocks = DIV_ROUND_UP(num_points, block_size); - - for (size_t block = 0; block < num_blocks; block++) { - size_t start_id = block * block_size; - size_t end_id = (std::min)((block + 1) * block_size, num_points); - size_t cur_blk_size = end_id - start_id; - - base_reader.read((char *)block_data_T.get(), - sizeof(T) * (cur_blk_size * dim)); - - for (size_t p = 0; p < cur_blk_size; p++) { - uint32_t original_point_map_id = (uint32_t)(start_id + p); - if (cur_pos == shard_size) break; - if (original_point_map_id == shard_ids[cur_pos]) { - cur_pos++; - shard_data_writer.write((char *)(block_data_T.get() + p * dim), - sizeof(T) * dim); - num_written++; - } + std::unique_ptr shard_counts = std::make_unique(num_centers); + std::vector shard_data_writer(num_centers); + std::vector shard_idmap_writer(num_centers); + uint32_t dummy_size = 0; + uint32_t const_one = 1; + + for (size_t i = 0; i < num_centers; i++) + { + std::string data_filename = prefix_path + "_subshard-" + std::to_string(i) + ".bin"; + std::string idmap_filename = prefix_path + "_subshard-" + std::to_string(i) + "_ids_uint32.bin"; + shard_data_writer[i] = std::ofstream(data_filename.c_str(), std::ios::binary); + shard_idmap_writer[i] = std::ofstream(idmap_filename.c_str(), std::ios::binary); + shard_data_writer[i].write((char *)&dummy_size, sizeof(uint32_t)); + shard_data_writer[i].write((char *)&basedim32, sizeof(uint32_t)); + shard_idmap_writer[i].write((char *)&dummy_size, sizeof(uint32_t)); + shard_idmap_writer[i].write((char *)&const_one, sizeof(uint32_t)); + shard_counts[i] = 0; + } + + size_t block_size = num_points <= BLOCK_SIZE ? num_points : BLOCK_SIZE; + std::unique_ptr block_closest_centers = std::make_unique(block_size * k_base); + std::unique_ptr block_data_T = std::make_unique(block_size * dim); + std::unique_ptr block_data_float = std::make_unique(block_size * dim); + + size_t num_blocks = DIV_ROUND_UP(num_points, block_size); + + for (size_t block = 0; block < num_blocks; block++) + { + size_t start_id = block * block_size; + size_t end_id = (std::min)((block + 1) * block_size, num_points); + size_t cur_blk_size = end_id - start_id; + + base_reader.read((char *)block_data_T.get(), sizeof(T) * (cur_blk_size * dim)); + diskann::convert_types(block_data_T.get(), block_data_float.get(), cur_blk_size, dim); + + math_utils::compute_closest_centers(block_data_float.get(), cur_blk_size, dim, pivots, num_centers, k_base, + block_closest_centers.get()); + + for (size_t p = 0; p < cur_blk_size; p++) + { + for (size_t p1 = 0; p1 < k_base; p1++) + { + size_t shard_id = block_closest_centers[p * k_base + p1]; + uint32_t original_point_map_id = (uint32_t)(start_id + p); + shard_data_writer[shard_id].write((char *)(block_data_T.get() + p * dim), sizeof(T) * dim); + shard_idmap_writer[shard_id].write((char *)&original_point_map_id, sizeof(uint32_t)); + shard_counts[shard_id]++; + } + } } - if (cur_pos == shard_size) break; - } - diskann::cout << "Written file with " << num_written << " points" - << std::endl; + size_t total_count = 0; + diskann::cout << "Actual shard sizes: " << std::flush; + for (size_t i = 0; i < num_centers; i++) + { + uint32_t cur_shard_count = (uint32_t)shard_counts[i]; + total_count += cur_shard_count; + diskann::cout << cur_shard_count << " "; + shard_data_writer[i].seekp(0); + shard_data_writer[i].write((char *)&cur_shard_count, sizeof(uint32_t)); + shard_data_writer[i].close(); + shard_idmap_writer[i].seekp(0); + shard_idmap_writer[i].write((char *)&cur_shard_count, sizeof(uint32_t)); + shard_idmap_writer[i].close(); + } - shard_data_writer.seekp(0); - shard_data_writer.write((char *)&num_written, sizeof(uint32_t)); - shard_data_writer.close(); - delete[] shard_ids; - return 0; + diskann::cout << "\n Partitioned " << num_points << " with replication factor " << k_base << " to get " + << total_count << " points across " << num_centers << " shards " << std::endl; + return 0; } -// partitions a large base file into many shards using k-means hueristic -// on a random sample generated using sampling_rate probability. After this, it -// assignes each base point to the closest k_base nearest centers and creates -// the shards. -// The total number of points across all shards will be k_base * num_points. - +// useful for partitioning large dataset. we first generate only the IDS for +// each shard, and retrieve the actual vectors on demand. template -int partition(const std::string data_file, const float sampling_rate, - size_t num_parts, size_t max_k_means_reps, - const std::string prefix_path, size_t k_base) { - size_t train_dim; - size_t num_train; - float *train_data_float; - - gen_random_slice(data_file, sampling_rate, train_data_float, num_train, - train_dim); - - float *pivot_data; - - std::string cur_file = std::string(prefix_path); - std::string output_file; - - // kmeans_partitioning on training data - - // cur_file = cur_file + "_kmeans_partitioning-" + - // std::to_string(num_parts); - output_file = cur_file + "_centroids.bin"; +int shard_data_into_clusters_only_ids(const std::string data_file, float *pivots, const size_t num_centers, + const size_t dim, const size_t k_base, std::string prefix_path) +{ + size_t read_blk_size = 64 * 1024 * 1024; + // uint64_t write_blk_size = 64 * 1024 * 1024; + // create cached reader + writer + cached_ifstream base_reader(data_file, read_blk_size); + uint32_t npts32; + uint32_t basedim32; + base_reader.read((char *)&npts32, sizeof(uint32_t)); + base_reader.read((char *)&basedim32, sizeof(uint32_t)); + size_t num_points = npts32; + if (basedim32 != dim) + { + diskann::cout << "Error. dimensions dont match for train set and base set" << std::endl; + return -1; + } - pivot_data = new float[num_parts * train_dim]; + std::unique_ptr shard_counts = std::make_unique(num_centers); - // Process Global k-means for kmeans_partitioning Step - diskann::cout << "Processing global k-means (kmeans_partitioning Step)" - << std::endl; - kmeans::kmeanspp_selecting_pivots(train_data_float, num_train, train_dim, - pivot_data, num_parts); + std::vector shard_idmap_writer(num_centers); + uint32_t dummy_size = 0; + uint32_t const_one = 1; - kmeans::run_lloyds(train_data_float, num_train, train_dim, pivot_data, - num_parts, max_k_means_reps, NULL, NULL); + for (size_t i = 0; i < num_centers; i++) + { + std::string idmap_filename = prefix_path + "_subshard-" + std::to_string(i) + "_ids_uint32.bin"; + shard_idmap_writer[i] = std::ofstream(idmap_filename.c_str(), std::ios::binary); + shard_idmap_writer[i].write((char *)&dummy_size, sizeof(uint32_t)); + shard_idmap_writer[i].write((char *)&const_one, sizeof(uint32_t)); + shard_counts[i] = 0; + } - diskann::cout << "Saving global k-center pivots" << std::endl; - diskann::save_bin(output_file.c_str(), pivot_data, (size_t)num_parts, - train_dim); + size_t block_size = num_points <= BLOCK_SIZE ? num_points : BLOCK_SIZE; + std::unique_ptr block_closest_centers = std::make_unique(block_size * k_base); + std::unique_ptr block_data_T = std::make_unique(block_size * dim); + std::unique_ptr block_data_float = std::make_unique(block_size * dim); + + size_t num_blocks = DIV_ROUND_UP(num_points, block_size); + + for (size_t block = 0; block < num_blocks; block++) + { + size_t start_id = block * block_size; + size_t end_id = (std::min)((block + 1) * block_size, num_points); + size_t cur_blk_size = end_id - start_id; + + base_reader.read((char *)block_data_T.get(), sizeof(T) * (cur_blk_size * dim)); + diskann::convert_types(block_data_T.get(), block_data_float.get(), cur_blk_size, dim); + + math_utils::compute_closest_centers(block_data_float.get(), cur_blk_size, dim, pivots, num_centers, k_base, + block_closest_centers.get()); + + for (size_t p = 0; p < cur_blk_size; p++) + { + for (size_t p1 = 0; p1 < k_base; p1++) + { + size_t shard_id = block_closest_centers[p * k_base + p1]; + uint32_t original_point_map_id = (uint32_t)(start_id + p); + shard_idmap_writer[shard_id].write((char *)&original_point_map_id, sizeof(uint32_t)); + shard_counts[shard_id]++; + } + } + } - // now pivots are ready. need to stream base points and assign them to - // closest clusters. + size_t total_count = 0; + diskann::cout << "Actual shard sizes: " << std::flush; + for (size_t i = 0; i < num_centers; i++) + { + uint32_t cur_shard_count = (uint32_t)shard_counts[i]; + total_count += cur_shard_count; + diskann::cout << cur_shard_count << " "; + shard_idmap_writer[i].seekp(0); + shard_idmap_writer[i].write((char *)&cur_shard_count, sizeof(uint32_t)); + shard_idmap_writer[i].close(); + } - shard_data_into_clusters(data_file, pivot_data, num_parts, train_dim, - k_base, prefix_path); - delete[] pivot_data; - delete[] train_data_float; - return 0; + diskann::cout << "\n Partitioned " << num_points << " with replication factor " << k_base << " to get " + << total_count << " points across " << num_centers << " shards " << std::endl; + return 0; } template -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) { - size_t train_dim; - size_t num_train; - float *train_data_float; - size_t max_k_means_reps = 10; +int retrieve_shard_data_from_ids(const std::string data_file, std::string idmap_filename, std::string data_filename) +{ + size_t read_blk_size = 64 * 1024 * 1024; + // uint64_t write_blk_size = 64 * 1024 * 1024; + // create cached reader + writer + cached_ifstream base_reader(data_file, read_blk_size); + uint32_t npts32; + uint32_t basedim32; + base_reader.read((char *)&npts32, sizeof(uint32_t)); + base_reader.read((char *)&basedim32, sizeof(uint32_t)); + size_t num_points = npts32; + size_t dim = basedim32; + + uint32_t dummy_size = 0; + + std::ofstream shard_data_writer(data_filename.c_str(), std::ios::binary); + shard_data_writer.write((char *)&dummy_size, sizeof(uint32_t)); + shard_data_writer.write((char *)&basedim32, sizeof(uint32_t)); + + uint32_t *shard_ids; + uint64_t shard_size, tmp; + diskann::load_bin(idmap_filename, shard_ids, shard_size, tmp); + + uint32_t cur_pos = 0; + uint32_t num_written = 0; + std::cout << "Shard has " << shard_size << " points" << std::endl; + + size_t block_size = num_points <= BLOCK_SIZE ? num_points : BLOCK_SIZE; + std::unique_ptr block_data_T = std::make_unique(block_size * dim); + + size_t num_blocks = DIV_ROUND_UP(num_points, block_size); + + for (size_t block = 0; block < num_blocks; block++) + { + size_t start_id = block * block_size; + size_t end_id = (std::min)((block + 1) * block_size, num_points); + size_t cur_blk_size = end_id - start_id; + + base_reader.read((char *)block_data_T.get(), sizeof(T) * (cur_blk_size * dim)); + + for (size_t p = 0; p < cur_blk_size; p++) + { + uint32_t original_point_map_id = (uint32_t)(start_id + p); + if (cur_pos == shard_size) + break; + if (original_point_map_id == shard_ids[cur_pos]) + { + cur_pos++; + shard_data_writer.write((char *)(block_data_T.get() + p * dim), sizeof(T) * dim); + num_written++; + } + } + if (cur_pos == shard_size) + break; + } - int num_parts = 3; - bool fit_in_ram = false; + diskann::cout << "Written file with " << num_written << " points" << std::endl; - gen_random_slice(data_file, sampling_rate, train_data_float, num_train, - train_dim); + shard_data_writer.seekp(0); + shard_data_writer.write((char *)&num_written, sizeof(uint32_t)); + shard_data_writer.close(); + delete[] shard_ids; + return 0; +} - size_t test_dim; - size_t num_test; - float *test_data_float; - gen_random_slice(data_file, sampling_rate, test_data_float, num_test, - test_dim); +// partitions a large base file into many shards using k-means hueristic +// on a random sample generated using sampling_rate probability. After this, it +// assignes each base point to the closest k_base nearest centers and creates +// the shards. +// The total number of points across all shards will be k_base * num_points. - float *pivot_data = nullptr; +template +int partition(const std::string data_file, const float sampling_rate, size_t num_parts, size_t max_k_means_reps, + const std::string prefix_path, size_t k_base) +{ + size_t train_dim; + size_t num_train; + float *train_data_float; - std::string cur_file = std::string(prefix_path); - std::string output_file; + gen_random_slice(data_file, sampling_rate, train_data_float, num_train, train_dim); - // kmeans_partitioning on training data + float *pivot_data; - // cur_file = cur_file + "_kmeans_partitioning-" + - // std::to_string(num_parts); - output_file = cur_file + "_centroids.bin"; + std::string cur_file = std::string(prefix_path); + std::string output_file; - while (!fit_in_ram) { - fit_in_ram = true; + // kmeans_partitioning on training data - double max_ram_usage = 0; - if (pivot_data != nullptr) delete[] pivot_data; + // cur_file = cur_file + "_kmeans_partitioning-" + + // std::to_string(num_parts); + output_file = cur_file + "_centroids.bin"; pivot_data = new float[num_parts * train_dim]; + // Process Global k-means for kmeans_partitioning Step - diskann::cout << "Processing global k-means (kmeans_partitioning Step)" - << std::endl; - kmeans::kmeanspp_selecting_pivots(train_data_float, num_train, train_dim, - pivot_data, num_parts); + diskann::cout << "Processing global k-means (kmeans_partitioning Step)" << std::endl; + kmeans::kmeanspp_selecting_pivots(train_data_float, num_train, train_dim, pivot_data, num_parts); + + kmeans::run_lloyds(train_data_float, num_train, train_dim, pivot_data, num_parts, max_k_means_reps, NULL, NULL); - kmeans::run_lloyds(train_data_float, num_train, train_dim, pivot_data, - num_parts, max_k_means_reps, NULL, NULL); + diskann::cout << "Saving global k-center pivots" << std::endl; + diskann::save_bin(output_file.c_str(), pivot_data, (size_t)num_parts, train_dim); // now pivots are ready. need to stream base points and assign them to // closest clusters. - std::vector cluster_sizes; - estimate_cluster_sizes(test_data_float, num_test, pivot_data, num_parts, - train_dim, k_base, cluster_sizes); - - for (auto &p : cluster_sizes) { - // to account for the fact that p is the size of the shard over the - // testing sample. - p = (uint64_t)(p / sampling_rate); - double cur_shard_ram_estimate = - diskann::estimate_ram_usage(p, train_dim, sizeof(T), graph_degree); + shard_data_into_clusters(data_file, pivot_data, num_parts, train_dim, k_base, prefix_path); + delete[] pivot_data; + delete[] train_data_float; + return 0; +} - if (cur_shard_ram_estimate > max_ram_usage) - max_ram_usage = cur_shard_ram_estimate; - } - diskann::cout << "With " << num_parts << " parts, max estimated RAM usage: " - << max_ram_usage / (1024 * 1024 * 1024) - << "GB, budget given is " << ram_budget << std::endl; - if (max_ram_usage > 1024 * 1024 * 1024 * ram_budget) { - fit_in_ram = false; - num_parts += 2; +template +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) +{ + size_t train_dim; + size_t num_train; + float *train_data_float; + size_t max_k_means_reps = 10; + + int num_parts = 3; + bool fit_in_ram = false; + + gen_random_slice(data_file, sampling_rate, train_data_float, num_train, train_dim); + + size_t test_dim; + size_t num_test; + float *test_data_float; + gen_random_slice(data_file, sampling_rate, test_data_float, num_test, test_dim); + + float *pivot_data = nullptr; + + std::string cur_file = std::string(prefix_path); + std::string output_file; + + // kmeans_partitioning on training data + + // cur_file = cur_file + "_kmeans_partitioning-" + + // std::to_string(num_parts); + output_file = cur_file + "_centroids.bin"; + + while (!fit_in_ram) + { + fit_in_ram = true; + + double max_ram_usage = 0; + if (pivot_data != nullptr) + delete[] pivot_data; + + pivot_data = new float[num_parts * train_dim]; + // Process Global k-means for kmeans_partitioning Step + diskann::cout << "Processing global k-means (kmeans_partitioning Step)" << std::endl; + kmeans::kmeanspp_selecting_pivots(train_data_float, num_train, train_dim, pivot_data, num_parts); + + kmeans::run_lloyds(train_data_float, num_train, train_dim, pivot_data, num_parts, max_k_means_reps, NULL, NULL); + + // now pivots are ready. need to stream base points and assign them to + // closest clusters. + + std::vector cluster_sizes; + estimate_cluster_sizes(test_data_float, num_test, pivot_data, num_parts, train_dim, k_base, cluster_sizes); + + for (auto &p : cluster_sizes) + { + // to account for the fact that p is the size of the shard over the + // testing sample. + p = (uint64_t)(p / sampling_rate); + double cur_shard_ram_estimate = diskann::estimate_ram_usage(p, train_dim, sizeof(T), graph_degree); + + if (cur_shard_ram_estimate > max_ram_usage) + max_ram_usage = cur_shard_ram_estimate; + } + diskann::cout << "With " << num_parts + << " parts, max estimated RAM usage: " << max_ram_usage / (1024 * 1024 * 1024) + << "GB, budget given is " << ram_budget << std::endl; + if (max_ram_usage > 1024 * 1024 * 1024 * ram_budget) + { + fit_in_ram = false; + num_parts += 2; + } } - } - - diskann::cout << "Saving global k-center pivots" << std::endl; - diskann::save_bin(output_file.c_str(), pivot_data, (size_t)num_parts, - train_dim); - - shard_data_into_clusters_only_ids(data_file, pivot_data, num_parts, - train_dim, k_base, prefix_path); - delete[] pivot_data; - delete[] train_data_float; - delete[] test_data_float; - return num_parts; + + diskann::cout << "Saving global k-center pivots" << std::endl; + diskann::save_bin(output_file.c_str(), pivot_data, (size_t)num_parts, train_dim); + + shard_data_into_clusters_only_ids(data_file, pivot_data, num_parts, train_dim, k_base, prefix_path); + delete[] pivot_data; + delete[] train_data_float; + delete[] test_data_float; + return num_parts; } // Instantations of supported templates -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); -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 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, - float *&sampled_data, size_t &slice_size); -template void DISKANN_DLLEXPORT gen_random_slice( - const int8_t *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 std::string data_file, double p_val, float *&sampled_data, - size_t &slice_size, size_t &ndims); -template void DISKANN_DLLEXPORT gen_random_slice( - const std::string data_file, double p_val, float *&sampled_data, - size_t &slice_size, size_t &ndims); -template void DISKANN_DLLEXPORT gen_random_slice( - const std::string data_file, double p_val, float *&sampled_data, - size_t &slice_size, size_t &ndims); - -template DISKANN_DLLEXPORT int partition( - const std::string data_file, const float sampling_rate, size_t num_centers, - size_t max_k_means_reps, const std::string prefix_path, size_t k_base); -template DISKANN_DLLEXPORT int partition( - const std::string data_file, const float sampling_rate, size_t num_centers, - size_t max_k_means_reps, const std::string prefix_path, size_t k_base); -template DISKANN_DLLEXPORT int partition( - const std::string data_file, const float sampling_rate, size_t num_centers, - size_t max_k_means_reps, const std::string prefix_path, size_t k_base); - -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 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 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 retrieve_shard_data_from_ids( - const std::string data_file, std::string idmap_filename, - std::string data_filename); -template DISKANN_DLLEXPORT int retrieve_shard_data_from_ids( - const std::string data_file, std::string idmap_filename, - std::string data_filename); -template DISKANN_DLLEXPORT int retrieve_shard_data_from_ids( - const std::string data_file, std::string idmap_filename, - std::string data_filename); \ No newline at end of file +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); +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 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, float *&sampled_data, size_t &slice_size); +template void DISKANN_DLLEXPORT gen_random_slice(const int8_t *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 std::string data_file, double p_val, float *&sampled_data, + size_t &slice_size, size_t &ndims); +template void DISKANN_DLLEXPORT gen_random_slice(const std::string data_file, double p_val, + float *&sampled_data, size_t &slice_size, size_t &ndims); +template void DISKANN_DLLEXPORT gen_random_slice(const std::string data_file, double p_val, + float *&sampled_data, size_t &slice_size, size_t &ndims); + +template DISKANN_DLLEXPORT int partition(const std::string data_file, const float sampling_rate, + size_t num_centers, size_t max_k_means_reps, + const std::string prefix_path, size_t k_base); +template DISKANN_DLLEXPORT int partition(const std::string data_file, const float sampling_rate, + size_t num_centers, size_t max_k_means_reps, + const std::string prefix_path, size_t k_base); +template DISKANN_DLLEXPORT int partition(const std::string data_file, const float sampling_rate, + size_t num_centers, size_t max_k_means_reps, + const std::string prefix_path, size_t k_base); + +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 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 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 retrieve_shard_data_from_ids(const std::string data_file, + std::string idmap_filename, + std::string data_filename); +template DISKANN_DLLEXPORT int retrieve_shard_data_from_ids(const std::string data_file, + std::string idmap_filename, + std::string data_filename); +template DISKANN_DLLEXPORT int retrieve_shard_data_from_ids(const std::string data_file, + std::string idmap_filename, + std::string data_filename); \ No newline at end of file diff --git a/src/pq.cpp b/src/pq.cpp index 3d453c046..3f976d66d 100644 --- a/src/pq.cpp +++ b/src/pq.cpp @@ -11,307 +11,333 @@ // block size for reading/processing large files and matrices in blocks #define BLOCK_SIZE 5000000 -namespace diskann { -FixedChunkPQTable::FixedChunkPQTable() {} +namespace diskann +{ +FixedChunkPQTable::FixedChunkPQTable() +{ +} -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; + 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) { +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) { +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"; + 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); + 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); + 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__); - } + 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]); + 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]); + 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__); - } + 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; + this->ndims = nc; #ifdef EXEC_ENV_OLS - diskann::load_bin(files, pq_table_file, centroid, nr, nc, - file_offset_data[1]); + 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]); + 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; - } + 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]); + 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]); + 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; - - if (file_exists(rotmat_file)) { + 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; + + if (file_exists(rotmat_file)) + { #ifdef EXEC_ENV_OLS - diskann::load_bin(files, rotmat_file, (float *&)rotmat_tr, nr, nc); + diskann::load_bin(files, rotmat_file, (float *&)rotmat_tr, nr, nc); #else - diskann::load_bin(rotmat_file, rotmat_tr, nr, nc); + 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__); + 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; } - 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]; + + // 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); +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]; - } +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)); } - 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); - } +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; +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; + 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; +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) + 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::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 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++) { - memcpy(out + i * ndims, all_coords + ids[i] * ndims, - ndims * sizeof(uint8_t)); - } +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++) + { + memcpy(out + i * ndims, all_coords + ids[i] * ndims, ndims * sizeof(uint8_t)); + } } -void pq_dist_lookup(const uint8_t *pq_ids, const size_t n_pts, - const size_t pq_nchunks, const float *pq_dists, - std::vector &dists_out) { - //_mm_prefetch((char*) dists_out, _MM_HINT_T0); - _mm_prefetch((char *)pq_ids, _MM_HINT_T0); - _mm_prefetch((char *)(pq_ids + 64), _MM_HINT_T0); - _mm_prefetch((char *)(pq_ids + 128), _MM_HINT_T0); - dists_out.clear(); - dists_out.resize(n_pts, 0); - for (size_t chunk = 0; chunk < pq_nchunks; chunk++) { - const float *chunk_dists = pq_dists + 256 * chunk; - if (chunk < pq_nchunks - 1) { - _mm_prefetch((char *)(chunk_dists + 256), _MM_HINT_T0); - } - for (size_t idx = 0; idx < n_pts; idx++) { - uint8_t pq_centerid = pq_ids[pq_nchunks * idx + chunk]; - dists_out[idx] += chunk_dists[pq_centerid]; +void pq_dist_lookup(const uint8_t *pq_ids, const size_t n_pts, const size_t pq_nchunks, const float *pq_dists, + std::vector &dists_out) +{ + //_mm_prefetch((char*) dists_out, _MM_HINT_T0); + _mm_prefetch((char *)pq_ids, _MM_HINT_T0); + _mm_prefetch((char *)(pq_ids + 64), _MM_HINT_T0); + _mm_prefetch((char *)(pq_ids + 128), _MM_HINT_T0); + dists_out.clear(); + dists_out.resize(n_pts, 0); + for (size_t chunk = 0; chunk < pq_nchunks; chunk++) + { + const float *chunk_dists = pq_dists + 256 * chunk; + if (chunk < pq_nchunks - 1) + { + _mm_prefetch((char *)(chunk_dists + 256), _MM_HINT_T0); + } + for (size_t idx = 0; idx < n_pts; idx++) + { + uint8_t pq_centerid = pq_ids[pq_nchunks * idx + chunk]; + dists_out[idx] += chunk_dists[pq_centerid]; + } } - } } // Need to replace calls to these functions with calls to vector& based // functions above -void aggregate_coords(const uint32_t *ids, const size_t n_ids, - const uint8_t *all_coords, const size_t ndims, - uint8_t *out) { - for (size_t i = 0; i < n_ids; i++) { - memcpy(out + i * ndims, all_coords + ids[i] * ndims, - ndims * sizeof(uint8_t)); - } +void aggregate_coords(const uint32_t *ids, const size_t n_ids, const uint8_t *all_coords, const size_t ndims, + uint8_t *out) +{ + for (size_t i = 0; i < n_ids; i++) + { + memcpy(out + i * ndims, all_coords + ids[i] * ndims, ndims * sizeof(uint8_t)); + } } -void pq_dist_lookup(const uint8_t *pq_ids, const size_t n_pts, - const size_t pq_nchunks, const float *pq_dists, - float *dists_out) { - _mm_prefetch((char *)dists_out, _MM_HINT_T0); - _mm_prefetch((char *)pq_ids, _MM_HINT_T0); - _mm_prefetch((char *)(pq_ids + 64), _MM_HINT_T0); - _mm_prefetch((char *)(pq_ids + 128), _MM_HINT_T0); - memset(dists_out, 0, n_pts * sizeof(float)); - for (size_t chunk = 0; chunk < pq_nchunks; chunk++) { - const float *chunk_dists = pq_dists + 256 * chunk; - if (chunk < pq_nchunks - 1) { - _mm_prefetch((char *)(chunk_dists + 256), _MM_HINT_T0); - } - for (size_t idx = 0; idx < n_pts; idx++) { - uint8_t pq_centerid = pq_ids[pq_nchunks * idx + chunk]; - dists_out[idx] += chunk_dists[pq_centerid]; +void pq_dist_lookup(const uint8_t *pq_ids, const size_t n_pts, const size_t pq_nchunks, const float *pq_dists, + float *dists_out) +{ + _mm_prefetch((char *)dists_out, _MM_HINT_T0); + _mm_prefetch((char *)pq_ids, _MM_HINT_T0); + _mm_prefetch((char *)(pq_ids + 64), _MM_HINT_T0); + _mm_prefetch((char *)(pq_ids + 128), _MM_HINT_T0); + memset(dists_out, 0, n_pts * sizeof(float)); + for (size_t chunk = 0; chunk < pq_nchunks; chunk++) + { + const float *chunk_dists = pq_dists + 256 * chunk; + if (chunk < pq_nchunks - 1) + { + _mm_prefetch((char *)(chunk_dists + 256), _MM_HINT_T0); + } + for (size_t idx = 0; idx < n_pts; idx++) + { + uint8_t pq_centerid = pq_ids[pq_nchunks * idx + chunk]; + dists_out[idx] += chunk_dists[pq_centerid]; + } } - } } // given training data in train_data of dimensions num_train * dim, generate @@ -319,372 +345,367 @@ void pq_dist_lookup(const uint8_t *pq_ids, const size_t n_pts, // 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 -int generate_pq_pivots(const float *const passed_train_data, size_t num_train, - uint32_t dim, uint32_t num_centers, - uint32_t num_pq_chunks, uint32_t max_k_means_reps, - std::string pq_pivots_path, bool make_zero_mean) { - if (num_pq_chunks > dim) { - diskann::cout << " Error: number of chunks more than dimension" - << std::endl; - return -1; - } - - std::unique_ptr train_data = - std::make_unique(num_train * dim); - std::memcpy(train_data.get(), passed_train_data, - num_train * dim * sizeof(float)); - - 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, METADATA_SIZE); - if (file_dim == dim && file_num_centers == num_centers) { - diskann::cout << "PQ pivot file exists. Not generating again" - << std::endl; - return -1; +int generate_pq_pivots(const float *const passed_train_data, size_t num_train, uint32_t dim, uint32_t num_centers, + uint32_t num_pq_chunks, uint32_t max_k_means_reps, std::string pq_pivots_path, + bool make_zero_mean) +{ + if (num_pq_chunks > dim) + { + diskann::cout << " Error: number of chunks more than dimension" << 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++) { - centroid[d] = 0; - } - if (make_zero_mean) { // If we use L2 distance, there is an option to - // translate all vectors to make them centered and - // then compute PQ. This needs to be set to false - // when using PQ for MIPS as such translations dont - // preserve inner products. - for (uint64_t d = 0; d < dim; d++) { - for (uint64_t p = 0; p < num_train; p++) { - centroid[d] += train_data[p * dim + d]; - } - centroid[d] /= num_train; + + std::unique_ptr train_data = std::make_unique(num_train * dim); + std::memcpy(train_data.get(), passed_train_data, num_train * dim * sizeof(float)); + + 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, METADATA_SIZE); + if (file_dim == dim && file_num_centers == num_centers) + { + diskann::cout << "PQ pivot file exists. Not generating again" << std::endl; + return -1; + } } - for (uint64_t d = 0; d < dim; d++) { - for (uint64_t p = 0; p < num_train; p++) { - train_data[p * dim + d] -= centroid[d]; - } + // Calculate centroid and center the training data + std::unique_ptr centroid = std::make_unique(dim); + for (uint64_t d = 0; d < dim; d++) + { + centroid[d] = 0; } - } - - std::vector chunk_offsets; - - size_t low_val = (size_t)std::floor((double)dim / (double)num_pq_chunks); - size_t high_val = (size_t)std::ceil((double)dim / (double)num_pq_chunks); - size_t max_num_high = dim - (low_val * num_pq_chunks); - size_t cur_num_high = 0; - size_t cur_bin_threshold = high_val; - - std::vector> bin_to_dims(num_pq_chunks); - tsl::robin_map dim_to_bin; - std::vector bin_loads(num_pq_chunks, 0); - - // Process dimensions not inserted by previous loop - for (uint32_t d = 0; d < dim; d++) { - if (dim_to_bin.find(d) != dim_to_bin.end()) continue; - auto cur_best = num_pq_chunks + 1; - float cur_best_load = std::numeric_limits::max(); - for (uint32_t b = 0; b < num_pq_chunks; b++) { - if (bin_loads[b] < cur_best_load && - bin_to_dims[b].size() < cur_bin_threshold) { - cur_best = b; - cur_best_load = bin_loads[b]; - } + if (make_zero_mean) + { // If we use L2 distance, there is an option to + // translate all vectors to make them centered and + // then compute PQ. This needs to be set to false + // when using PQ for MIPS as such translations dont + // preserve inner products. + for (uint64_t d = 0; d < dim; d++) + { + for (uint64_t p = 0; p < num_train; p++) + { + centroid[d] += train_data[p * dim + d]; + } + centroid[d] /= num_train; + } + + for (uint64_t d = 0; d < dim; d++) + { + for (uint64_t p = 0; p < num_train; p++) + { + train_data[p * dim + d] -= centroid[d]; + } + } } - bin_to_dims[cur_best].push_back(d); - if (bin_to_dims[cur_best].size() == high_val) { - cur_num_high++; - if (cur_num_high == max_num_high) cur_bin_threshold = low_val; + + std::vector chunk_offsets; + + size_t low_val = (size_t)std::floor((double)dim / (double)num_pq_chunks); + size_t high_val = (size_t)std::ceil((double)dim / (double)num_pq_chunks); + size_t max_num_high = dim - (low_val * num_pq_chunks); + size_t cur_num_high = 0; + size_t cur_bin_threshold = high_val; + + std::vector> bin_to_dims(num_pq_chunks); + tsl::robin_map dim_to_bin; + std::vector bin_loads(num_pq_chunks, 0); + + // Process dimensions not inserted by previous loop + for (uint32_t d = 0; d < dim; d++) + { + if (dim_to_bin.find(d) != dim_to_bin.end()) + continue; + auto cur_best = num_pq_chunks + 1; + float cur_best_load = std::numeric_limits::max(); + for (uint32_t b = 0; b < num_pq_chunks; b++) + { + if (bin_loads[b] < cur_best_load && bin_to_dims[b].size() < cur_bin_threshold) + { + cur_best = b; + cur_best_load = bin_loads[b]; + } + } + bin_to_dims[cur_best].push_back(d); + if (bin_to_dims[cur_best].size() == high_val) + { + cur_num_high++; + if (cur_num_high == max_num_high) + cur_bin_threshold = low_val; + } } - } - chunk_offsets.clear(); - chunk_offsets.push_back(0); + chunk_offsets.clear(); + chunk_offsets.push_back(0); - for (uint32_t b = 0; b < num_pq_chunks; b++) { - if (b > 0) - chunk_offsets.push_back(chunk_offsets[b - 1] + - (uint32_t)bin_to_dims[b - 1].size()); - } - chunk_offsets.push_back(dim); + for (uint32_t b = 0; b < num_pq_chunks; b++) + { + if (b > 0) + chunk_offsets.push_back(chunk_offsets[b - 1] + (uint32_t)bin_to_dims[b - 1].size()); + } + chunk_offsets.push_back(dim); - full_pivot_data.reset(new float[num_centers * dim]); + full_pivot_data.reset(new float[num_centers * dim]); - for (size_t i = 0; i < num_pq_chunks; i++) { - size_t cur_chunk_size = chunk_offsets[i + 1] - chunk_offsets[i]; + for (size_t i = 0; i < num_pq_chunks; i++) + { + size_t cur_chunk_size = chunk_offsets[i + 1] - chunk_offsets[i]; - if (cur_chunk_size == 0) continue; - std::unique_ptr cur_pivot_data = - std::make_unique(num_centers * cur_chunk_size); - std::unique_ptr cur_data = - std::make_unique(num_train * cur_chunk_size); - std::unique_ptr closest_center = - std::make_unique(num_train); + if (cur_chunk_size == 0) + continue; + std::unique_ptr cur_pivot_data = std::make_unique(num_centers * cur_chunk_size); + std::unique_ptr cur_data = std::make_unique(num_train * cur_chunk_size); + 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; + diskann::cout << "Processing chunk " << i << " with dimensions [" << chunk_offsets[i] << ", " + << chunk_offsets[i + 1] << ")" << std::endl; #pragma omp parallel for schedule(static, 65536) - for (int64_t j = 0; j < (int64_t)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)); - } + for (int64_t j = 0; j < (int64_t)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)); + } - kmeans::kmeanspp_selecting_pivots(cur_data.get(), num_train, cur_chunk_size, - cur_pivot_data.get(), num_centers); + kmeans::kmeanspp_selecting_pivots(cur_data.get(), num_train, cur_chunk_size, cur_pivot_data.get(), 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()); + 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()); - 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)); + 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)); + } } - } - - std::vector cumul_bytes(4, 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(), chunk_offsets.data(), - chunk_offsets.size(), 1, cumul_bytes[2]); - diskann::save_bin(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; - - return 0; + + std::vector cumul_bytes(4, 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(), chunk_offsets.data(), + chunk_offsets.size(), 1, cumul_bytes[2]); + diskann::save_bin(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; + + return 0; } -int generate_opq_pivots(const float *passed_train_data, size_t num_train, - uint32_t dim, uint32_t num_centers, - uint32_t num_pq_chunks, std::string opq_pivots_path, - bool make_zero_mean) { - if (num_pq_chunks > dim) { - diskann::cout << " Error: number of chunks more than dimension" - << std::endl; - return -1; - } - - std::unique_ptr train_data = - std::make_unique(num_train * dim); - std::memcpy(train_data.get(), passed_train_data, - num_train * dim * sizeof(float)); - - std::unique_ptr rotated_train_data = - std::make_unique(num_train * dim); - std::unique_ptr rotated_and_quantized_train_data = - std::make_unique(num_train * dim); - - std::unique_ptr full_pivot_data; - - // rotation matrix for OPQ - std::unique_ptr rotmat_tr; - - // matrices for SVD - std::unique_ptr Umat = std::make_unique(dim * dim); - std::unique_ptr Vmat_T = std::make_unique(dim * dim); - std::unique_ptr singular_values = std::make_unique(dim); - std::unique_ptr correlation_matrix = - std::make_unique(dim * dim); - - // Calculate centroid and center the training data - std::unique_ptr centroid = std::make_unique(dim); - for (uint64_t d = 0; d < dim; d++) { - centroid[d] = 0; - } - if (make_zero_mean) { // If we use L2 distance, there is an option to - // translate all vectors to make them centered and - // then compute PQ. This needs to be set to false - // when using PQ for MIPS as such translations dont - // preserve inner products. - for (uint64_t d = 0; d < dim; d++) { - for (uint64_t p = 0; p < num_train; p++) { - centroid[d] += train_data[p * dim + d]; - } - centroid[d] /= num_train; +int generate_opq_pivots(const float *passed_train_data, size_t num_train, uint32_t dim, uint32_t num_centers, + uint32_t num_pq_chunks, std::string opq_pivots_path, bool make_zero_mean) +{ + if (num_pq_chunks > dim) + { + diskann::cout << " Error: number of chunks more than dimension" << std::endl; + return -1; + } + + std::unique_ptr train_data = std::make_unique(num_train * dim); + std::memcpy(train_data.get(), passed_train_data, num_train * dim * sizeof(float)); + + std::unique_ptr rotated_train_data = std::make_unique(num_train * dim); + std::unique_ptr rotated_and_quantized_train_data = std::make_unique(num_train * dim); + + std::unique_ptr full_pivot_data; + + // rotation matrix for OPQ + std::unique_ptr rotmat_tr; + + // matrices for SVD + std::unique_ptr Umat = std::make_unique(dim * dim); + std::unique_ptr Vmat_T = std::make_unique(dim * dim); + std::unique_ptr singular_values = std::make_unique(dim); + std::unique_ptr correlation_matrix = std::make_unique(dim * dim); + + // Calculate centroid and center the training data + std::unique_ptr centroid = std::make_unique(dim); + for (uint64_t d = 0; d < dim; d++) + { + centroid[d] = 0; } - for (uint64_t d = 0; d < dim; d++) { - for (uint64_t p = 0; p < num_train; p++) { - train_data[p * dim + d] -= centroid[d]; - } + if (make_zero_mean) + { // If we use L2 distance, there is an option to + // translate all vectors to make them centered and + // then compute PQ. This needs to be set to false + // when using PQ for MIPS as such translations dont + // preserve inner products. + for (uint64_t d = 0; d < dim; d++) + { + for (uint64_t p = 0; p < num_train; p++) + { + centroid[d] += train_data[p * dim + d]; + } + centroid[d] /= num_train; + } + for (uint64_t d = 0; d < dim; d++) + { + for (uint64_t p = 0; p < num_train; p++) + { + train_data[p * dim + d] -= centroid[d]; + } + } } - } - - std::vector chunk_offsets; - - size_t low_val = (size_t)std::floor((double)dim / (double)num_pq_chunks); - size_t high_val = (size_t)std::ceil((double)dim / (double)num_pq_chunks); - size_t max_num_high = dim - (low_val * num_pq_chunks); - size_t cur_num_high = 0; - size_t cur_bin_threshold = high_val; - - std::vector> bin_to_dims(num_pq_chunks); - tsl::robin_map dim_to_bin; - std::vector bin_loads(num_pq_chunks, 0); - - // Process dimensions not inserted by previous loop - for (uint32_t d = 0; d < dim; d++) { - if (dim_to_bin.find(d) != dim_to_bin.end()) continue; - auto cur_best = num_pq_chunks + 1; - float cur_best_load = std::numeric_limits::max(); - for (uint32_t b = 0; b < num_pq_chunks; b++) { - if (bin_loads[b] < cur_best_load && - bin_to_dims[b].size() < cur_bin_threshold) { - cur_best = b; - cur_best_load = bin_loads[b]; - } + + std::vector chunk_offsets; + + size_t low_val = (size_t)std::floor((double)dim / (double)num_pq_chunks); + size_t high_val = (size_t)std::ceil((double)dim / (double)num_pq_chunks); + size_t max_num_high = dim - (low_val * num_pq_chunks); + size_t cur_num_high = 0; + size_t cur_bin_threshold = high_val; + + std::vector> bin_to_dims(num_pq_chunks); + tsl::robin_map dim_to_bin; + std::vector bin_loads(num_pq_chunks, 0); + + // Process dimensions not inserted by previous loop + for (uint32_t d = 0; d < dim; d++) + { + if (dim_to_bin.find(d) != dim_to_bin.end()) + continue; + auto cur_best = num_pq_chunks + 1; + float cur_best_load = std::numeric_limits::max(); + for (uint32_t b = 0; b < num_pq_chunks; b++) + { + if (bin_loads[b] < cur_best_load && bin_to_dims[b].size() < cur_bin_threshold) + { + cur_best = b; + cur_best_load = bin_loads[b]; + } + } + bin_to_dims[cur_best].push_back(d); + if (bin_to_dims[cur_best].size() == high_val) + { + cur_num_high++; + if (cur_num_high == max_num_high) + cur_bin_threshold = low_val; + } } - bin_to_dims[cur_best].push_back(d); - if (bin_to_dims[cur_best].size() == high_val) { - cur_num_high++; - if (cur_num_high == max_num_high) cur_bin_threshold = low_val; + + chunk_offsets.clear(); + chunk_offsets.push_back(0); + + for (uint32_t b = 0; b < num_pq_chunks; b++) + { + if (b > 0) + chunk_offsets.push_back(chunk_offsets[b - 1] + (uint32_t)bin_to_dims[b - 1].size()); } - } - - chunk_offsets.clear(); - chunk_offsets.push_back(0); - - for (uint32_t b = 0; b < num_pq_chunks; b++) { - if (b > 0) - chunk_offsets.push_back(chunk_offsets[b - 1] + - (uint32_t)bin_to_dims[b - 1].size()); - } - chunk_offsets.push_back(dim); - - full_pivot_data.reset(new float[num_centers * dim]); - rotmat_tr.reset(new float[dim * dim]); - - std::memset(rotmat_tr.get(), 0, dim * dim * sizeof(float)); - for (uint32_t d1 = 0; d1 < dim; d1++) *(rotmat_tr.get() + d1 * dim + d1) = 1; - - for (uint32_t rnd = 0; rnd < MAX_OPQ_ITERS; rnd++) { - // rotate the training data using the current rotation matrix - cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, (MKL_INT)num_train, - (MKL_INT)dim, (MKL_INT)dim, 1.0f, train_data.get(), - (MKL_INT)dim, rotmat_tr.get(), (MKL_INT)dim, 0.0f, - rotated_train_data.get(), (MKL_INT)dim); - - // compute the PQ pivots on the rotated space - for (size_t i = 0; i < num_pq_chunks; i++) { - size_t cur_chunk_size = chunk_offsets[i + 1] - chunk_offsets[i]; - - if (cur_chunk_size == 0) continue; - std::unique_ptr cur_pivot_data = - std::make_unique(num_centers * cur_chunk_size); - std::unique_ptr cur_data = - std::make_unique(num_train * cur_chunk_size); - 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; + chunk_offsets.push_back(dim); + + full_pivot_data.reset(new float[num_centers * dim]); + rotmat_tr.reset(new float[dim * dim]); + + std::memset(rotmat_tr.get(), 0, dim * dim * sizeof(float)); + for (uint32_t d1 = 0; d1 < dim; d1++) + *(rotmat_tr.get() + d1 * dim + d1) = 1; + + for (uint32_t rnd = 0; rnd < MAX_OPQ_ITERS; rnd++) + { + // rotate the training data using the current rotation matrix + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, (MKL_INT)num_train, (MKL_INT)dim, (MKL_INT)dim, 1.0f, + train_data.get(), (MKL_INT)dim, rotmat_tr.get(), (MKL_INT)dim, 0.0f, rotated_train_data.get(), + (MKL_INT)dim); + + // compute the PQ pivots on the rotated space + for (size_t i = 0; i < num_pq_chunks; i++) + { + size_t cur_chunk_size = chunk_offsets[i + 1] - chunk_offsets[i]; + + if (cur_chunk_size == 0) + continue; + std::unique_ptr cur_pivot_data = std::make_unique(num_centers * cur_chunk_size); + std::unique_ptr cur_data = std::make_unique(num_train * cur_chunk_size); + 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; #pragma omp parallel for schedule(static, 65536) - for (int64_t j = 0; j < (int64_t)num_train; j++) { - std::memcpy(cur_data.get() + j * cur_chunk_size, - rotated_train_data.get() + j * dim + chunk_offsets[i], - cur_chunk_size * sizeof(float)); - } - - if (rnd == 0) { - kmeans::kmeanspp_selecting_pivots(cur_data.get(), num_train, - cur_chunk_size, cur_pivot_data.get(), - num_centers); - } else { - for (uint64_t j = 0; j < num_centers; j++) { - std::memcpy(cur_pivot_data.get() + j * cur_chunk_size, - full_pivot_data.get() + j * dim + chunk_offsets[i], - cur_chunk_size * sizeof(float)); + for (int64_t j = 0; j < (int64_t)num_train; j++) + { + std::memcpy(cur_data.get() + j * cur_chunk_size, rotated_train_data.get() + j * dim + chunk_offsets[i], + cur_chunk_size * sizeof(float)); + } + + if (rnd == 0) + { + kmeans::kmeanspp_selecting_pivots(cur_data.get(), num_train, cur_chunk_size, cur_pivot_data.get(), + num_centers); + } + else + { + for (uint64_t j = 0; j < num_centers; j++) + { + std::memcpy(cur_pivot_data.get() + j * cur_chunk_size, + full_pivot_data.get() + j * dim + chunk_offsets[i], cur_chunk_size * sizeof(float)); + } + } + + uint32_t num_lloyds_iters = 8; + kmeans::run_lloyds(cur_data.get(), num_train, cur_chunk_size, cur_pivot_data.get(), num_centers, + num_lloyds_iters, NULL, closest_center.get()); + + 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)); + } + + for (size_t j = 0; j < num_train; j++) + { + std::memcpy(rotated_and_quantized_train_data.get() + j * dim + chunk_offsets[i], + cur_pivot_data.get() + (size_t)closest_center[j] * cur_chunk_size, + cur_chunk_size * sizeof(float)); + } } - } - - uint32_t num_lloyds_iters = 8; - kmeans::run_lloyds(cur_data.get(), num_train, cur_chunk_size, - cur_pivot_data.get(), num_centers, num_lloyds_iters, - NULL, closest_center.get()); - - 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)); - } - - for (size_t j = 0; j < num_train; j++) { - std::memcpy( - rotated_and_quantized_train_data.get() + j * dim + chunk_offsets[i], - cur_pivot_data.get() + (size_t)closest_center[j] * cur_chunk_size, - cur_chunk_size * sizeof(float)); - } - } - // compute the correlation matrix between the original data and the - // quantized data to compute the new rotation - cblas_sgemm(CblasRowMajor, CblasTrans, CblasNoTrans, (MKL_INT)dim, - (MKL_INT)dim, (MKL_INT)num_train, 1.0f, train_data.get(), - (MKL_INT)dim, rotated_and_quantized_train_data.get(), - (MKL_INT)dim, 0.0f, correlation_matrix.get(), (MKL_INT)dim); - - // compute the SVD of the correlation matrix to help determine the new - // rotation matrix - uint32_t errcode = LAPACKE_sgesdd( - LAPACK_ROW_MAJOR, 'A', (MKL_INT)dim, (MKL_INT)dim, - correlation_matrix.get(), (MKL_INT)dim, singular_values.get(), - Umat.get(), (MKL_INT)dim, Vmat_T.get(), (MKL_INT)dim); - - if (errcode > 0) { - std::cout << "SVD failed to converge." << std::endl; - exit(-1); + // compute the correlation matrix between the original data and the + // quantized data to compute the new rotation + cblas_sgemm(CblasRowMajor, CblasTrans, CblasNoTrans, (MKL_INT)dim, (MKL_INT)dim, (MKL_INT)num_train, 1.0f, + train_data.get(), (MKL_INT)dim, rotated_and_quantized_train_data.get(), (MKL_INT)dim, 0.0f, + correlation_matrix.get(), (MKL_INT)dim); + + // compute the SVD of the correlation matrix to help determine the new + // rotation matrix + uint32_t errcode = + LAPACKE_sgesdd(LAPACK_ROW_MAJOR, 'A', (MKL_INT)dim, (MKL_INT)dim, correlation_matrix.get(), (MKL_INT)dim, + singular_values.get(), Umat.get(), (MKL_INT)dim, Vmat_T.get(), (MKL_INT)dim); + + if (errcode > 0) + { + std::cout << "SVD failed to converge." << std::endl; + exit(-1); + } + + // compute the new rotation matrix from the singular vectors as R^T = U + // V^T + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, (MKL_INT)dim, (MKL_INT)dim, (MKL_INT)dim, 1.0f, + Umat.get(), (MKL_INT)dim, Vmat_T.get(), (MKL_INT)dim, 0.0f, rotmat_tr.get(), (MKL_INT)dim); } - // compute the new rotation matrix from the singular vectors as R^T = U - // V^T - cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, (MKL_INT)dim, - (MKL_INT)dim, (MKL_INT)dim, 1.0f, Umat.get(), (MKL_INT)dim, - Vmat_T.get(), (MKL_INT)dim, 0.0f, rotmat_tr.get(), - (MKL_INT)dim); - } - - std::vector cumul_bytes(4, 0); - cumul_bytes[0] = METADATA_SIZE; - cumul_bytes[1] = - cumul_bytes[0] + - diskann::save_bin(opq_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( - opq_pivots_path.c_str(), centroid.get(), - (size_t)dim, 1, cumul_bytes[1]); - cumul_bytes[3] = - cumul_bytes[2] + - diskann::save_bin(opq_pivots_path.c_str(), chunk_offsets.data(), - chunk_offsets.size(), 1, cumul_bytes[2]); - diskann::save_bin(opq_pivots_path.c_str(), cumul_bytes.data(), - cumul_bytes.size(), 1, 0); - - diskann::cout << "Saved opq pivot data to " << opq_pivots_path << " of size " - << cumul_bytes[cumul_bytes.size() - 1] << "B." << std::endl; - - std::string rotmat_path = opq_pivots_path + "_rotation_matrix.bin"; - diskann::save_bin(rotmat_path.c_str(), rotmat_tr.get(), dim, dim); - - return 0; + std::vector cumul_bytes(4, 0); + cumul_bytes[0] = METADATA_SIZE; + cumul_bytes[1] = cumul_bytes[0] + diskann::save_bin(opq_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(opq_pivots_path.c_str(), centroid.get(), (size_t)dim, 1, cumul_bytes[1]); + cumul_bytes[3] = cumul_bytes[2] + diskann::save_bin(opq_pivots_path.c_str(), chunk_offsets.data(), + chunk_offsets.size(), 1, cumul_bytes[2]); + diskann::save_bin(opq_pivots_path.c_str(), cumul_bytes.data(), cumul_bytes.size(), 1, 0); + + diskann::cout << "Saved opq pivot data to " << opq_pivots_path << " of size " << cumul_bytes[cumul_bytes.size() - 1] + << "B." << std::endl; + + std::string rotmat_path = opq_pivots_path + "_rotation_matrix.bin"; + diskann::save_bin(rotmat_path.c_str(), rotmat_tr.get(), dim, dim); + + return 0; } // streams the base file (data_file), and computes the closest centers in each @@ -693,358 +714,344 @@ int generate_opq_pivots(const float *passed_train_data, size_t num_train, // If the numbber of centers is < 256, it stores as byte vector, else as // 4-byte vector in binary format. template -int generate_pq_data_from_pivots(const std::string &data_file, - uint32_t num_centers, uint32_t num_pq_chunks, - const std::string &pq_pivots_path, - const std::string &pq_compressed_vectors_path, - bool use_opq) { - size_t read_blk_size = 64 * 1024 * 1024; - cached_ifstream base_reader(data_file, read_blk_size); - uint32_t npts32; - uint32_t basedim32; - base_reader.read((char *)&npts32, sizeof(uint32_t)); - base_reader.read((char *)&basedim32, sizeof(uint32_t)); - size_t num_points = npts32; - size_t dim = basedim32; - - std::unique_ptr full_pivot_data; - std::unique_ptr rotmat_tr; - std::unique_ptr centroid; - std::unique_ptr chunk_offsets; - - std::string inflated_pq_file = pq_compressed_vectors_path + "_inflated.bin"; - - if (!file_exists(pq_pivots_path)) { - std::cout << "ERROR: PQ k-means pivot file not found" << std::endl; - throw diskann::ANNException("PQ k-means pivot file not found", -1); - } else { - size_t nr, nc; - std::unique_ptr file_offset_data; +int generate_pq_data_from_pivots(const std::string &data_file, uint32_t num_centers, uint32_t num_pq_chunks, + const std::string &pq_pivots_path, const std::string &pq_compressed_vectors_path, + bool use_opq) +{ + size_t read_blk_size = 64 * 1024 * 1024; + cached_ifstream base_reader(data_file, read_blk_size); + uint32_t npts32; + uint32_t basedim32; + base_reader.read((char *)&npts32, sizeof(uint32_t)); + base_reader.read((char *)&basedim32, sizeof(uint32_t)); + size_t num_points = npts32; + size_t dim = basedim32; + + std::unique_ptr full_pivot_data; + std::unique_ptr rotmat_tr; + std::unique_ptr centroid; + std::unique_ptr chunk_offsets; + + std::string inflated_pq_file = pq_compressed_vectors_path + "_inflated.bin"; + + if (!file_exists(pq_pivots_path)) + { + std::cout << "ERROR: PQ k-means pivot file not found" << std::endl; + throw diskann::ANNException("PQ k-means pivot file not found", -1); + } + else + { + size_t nr, nc; + std::unique_ptr file_offset_data; + + diskann::load_bin(pq_pivots_path.c_str(), file_offset_data, nr, nc, 0); + + if (nr != 4) + { + diskann::cout << "Error reading pq_pivots file " << pq_pivots_path + << ". Offsets dont contain correct metadata, # offsets = " << nr << ", but expecting 4."; + throw diskann::ANNException("Error reading pq_pivots file at offsets data.", -1, __FUNCSIG__, __FILE__, + __LINE__); + } - diskann::load_bin(pq_pivots_path.c_str(), file_offset_data, nr, nc, - 0); + diskann::load_bin(pq_pivots_path.c_str(), full_pivot_data, nr, nc, file_offset_data[0]); - if (nr != 4) { - diskann::cout << "Error reading pq_pivots file " << pq_pivots_path - << ". Offsets dont contain correct metadata, # offsets = " - << nr << ", but expecting 4."; - throw diskann::ANNException( - "Error reading pq_pivots file at offsets data.", -1, __FUNCSIG__, - __FILE__, __LINE__); - } + 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__); + } - 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__); - } + diskann::load_bin(pq_pivots_path.c_str(), centroid, nr, nc, file_offset_data[1]); - 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__); + } - 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__); - } + diskann::load_bin(pq_pivots_path.c_str(), chunk_offsets, nr, nc, file_offset_data[2]); - diskann::load_bin(pq_pivots_path.c_str(), chunk_offsets, nr, nc, - file_offset_data[2]); - - 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__); - } + 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__); + } - if (use_opq) { - std::string rotmat_path = pq_pivots_path + "_rotation_matrix.bin"; - diskann::load_bin(rotmat_path.c_str(), rotmat_tr, nr, nc); - if (nr != (uint64_t)dim || nc != dim) { - diskann::cout << "Error reading rotation matrix file." << std::endl; - throw diskann::ANNException("Error reading rotation matrix file.", -1, - __FUNCSIG__, __FILE__, __LINE__); - } - } + if (use_opq) + { + std::string rotmat_path = pq_pivots_path + "_rotation_matrix.bin"; + diskann::load_bin(rotmat_path.c_str(), rotmat_tr, nr, nc); + if (nr != (uint64_t)dim || nc != dim) + { + diskann::cout << "Error reading rotation matrix file." << std::endl; + throw diskann::ANNException("Error reading rotation matrix file.", -1, __FUNCSIG__, __FILE__, __LINE__); + } + } - diskann::cout << "Loaded PQ pivot information" << std::endl; - } + diskann::cout << "Loaded PQ pivot information" << std::endl; + } - std::ofstream compressed_file_writer(pq_compressed_vectors_path, - std::ios::binary); - uint32_t num_pq_chunks_u32 = num_pq_chunks; + std::ofstream compressed_file_writer(pq_compressed_vectors_path, std::ios::binary); + uint32_t 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)); + compressed_file_writer.write((char *)&num_points, sizeof(uint32_t)); + compressed_file_writer.write((char *)&num_pq_chunks_u32, sizeof(uint32_t)); - size_t block_size = num_points <= BLOCK_SIZE ? num_points : BLOCK_SIZE; + size_t block_size = num_points <= BLOCK_SIZE ? num_points : BLOCK_SIZE; #ifdef SAVE_INFLATED_PQ - std::ofstream inflated_file_writer(inflated_pq_file, std::ios::binary); - inflated_file_writer.write((char *)&num_points, sizeof(uint32_t)); - inflated_file_writer.write((char *)&basedim32, sizeof(uint32_t)); + std::ofstream inflated_file_writer(inflated_pq_file, std::ios::binary); + inflated_file_writer.write((char *)&num_points, sizeof(uint32_t)); + inflated_file_writer.write((char *)&basedim32, sizeof(uint32_t)); - std::unique_ptr block_inflated_base = - std::make_unique(block_size * dim); - std::memset(block_inflated_base.get(), 0, block_size * dim * sizeof(float)); + std::unique_ptr block_inflated_base = std::make_unique(block_size * dim); + std::memset(block_inflated_base.get(), 0, block_size * dim * sizeof(float)); #endif - std::unique_ptr block_compressed_base = - std::make_unique(block_size * (size_t)num_pq_chunks); - std::memset(block_compressed_base.get(), 0, - block_size * (size_t)num_pq_chunks * sizeof(uint32_t)); + std::unique_ptr block_compressed_base = + std::make_unique(block_size * (size_t)num_pq_chunks); + std::memset(block_compressed_base.get(), 0, block_size * (size_t)num_pq_chunks * sizeof(uint32_t)); - std::unique_ptr block_data_T = std::make_unique(block_size * dim); - std::unique_ptr block_data_float = - std::make_unique(block_size * dim); - std::unique_ptr block_data_tmp = - std::make_unique(block_size * dim); + std::unique_ptr block_data_T = std::make_unique(block_size * dim); + std::unique_ptr block_data_float = std::make_unique(block_size * dim); + std::unique_ptr block_data_tmp = std::make_unique(block_size * dim); - size_t num_blocks = DIV_ROUND_UP(num_points, block_size); + size_t num_blocks = DIV_ROUND_UP(num_points, block_size); - for (size_t block = 0; block < num_blocks; block++) { - size_t start_id = block * block_size; - size_t end_id = (std::min)((block + 1) * block_size, num_points); - size_t cur_blk_size = end_id - start_id; + for (size_t block = 0; block < num_blocks; block++) + { + size_t start_id = block * block_size; + size_t end_id = (std::min)((block + 1) * block_size, num_points); + size_t cur_blk_size = end_id - start_id; - base_reader.read((char *)(block_data_T.get()), - sizeof(T) * (cur_blk_size * dim)); - diskann::convert_types(block_data_T.get(), block_data_tmp.get(), - cur_blk_size, dim); + base_reader.read((char *)(block_data_T.get()), sizeof(T) * (cur_blk_size * dim)); + 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 << ").." << std::flush; - for (size_t p = 0; p < cur_blk_size; p++) { - for (uint64_t d = 0; d < dim; d++) { - block_data_tmp[p * dim + d] -= centroid[d]; - } - } + for (size_t p = 0; p < cur_blk_size; p++) + { + for (uint64_t d = 0; d < dim; d++) + { + block_data_tmp[p * dim + d] -= centroid[d]; + } + } - for (size_t p = 0; p < cur_blk_size; p++) { - for (uint64_t d = 0; d < dim; d++) { - block_data_float[p * dim + d] = block_data_tmp[p * dim + d]; - } - } + for (size_t p = 0; p < cur_blk_size; p++) + { + for (uint64_t d = 0; d < dim; d++) + { + block_data_float[p * dim + d] = block_data_tmp[p * dim + d]; + } + } - if (use_opq) { - // rotate the current block with the trained rotation matrix before - // PQ - cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, - (MKL_INT)cur_blk_size, (MKL_INT)dim, (MKL_INT)dim, 1.0f, - block_data_float.get(), (MKL_INT)dim, rotmat_tr.get(), - (MKL_INT)dim, 0.0f, block_data_tmp.get(), (MKL_INT)dim); - std::memcpy(block_data_float.get(), block_data_tmp.get(), - cur_blk_size * dim * sizeof(float)); - } + if (use_opq) + { + // rotate the current block with the trained rotation matrix before + // PQ + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, (MKL_INT)cur_blk_size, (MKL_INT)dim, (MKL_INT)dim, + 1.0f, block_data_float.get(), (MKL_INT)dim, rotmat_tr.get(), (MKL_INT)dim, 0.0f, + block_data_tmp.get(), (MKL_INT)dim); + std::memcpy(block_data_float.get(), block_data_tmp.get(), cur_blk_size * dim * sizeof(float)); + } - for (size_t i = 0; i < num_pq_chunks; i++) { - size_t cur_chunk_size = chunk_offsets[i + 1] - chunk_offsets[i]; - if (cur_chunk_size == 0) continue; + for (size_t i = 0; i < num_pq_chunks; i++) + { + size_t cur_chunk_size = chunk_offsets[i + 1] - chunk_offsets[i]; + if (cur_chunk_size == 0) + continue; - std::unique_ptr cur_pivot_data = - std::make_unique(num_centers * cur_chunk_size); - std::unique_ptr cur_data = - std::make_unique(cur_blk_size * cur_chunk_size); - std::unique_ptr closest_center = - std::make_unique(cur_blk_size); + std::unique_ptr cur_pivot_data = std::make_unique(num_centers * cur_chunk_size); + std::unique_ptr cur_data = std::make_unique(cur_blk_size * cur_chunk_size); + std::unique_ptr closest_center = std::make_unique(cur_blk_size); #pragma omp parallel for schedule(static, 8192) - for (int64_t j = 0; j < (int64_t)cur_blk_size; j++) { - for (size_t k = 0; k < cur_chunk_size; k++) - cur_data[j * cur_chunk_size + k] = - block_data_float[j * dim + chunk_offsets[i] + k]; - } + for (int64_t j = 0; j < (int64_t)cur_blk_size; j++) + { + for (size_t k = 0; k < cur_chunk_size; k++) + cur_data[j * cur_chunk_size + k] = block_data_float[j * dim + chunk_offsets[i] + k]; + } #pragma omp parallel for schedule(static, 1) - for (int64_t j = 0; j < (int64_t)num_centers; j++) { - std::memcpy(cur_pivot_data.get() + j * cur_chunk_size, - full_pivot_data.get() + j * dim + chunk_offsets[i], - cur_chunk_size * sizeof(float)); - } + for (int64_t j = 0; j < (int64_t)num_centers; j++) + { + std::memcpy(cur_pivot_data.get() + j * cur_chunk_size, + full_pivot_data.get() + j * dim + chunk_offsets[i], cur_chunk_size * sizeof(float)); + } - math_utils::compute_closest_centers(cur_data.get(), cur_blk_size, - cur_chunk_size, cur_pivot_data.get(), - num_centers, 1, closest_center.get()); + 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 < (int64_t)cur_blk_size; j++) { - block_compressed_base[j * num_pq_chunks + i] = closest_center[j]; + for (int64_t j = 0; j < (int64_t)cur_blk_size; j++) + { + block_compressed_base[j * num_pq_chunks + i] = closest_center[j]; #ifdef SAVE_INFLATED_PQ - for (size_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]; + for (size_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 - } - } + } + } - if (num_centers > 256) { - compressed_file_writer.write( - (char *)(block_compressed_base.get()), - cur_blk_size * num_pq_chunks * sizeof(uint32_t)); - } else { - std::unique_ptr pVec = - std::make_unique(cur_blk_size * num_pq_chunks); - diskann::convert_types( - block_compressed_base.get(), pVec.get(), cur_blk_size, num_pq_chunks); - compressed_file_writer.write( - (char *)(pVec.get()), cur_blk_size * num_pq_chunks * sizeof(uint8_t)); - } + if (num_centers > 256) + { + compressed_file_writer.write((char *)(block_compressed_base.get()), + cur_blk_size * num_pq_chunks * sizeof(uint32_t)); + } + else + { + std::unique_ptr pVec = std::make_unique(cur_blk_size * num_pq_chunks); + diskann::convert_types(block_compressed_base.get(), pVec.get(), cur_blk_size, + num_pq_chunks); + compressed_file_writer.write((char *)(pVec.get()), cur_blk_size * num_pq_chunks * sizeof(uint8_t)); + } #ifdef SAVE_INFLATED_PQ - inflated_file_writer.write((char *)(block_inflated_base.get()), - cur_blk_size * dim * sizeof(float)); + inflated_file_writer.write((char *)(block_inflated_base.get()), cur_blk_size * dim * sizeof(float)); #endif - diskann::cout << ".done." << std::endl; - } + diskann::cout << ".done." << std::endl; + } // Gopal. Splitting diskann_dll into separate DLLs for search and build. // This code should only be available in the "build" DLL. -#if defined(RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS) && \ - defined(DISKANN_BUILD) - MallocExtension::instance()->ReleaseFreeMemory(); +#if defined(RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS) && defined(DISKANN_BUILD) + MallocExtension::instance()->ReleaseFreeMemory(); #endif - compressed_file_writer.close(); + compressed_file_writer.close(); #ifdef SAVE_INFLATED_PQ - inflated_file_writer.close(); + inflated_file_writer.close(); #endif - return 0; + return 0; } template -void generate_disk_quantized_data( - const std::string &data_file_to_use, const std::string &disk_pq_pivots_path, - const std::string &disk_pq_compressed_vectors_path, - diskann::Metric compareMetric, const double p_val, size_t &disk_pq_dims) { - size_t train_size, train_dim; - float *train_data; - - // instantiates train_data with random sample updates train_size - gen_random_slice(data_file_to_use.c_str(), p_val, train_data, train_size, - train_dim); - diskann::cout << "Training data with " << train_size << " samples loaded." - << std::endl; - - if (disk_pq_dims > train_dim) disk_pq_dims = train_dim; - - std::cout << "Compressing base for disk-PQ into " << disk_pq_dims - << " chunks " << std::endl; - generate_pq_pivots(train_data, train_size, (uint32_t)train_dim, 256, - (uint32_t)disk_pq_dims, NUM_KMEANS_REPS_PQ, - disk_pq_pivots_path, false); - if (compareMetric == diskann::Metric::INNER_PRODUCT) - generate_pq_data_from_pivots( - data_file_to_use, 256, (uint32_t)disk_pq_dims, disk_pq_pivots_path, - disk_pq_compressed_vectors_path); - else - generate_pq_data_from_pivots(data_file_to_use, 256, - (uint32_t)disk_pq_dims, disk_pq_pivots_path, - disk_pq_compressed_vectors_path); - - delete[] train_data; +void generate_disk_quantized_data(const std::string &data_file_to_use, const std::string &disk_pq_pivots_path, + const std::string &disk_pq_compressed_vectors_path, diskann::Metric compareMetric, + const double p_val, size_t &disk_pq_dims) +{ + size_t train_size, train_dim; + float *train_data; + + // instantiates train_data with random sample updates train_size + gen_random_slice(data_file_to_use.c_str(), p_val, train_data, train_size, train_dim); + diskann::cout << "Training data with " << train_size << " samples loaded." << std::endl; + + if (disk_pq_dims > train_dim) + disk_pq_dims = train_dim; + + std::cout << "Compressing base for disk-PQ into " << disk_pq_dims << " chunks " << std::endl; + generate_pq_pivots(train_data, train_size, (uint32_t)train_dim, 256, (uint32_t)disk_pq_dims, NUM_KMEANS_REPS_PQ, + disk_pq_pivots_path, false); + if (compareMetric == diskann::Metric::INNER_PRODUCT) + generate_pq_data_from_pivots(data_file_to_use, 256, (uint32_t)disk_pq_dims, disk_pq_pivots_path, + disk_pq_compressed_vectors_path); + else + generate_pq_data_from_pivots(data_file_to_use, 256, (uint32_t)disk_pq_dims, disk_pq_pivots_path, + disk_pq_compressed_vectors_path); + + delete[] train_data; } template -void generate_quantized_data(const std::string &data_file_to_use, - const std::string &pq_pivots_path, - const std::string &pq_compressed_vectors_path, - diskann::Metric compareMetric, const double p_val, - const size_t num_pq_chunks, const bool use_opq, - const std::string &codebook_prefix) { - size_t train_size, train_dim; - float *train_data; - if (!file_exists(codebook_prefix)) { - // instantiates train_data with random sample updates train_size - gen_random_slice(data_file_to_use.c_str(), p_val, train_data, train_size, - train_dim); - diskann::cout << "Training data with " << train_size << " samples loaded." - << std::endl; - - bool make_zero_mean = true; - if (compareMetric == diskann::Metric::INNER_PRODUCT) make_zero_mean = false; - if (use_opq) // we also do not center the data for OPQ - make_zero_mean = false; - - if (!use_opq) { - generate_pq_pivots(train_data, train_size, (uint32_t)train_dim, - NUM_PQ_CENTROIDS, (uint32_t)num_pq_chunks, - NUM_KMEANS_REPS_PQ, pq_pivots_path, make_zero_mean); - } else { - generate_opq_pivots(train_data, train_size, (uint32_t)train_dim, - NUM_PQ_CENTROIDS, (uint32_t)num_pq_chunks, - pq_pivots_path, make_zero_mean); +void generate_quantized_data(const std::string &data_file_to_use, const std::string &pq_pivots_path, + const std::string &pq_compressed_vectors_path, diskann::Metric compareMetric, + const double p_val, const size_t num_pq_chunks, const bool use_opq, + const std::string &codebook_prefix) +{ + size_t train_size, train_dim; + float *train_data; + if (!file_exists(codebook_prefix)) + { + // instantiates train_data with random sample updates train_size + gen_random_slice(data_file_to_use.c_str(), p_val, train_data, train_size, train_dim); + diskann::cout << "Training data with " << train_size << " samples loaded." << std::endl; + + bool make_zero_mean = true; + if (compareMetric == diskann::Metric::INNER_PRODUCT) + make_zero_mean = false; + if (use_opq) // we also do not center the data for OPQ + make_zero_mean = false; + + if (!use_opq) + { + generate_pq_pivots(train_data, train_size, (uint32_t)train_dim, NUM_PQ_CENTROIDS, (uint32_t)num_pq_chunks, + NUM_KMEANS_REPS_PQ, pq_pivots_path, make_zero_mean); + } + else + { + generate_opq_pivots(train_data, train_size, (uint32_t)train_dim, NUM_PQ_CENTROIDS, (uint32_t)num_pq_chunks, + pq_pivots_path, make_zero_mean); + } + delete[] train_data; } - delete[] train_data; - } else { - diskann::cout << "Skip Training with predefined pivots in: " - << pq_pivots_path << std::endl; - } - generate_pq_data_from_pivots(data_file_to_use, NUM_PQ_CENTROIDS, - (uint32_t)num_pq_chunks, pq_pivots_path, - pq_compressed_vectors_path, use_opq); + else + { + diskann::cout << "Skip Training with predefined pivots in: " << pq_pivots_path << std::endl; + } + generate_pq_data_from_pivots(data_file_to_use, NUM_PQ_CENTROIDS, (uint32_t)num_pq_chunks, pq_pivots_path, + pq_compressed_vectors_path, use_opq); } // Instantations of supported templates -template DISKANN_DLLEXPORT int generate_pq_data_from_pivots( - const std::string &data_file, uint32_t num_centers, uint32_t num_pq_chunks, - const std::string &pq_pivots_path, - const std::string &pq_compressed_vectors_path, bool use_opq); -template DISKANN_DLLEXPORT int generate_pq_data_from_pivots( - const std::string &data_file, uint32_t num_centers, uint32_t num_pq_chunks, - const std::string &pq_pivots_path, - const std::string &pq_compressed_vectors_path, bool use_opq); -template DISKANN_DLLEXPORT int generate_pq_data_from_pivots( - const std::string &data_file, uint32_t num_centers, uint32_t num_pq_chunks, - const std::string &pq_pivots_path, - const std::string &pq_compressed_vectors_path, bool use_opq); - -template DISKANN_DLLEXPORT void generate_disk_quantized_data( - const std::string &data_file_to_use, const std::string &disk_pq_pivots_path, - const std::string &disk_pq_compressed_vectors_path, - diskann::Metric compareMetric, const double p_val, size_t &disk_pq_dims); +template DISKANN_DLLEXPORT int generate_pq_data_from_pivots(const std::string &data_file, uint32_t num_centers, + uint32_t num_pq_chunks, + const std::string &pq_pivots_path, + const std::string &pq_compressed_vectors_path, + bool use_opq); +template DISKANN_DLLEXPORT int generate_pq_data_from_pivots(const std::string &data_file, uint32_t num_centers, + uint32_t num_pq_chunks, + const std::string &pq_pivots_path, + const std::string &pq_compressed_vectors_path, + bool use_opq); +template DISKANN_DLLEXPORT int generate_pq_data_from_pivots(const std::string &data_file, uint32_t num_centers, + uint32_t num_pq_chunks, + const std::string &pq_pivots_path, + const std::string &pq_compressed_vectors_path, + bool use_opq); + +template DISKANN_DLLEXPORT void generate_disk_quantized_data(const std::string &data_file_to_use, + const std::string &disk_pq_pivots_path, + const std::string &disk_pq_compressed_vectors_path, + diskann::Metric compareMetric, const double p_val, + size_t &disk_pq_dims); template DISKANN_DLLEXPORT void generate_disk_quantized_data( const std::string &data_file_to_use, const std::string &disk_pq_pivots_path, - const std::string &disk_pq_compressed_vectors_path, - diskann::Metric compareMetric, const double p_val, size_t &disk_pq_dims); - -template DISKANN_DLLEXPORT void generate_disk_quantized_data( - const std::string &data_file_to_use, const std::string &disk_pq_pivots_path, - const std::string &disk_pq_compressed_vectors_path, - diskann::Metric compareMetric, const double p_val, size_t &disk_pq_dims); - -template DISKANN_DLLEXPORT void generate_quantized_data( - const std::string &data_file_to_use, const std::string &pq_pivots_path, - const std::string &pq_compressed_vectors_path, - diskann::Metric compareMetric, const double p_val, - const size_t num_pq_chunks, const bool use_opq, - const std::string &codebook_prefix); - -template DISKANN_DLLEXPORT void generate_quantized_data( - const std::string &data_file_to_use, const std::string &pq_pivots_path, - const std::string &pq_compressed_vectors_path, - diskann::Metric compareMetric, const double p_val, - const size_t num_pq_chunks, const bool use_opq, - const std::string &codebook_prefix); - -template DISKANN_DLLEXPORT void generate_quantized_data( - const std::string &data_file_to_use, const std::string &pq_pivots_path, - const std::string &pq_compressed_vectors_path, - diskann::Metric compareMetric, const double p_val, - const size_t num_pq_chunks, const bool use_opq, - const std::string &codebook_prefix); -} // namespace diskann + const std::string &disk_pq_compressed_vectors_path, diskann::Metric compareMetric, const double p_val, + size_t &disk_pq_dims); + +template DISKANN_DLLEXPORT void generate_disk_quantized_data(const std::string &data_file_to_use, + const std::string &disk_pq_pivots_path, + const std::string &disk_pq_compressed_vectors_path, + diskann::Metric compareMetric, const double p_val, + size_t &disk_pq_dims); + +template DISKANN_DLLEXPORT void generate_quantized_data(const std::string &data_file_to_use, + const std::string &pq_pivots_path, + const std::string &pq_compressed_vectors_path, + diskann::Metric compareMetric, const double p_val, + const size_t num_pq_chunks, const bool use_opq, + const std::string &codebook_prefix); + +template DISKANN_DLLEXPORT void generate_quantized_data(const std::string &data_file_to_use, + const std::string &pq_pivots_path, + const std::string &pq_compressed_vectors_path, + diskann::Metric compareMetric, const double p_val, + const size_t num_pq_chunks, const bool use_opq, + const std::string &codebook_prefix); + +template DISKANN_DLLEXPORT void generate_quantized_data(const std::string &data_file_to_use, + const std::string &pq_pivots_path, + const std::string &pq_compressed_vectors_path, + diskann::Metric compareMetric, const double p_val, + const size_t num_pq_chunks, const bool use_opq, + const std::string &codebook_prefix); +} // namespace diskann diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index fa78f689b..4d137e560 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -21,1457 +21,1536 @@ #define NODE_SECTOR_NO(node_id) (((uint64_t)(node_id)) / nnodes_per_sector + 1) // obtains region of sector containing node -#define OFFSET_TO_NODE(sector_buf, node_id) \ - ((char *)sector_buf + \ - (((uint64_t)node_id) % nnodes_per_sector) * max_node_len) +#define OFFSET_TO_NODE(sector_buf, node_id) \ + ((char *)sector_buf + (((uint64_t)node_id) % nnodes_per_sector) * max_node_len) // returns region of `node_buf` containing [NNBRS][NBR_ID(uint32_t)] -#define OFFSET_TO_NODE_NHOOD(node_buf) \ - (unsigned *)((char *)node_buf + disk_bytes_per_point) +#define OFFSET_TO_NODE_NHOOD(node_buf) (unsigned *)((char *)node_buf + disk_bytes_per_point) // returns region of `node_buf` containing [COORD(T)] #define OFFSET_TO_NODE_COORDS(node_buf) (T *)(node_buf) // sector # beyond the end of graph where data for id is present for reordering -#define VECTOR_SECTOR_NO(id) \ - (((uint64_t)(id)) / nvecs_per_sector + reorder_data_start_sector) +#define VECTOR_SECTOR_NO(id) (((uint64_t)(id)) / nvecs_per_sector + reorder_data_start_sector) // sector # beyond the end of graph where data for id is present for reordering -#define VECTOR_SECTOR_OFFSET(id) \ - ((((uint64_t)(id)) % nvecs_per_sector) * data_dim * sizeof(float)) +#define VECTOR_SECTOR_OFFSET(id) ((((uint64_t)(id)) % nvecs_per_sector) * data_dim * sizeof(float)) -namespace diskann { +namespace diskann +{ template -PQFlashIndex::PQFlashIndex( - std::shared_ptr &fileReader, diskann::Metric m) - : reader(fileReader), metric(m) { - if (m == diskann::Metric::COSINE || m == diskann::Metric::INNER_PRODUCT) { - if (std::is_floating_point::value) { - diskann::cout << "Cosine metric chosen for (normalized) float data." - "Changing distance to L2 to boost accuracy." - << std::endl; - metric = diskann::Metric::L2; - } 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; +PQFlashIndex::PQFlashIndex(std::shared_ptr &fileReader, diskann::Metric m) + : reader(fileReader), metric(m) +{ + if (m == diskann::Metric::COSINE || m == diskann::Metric::INNER_PRODUCT) + { + if (std::is_floating_point::value) + { + diskann::cout << "Cosine metric chosen for (normalized) float data." + "Changing distance to L2 to boost accuracy." + << std::endl; + metric = diskann::Metric::L2; + } + 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(metric)); - this->dist_cmp_float.reset(diskann::get_distance_function(metric)); + this->dist_cmp.reset(diskann::get_distance_function(metric)); + this->dist_cmp_float.reset(diskann::get_distance_function(metric)); } -template -PQFlashIndex::~PQFlashIndex() { +template PQFlashIndex::~PQFlashIndex() +{ #ifndef EXEC_ENV_OLS - if (data != nullptr) { - delete[] data; - } + if (data != nullptr) + { + delete[] data; + } #endif - if (centroid_data != nullptr) aligned_free(centroid_data); - // delete backing bufs for nhood and coord cache - if (nhood_cache_buf != nullptr) { - delete[] nhood_cache_buf; - diskann::aligned_free(coord_cache_buf); - } + if (centroid_data != nullptr) + aligned_free(centroid_data); + // delete backing bufs for nhood and coord cache + if (nhood_cache_buf != nullptr) + { + delete[] nhood_cache_buf; + diskann::aligned_free(coord_cache_buf); + } - if (load_flag) { - diskann::cout << "Clearing scratch" << std::endl; - ScratchStoreManager> manager(this->thread_data); - manager.destroy(); - this->reader->deregister_all_threads(); - reader->close(); - } - if (_pts_to_label_offsets != nullptr) { - delete[] _pts_to_label_offsets; - } - - if (_pts_to_labels != nullptr) { - delete[] _pts_to_labels; - } + if (load_flag) + { + diskann::cout << "Clearing scratch" << std::endl; + ScratchStoreManager> manager(this->thread_data); + manager.destroy(); + this->reader->deregister_all_threads(); + reader->close(); + } + if (_pts_to_label_offsets != nullptr) + { + delete[] _pts_to_label_offsets; + } + + if (_pts_to_labels != nullptr) + { + delete[] _pts_to_labels; + } } template -void PQFlashIndex::setup_thread_data(uint64_t nthreads, - uint64_t visited_reserve) { - diskann::cout << "Setting up thread-specific contexts for nthreads: " - << nthreads << std::endl; +void PQFlashIndex::setup_thread_data(uint64_t nthreads, uint64_t visited_reserve) +{ + diskann::cout << "Setting up thread-specific contexts for nthreads: " << nthreads << std::endl; // omp parallel for to generate unique thread IDs #pragma omp parallel for num_threads((int)nthreads) - for (int64_t thread = 0; thread < (int64_t)nthreads; thread++) { -#pragma omp critical + for (int64_t thread = 0; thread < (int64_t)nthreads; thread++) { - SSDThreadData *data = - new SSDThreadData(this->aligned_dim, visited_reserve); - this->reader->register_thread(); - data->ctx = this->reader->get_ctx(); - this->thread_data.push(data); +#pragma omp critical + { + SSDThreadData *data = new SSDThreadData(this->aligned_dim, visited_reserve); + this->reader->register_thread(); + data->ctx = this->reader->get_ctx(); + this->thread_data.push(data); + } } - } - load_flag = true; + load_flag = true; } -template -void PQFlashIndex::load_cache_list( - std::vector &node_list) { - diskann::cout << "Loading the cache list into memory.." << std::flush; - size_t num_cached_nodes = node_list.size(); - - // borrow thread data - ScratchStoreManager> manager(this->thread_data); - auto this_thread_data = manager.scratch_space(); - IOContext &ctx = this_thread_data->ctx; - - nhood_cache_buf = new uint32_t[num_cached_nodes * (max_degree + 1)]; - memset(nhood_cache_buf, 0, num_cached_nodes * (max_degree + 1)); - - size_t coord_cache_buf_len = num_cached_nodes * aligned_dim; - diskann::alloc_aligned((void **)&coord_cache_buf, - coord_cache_buf_len * sizeof(T), 8 * sizeof(T)); - memset(coord_cache_buf, 0, coord_cache_buf_len * sizeof(T)); - - size_t BLOCK_SIZE = 8; - size_t num_blocks = DIV_ROUND_UP(num_cached_nodes, BLOCK_SIZE); - - for (size_t block = 0; block < num_blocks; block++) { - size_t start_idx = block * BLOCK_SIZE; - size_t end_idx = (std::min)(num_cached_nodes, (block + 1) * BLOCK_SIZE); - std::vector read_reqs; - std::vector> nhoods; - for (size_t node_idx = start_idx; node_idx < end_idx; node_idx++) { - AlignedRead read; - char *buf = nullptr; - alloc_aligned((void **)&buf, SECTOR_LEN, SECTOR_LEN); - nhoods.push_back(std::make_pair(node_list[node_idx], buf)); - read.len = SECTOR_LEN; - read.buf = buf; - read.offset = NODE_SECTOR_NO(node_list[node_idx]) * SECTOR_LEN; - read_reqs.push_back(read); - } +template void PQFlashIndex::load_cache_list(std::vector &node_list) +{ + diskann::cout << "Loading the cache list into memory.." << std::flush; + size_t num_cached_nodes = node_list.size(); + + // borrow thread data + ScratchStoreManager> manager(this->thread_data); + auto this_thread_data = manager.scratch_space(); + IOContext &ctx = this_thread_data->ctx; + + nhood_cache_buf = new uint32_t[num_cached_nodes * (max_degree + 1)]; + memset(nhood_cache_buf, 0, num_cached_nodes * (max_degree + 1)); - reader->read(read_reqs, ctx); + size_t coord_cache_buf_len = num_cached_nodes * aligned_dim; + diskann::alloc_aligned((void **)&coord_cache_buf, coord_cache_buf_len * sizeof(T), 8 * sizeof(T)); + memset(coord_cache_buf, 0, coord_cache_buf_len * sizeof(T)); - size_t node_idx = start_idx; - for (uint32_t i = 0; i < read_reqs.size(); i++) { -#if defined(_WINDOWS) && \ - defined(USE_BING_INFRA) // this block is to handle failed reads in - // production settings - if ((*ctx.m_pRequestsStatus)[i] != IOContext::READ_SUCCESS) { - continue; - } + size_t BLOCK_SIZE = 8; + size_t num_blocks = DIV_ROUND_UP(num_cached_nodes, BLOCK_SIZE); + + for (size_t block = 0; block < num_blocks; block++) + { + size_t start_idx = block * BLOCK_SIZE; + size_t end_idx = (std::min)(num_cached_nodes, (block + 1) * BLOCK_SIZE); + std::vector read_reqs; + std::vector> nhoods; + for (size_t node_idx = start_idx; node_idx < end_idx; node_idx++) + { + AlignedRead read; + char *buf = nullptr; + alloc_aligned((void **)&buf, SECTOR_LEN, SECTOR_LEN); + nhoods.push_back(std::make_pair(node_list[node_idx], buf)); + read.len = SECTOR_LEN; + read.buf = buf; + read.offset = NODE_SECTOR_NO(node_list[node_idx]) * SECTOR_LEN; + read_reqs.push_back(read); + } + + reader->read(read_reqs, ctx); + + size_t node_idx = start_idx; + for (uint32_t i = 0; i < read_reqs.size(); i++) + { +#if defined(_WINDOWS) && defined(USE_BING_INFRA) // this block is to handle failed reads in + // production settings + if ((*ctx.m_pRequestsStatus)[i] != IOContext::READ_SUCCESS) + { + continue; + } #endif - auto &nhood = nhoods[i]; - char *node_buf = OFFSET_TO_NODE(nhood.second, nhood.first); - T *node_coords = OFFSET_TO_NODE_COORDS(node_buf); - T *cached_coords = coord_cache_buf + node_idx * aligned_dim; - memcpy(cached_coords, node_coords, disk_bytes_per_point); - coord_cache.insert(std::make_pair(nhood.first, cached_coords)); - - // insert node nhood into nhood_cache - uint32_t *node_nhood = OFFSET_TO_NODE_NHOOD(node_buf); - - auto nnbrs = *node_nhood; - uint32_t *nbrs = node_nhood + 1; - std::pair cnhood; - cnhood.first = nnbrs; - cnhood.second = nhood_cache_buf + node_idx * (max_degree + 1); - memcpy(cnhood.second, nbrs, nnbrs * sizeof(uint32_t)); - nhood_cache.insert(std::make_pair(nhood.first, cnhood)); - aligned_free(nhood.second); - node_idx++; + auto &nhood = nhoods[i]; + char *node_buf = OFFSET_TO_NODE(nhood.second, nhood.first); + T *node_coords = OFFSET_TO_NODE_COORDS(node_buf); + T *cached_coords = coord_cache_buf + node_idx * aligned_dim; + memcpy(cached_coords, node_coords, disk_bytes_per_point); + coord_cache.insert(std::make_pair(nhood.first, cached_coords)); + + // insert node nhood into nhood_cache + uint32_t *node_nhood = OFFSET_TO_NODE_NHOOD(node_buf); + + auto nnbrs = *node_nhood; + uint32_t *nbrs = node_nhood + 1; + std::pair cnhood; + cnhood.first = nnbrs; + cnhood.second = nhood_cache_buf + node_idx * (max_degree + 1); + memcpy(cnhood.second, nbrs, nnbrs * sizeof(uint32_t)); + nhood_cache.insert(std::make_pair(nhood.first, cnhood)); + aligned_free(nhood.second); + node_idx++; + } } - } - diskann::cout << "..done." << std::endl; + diskann::cout << "..done." << std::endl; } #ifdef EXEC_ENV_OLS template -void PQFlashIndex::generate_cache_list_from_sample_queries( - MemoryMappedFiles &files, std::string sample_bin, uint64_t l_search, - uint64_t beamwidth, uint64_t num_nodes_to_cache, uint32_t nthreads, - std::vector &node_list) { +void PQFlashIndex::generate_cache_list_from_sample_queries(MemoryMappedFiles &files, std::string sample_bin, + uint64_t l_search, uint64_t beamwidth, + uint64_t num_nodes_to_cache, uint32_t nthreads, + std::vector &node_list) +{ #else template -void PQFlashIndex::generate_cache_list_from_sample_queries( - std::string sample_bin, uint64_t l_search, uint64_t beamwidth, - uint64_t num_nodes_to_cache, uint32_t nthreads, - std::vector &node_list) { +void PQFlashIndex::generate_cache_list_from_sample_queries(std::string sample_bin, uint64_t l_search, + uint64_t beamwidth, uint64_t num_nodes_to_cache, + uint32_t nthreads, + std::vector &node_list) +{ #endif - this->count_visited_nodes = true; - this->node_visit_counter.clear(); - this->node_visit_counter.resize(this->num_points); - for (uint32_t i = 0; i < node_visit_counter.size(); i++) { - this->node_visit_counter[i].first = i; - this->node_visit_counter[i].second = 0; - } + this->count_visited_nodes = true; + this->node_visit_counter.clear(); + this->node_visit_counter.resize(this->num_points); + for (uint32_t i = 0; i < node_visit_counter.size(); i++) + { + this->node_visit_counter[i].first = i; + this->node_visit_counter[i].second = 0; + } - uint64_t sample_num, sample_dim, sample_aligned_dim; - T *samples; + uint64_t sample_num, sample_dim, sample_aligned_dim; + T *samples; #ifdef EXEC_ENV_OLS - if (files.fileExists(sample_bin)) { - diskann::load_aligned_bin(files, sample_bin, samples, sample_num, - sample_dim, sample_aligned_dim); - } + if (files.fileExists(sample_bin)) + { + diskann::load_aligned_bin(files, sample_bin, samples, sample_num, sample_dim, sample_aligned_dim); + } #else - if (file_exists(sample_bin)) { - diskann::load_aligned_bin(sample_bin, samples, sample_num, sample_dim, - sample_aligned_dim); - } + if (file_exists(sample_bin)) + { + diskann::load_aligned_bin(sample_bin, samples, sample_num, sample_dim, sample_aligned_dim); + } #endif - else { - diskann::cerr << "Sample bin file not found. Not generating cache." - << std::endl; - return; - } + else + { + diskann::cerr << "Sample bin file not found. Not generating cache." << std::endl; + return; + } - std::vector tmp_result_ids_64(sample_num, 0); - std::vector tmp_result_dists(sample_num, 0); + std::vector tmp_result_ids_64(sample_num, 0); + std::vector tmp_result_dists(sample_num, 0); #pragma omp parallel for schedule(dynamic, 1) num_threads(nthreads) - for (int64_t 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_dists.data() + (i * 1), beamwidth); - } - - std::sort(this->node_visit_counter.begin(), node_visit_counter.end(), - [](std::pair &left, - std::pair &right) { - return left.second > right.second; - }); - node_list.clear(); - node_list.shrink_to_fit(); - node_list.reserve(num_nodes_to_cache); - for (uint64_t i = 0; i < num_nodes_to_cache; i++) { - node_list.push_back(this->node_visit_counter[i].first); - } - this->count_visited_nodes = false; - - diskann::aligned_free(samples); + for (int64_t 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_dists.data() + (i * 1), beamwidth); + } + + std::sort(this->node_visit_counter.begin(), node_visit_counter.end(), + [](std::pair &left, std::pair &right) { + return left.second > right.second; + }); + node_list.clear(); + node_list.shrink_to_fit(); + node_list.reserve(num_nodes_to_cache); + for (uint64_t i = 0; i < num_nodes_to_cache; i++) + { + node_list.push_back(this->node_visit_counter[i].first); + } + this->count_visited_nodes = false; + + diskann::aligned_free(samples); } template -void PQFlashIndex::cache_bfs_levels(uint64_t num_nodes_to_cache, - std::vector &node_list, - const bool shuffle) { - std::random_device rng; - std::mt19937 urng(rng()); - - tsl::robin_set node_set; - - // Do not cache more than 10% of the nodes in the index - uint64_t tenp_nodes = (uint64_t)(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 - ScratchStoreManager> manager(this->thread_data); - auto this_thread_data = manager.scratch_space(); - IOContext &ctx = this_thread_data->ctx; - - std::unique_ptr> cur_level, prev_level; - cur_level = std::make_unique>(); - prev_level = std::make_unique>(); - - for (uint64_t miter = 0; miter < num_medoids; miter++) { - cur_level->insert(medoids[miter]); - } - - uint64_t lvl = 1; - uint64_t prev_node_set_size = 0; - while ((node_set.size() + cur_level->size() < num_nodes_to_cache) && - cur_level->size() != 0) { - // swap prev_level and cur_level - std::swap(prev_level, cur_level); - // clear cur_level - cur_level->clear(); - - std::vector nodes_to_expand; - - for (const uint32_t &id : *prev_level) { - if (node_set.find(id) != node_set.end()) { - continue; - } - node_set.insert(id); - nodes_to_expand.push_back(id); +void PQFlashIndex::cache_bfs_levels(uint64_t num_nodes_to_cache, std::vector &node_list, + const bool shuffle) +{ + std::random_device rng; + std::mt19937 urng(rng()); + + tsl::robin_set node_set; + + // Do not cache more than 10% of the nodes in the index + uint64_t tenp_nodes = (uint64_t)(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; - if (shuffle) - std::shuffle(nodes_to_expand.begin(), nodes_to_expand.end(), urng); - else - std::sort(nodes_to_expand.begin(), nodes_to_expand.end()); + // borrow thread data + ScratchStoreManager> manager(this->thread_data); + auto this_thread_data = manager.scratch_space(); + IOContext &ctx = this_thread_data->ctx; - diskann::cout << "Level: " << lvl << std::flush; - bool finish_flag = false; - - uint64_t BLOCK_SIZE = 1024; - uint64_t nblocks = DIV_ROUND_UP(nodes_to_expand.size(), BLOCK_SIZE); - for (size_t block = 0; block < nblocks && !finish_flag; block++) { - diskann::cout << "." << std::flush; - 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> nhoods; - for (size_t cur_pt = start; cur_pt < end; cur_pt++) { - char *buf = nullptr; - alloc_aligned((void **)&buf, SECTOR_LEN, SECTOR_LEN); - nhoods.emplace_back(nodes_to_expand[cur_pt], buf); - AlignedRead read; - read.len = SECTOR_LEN; - read.buf = buf; - 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 (uint32_t i = 0; i < read_reqs.size(); i++) { -#if defined(_WINDOWS) && \ - defined(USE_BING_INFRA) // this block is to handle read failures in - // production settings - if ((*ctx.m_pRequestsStatus)[i] != IOContext::READ_SUCCESS) { - continue; + std::unique_ptr> cur_level, prev_level; + cur_level = std::make_unique>(); + prev_level = std::make_unique>(); + + for (uint64_t miter = 0; miter < num_medoids; miter++) + { + cur_level->insert(medoids[miter]); + } + + uint64_t lvl = 1; + uint64_t prev_node_set_size = 0; + while ((node_set.size() + cur_level->size() < num_nodes_to_cache) && cur_level->size() != 0) + { + // swap prev_level and cur_level + std::swap(prev_level, cur_level); + // clear cur_level + cur_level->clear(); + + std::vector nodes_to_expand; + + for (const uint32_t &id : *prev_level) + { + if (node_set.find(id) != node_set.end()) + { + continue; + } + node_set.insert(id); + nodes_to_expand.push_back(id); } + + if (shuffle) + std::shuffle(nodes_to_expand.begin(), nodes_to_expand.end(), urng); + else + std::sort(nodes_to_expand.begin(), nodes_to_expand.end()); + + diskann::cout << "Level: " << lvl << std::flush; + bool finish_flag = false; + + uint64_t BLOCK_SIZE = 1024; + uint64_t nblocks = DIV_ROUND_UP(nodes_to_expand.size(), BLOCK_SIZE); + for (size_t block = 0; block < nblocks && !finish_flag; block++) + { + diskann::cout << "." << std::flush; + 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> nhoods; + for (size_t cur_pt = start; cur_pt < end; cur_pt++) + { + char *buf = nullptr; + alloc_aligned((void **)&buf, SECTOR_LEN, SECTOR_LEN); + nhoods.emplace_back(nodes_to_expand[cur_pt], buf); + AlignedRead read; + read.len = SECTOR_LEN; + read.buf = buf; + 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 (uint32_t i = 0; i < read_reqs.size(); i++) + { +#if defined(_WINDOWS) && defined(USE_BING_INFRA) // this block is to handle read failures in + // production settings + if ((*ctx.m_pRequestsStatus)[i] != IOContext::READ_SUCCESS) + { + continue; + } #endif - auto &nhood = nhoods[i]; - - // insert node coord into coord_cache - char *node_buf = OFFSET_TO_NODE(nhood.second, nhood.first); - uint32_t *node_nhood = OFFSET_TO_NODE_NHOOD(node_buf); - uint64_t nnbrs = (uint64_t)*node_nhood; - uint32_t *nbrs = node_nhood + 1; - // explore next level - for (uint64_t j = 0; j < nnbrs && !finish_flag; j++) { - if (node_set.find(nbrs[j]) == node_set.end()) { - cur_level->insert(nbrs[j]); - } - if (cur_level->size() + node_set.size() >= num_nodes_to_cache) { - finish_flag = true; - } + auto &nhood = nhoods[i]; + + // insert node coord into coord_cache + char *node_buf = OFFSET_TO_NODE(nhood.second, nhood.first); + uint32_t *node_nhood = OFFSET_TO_NODE_NHOOD(node_buf); + uint64_t nnbrs = (uint64_t)*node_nhood; + uint32_t *nbrs = node_nhood + 1; + // explore next level + for (uint64_t j = 0; j < nnbrs && !finish_flag; j++) + { + if (node_set.find(nbrs[j]) == node_set.end()) + { + cur_level->insert(nbrs[j]); + } + if (cur_level->size() + node_set.size() >= num_nodes_to_cache) + { + finish_flag = true; + } + } + aligned_free(nhood.second); + } } - aligned_free(nhood.second); - } - } - diskann::cout << ". #nodes: " << node_set.size() - prev_node_set_size - << ", #nodes thus far: " << node_list.size() << std::endl; - prev_node_set_size = node_set.size(); - lvl++; - } + diskann::cout << ". #nodes: " << node_set.size() - prev_node_set_size + << ", #nodes thus far: " << node_list.size() << std::endl; + prev_node_set_size = node_set.size(); + lvl++; + } - assert(node_set.size() + cur_level->size() == num_nodes_to_cache || - cur_level->size() == 0); + assert(node_set.size() + cur_level->size() == num_nodes_to_cache || cur_level->size() == 0); - node_list.clear(); - node_list.reserve(node_set.size() + cur_level->size()); - for (auto node : node_set) node_list.push_back(node); - for (auto node : *cur_level) node_list.push_back(node); + node_list.clear(); + node_list.reserve(node_set.size() + cur_level->size()); + for (auto node : node_set) + node_list.push_back(node); + for (auto node : *cur_level) + node_list.push_back(node); - diskann::cout << "Level: " << lvl << std::flush; - diskann::cout << ". #nodes: " << node_list.size() - prev_node_set_size - << ", #nodes thus far: " << node_list.size() << std::endl; - diskann::cout << "done" << std::endl; + diskann::cout << "Level: " << lvl << std::flush; + diskann::cout << ". #nodes: " << node_list.size() - prev_node_set_size << ", #nodes thus far: " << node_list.size() + << std::endl; + diskann::cout << "done" << std::endl; } -template -void PQFlashIndex::use_medoids_data_as_centroids() { - if (centroid_data != nullptr) aligned_free(centroid_data); - alloc_aligned(((void **)¢roid_data), - num_medoids * aligned_dim * sizeof(float), 32); - std::memset(centroid_data, 0, num_medoids * aligned_dim * sizeof(float)); - - // borrow ctx - ScratchStoreManager> manager(this->thread_data); - auto data = manager.scratch_space(); - IOContext &ctx = data->ctx; - diskann::cout << "Loading centroid data from medoids vector data of " - << num_medoids << " medoid(s)" << std::endl; - for (uint64_t cur_m = 0; cur_m < num_medoids; cur_m++) { - auto medoid = medoids[cur_m]; - // read medoid nhood - char *medoid_buf = nullptr; - alloc_aligned((void **)&medoid_buf, SECTOR_LEN, SECTOR_LEN); - std::vector medoid_read(1); - medoid_read[0].len = SECTOR_LEN; - medoid_read[0].buf = medoid_buf; - medoid_read[0].offset = NODE_SECTOR_NO(medoid) * SECTOR_LEN; - reader->read(medoid_read, ctx); - - // all data about medoid - char *medoid_node_buf = OFFSET_TO_NODE(medoid_buf, medoid); - - // add medoid coords to `coord_cache` - T *medoid_coords = new T[data_dim]; - T *medoid_disk_coords = OFFSET_TO_NODE_COORDS(medoid_node_buf); - memcpy(medoid_coords, medoid_disk_coords, disk_bytes_per_point); - - if (!use_disk_index_pq) { - for (uint32_t i = 0; i < data_dim; i++) - centroid_data[cur_m * aligned_dim + i] = medoid_coords[i]; - } else { - disk_pq_table.inflate_vector((uint8_t *)medoid_coords, - (centroid_data + cur_m * aligned_dim)); - } +template void PQFlashIndex::use_medoids_data_as_centroids() +{ + if (centroid_data != nullptr) + aligned_free(centroid_data); + alloc_aligned(((void **)¢roid_data), num_medoids * aligned_dim * sizeof(float), 32); + std::memset(centroid_data, 0, num_medoids * aligned_dim * sizeof(float)); + + // borrow ctx + ScratchStoreManager> manager(this->thread_data); + auto data = manager.scratch_space(); + IOContext &ctx = data->ctx; + diskann::cout << "Loading centroid data from medoids vector data of " << num_medoids << " medoid(s)" << std::endl; + for (uint64_t cur_m = 0; cur_m < num_medoids; cur_m++) + { + auto medoid = medoids[cur_m]; + // read medoid nhood + char *medoid_buf = nullptr; + alloc_aligned((void **)&medoid_buf, SECTOR_LEN, SECTOR_LEN); + std::vector medoid_read(1); + medoid_read[0].len = SECTOR_LEN; + medoid_read[0].buf = medoid_buf; + medoid_read[0].offset = NODE_SECTOR_NO(medoid) * SECTOR_LEN; + reader->read(medoid_read, ctx); + + // all data about medoid + char *medoid_node_buf = OFFSET_TO_NODE(medoid_buf, medoid); + + // add medoid coords to `coord_cache` + T *medoid_coords = new T[data_dim]; + T *medoid_disk_coords = OFFSET_TO_NODE_COORDS(medoid_node_buf); + memcpy(medoid_coords, medoid_disk_coords, disk_bytes_per_point); + + if (!use_disk_index_pq) + { + for (uint32_t i = 0; i < data_dim; i++) + centroid_data[cur_m * aligned_dim + i] = medoid_coords[i]; + } + else + { + disk_pq_table.inflate_vector((uint8_t *)medoid_coords, (centroid_data + cur_m * aligned_dim)); + } - aligned_free(medoid_buf); - delete[] medoid_coords; - } + aligned_free(medoid_buf); + delete[] medoid_coords; + } } template -inline int32_t PQFlashIndex::get_filter_number( - const LabelT &filter_label) { - int idx = -1; - for (uint32_t i = 0; i < _filter_list.size(); i++) { - if (_filter_list[i] == filter_label) { - idx = i; - break; +inline int32_t PQFlashIndex::get_filter_number(const LabelT &filter_label) +{ + int idx = -1; + for (uint32_t i = 0; i < _filter_list.size(); i++) + { + if (_filter_list[i] == filter_label) + { + idx = i; + break; + } } - } - return idx; + return idx; } template -std::unordered_map PQFlashIndex::load_label_map( - const std::string &labels_map_file) { - std::unordered_map string_to_int_mp; - std::ifstream map_reader(labels_map_file); - std::string line, token; - LabelT token_as_num; - std::string label_str; - while (std::getline(map_reader, line)) { - std::istringstream iss(line); - getline(iss, token, '\t'); - label_str = token; - getline(iss, token, '\t'); - token_as_num = std::stoul(token); - string_to_int_mp[label_str] = token_as_num; - } - return string_to_int_mp; +std::unordered_map PQFlashIndex::load_label_map(const std::string &labels_map_file) +{ + std::unordered_map string_to_int_mp; + std::ifstream map_reader(labels_map_file); + std::string line, token; + LabelT token_as_num; + std::string label_str; + while (std::getline(map_reader, line)) + { + std::istringstream iss(line); + getline(iss, token, '\t'); + label_str = token; + getline(iss, token, '\t'); + token_as_num = std::stoul(token); + string_to_int_mp[label_str] = token_as_num; + } + return string_to_int_mp; } template -LabelT PQFlashIndex::get_converted_label( - const std::string &filter_label) { - if (_label_map.find(filter_label) != _label_map.end()) { - return _label_map[filter_label]; - } - std::stringstream stream; - stream << "Unable to find label in the Label Map"; - diskann::cerr << stream.str() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); +LabelT PQFlashIndex::get_converted_label(const std::string &filter_label) +{ + if (_label_map.find(filter_label) != _label_map.end()) + { + return _label_map[filter_label]; + } + std::stringstream stream; + stream << "Unable to find label in the Label Map"; + diskann::cerr << stream.str() << std::endl; + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } template -void PQFlashIndex::get_label_file_metadata( - std::string map_file, uint32_t &num_pts, uint32_t &num_total_labels) { - std::ifstream infile(map_file); - std::string line, token; - num_pts = 0; - num_total_labels = 0; - - while (std::getline(infile, line)) { - std::istringstream iss(line); - while (getline(iss, token, ',')) { - token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); - token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); - num_total_labels++; +void PQFlashIndex::get_label_file_metadata(std::string map_file, uint32_t &num_pts, + uint32_t &num_total_labels) +{ + std::ifstream infile(map_file); + std::string line, token; + num_pts = 0; + num_total_labels = 0; + + while (std::getline(infile, line)) + { + std::istringstream iss(line); + while (getline(iss, token, ',')) + { + token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); + token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); + num_total_labels++; + } + num_pts++; } - num_pts++; - } - diskann::cout << "Labels file metadata: num_points: " << num_pts - << ", #total_labels: " << num_total_labels << std::endl; - infile.close(); + diskann::cout << "Labels file metadata: num_points: " << num_pts << ", #total_labels: " << num_total_labels + << std::endl; + infile.close(); } template -inline bool PQFlashIndex::point_has_label(uint32_t point_id, - uint32_t label_id) { - uint32_t start_vec = _pts_to_label_offsets[point_id]; - uint32_t num_lbls = _pts_to_labels[start_vec]; - bool ret_val = false; - for (uint32_t i = 0; i < num_lbls; i++) { - if (_pts_to_labels[start_vec + 1 + i] == label_id) { - ret_val = true; - break; +inline bool PQFlashIndex::point_has_label(uint32_t point_id, uint32_t label_id) +{ + uint32_t start_vec = _pts_to_label_offsets[point_id]; + uint32_t num_lbls = _pts_to_labels[start_vec]; + bool ret_val = false; + for (uint32_t i = 0; i < num_lbls; i++) + { + if (_pts_to_labels[start_vec + 1 + i] == label_id) + { + ret_val = true; + break; + } } - } - return ret_val; + return ret_val; } template -void PQFlashIndex::parse_label_file(const std::string &label_file, - size_t &num_points_labels) { - std::ifstream infile(label_file); - if (infile.fail()) { - throw diskann::ANNException( - std::string("Failed to open file ") + label_file, -1); - } - - std::string line, token; - uint32_t line_cnt = 0; - - uint32_t num_pts_in_label_file; - uint32_t num_total_labels; - get_label_file_metadata(label_file, num_pts_in_label_file, num_total_labels); - - _pts_to_label_offsets = new uint32_t[num_pts_in_label_file]; - _pts_to_labels = new uint32_t[num_pts_in_label_file + num_total_labels]; - uint32_t counter = 0; - - while (std::getline(infile, line)) { - std::istringstream iss(line); - std::vector lbls(0); - - _pts_to_label_offsets[line_cnt] = counter; - uint32_t &num_lbls_in_cur_pt = _pts_to_labels[counter]; - num_lbls_in_cur_pt = 0; - counter++; - getline(iss, token, '\t'); - std::istringstream new_iss(token); - while (getline(new_iss, token, ',')) { - token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); - token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); - LabelT token_as_num = std::stoul(token); - if (_labels.find(token_as_num) == _labels.end()) { - _filter_list.emplace_back(token_as_num); - } - int32_t filter_num = get_filter_number(token_as_num); - if (filter_num == -1) { - diskann::cout << "Error!! " << std::endl; - exit(-1); - } - _pts_to_labels[counter++] = filter_num; - num_lbls_in_cur_pt++; - _labels.insert(token_as_num); +void PQFlashIndex::parse_label_file(const std::string &label_file, size_t &num_points_labels) +{ + std::ifstream infile(label_file); + if (infile.fail()) + { + throw diskann::ANNException(std::string("Failed to open file ") + label_file, -1); } - if (num_lbls_in_cur_pt == 0) { - diskann::cout << "No label found for point " << line_cnt << std::endl; - exit(-1); + std::string line, token; + uint32_t line_cnt = 0; + + uint32_t num_pts_in_label_file; + uint32_t num_total_labels; + get_label_file_metadata(label_file, num_pts_in_label_file, num_total_labels); + + _pts_to_label_offsets = new uint32_t[num_pts_in_label_file]; + _pts_to_labels = new uint32_t[num_pts_in_label_file + num_total_labels]; + uint32_t counter = 0; + + while (std::getline(infile, line)) + { + std::istringstream iss(line); + std::vector lbls(0); + + _pts_to_label_offsets[line_cnt] = counter; + uint32_t &num_lbls_in_cur_pt = _pts_to_labels[counter]; + num_lbls_in_cur_pt = 0; + counter++; + getline(iss, token, '\t'); + std::istringstream new_iss(token); + while (getline(new_iss, token, ',')) + { + token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); + token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); + LabelT token_as_num = std::stoul(token); + if (_labels.find(token_as_num) == _labels.end()) + { + _filter_list.emplace_back(token_as_num); + } + int32_t filter_num = get_filter_number(token_as_num); + if (filter_num == -1) + { + diskann::cout << "Error!! " << std::endl; + exit(-1); + } + _pts_to_labels[counter++] = filter_num; + num_lbls_in_cur_pt++; + _labels.insert(token_as_num); + } + + if (num_lbls_in_cur_pt == 0) + { + diskann::cout << "No label found for point " << line_cnt << std::endl; + exit(-1); + } + line_cnt++; } - line_cnt++; - } - infile.close(); - num_points_labels = line_cnt; + infile.close(); + num_points_labels = line_cnt; } -template -void PQFlashIndex::set_universal_label(const LabelT &label) { - int32_t temp_filter_num = get_filter_number(label); - if (temp_filter_num == -1) { - diskann::cout << "Error, could not find universal label." << std::endl; - } else { - _use_universal_label = true; - _universal_filter_num = (uint32_t)temp_filter_num; - } +template void PQFlashIndex::set_universal_label(const LabelT &label) +{ + int32_t temp_filter_num = get_filter_number(label); + if (temp_filter_num == -1) + { + diskann::cout << "Error, could not find universal label." << std::endl; + } + else + { + _use_universal_label = true; + _universal_filter_num = (uint32_t)temp_filter_num; + } } #ifdef EXEC_ENV_OLS template -int PQFlashIndex::load(MemoryMappedFiles &files, - uint32_t num_threads, - const char *index_prefix) { +int PQFlashIndex::load(MemoryMappedFiles &files, uint32_t num_threads, const char *index_prefix) +{ #else -template -int PQFlashIndex::load(uint32_t num_threads, - const char *index_prefix) { +template int PQFlashIndex::load(uint32_t num_threads, const char *index_prefix) +{ #endif - std::string pq_table_bin = std::string(index_prefix) + "_pq_pivots.bin"; - std::string pq_compressed_vectors = - std::string(index_prefix) + "_pq_compressed.bin"; - std::string disk_index_file = std::string(index_prefix) + "_disk.index"; + std::string pq_table_bin = std::string(index_prefix) + "_pq_pivots.bin"; + std::string pq_compressed_vectors = std::string(index_prefix) + "_pq_compressed.bin"; + std::string disk_index_file = std::string(index_prefix) + "_disk.index"; #ifdef EXEC_ENV_OLS - return load_from_separate_paths(files, num_threads, disk_index_file.c_str(), - pq_table_bin.c_str(), - pq_compressed_vectors.c_str()); + return load_from_separate_paths(files, num_threads, disk_index_file.c_str(), pq_table_bin.c_str(), + pq_compressed_vectors.c_str()); #else - return load_from_separate_paths(num_threads, disk_index_file.c_str(), - pq_table_bin.c_str(), - pq_compressed_vectors.c_str()); + return load_from_separate_paths(num_threads, disk_index_file.c_str(), pq_table_bin.c_str(), + pq_compressed_vectors.c_str()); #endif } #ifdef EXEC_ENV_OLS template -int PQFlashIndex::load_from_separate_paths( - diskann::MemoryMappedFiles &files, uint32_t num_threads, - const char *index_filepath, const char *pivots_filepath, - const char *compressed_filepath) { +int PQFlashIndex::load_from_separate_paths(diskann::MemoryMappedFiles &files, uint32_t num_threads, + const char *index_filepath, const char *pivots_filepath, + const char *compressed_filepath) +{ #else template -int PQFlashIndex::load_from_separate_paths( - uint32_t num_threads, const char *index_filepath, - const char *pivots_filepath, const char *compressed_filepath) { +int PQFlashIndex::load_from_separate_paths(uint32_t num_threads, const char *index_filepath, + const char *pivots_filepath, const char *compressed_filepath) +{ #endif - std::string pq_table_bin = pivots_filepath; - std::string pq_compressed_vectors = compressed_filepath; - std::string disk_index_file = index_filepath; - std::string medoids_file = std::string(disk_index_file) + "_medoids.bin"; - std::string centroids_file = std::string(disk_index_file) + "_centroids.bin"; - - std::string labels_file = std ::string(disk_index_file) + "_labels.txt"; - std::string labels_to_medoids = - std ::string(disk_index_file) + "_labels_to_medoids.txt"; - std::string dummy_map_file = std ::string(disk_index_file) + "_dummy_map.txt"; - std::string labels_map_file = - std ::string(disk_index_file) + "_labels_map.txt"; - size_t num_pts_in_label_file = 0; - - size_t pq_file_dim, pq_file_num_centroids; + std::string pq_table_bin = pivots_filepath; + std::string pq_compressed_vectors = compressed_filepath; + std::string disk_index_file = index_filepath; + std::string medoids_file = std::string(disk_index_file) + "_medoids.bin"; + std::string centroids_file = std::string(disk_index_file) + "_centroids.bin"; + + std::string labels_file = std ::string(disk_index_file) + "_labels.txt"; + std::string labels_to_medoids = std ::string(disk_index_file) + "_labels_to_medoids.txt"; + std::string dummy_map_file = std ::string(disk_index_file) + "_dummy_map.txt"; + std::string labels_map_file = std ::string(disk_index_file) + "_labels_map.txt"; + size_t num_pts_in_label_file = 0; + + size_t pq_file_dim, pq_file_num_centroids; #ifdef EXEC_ENV_OLS - get_bin_metadata(files, pq_table_bin, pq_file_num_centroids, pq_file_dim, - METADATA_SIZE); + get_bin_metadata(files, pq_table_bin, pq_file_num_centroids, pq_file_dim, METADATA_SIZE); #else - get_bin_metadata(pq_table_bin, pq_file_num_centroids, pq_file_dim, - METADATA_SIZE); + get_bin_metadata(pq_table_bin, pq_file_num_centroids, pq_file_dim, METADATA_SIZE); #endif - this->disk_index_file = disk_index_file; + this->disk_index_file = disk_index_file; - if (pq_file_num_centroids != 256) { - diskann::cout << "Error. Number of PQ centroids is not 256. Exiting." - << std::endl; - return -1; - } - - this->data_dim = pq_file_dim; - // will reset later if we use PQ on disk - this->disk_data_dim = this->data_dim; - // will change later if we use PQ on disk or if we are using - // inner product without PQ - this->disk_bytes_per_point = this->data_dim * sizeof(T); - this->aligned_dim = ROUND_UP(pq_file_dim, 8); - - size_t npts_u64, nchunks_u64; + if (pq_file_num_centroids != 256) + { + diskann::cout << "Error. Number of PQ centroids is not 256. Exiting." << std::endl; + return -1; + } + + this->data_dim = pq_file_dim; + // will reset later if we use PQ on disk + this->disk_data_dim = this->data_dim; + // will change later if we use PQ on disk or if we are using + // inner product without PQ + this->disk_bytes_per_point = this->data_dim * sizeof(T); + this->aligned_dim = ROUND_UP(pq_file_dim, 8); + + size_t npts_u64, nchunks_u64; #ifdef EXEC_ENV_OLS - diskann::load_bin(files, pq_compressed_vectors, this->data, npts_u64, - nchunks_u64); + diskann::load_bin(files, pq_compressed_vectors, this->data, npts_u64, nchunks_u64); #else - diskann::load_bin(pq_compressed_vectors, this->data, npts_u64, - nchunks_u64); + diskann::load_bin(pq_compressed_vectors, this->data, npts_u64, nchunks_u64); #endif - this->num_points = npts_u64; - this->n_chunks = nchunks_u64; - if (file_exists(labels_file)) { - parse_label_file(labels_file, num_pts_in_label_file); - assert(num_pts_in_label_file == this->num_points); - _label_map = load_label_map(labels_map_file); - if (file_exists(labels_to_medoids)) { - std::ifstream medoid_stream(labels_to_medoids); - assert(medoid_stream.is_open()); - std::string line, token; - - _filter_to_medoid_id.clear(); - try { - while (std::getline(medoid_stream, line)) { - std::istringstream iss(line); - uint32_t cnt = 0; - uint32_t medoid = 0; - LabelT label; - while (std::getline(iss, token, ',')) { - if (cnt == 0) - label = std::stoul(token); - else - medoid = (uint32_t)stoul(token); - cnt++; - } - _filter_to_medoid_id[label] = medoid; + this->num_points = npts_u64; + this->n_chunks = nchunks_u64; + if (file_exists(labels_file)) + { + parse_label_file(labels_file, num_pts_in_label_file); + assert(num_pts_in_label_file == this->num_points); + _label_map = load_label_map(labels_map_file); + if (file_exists(labels_to_medoids)) + { + std::ifstream medoid_stream(labels_to_medoids); + assert(medoid_stream.is_open()); + std::string line, token; + + _filter_to_medoid_id.clear(); + try + { + while (std::getline(medoid_stream, line)) + { + std::istringstream iss(line); + uint32_t cnt = 0; + uint32_t medoid = 0; + LabelT label; + while (std::getline(iss, token, ',')) + { + if (cnt == 0) + label = std::stoul(token); + else + medoid = (uint32_t)stoul(token); + cnt++; + } + _filter_to_medoid_id[label] = medoid; + } + } + catch (std::system_error &e) + { + throw FileException(labels_to_medoids, e, __FUNCSIG__, __FILE__, __LINE__); + } } - } catch (std::system_error &e) { - throw FileException(labels_to_medoids, e, __FUNCSIG__, __FILE__, - __LINE__); - } - } - std::string univ_label_file = - std ::string(disk_index_file) + "_universal_label.txt"; - if (file_exists(univ_label_file)) { - std::ifstream universal_label_reader(univ_label_file); - assert(universal_label_reader.is_open()); - std::string univ_label; - universal_label_reader >> univ_label; - universal_label_reader.close(); - LabelT label_as_num = std::stoul(univ_label); - set_universal_label(label_as_num); - } - if (file_exists(dummy_map_file)) { - std::ifstream dummy_map_stream(dummy_map_file); - assert(dummy_map_stream.is_open()); - std::string line, token; - - while (std::getline(dummy_map_stream, line)) { - std::istringstream iss(line); - uint32_t cnt = 0; - uint32_t dummy_id; - uint32_t real_id; - while (std::getline(iss, token, ',')) { - if (cnt == 0) - dummy_id = (uint32_t)stoul(token); - else - real_id = (uint32_t)stoul(token); - cnt++; + std::string univ_label_file = std ::string(disk_index_file) + "_universal_label.txt"; + if (file_exists(univ_label_file)) + { + std::ifstream universal_label_reader(univ_label_file); + assert(universal_label_reader.is_open()); + std::string univ_label; + universal_label_reader >> univ_label; + universal_label_reader.close(); + LabelT label_as_num = std::stoul(univ_label); + set_universal_label(label_as_num); + } + if (file_exists(dummy_map_file)) + { + std::ifstream dummy_map_stream(dummy_map_file); + assert(dummy_map_stream.is_open()); + std::string line, token; + + while (std::getline(dummy_map_stream, line)) + { + std::istringstream iss(line); + uint32_t cnt = 0; + uint32_t dummy_id; + uint32_t real_id; + while (std::getline(iss, token, ',')) + { + if (cnt == 0) + dummy_id = (uint32_t)stoul(token); + else + real_id = (uint32_t)stoul(token); + cnt++; + } + _dummy_pts.insert(dummy_id); + _has_dummy_pts.insert(real_id); + _dummy_to_real_map[dummy_id] = real_id; + + if (_real_to_dummy_map.find(real_id) == _real_to_dummy_map.end()) + _real_to_dummy_map[real_id] = std::vector(); + + _real_to_dummy_map[real_id].emplace_back(dummy_id); + } + dummy_map_stream.close(); + diskann::cout << "Loaded dummy map" << std::endl; } - _dummy_pts.insert(dummy_id); - _has_dummy_pts.insert(real_id); - _dummy_to_real_map[dummy_id] = real_id; - - if (_real_to_dummy_map.find(real_id) == _real_to_dummy_map.end()) - _real_to_dummy_map[real_id] = std::vector(); - - _real_to_dummy_map[real_id].emplace_back(dummy_id); - } - dummy_map_stream.close(); - diskann::cout << "Loaded dummy map" << std::endl; } - } #ifdef EXEC_ENV_OLS - pq_table.load_pq_centroid_bin(files, pq_table_bin.c_str(), nchunks_u64); + 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); + 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 << "Loaded PQ centroids and in-memory compressed vectors. #points: " << num_points + << " #dim: " << data_dim << " #aligned_dim: " << aligned_dim << " #chunks: " << n_chunks << std::endl; - if (n_chunks > MAX_PQ_CHUNKS) { - std::stringstream stream; - stream << "Error loading index. Ensure that max PQ bytes for in-memory " - "PQ data does not exceed " - << MAX_PQ_CHUNKS << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); - } - - std::string disk_pq_pivots_path = this->disk_index_file + "_pq_pivots.bin"; - if (file_exists(disk_pq_pivots_path)) { - use_disk_index_pq = true; + if (n_chunks > MAX_PQ_CHUNKS) + { + std::stringstream stream; + stream << "Error loading index. Ensure that max PQ bytes for in-memory " + "PQ data does not exceed " + << MAX_PQ_CHUNKS << std::endl; + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } + + std::string disk_pq_pivots_path = this->disk_index_file + "_pq_pivots.bin"; + if (file_exists(disk_pq_pivots_path)) + { + use_disk_index_pq = true; #ifdef EXEC_ENV_OLS - // giving 0 chunks to make the pq_table infer from the - // chunk_offsets file the correct value - disk_pq_table.load_pq_centroid_bin(files, disk_pq_pivots_path.c_str(), 0); + // giving 0 chunks to make the pq_table infer from the + // chunk_offsets file the correct value + disk_pq_table.load_pq_centroid_bin(files, disk_pq_pivots_path.c_str(), 0); #else - // giving 0 chunks to make the pq_table infer from the - // chunk_offsets file the correct value - disk_pq_table.load_pq_centroid_bin(disk_pq_pivots_path.c_str(), 0); + // giving 0 chunks to make the pq_table infer from the + // 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_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." << std::endl; - } + disk_pq_n_chunks = disk_pq_table.get_num_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." + << 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); + // 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); #else - std::ifstream index_metadata(disk_index_file, std::ios::binary); + std::ifstream index_metadata(disk_index_file, std::ios::binary); #endif - uint32_t nr, nc; // metadata itself is stored as bin format (nr is number of - // metadata, nc should be 1) - READ_U32(index_metadata, nr); - READ_U32(index_metadata, nc); - - uint64_t disk_nnodes; - uint64_t disk_ndims; // can be disk PQ dim if disk_PQ is set to true - READ_U64(index_metadata, disk_nnodes); - READ_U64(index_metadata, disk_ndims); - - if (disk_nnodes != num_points) { - diskann::cout << "Mismatch in #points for compressed data file and disk " - "index file: " - << disk_nnodes << " vs " << num_points << std::endl; - 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 - disk_bytes_per_point) / sizeof(uint32_t)) - 1; - - if (max_degree > MAX_GRAPH_DEGREE) { - std::stringstream stream; - stream << "Error loading index. Ensure that max graph degree (R) does " - "not exceed " - << MAX_GRAPH_DEGREE << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); - } - - // setting up concept of frozen points in disk index for streaming-DiskANN - READ_U64(index_metadata, this->num_frozen_points); - uint64_t file_frozen_id; - 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, this->reorder_data_exists); - if (this->reorder_data_exists) { - if (this->use_disk_index_pq == false) { - throw ANNException( - "Reordering is designed for used with disk PQ " - "compression option", - -1, __FUNCSIG__, __FILE__, __LINE__); + uint32_t nr, nc; // metadata itself is stored as bin format (nr is number of + // metadata, nc should be 1) + READ_U32(index_metadata, nr); + READ_U32(index_metadata, nc); + + uint64_t disk_nnodes; + uint64_t disk_ndims; // can be disk PQ dim if disk_PQ is set to true + READ_U64(index_metadata, disk_nnodes); + READ_U64(index_metadata, disk_ndims); + + if (disk_nnodes != num_points) + { + diskann::cout << "Mismatch in #points for compressed data file and disk " + "index file: " + << disk_nnodes << " vs " << num_points << std::endl; + 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 - disk_bytes_per_point) / sizeof(uint32_t)) - 1; + + if (max_degree > MAX_GRAPH_DEGREE) + { + std::stringstream stream; + stream << "Error loading index. Ensure that max graph degree (R) does " + "not exceed " + << MAX_GRAPH_DEGREE << std::endl; + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } - READ_U64(index_metadata, this->reorder_data_start_sector); - READ_U64(index_metadata, this->ndims_reorder_vecs); - READ_U64(index_metadata, this->nvecs_per_sector); - } - 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; + // setting up concept of frozen points in disk index for streaming-DiskANN + READ_U64(index_metadata, this->num_frozen_points); + uint64_t file_frozen_id; + 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, this->reorder_data_exists); + if (this->reorder_data_exists) + { + if (this->use_disk_index_pq == false) + { + throw ANNException("Reordering is designed for used with disk PQ " + "compression option", + -1, __FUNCSIG__, __FILE__, __LINE__); + } + READ_U64(index_metadata, this->reorder_data_start_sector); + READ_U64(index_metadata, this->ndims_reorder_vecs); + READ_U64(index_metadata, this->nvecs_per_sector); + } + + 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; + delete[] bytes; #else - index_metadata.close(); + index_metadata.close(); #endif #ifndef EXEC_ENV_OLS - // open AlignedFileReader handle to index_file - std::string index_fname(disk_index_file); - reader->open(index_fname); - this->setup_thread_data(num_threads); - this->max_nthreads = num_threads; + // open AlignedFileReader handle to index_file + std::string index_fname(disk_index_file); + reader->open(index_fname); + 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; - diskann::load_bin(files, medoids_file, medoids, num_medoids, - tmp_dim); + if (files.fileExists(medoids_file)) + { + size_t tmp_dim; + diskann::load_bin(files, medoids_file, medoids, num_medoids, tmp_dim); #else - if (file_exists(medoids_file)) { - size_t tmp_dim; - diskann::load_bin(medoids_file, medoids, num_medoids, tmp_dim); + if (file_exists(medoids_file)) + { + size_t tmp_dim; + diskann::load_bin(medoids_file, medoids, num_medoids, tmp_dim); #endif - if (tmp_dim != 1) { - std::stringstream stream; - stream << "Error loading medoids file. Expected bin format of m times " - "1 vector of uint32_t." - << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); - } + if (tmp_dim != 1) + { + std::stringstream stream; + stream << "Error loading medoids file. Expected bin format of m times " + "1 vector of uint32_t." + << std::endl; + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } #ifdef EXEC_ENV_OLS - if (!files.fileExists(centroids_file)) { + if (!files.fileExists(centroids_file)) + { #else - if (!file_exists(centroids_file)) { + if (!file_exists(centroids_file)) + { #endif - diskann::cout - << "Centroid data file not found. Using corresponding vectors " - "for the medoids " - << std::endl; - use_medoids_data_as_centroids(); - } else { - size_t num_centroids, aligned_tmp_dim; + diskann::cout << "Centroid data file not found. Using corresponding vectors " + "for the medoids " + << std::endl; + use_medoids_data_as_centroids(); + } + else + { + size_t num_centroids, aligned_tmp_dim; #ifdef EXEC_ENV_OLS - diskann::load_aligned_bin(files, centroids_file, centroid_data, - num_centroids, tmp_dim, aligned_tmp_dim); + diskann::load_aligned_bin(files, centroids_file, centroid_data, num_centroids, tmp_dim, + aligned_tmp_dim); #else - diskann::load_aligned_bin(centroids_file, centroid_data, - num_centroids, tmp_dim, aligned_tmp_dim); + diskann::load_aligned_bin(centroids_file, centroid_data, num_centroids, tmp_dim, aligned_tmp_dim); #endif - if (aligned_tmp_dim != aligned_dim || num_centroids != num_medoids) { - std::stringstream stream; - 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."; - diskann::cerr << stream.str() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); - } + if (aligned_tmp_dim != aligned_dim || num_centroids != num_medoids) + { + std::stringstream stream; + 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."; + diskann::cerr << stream.str() << std::endl; + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } + } + } + else + { + num_medoids = 1; + medoids = new uint32_t[1]; + medoids[0] = (uint32_t)(medoid_id_on_file); + use_medoids_data_as_centroids(); } - } else { - num_medoids = 1; - medoids = new uint32_t[1]; - medoids[0] = (uint32_t)(medoid_id_on_file); - use_medoids_data_as_centroids(); - } - - std::string norm_file = std::string(disk_index_file) + "_max_base_norm.bin"; - - if (file_exists(norm_file) && metric == diskann::Metric::INNER_PRODUCT) { - uint64_t dumr, dumc; - float *norm_val; - diskann::load_bin(norm_file, norm_val, dumr, dumc); - this->max_base_norm = norm_val[0]; - diskann::cout << "Setting re-scaling factor of base vectors to " - << this->max_base_norm << std::endl; - delete[] norm_val; - } - diskann::cout << "done.." << std::endl; - return 0; + + std::string norm_file = std::string(disk_index_file) + "_max_base_norm.bin"; + + if (file_exists(norm_file) && metric == diskann::Metric::INNER_PRODUCT) + { + uint64_t dumr, dumc; + float *norm_val; + diskann::load_bin(norm_file, norm_val, dumr, dumc); + this->max_base_norm = norm_val[0]; + diskann::cout << "Setting re-scaling factor of base vectors to " << this->max_base_norm << std::endl; + delete[] norm_val; + } + diskann::cout << "done.." << std::endl; + return 0; } #ifdef USE_BING_INFRA -bool getNextCompletedRequest(const IOContext &ctx, size_t size, - int &completedIndex) { - bool waitsRemaining = false; - long completeCount = ctx.m_completeCount; - do { - for (int i = 0; i < size; i++) { - auto ithStatus = (*ctx.m_pRequestsStatus)[i]; - if (ithStatus == IOContext::Status::READ_SUCCESS) { - completedIndex = i; - return true; - } else if (ithStatus == IOContext::Status::READ_WAIT) { - waitsRemaining = true; - } - } +bool getNextCompletedRequest(const IOContext &ctx, size_t size, int &completedIndex) +{ + bool waitsRemaining = false; + long completeCount = ctx.m_completeCount; + do + { + for (int i = 0; i < size; i++) + { + auto ithStatus = (*ctx.m_pRequestsStatus)[i]; + if (ithStatus == IOContext::Status::READ_SUCCESS) + { + completedIndex = i; + return true; + } + else if (ithStatus == IOContext::Status::READ_WAIT) + { + waitsRemaining = true; + } + } - // if we didn't find one in READ_SUCCESS, wait for one to complete. - if (waitsRemaining) { - WaitOnAddress(&ctx.m_completeCount, &completeCount, sizeof(completeCount), - 100); - // this assumes the knowledge of the reader behavior (implicit - // contract). need better factoring? - } - } while (waitsRemaining); + // if we didn't find one in READ_SUCCESS, wait for one to complete. + if (waitsRemaining) + { + WaitOnAddress(&ctx.m_completeCount, &completeCount, sizeof(completeCount), 100); + // this assumes the knowledge of the reader behavior (implicit + // contract). need better factoring? + } + } while (waitsRemaining); - completedIndex = -1; - return false; + completedIndex = -1; + return false; } #endif template -void PQFlashIndex::cached_beam_search( - const T *query1, const uint64_t k_search, const uint64_t l_search, - uint64_t *indices, float *distances, const uint64_t beam_width, - const bool use_reorder_data, QueryStats *stats) { - cached_beam_search(query1, k_search, l_search, indices, distances, beam_width, - std::numeric_limits::max(), use_reorder_data, - stats); +void PQFlashIndex::cached_beam_search(const T *query1, const uint64_t k_search, const uint64_t l_search, + uint64_t *indices, float *distances, const uint64_t beam_width, + const bool use_reorder_data, QueryStats *stats) +{ + cached_beam_search(query1, k_search, l_search, indices, distances, beam_width, std::numeric_limits::max(), + use_reorder_data, stats); } template -void PQFlashIndex::cached_beam_search( - const T *query1, const uint64_t k_search, const uint64_t l_search, - uint64_t *indices, float *distances, const uint64_t beam_width, - const bool use_filter, const LabelT &filter_label, - const bool use_reorder_data, QueryStats *stats) { - cached_beam_search(query1, k_search, l_search, indices, distances, beam_width, - use_filter, filter_label, - std::numeric_limits::max(), use_reorder_data, - stats); +void PQFlashIndex::cached_beam_search(const T *query1, const uint64_t k_search, const uint64_t l_search, + uint64_t *indices, float *distances, const uint64_t beam_width, + const bool use_filter, const LabelT &filter_label, + const bool use_reorder_data, QueryStats *stats) +{ + cached_beam_search(query1, k_search, l_search, indices, distances, beam_width, use_filter, filter_label, + std::numeric_limits::max(), use_reorder_data, stats); } template -void PQFlashIndex::cached_beam_search( - const T *query1, const uint64_t k_search, const uint64_t l_search, - uint64_t *indices, float *distances, const uint64_t beam_width, - const uint32_t io_limit, const bool use_reorder_data, QueryStats *stats) { - LabelT dummy_filter = 0; - cached_beam_search(query1, k_search, l_search, indices, distances, beam_width, - false, dummy_filter, std::numeric_limits::max(), - use_reorder_data, stats); +void PQFlashIndex::cached_beam_search(const T *query1, const uint64_t k_search, const uint64_t l_search, + uint64_t *indices, float *distances, const uint64_t beam_width, + const uint32_t io_limit, const bool use_reorder_data, + QueryStats *stats) +{ + LabelT dummy_filter = 0; + cached_beam_search(query1, k_search, l_search, indices, distances, beam_width, false, dummy_filter, + std::numeric_limits::max(), use_reorder_data, stats); } template -void PQFlashIndex::cached_beam_search( - const T *query1, const uint64_t k_search, const uint64_t l_search, - uint64_t *indices, float *distances, const uint64_t beam_width, - const bool use_filter, const LabelT &filter_label, const uint32_t io_limit, - const bool use_reorder_data, QueryStats *stats) { - int32_t filter_num = 0; - if (use_filter) { - filter_num = get_filter_number(filter_label); - if (filter_num < 0) { - if (!_use_universal_label) { - return; - } else { - filter_num = _universal_filter_num; - } - } - } - - if (beam_width > MAX_N_SECTOR_READS) - throw ANNException("Beamwidth can not be higher than MAX_N_SECTOR_READS", - -1, __FUNCSIG__, __FILE__, __LINE__); - - ScratchStoreManager> manager(this->thread_data); - auto data = manager.scratch_space(); - IOContext &ctx = data->ctx; - auto query_scratch = &(data->scratch); - auto pq_query_scratch = query_scratch->_pq_scratch; - - // reset query scratch - query_scratch->reset(); - - // copy query to thread specific aligned and allocated memory (for distance - // calculations we need aligned data) - float query_norm = 0; - T *aligned_query_T = query_scratch->aligned_query_T; - float *query_float = pq_query_scratch->aligned_query_float; - float *query_rotated = pq_query_scratch->rotated_query; - - // if inner product, we laso normalize the query and set the last coordinate - // to 0 (this is the extra coordindate used to convert MIPS to L2 search) - if (metric == diskann::Metric::INNER_PRODUCT) { - for (size_t i = 0; i < this->data_dim - 1; i++) { - aligned_query_T[i] = query1[i]; - query_norm += query1[i] * query1[i]; +void PQFlashIndex::cached_beam_search(const T *query1, const uint64_t k_search, const uint64_t l_search, + uint64_t *indices, float *distances, const uint64_t beam_width, + const bool use_filter, const LabelT &filter_label, + const uint32_t io_limit, const bool use_reorder_data, + QueryStats *stats) +{ + int32_t filter_num = 0; + if (use_filter) + { + filter_num = get_filter_number(filter_label); + if (filter_num < 0) + { + if (!_use_universal_label) + { + return; + } + else + { + filter_num = _universal_filter_num; + } + } } - aligned_query_T[this->data_dim - 1] = 0; - query_norm = std::sqrt(query_norm); + if (beam_width > MAX_N_SECTOR_READS) + throw ANNException("Beamwidth can not be higher than MAX_N_SECTOR_READS", -1, __FUNCSIG__, __FILE__, __LINE__); - for (size_t i = 0; i < this->data_dim - 1; i++) { - aligned_query_T[i] /= query_norm; - } - pq_query_scratch->set(this->data_dim, aligned_query_T); - } else { - for (size_t i = 0; i < this->data_dim; i++) { - aligned_query_T[i] = query1[i]; + ScratchStoreManager> manager(this->thread_data); + auto data = manager.scratch_space(); + IOContext &ctx = data->ctx; + auto query_scratch = &(data->scratch); + auto pq_query_scratch = query_scratch->_pq_scratch; + + // reset query scratch + query_scratch->reset(); + + // copy query to thread specific aligned and allocated memory (for distance + // calculations we need aligned data) + float query_norm = 0; + T *aligned_query_T = query_scratch->aligned_query_T; + float *query_float = pq_query_scratch->aligned_query_float; + float *query_rotated = pq_query_scratch->rotated_query; + + // if inner product, we laso normalize the query and set the last coordinate + // to 0 (this is the extra coordindate used to convert MIPS to L2 search) + if (metric == diskann::Metric::INNER_PRODUCT) + { + for (size_t i = 0; i < this->data_dim - 1; i++) + { + aligned_query_T[i] = query1[i]; + query_norm += query1[i] * query1[i]; + } + aligned_query_T[this->data_dim - 1] = 0; + + query_norm = std::sqrt(query_norm); + + for (size_t i = 0; i < this->data_dim - 1; i++) + { + aligned_query_T[i] /= query_norm; + } + pq_query_scratch->set(this->data_dim, aligned_query_T); } - pq_query_scratch->set(this->data_dim, aligned_query_T); - } - - // pointers to buffers for data - T *data_buf = query_scratch->coord_scratch; - uint64_t &data_buf_idx = query_scratch->coord_idx; - _mm_prefetch((char *)data_buf, _MM_HINT_T1); - - // sector scratch - char *sector_scratch = query_scratch->sector_scratch; - uint64_t §or_scratch_idx = query_scratch->sector_idx; - - // query <-> PQ chunk centers distances - pq_table.preprocess_query(query_rotated); // center the query and rotate if - // we have a rotation matrix - float *pq_dists = pq_query_scratch->aligned_pqtable_dist_scratch; - pq_table.populate_chunk_distances(query_rotated, pq_dists); - - // query <-> neighbor list - float *dist_scratch = pq_query_scratch->aligned_dist_scratch; - uint8_t *pq_coord_scratch = pq_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 uint32_t *ids, - const uint64_t n_ids, - float *dists_out) { - diskann::aggregate_coords(ids, n_ids, this->data, this->n_chunks, - pq_coord_scratch); - diskann::pq_dist_lookup(pq_coord_scratch, n_ids, this->n_chunks, pq_dists, - dists_out); - }; - Timer query_timer, io_timer, cpu_timer; - - tsl::robin_set &visited = query_scratch->visited; - NeighborPriorityQueue &retset = query_scratch->retset; - retset.reserve(l_search); - std::vector &full_retset = query_scratch->full_retset; - - uint32_t best_medoid = 0; - float best_dist = (std::numeric_limits::max)(); - if (!use_filter) { - for (uint64_t cur_m = 0; cur_m < num_medoids; cur_m++) { - float cur_expanded_dist = dist_cmp_float->compare( - query_float, centroid_data + aligned_dim * cur_m, - (uint32_t)aligned_dim); - if (cur_expanded_dist < best_dist) { - best_medoid = medoids[cur_m]; - best_dist = cur_expanded_dist; - } + else + { + for (size_t i = 0; i < this->data_dim; i++) + { + aligned_query_T[i] = query1[i]; + } + pq_query_scratch->set(this->data_dim, aligned_query_T); } - } else if (_filter_to_medoid_id.find(filter_label) != - _filter_to_medoid_id.end()) { - best_medoid = _filter_to_medoid_id[filter_label]; - } else { - throw ANNException("Cannot find medoid for specified filter.", -1, - __FUNCSIG__, __FILE__, __LINE__); - } - - compute_dists(&best_medoid, 1, dist_scratch); - retset.insert(Neighbor(best_medoid, dist_scratch[0])); - visited.insert(best_medoid); - - uint32_t cmps = 0; - uint32_t hops = 0; - uint32_t num_ios = 0; - - // cleared every iteration - std::vector frontier; - frontier.reserve(2 * beam_width); - std::vector> frontier_nhoods; - frontier_nhoods.reserve(2 * beam_width); - std::vector frontier_read_reqs; - frontier_read_reqs.reserve(2 * beam_width); - std::vector>> - cached_nhoods; - cached_nhoods.reserve(2 * beam_width); - - while (retset.has_unexpanded_node() && num_ios < io_limit) { - // clear iteration state - frontier.clear(); - frontier_nhoods.clear(); - frontier_read_reqs.clear(); - cached_nhoods.clear(); - sector_scratch_idx = 0; - // find new beam - uint32_t num_seen = 0; - while (retset.has_unexpanded_node() && frontier.size() < beam_width && - num_seen < beam_width) { - auto nbr = retset.closest_unexpanded(); - num_seen++; - auto iter = nhood_cache.find(nbr.id); - if (iter != nhood_cache.end()) { - cached_nhoods.push_back(std::make_pair(nbr.id, iter->second)); - if (stats != nullptr) { - stats->n_cache_hits++; + + // pointers to buffers for data + T *data_buf = query_scratch->coord_scratch; + uint64_t &data_buf_idx = query_scratch->coord_idx; + _mm_prefetch((char *)data_buf, _MM_HINT_T1); + + // sector scratch + char *sector_scratch = query_scratch->sector_scratch; + uint64_t §or_scratch_idx = query_scratch->sector_idx; + + // query <-> PQ chunk centers distances + pq_table.preprocess_query(query_rotated); // center the query and rotate if + // we have a rotation matrix + float *pq_dists = pq_query_scratch->aligned_pqtable_dist_scratch; + pq_table.populate_chunk_distances(query_rotated, pq_dists); + + // query <-> neighbor list + float *dist_scratch = pq_query_scratch->aligned_dist_scratch; + uint8_t *pq_coord_scratch = pq_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 uint32_t *ids, const uint64_t n_ids, + float *dists_out) { + diskann::aggregate_coords(ids, n_ids, this->data, this->n_chunks, pq_coord_scratch); + diskann::pq_dist_lookup(pq_coord_scratch, n_ids, this->n_chunks, pq_dists, dists_out); + }; + Timer query_timer, io_timer, cpu_timer; + + tsl::robin_set &visited = query_scratch->visited; + NeighborPriorityQueue &retset = query_scratch->retset; + retset.reserve(l_search); + std::vector &full_retset = query_scratch->full_retset; + + uint32_t best_medoid = 0; + float best_dist = (std::numeric_limits::max)(); + if (!use_filter) + { + for (uint64_t cur_m = 0; cur_m < num_medoids; cur_m++) + { + float cur_expanded_dist = + dist_cmp_float->compare(query_float, centroid_data + aligned_dim * cur_m, (uint32_t)aligned_dim); + if (cur_expanded_dist < best_dist) + { + best_medoid = medoids[cur_m]; + best_dist = cur_expanded_dist; + } } - } else { - frontier.push_back(nbr.id); - } - if (this->count_visited_nodes) { - reinterpret_cast &>( - this->node_visit_counter[nbr.id].second) - .fetch_add(1); - } + } + else if (_filter_to_medoid_id.find(filter_label) != _filter_to_medoid_id.end()) + { + best_medoid = _filter_to_medoid_id[filter_label]; + } + else + { + throw ANNException("Cannot find medoid for specified filter.", -1, __FUNCSIG__, __FILE__, __LINE__); } - // read nhoods of frontier ids - if (!frontier.empty()) { - if (stats != nullptr) stats->n_hops++; - for (uint64_t i = 0; i < frontier.size(); i++) { - auto id = frontier[i]; - std::pair fnhood; - fnhood.first = id; - fnhood.second = sector_scratch + sector_scratch_idx * SECTOR_LEN; - sector_scratch_idx++; - frontier_nhoods.push_back(fnhood); - frontier_read_reqs.emplace_back( - NODE_SECTOR_NO(((size_t)id)) * SECTOR_LEN, SECTOR_LEN, - fnhood.second); - if (stats != nullptr) { - stats->n_4k++; - stats->n_ios++; + compute_dists(&best_medoid, 1, dist_scratch); + retset.insert(Neighbor(best_medoid, dist_scratch[0])); + visited.insert(best_medoid); + + uint32_t cmps = 0; + uint32_t hops = 0; + uint32_t num_ios = 0; + + // cleared every iteration + std::vector frontier; + frontier.reserve(2 * beam_width); + std::vector> frontier_nhoods; + frontier_nhoods.reserve(2 * beam_width); + std::vector frontier_read_reqs; + frontier_read_reqs.reserve(2 * beam_width); + std::vector>> cached_nhoods; + cached_nhoods.reserve(2 * beam_width); + + while (retset.has_unexpanded_node() && num_ios < io_limit) + { + // clear iteration state + frontier.clear(); + frontier_nhoods.clear(); + frontier_read_reqs.clear(); + cached_nhoods.clear(); + sector_scratch_idx = 0; + // find new beam + uint32_t num_seen = 0; + while (retset.has_unexpanded_node() && frontier.size() < beam_width && num_seen < beam_width) + { + auto nbr = retset.closest_unexpanded(); + num_seen++; + auto iter = nhood_cache.find(nbr.id); + if (iter != nhood_cache.end()) + { + cached_nhoods.push_back(std::make_pair(nbr.id, iter->second)); + if (stats != nullptr) + { + stats->n_cache_hits++; + } + } + else + { + frontier.push_back(nbr.id); + } + if (this->count_visited_nodes) + { + reinterpret_cast &>(this->node_visit_counter[nbr.id].second).fetch_add(1); + } } - num_ios++; - } - io_timer.reset(); + + // read nhoods of frontier ids + if (!frontier.empty()) + { + if (stats != nullptr) + stats->n_hops++; + for (uint64_t i = 0; i < frontier.size(); i++) + { + auto id = frontier[i]; + std::pair fnhood; + fnhood.first = id; + fnhood.second = sector_scratch + sector_scratch_idx * SECTOR_LEN; + sector_scratch_idx++; + frontier_nhoods.push_back(fnhood); + frontier_read_reqs.emplace_back(NODE_SECTOR_NO(((size_t)id)) * SECTOR_LEN, SECTOR_LEN, fnhood.second); + if (stats != nullptr) + { + stats->n_4k++; + stats->n_ios++; + } + num_ios++; + } + io_timer.reset(); #ifdef USE_BING_INFRA - reader->read(frontier_read_reqs, ctx, - true); // async reader windows. + 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); // synchronous IO linux #endif - if (stats != nullptr) { - stats->io_us += (double)io_timer.elapsed(); - } - } + if (stats != nullptr) + { + 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; - float cur_expanded_dist; - if (!use_disk_index_pq) { - cur_expanded_dist = dist_cmp->compare( - aligned_query_T, node_fp_coords_copy, (uint32_t)aligned_dim); - } else { - if (metric == diskann::Metric::INNER_PRODUCT) - cur_expanded_dist = disk_pq_table.inner_product( - query_float, (uint8_t *)node_fp_coords_copy); - else - cur_expanded_dist = - disk_pq_table.l2_distance( // disk_pq does not support OPQ yet - query_float, (uint8_t *)node_fp_coords_copy); - } - full_retset.push_back( - Neighbor((uint32_t)cached_nhood.first, cur_expanded_dist)); - - uint64_t nnbrs = cached_nhood.second.first; - uint32_t *node_nbrs = cached_nhood.second.second; - - // compute node_nbrs <-> query dists in PQ space - cpu_timer.reset(); - compute_dists(node_nbrs, nnbrs, dist_scratch); - if (stats != nullptr) { - stats->n_cmps += (double)nnbrs; - stats->cpu_us += (double)cpu_timer.elapsed(); - } - - // process prefetched nhood - for (uint64_t m = 0; m < nnbrs; ++m) { - uint32_t id = node_nbrs[m]; - if (visited.insert(id).second) { - if (!use_filter && _dummy_pts.find(id) != _dummy_pts.end()) continue; - - if (use_filter && !point_has_label(id, filter_num) && - !point_has_label(id, _universal_filter_num)) - continue; - cmps++; - float dist = dist_scratch[m]; - Neighbor nn(id, dist); - retset.insert(nn); + // 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; + float cur_expanded_dist; + if (!use_disk_index_pq) + { + cur_expanded_dist = dist_cmp->compare(aligned_query_T, node_fp_coords_copy, (uint32_t)aligned_dim); + } + else + { + if (metric == diskann::Metric::INNER_PRODUCT) + cur_expanded_dist = disk_pq_table.inner_product(query_float, (uint8_t *)node_fp_coords_copy); + else + cur_expanded_dist = disk_pq_table.l2_distance( // disk_pq does not support OPQ yet + query_float, (uint8_t *)node_fp_coords_copy); + } + full_retset.push_back(Neighbor((uint32_t)cached_nhood.first, cur_expanded_dist)); + + uint64_t nnbrs = cached_nhood.second.first; + uint32_t *node_nbrs = cached_nhood.second.second; + + // compute node_nbrs <-> query dists in PQ space + cpu_timer.reset(); + compute_dists(node_nbrs, nnbrs, dist_scratch); + if (stats != nullptr) + { + stats->n_cmps += (double)nnbrs; + stats->cpu_us += (double)cpu_timer.elapsed(); + } + + // process prefetched nhood + for (uint64_t m = 0; m < nnbrs; ++m) + { + uint32_t id = node_nbrs[m]; + if (visited.insert(id).second) + { + if (!use_filter && _dummy_pts.find(id) != _dummy_pts.end()) + continue; + + if (use_filter && !point_has_label(id, filter_num) && !point_has_label(id, _universal_filter_num)) + continue; + cmps++; + float dist = dist_scratch[m]; + Neighbor nn(id, dist); + retset.insert(nn); + } + } } - } - } #ifdef USE_BING_INFRA - // process each frontier nhood - compute distances to unvisited nodes - int completedIndex = -1; - long requestCount = static_cast(frontier_read_reqs.size()); - // If we issued read requests and if a read is complete or there are - // reads in wait state, then enter the while loop. - while (requestCount > 0 && - getNextCompletedRequest(ctx, requestCount, completedIndex)) { - assert(completedIndex >= 0); - auto &frontier_nhood = frontier_nhoods[completedIndex]; - (*ctx.m_pRequestsStatus)[completedIndex] = IOContext::PROCESS_COMPLETE; + // process each frontier nhood - compute distances to unvisited nodes + int completedIndex = -1; + long requestCount = static_cast(frontier_read_reqs.size()); + // If we issued read requests and if a read is complete or there are + // reads in wait state, then enter the while loop. + while (requestCount > 0 && getNextCompletedRequest(ctx, requestCount, completedIndex)) + { + assert(completedIndex >= 0); + auto &frontier_nhood = frontier_nhoods[completedIndex]; + (*ctx.m_pRequestsStatus)[completedIndex] = IOContext::PROCESS_COMPLETE; #else - for (auto &frontier_nhood : frontier_nhoods) { + for (auto &frontier_nhood : frontier_nhoods) + { #endif - char *node_disk_buf = - OFFSET_TO_NODE(frontier_nhood.second, frontier_nhood.first); - uint32_t *node_buf = OFFSET_TO_NODE_NHOOD(node_disk_buf); - uint64_t nnbrs = (uint64_t)(*node_buf); - T *node_fp_coords = OFFSET_TO_NODE_COORDS(node_disk_buf); - // assert(data_buf_idx < MAX_N_CMPS); - if (data_buf_idx == MAX_N_CMPS) data_buf_idx = 0; - - T *node_fp_coords_copy = data_buf + (data_buf_idx * aligned_dim); - data_buf_idx++; - memcpy(node_fp_coords_copy, node_fp_coords, disk_bytes_per_point); - float cur_expanded_dist; - if (!use_disk_index_pq) { - cur_expanded_dist = dist_cmp->compare( - aligned_query_T, node_fp_coords_copy, (uint32_t)aligned_dim); - } else { - if (metric == diskann::Metric::INNER_PRODUCT) - cur_expanded_dist = disk_pq_table.inner_product( - query_float, (uint8_t *)node_fp_coords_copy); - else - cur_expanded_dist = disk_pq_table.l2_distance( - query_float, (uint8_t *)node_fp_coords_copy); - } - full_retset.push_back(Neighbor(frontier_nhood.first, cur_expanded_dist)); - uint32_t *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 += (double)nnbrs; - stats->cpu_us += (double)cpu_timer.elapsed(); - } - - cpu_timer.reset(); - // process prefetch-ed nhood - for (uint64_t m = 0; m < nnbrs; ++m) { - uint32_t id = node_nbrs[m]; - if (visited.insert(id).second) { - if (!use_filter && _dummy_pts.find(id) != _dummy_pts.end()) continue; - - if (use_filter && !point_has_label(id, filter_num) && - !point_has_label(id, _universal_filter_num)) - continue; - cmps++; - float dist = dist_scratch[m]; - if (stats != nullptr) { - stats->n_cmps++; - } - - Neighbor nn(id, dist); - retset.insert(nn); + char *node_disk_buf = OFFSET_TO_NODE(frontier_nhood.second, frontier_nhood.first); + uint32_t *node_buf = OFFSET_TO_NODE_NHOOD(node_disk_buf); + uint64_t nnbrs = (uint64_t)(*node_buf); + T *node_fp_coords = OFFSET_TO_NODE_COORDS(node_disk_buf); + // assert(data_buf_idx < MAX_N_CMPS); + if (data_buf_idx == MAX_N_CMPS) + data_buf_idx = 0; + + T *node_fp_coords_copy = data_buf + (data_buf_idx * aligned_dim); + data_buf_idx++; + memcpy(node_fp_coords_copy, node_fp_coords, disk_bytes_per_point); + float cur_expanded_dist; + if (!use_disk_index_pq) + { + cur_expanded_dist = dist_cmp->compare(aligned_query_T, node_fp_coords_copy, (uint32_t)aligned_dim); + } + else + { + if (metric == diskann::Metric::INNER_PRODUCT) + cur_expanded_dist = disk_pq_table.inner_product(query_float, (uint8_t *)node_fp_coords_copy); + else + cur_expanded_dist = disk_pq_table.l2_distance(query_float, (uint8_t *)node_fp_coords_copy); + } + full_retset.push_back(Neighbor(frontier_nhood.first, cur_expanded_dist)); + uint32_t *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 += (double)nnbrs; + stats->cpu_us += (double)cpu_timer.elapsed(); + } + + cpu_timer.reset(); + // process prefetch-ed nhood + for (uint64_t m = 0; m < nnbrs; ++m) + { + uint32_t id = node_nbrs[m]; + if (visited.insert(id).second) + { + if (!use_filter && _dummy_pts.find(id) != _dummy_pts.end()) + continue; + + if (use_filter && !point_has_label(id, filter_num) && !point_has_label(id, _universal_filter_num)) + continue; + cmps++; + float dist = dist_scratch[m]; + if (stats != nullptr) + { + stats->n_cmps++; + } + + Neighbor nn(id, dist); + retset.insert(nn); + } + } + + if (stats != nullptr) + { + stats->cpu_us += (double)cpu_timer.elapsed(); + } } - } - if (stats != nullptr) { - stats->cpu_us += (double)cpu_timer.elapsed(); - } + hops++; } - hops++; - } - - // re-sort by distance - std::sort(full_retset.begin(), full_retset.end()); + // re-sort by distance + std::sort(full_retset.begin(), full_retset.end()); - if (use_reorder_data) { - if (!(this->reorder_data_exists)) { - throw ANNException( - "Requested use of reordering data which does " - "not exist in index " - "file", - -1, __FUNCSIG__, __FILE__, __LINE__); - } + if (use_reorder_data) + { + if (!(this->reorder_data_exists)) + { + throw ANNException("Requested use of reordering data which does " + "not exist in index " + "file", + -1, __FUNCSIG__, __FILE__, __LINE__); + } - std::vector vec_read_reqs; + std::vector vec_read_reqs; - if (full_retset.size() > k_search * FULL_PRECISION_REORDER_MULTIPLIER) - full_retset.erase( - full_retset.begin() + k_search * FULL_PRECISION_REORDER_MULTIPLIER, - full_retset.end()); + if (full_retset.size() > k_search * FULL_PRECISION_REORDER_MULTIPLIER) + full_retset.erase(full_retset.begin() + k_search * FULL_PRECISION_REORDER_MULTIPLIER, full_retset.end()); - for (size_t i = 0; i < full_retset.size(); ++i) { - vec_read_reqs.emplace_back( - VECTOR_SECTOR_NO(((size_t)full_retset[i].id)) * SECTOR_LEN, - SECTOR_LEN, sector_scratch + i * SECTOR_LEN); + for (size_t i = 0; i < full_retset.size(); ++i) + { + vec_read_reqs.emplace_back(VECTOR_SECTOR_NO(((size_t)full_retset[i].id)) * SECTOR_LEN, SECTOR_LEN, + sector_scratch + i * SECTOR_LEN); - if (stats != nullptr) { - stats->n_4k++; - stats->n_ios++; - } - } + if (stats != nullptr) + { + stats->n_4k++; + stats->n_ios++; + } + } - io_timer.reset(); + io_timer.reset(); #ifdef USE_BING_INFRA - reader->read(vec_read_reqs, ctx, false); // sync reader windows. + reader->read(vec_read_reqs, ctx, false); // sync reader windows. #else - reader->read(vec_read_reqs, ctx); // synchronous IO linux + reader->read(vec_read_reqs, ctx); // synchronous IO linux #endif - if (stats != nullptr) { - stats->io_us += io_timer.elapsed(); - } + if (stats != nullptr) + { + stats->io_us += io_timer.elapsed(); + } - for (size_t i = 0; i < full_retset.size(); ++i) { - auto id = full_retset[i].id; - auto location = - (sector_scratch + i * SECTOR_LEN) + VECTOR_SECTOR_OFFSET(id); - full_retset[i].distance = - dist_cmp->compare(aligned_query_T, (T *)location, this->data_dim); - } + for (size_t i = 0; i < full_retset.size(); ++i) + { + auto id = full_retset[i].id; + auto location = (sector_scratch + i * SECTOR_LEN) + VECTOR_SECTOR_OFFSET(id); + full_retset[i].distance = dist_cmp->compare(aligned_query_T, (T *)location, this->data_dim); + } - std::sort(full_retset.begin(), full_retset.end()); - } + std::sort(full_retset.begin(), full_retset.end()); + } - // copy k_search values - for (uint64_t i = 0; i < k_search; i++) { - indices[i] = full_retset[i].id; + // copy k_search values + for (uint64_t i = 0; i < k_search; i++) + { + indices[i] = full_retset[i].id; - if (_dummy_pts.find(indices[i]) != _dummy_pts.end()) { - indices[i] = _dummy_to_real_map[indices[i]]; - } + if (_dummy_pts.find(indices[i]) != _dummy_pts.end()) + { + indices[i] = _dummy_to_real_map[indices[i]]; + } - if (distances != nullptr) { - distances[i] = full_retset[i].distance; - if (metric == diskann::Metric::INNER_PRODUCT) { - // flip the sign to convert min to max - distances[i] = (-distances[i]); - // rescale to revert back to original norms (cancelling the - // effect of base and query pre-processing) - if (max_base_norm != 0) distances[i] *= (max_base_norm * query_norm); - } + if (distances != nullptr) + { + distances[i] = full_retset[i].distance; + if (metric == diskann::Metric::INNER_PRODUCT) + { + // flip the sign to convert min to max + distances[i] = (-distances[i]); + // rescale to revert back to original norms (cancelling the + // effect of base and query pre-processing) + if (max_base_norm != 0) + distances[i] *= (max_base_norm * query_norm); + } + } } - } #ifdef USE_BING_INFRA - ctx.m_completeCount = 0; + ctx.m_completeCount = 0; #endif - if (stats != nullptr) { - stats->total_us = (double)query_timer.elapsed(); - } + if (stats != nullptr) + { + stats->total_us = (double)query_timer.elapsed(); + } } // range search returns results of all neighbors within distance of range. // indices and distances need to be pre-allocated of size l_search and the // return value is the number of matching hits. template -uint32_t PQFlashIndex::range_search( - const T *query1, const double range, const uint64_t min_l_search, - const uint64_t max_l_search, std::vector &indices, - std::vector &distances, const uint64_t min_beam_width, - QueryStats *stats) { - uint32_t res_count = 0; - - bool stop_flag = false; - - uint32_t l_search = min_l_search; // starting size of the candidate list - while (!stop_flag) { - indices.resize(l_search); - distances.resize(l_search); - uint64_t cur_bw = - min_beam_width > (l_search / 5) ? min_beam_width : l_search / 5; - cur_bw = (cur_bw > 100) ? 100 : cur_bw; - for (auto &x : distances) x = std::numeric_limits::max(); - this->cached_beam_search(query1, l_search, l_search, indices.data(), - distances.data(), cur_bw, false, stats); - for (uint32_t i = 0; i < l_search; i++) { - if (distances[i] > (float)range) { - res_count = i; - break; - } else if (i == l_search - 1) - res_count = l_search; +uint32_t PQFlashIndex::range_search(const T *query1, const double range, const uint64_t min_l_search, + const uint64_t max_l_search, std::vector &indices, + std::vector &distances, const uint64_t min_beam_width, + QueryStats *stats) +{ + uint32_t res_count = 0; + + bool stop_flag = false; + + uint32_t l_search = min_l_search; // starting size of the candidate list + while (!stop_flag) + { + indices.resize(l_search); + distances.resize(l_search); + uint64_t cur_bw = min_beam_width > (l_search / 5) ? min_beam_width : l_search / 5; + cur_bw = (cur_bw > 100) ? 100 : cur_bw; + for (auto &x : distances) + x = std::numeric_limits::max(); + this->cached_beam_search(query1, l_search, l_search, indices.data(), distances.data(), cur_bw, false, stats); + for (uint32_t i = 0; i < l_search; i++) + { + if (distances[i] > (float)range) + { + res_count = i; + break; + } + else if (i == l_search - 1) + res_count = l_search; + } + if (res_count < (uint32_t)(l_search / 2.0)) + stop_flag = true; + l_search = l_search * 2; + if (l_search > max_l_search) + stop_flag = true; } - if (res_count < (uint32_t)(l_search / 2.0)) stop_flag = true; - l_search = l_search * 2; - if (l_search > max_l_search) stop_flag = true; - } - indices.resize(res_count); - distances.resize(res_count); - return res_count; + indices.resize(res_count); + distances.resize(res_count); + return res_count; } -template -uint64_t PQFlashIndex::get_data_dim() { - return data_dim; +template uint64_t PQFlashIndex::get_data_dim() +{ + return data_dim; } -template -diskann::Metric PQFlashIndex::get_metric() { - return this->metric; +template diskann::Metric PQFlashIndex::get_metric() +{ + return this->metric; } #ifdef EXEC_ENV_OLS -template -char *PQFlashIndex::getHeaderBytes() { - IOContext &ctx = reader->get_ctx(); - AlignedRead readReq; - readReq.buf = new char[PQFlashIndex::HEADER_SIZE]; - readReq.len = PQFlashIndex::HEADER_SIZE; - readReq.offset = 0; +template char *PQFlashIndex::getHeaderBytes() +{ + IOContext &ctx = reader->get_ctx(); + AlignedRead readReq; + readReq.buf = new char[PQFlashIndex::HEADER_SIZE]; + readReq.len = PQFlashIndex::HEADER_SIZE; + readReq.offset = 0; - std::vector readReqs; - readReqs.push_back(readReq); + std::vector readReqs; + readReqs.push_back(readReq); - reader->read(readReqs, ctx, false); + reader->read(readReqs, ctx, false); - return (char *)readReq.buf; + return (char *)readReq.buf; } #endif @@ -1483,4 +1562,4 @@ template class PQFlashIndex; template class PQFlashIndex; template class PQFlashIndex; -} // namespace diskann +} // namespace diskann diff --git a/src/restapi/search_wrapper.cpp b/src/restapi/search_wrapper.cpp index 5fe049bf6..dc9f5734e 100644 --- a/src/restapi/search_wrapper.cpp +++ b/src/restapi/search_wrapper.cpp @@ -21,182 +21,189 @@ #endif #endif -namespace diskann { +namespace diskann +{ const unsigned int DEFAULT_W = 1; -SearchResult::SearchResult(unsigned int K, unsigned int elapsed_time_in_ms, - const unsigned *const indices, - const float *const distances, - const std::string *const tags, +SearchResult::SearchResult(unsigned int K, unsigned int elapsed_time_in_ms, const unsigned *const indices, + const float *const distances, const std::string *const tags, const unsigned *const partitions) - : _K(K), _search_time_in_ms(elapsed_time_in_ms) { - for (unsigned i = 0; i < K; ++i) { - this->_indices.push_back(indices[i]); - this->_distances.push_back(distances[i]); - if (tags != NULL) this->_tags.push_back(tags[i]); - if (partitions != NULL) this->_partitions.push_back(partitions[i]); - } - if (tags != nullptr) - this->_tags_enabled = true; - else - this->_tags_enabled = false; - - if (partitions != nullptr) - this->_partitions_enabled = true; - else - this->_partitions_enabled = false; + : _K(K), _search_time_in_ms(elapsed_time_in_ms) +{ + for (unsigned i = 0; i < K; ++i) + { + this->_indices.push_back(indices[i]); + this->_distances.push_back(distances[i]); + if (tags != NULL) + this->_tags.push_back(tags[i]); + if (partitions != NULL) + this->_partitions.push_back(partitions[i]); + } + if (tags != nullptr) + this->_tags_enabled = true; + else + this->_tags_enabled = false; + + if (partitions != nullptr) + this->_partitions_enabled = true; + else + this->_partitions_enabled = false; } -BaseSearch::BaseSearch(const std::string &tagsFile) { - if (tagsFile.size() != 0) { - std::ifstream in(tagsFile); +BaseSearch::BaseSearch(const std::string &tagsFile) +{ + if (tagsFile.size() != 0) + { + std::ifstream in(tagsFile); - if (!in.is_open()) { - std::cerr << "Could not open " << tagsFile << std::endl; - } + if (!in.is_open()) + { + std::cerr << "Could not open " << tagsFile << std::endl; + } - std::string tag; - while (std::getline(in, tag)) { - _tags_str.push_back(tag); - } + std::string tag; + while (std::getline(in, tag)) + { + _tags_str.push_back(tag); + } - _tags_enabled = true; + _tags_enabled = true; - std::cout << "Loaded " << _tags_str.size() << " tags from " << tagsFile - << std::endl; - } else { - _tags_enabled = false; - } + std::cout << "Loaded " << _tags_str.size() << " tags from " << tagsFile << std::endl; + } + else + { + _tags_enabled = false; + } } -void BaseSearch::lookup_tags(const unsigned K, const unsigned *indices, - std::string *ret_tags) { - if (_tags_enabled == false) - throw std::runtime_error("Can not look up tags as they are not enabled."); - else { - for (unsigned k = 0; k < K; ++k) { - if (indices[k] > _tags_str.size()) - throw std::runtime_error( - "In tag lookup, index exceeded the number of tags"); - else - ret_tags[k] = _tags_str[indices[k]]; +void BaseSearch::lookup_tags(const unsigned K, const unsigned *indices, std::string *ret_tags) +{ + if (_tags_enabled == false) + throw std::runtime_error("Can not look up tags as they are not enabled."); + else + { + for (unsigned k = 0; k < K; ++k) + { + if (indices[k] > _tags_str.size()) + throw std::runtime_error("In tag lookup, index exceeded the number of tags"); + else + ret_tags[k] = _tags_str[indices[k]]; + } } - } } template -InMemorySearch::InMemorySearch(const std::string &baseFile, - const std::string &indexFile, - const std::string &tagsFile, Metric m, - uint32_t num_threads, uint32_t search_l) - : BaseSearch(tagsFile) { - size_t dimensions, total_points = 0; - diskann::get_bin_metadata(baseFile, total_points, dimensions); - _index = std::unique_ptr>( - new diskann::Index(m, dimensions, total_points, false)); - - _index->load(indexFile.c_str(), num_threads, search_l); +InMemorySearch::InMemorySearch(const std::string &baseFile, const std::string &indexFile, + const std::string &tagsFile, Metric m, uint32_t num_threads, uint32_t search_l) + : BaseSearch(tagsFile) +{ + size_t dimensions, total_points = 0; + diskann::get_bin_metadata(baseFile, total_points, dimensions); + _index = std::unique_ptr>(new diskann::Index(m, dimensions, total_points, false)); + + _index->load(indexFile.c_str(), num_threads, search_l); } template -SearchResult InMemorySearch::search(const T *query, - const unsigned int dimensions, - const unsigned int K, - const unsigned int Ls) { - unsigned int *indices = new unsigned int[K]; - float *distances = new float[K]; - - auto startTime = std::chrono::high_resolution_clock::now(); - _index->search(query, K, Ls, indices, distances); - auto duration = std::chrono::duration_cast( - std::chrono::high_resolution_clock::now() - startTime) - .count(); - - std::string *tags = nullptr; - if (_tags_enabled) { - tags = new std::string[K]; - lookup_tags(K, indices, tags); - } - - SearchResult result(K, (unsigned int)duration, indices, distances, tags); - - delete[] indices; - delete[] distances; - return result; +SearchResult InMemorySearch::search(const T *query, const unsigned int dimensions, const unsigned int K, + const unsigned int Ls) +{ + unsigned int *indices = new unsigned int[K]; + float *distances = new float[K]; + + auto startTime = std::chrono::high_resolution_clock::now(); + _index->search(query, K, Ls, indices, distances); + auto duration = + std::chrono::duration_cast(std::chrono::high_resolution_clock::now() - startTime) + .count(); + + std::string *tags = nullptr; + if (_tags_enabled) + { + tags = new std::string[K]; + lookup_tags(K, indices, tags); + } + + SearchResult result(K, (unsigned int)duration, indices, distances, tags); + + delete[] indices; + delete[] distances; + return result; } -template -InMemorySearch::~InMemorySearch() {} +template InMemorySearch::~InMemorySearch() +{ +} template -PQFlashSearch::PQFlashSearch(const std::string &indexPrefix, - const unsigned num_nodes_to_cache, - const unsigned num_threads, - const std::string &tagsFile, Metric m) - : BaseSearch(tagsFile) { +PQFlashSearch::PQFlashSearch(const std::string &indexPrefix, const unsigned num_nodes_to_cache, + const unsigned num_threads, const std::string &tagsFile, Metric m) + : BaseSearch(tagsFile) +{ #ifdef _WINDOWS #ifndef USE_BING_INFRA - reader.reset(new WindowsAlignedFileReader()); + reader.reset(new WindowsAlignedFileReader()); #else - reader.reset(new diskann::BingAlignedFileReader()); + reader.reset(new diskann::BingAlignedFileReader()); #endif #else - auto ptr = new LinuxAlignedFileReader(); - reader.reset(ptr); + auto ptr = new LinuxAlignedFileReader(); + reader.reset(ptr); #endif - std::string index_prefix_path(indexPrefix); - std::string disk_index_file = index_prefix_path + "_disk.index"; - std::string warmup_query_file = index_prefix_path + "_sample_data.bin"; + std::string index_prefix_path(indexPrefix); + std::string disk_index_file = index_prefix_path + "_disk.index"; + std::string warmup_query_file = index_prefix_path + "_sample_data.bin"; - _index = std::unique_ptr>( - new diskann::PQFlashIndex(reader, m)); + _index = std::unique_ptr>(new diskann::PQFlashIndex(reader, m)); - int res = _index->load(num_threads, index_prefix_path.c_str()); + int res = _index->load(num_threads, index_prefix_path.c_str()); - if (res != 0) { - std::cerr << "Unable to load index. Status code: " << res << "." - << std::endl; - } + if (res != 0) + { + std::cerr << "Unable to load index. Status code: " << res << "." << std::endl; + } - std::vector node_list; - std::cout << "Caching " << num_nodes_to_cache << " BFS nodes around medoid(s)" - << std::endl; - _index->cache_bfs_levels(num_nodes_to_cache, node_list); - _index->load_cache_list(node_list); - omp_set_num_threads(num_threads); + std::vector node_list; + std::cout << "Caching " << num_nodes_to_cache << " BFS nodes around medoid(s)" << std::endl; + _index->cache_bfs_levels(num_nodes_to_cache, node_list); + _index->load_cache_list(node_list); + omp_set_num_threads(num_threads); } template -SearchResult PQFlashSearch::search(const T *query, - const unsigned int dimensions, - const unsigned int K, - const unsigned int Ls) { - uint64_t *indices_u64 = new uint64_t[K]; - unsigned *indices = new unsigned[K]; - float *distances = new float[K]; - - auto startTime = std::chrono::high_resolution_clock::now(); - _index->cached_beam_search(query, K, Ls, indices_u64, distances, DEFAULT_W); - auto duration = std::chrono::duration_cast( - std::chrono::high_resolution_clock::now() - startTime) - .count(); - for (unsigned k = 0; k < K; ++k) indices[k] = indices_u64[k]; - - std::string *tags = nullptr; - if (_tags_enabled) { - tags = new std::string[K]; - lookup_tags(K, indices, tags); - } - SearchResult result(K, (unsigned int)duration, indices, distances, tags); - delete[] indices_u64; - delete[] indices; - delete[] distances; - return result; +SearchResult PQFlashSearch::search(const T *query, const unsigned int dimensions, const unsigned int K, + const unsigned int Ls) +{ + uint64_t *indices_u64 = new uint64_t[K]; + unsigned *indices = new unsigned[K]; + float *distances = new float[K]; + + auto startTime = std::chrono::high_resolution_clock::now(); + _index->cached_beam_search(query, K, Ls, indices_u64, distances, DEFAULT_W); + auto duration = + std::chrono::duration_cast(std::chrono::high_resolution_clock::now() - startTime) + .count(); + for (unsigned k = 0; k < K; ++k) + indices[k] = indices_u64[k]; + + std::string *tags = nullptr; + if (_tags_enabled) + { + tags = new std::string[K]; + lookup_tags(K, indices, tags); + } + SearchResult result(K, (unsigned int)duration, indices, distances, tags); + delete[] indices_u64; + delete[] indices; + delete[] distances; + return result; } -template -PQFlashSearch::~PQFlashSearch() {} +template PQFlashSearch::~PQFlashSearch() +{ +} template class InMemorySearch; template class InMemorySearch; @@ -205,4 +212,4 @@ template class InMemorySearch; template class PQFlashSearch; template class PQFlashSearch; template class PQFlashSearch; -} // namespace diskann +} // namespace diskann diff --git a/src/restapi/server.cpp b/src/restapi/server.cpp index 6d365913c..f79b0affb 100644 --- a/src/restapi/server.cpp +++ b/src/restapi/server.cpp @@ -11,239 +11,261 @@ #include -namespace diskann { - -Server::Server( - web::uri &uri, - std::vector> &multi_searcher, - const std::string &typestring) - : _multi_search(multi_searcher.size() > 1 ? true : false) { - for (auto &searcher : multi_searcher) - _multi_searcher.push_back(std::move(searcher)); - - _listener = std::unique_ptr( - new web::http::experimental::listener::http_listener(uri)); - if (typestring == std::string("float")) { - _listener->support( - std::bind(&Server::handle_post, this, std::placeholders::_1)); - } else if (typestring == std::string("int8_t")) { - _listener->support( - web::http::methods::POST, - std::bind(&Server::handle_post, this, std::placeholders::_1)); - } else if (typestring == std::string("uint8_t")) { - _listener->support( - web::http::methods::POST, - std::bind(&Server::handle_post, this, std::placeholders::_1)); - } else { - throw "Unsupported type in server constuctor"; - } +namespace diskann +{ + +Server::Server(web::uri &uri, std::vector> &multi_searcher, + const std::string &typestring) + : _multi_search(multi_searcher.size() > 1 ? true : false) +{ + for (auto &searcher : multi_searcher) + _multi_searcher.push_back(std::move(searcher)); + + _listener = std::unique_ptr( + new web::http::experimental::listener::http_listener(uri)); + if (typestring == std::string("float")) + { + _listener->support(std::bind(&Server::handle_post, this, std::placeholders::_1)); + } + else if (typestring == std::string("int8_t")) + { + _listener->support(web::http::methods::POST, + std::bind(&Server::handle_post, this, std::placeholders::_1)); + } + else if (typestring == std::string("uint8_t")) + { + _listener->support(web::http::methods::POST, + std::bind(&Server::handle_post, this, std::placeholders::_1)); + } + else + { + throw "Unsupported type in server constuctor"; + } } -Server::~Server() {} +Server::~Server() +{ +} -pplx::task Server::open() { return _listener->open(); } -pplx::task Server::close() { return _listener->close(); } +pplx::task Server::open() +{ + return _listener->open(); +} +pplx::task Server::close() +{ + return _listener->close(); +} -diskann::SearchResult Server::aggregate_results( - const unsigned K, const std::vector &results) { - if (_multi_search) { - auto best_indices = new unsigned[K]; - auto best_distances = new float[K]; - auto best_partitions = new unsigned[K]; - auto best_tags = results[0].tags_enabled() ? new std::string[K] : nullptr; +diskann::SearchResult Server::aggregate_results(const unsigned K, const std::vector &results) +{ + if (_multi_search) + { + auto best_indices = new unsigned[K]; + auto best_distances = new float[K]; + auto best_partitions = new unsigned[K]; + auto best_tags = results[0].tags_enabled() ? new std::string[K] : nullptr; - auto numsearchers = _multi_searcher.size(); - std::vector pos(numsearchers, 0); + auto numsearchers = _multi_searcher.size(); + std::vector pos(numsearchers, 0); - for (size_t k = 0; k < K; ++k) { - float best_distance = std::numeric_limits::max(); - unsigned best_partition = 0; + for (size_t k = 0; k < K; ++k) + { + float best_distance = std::numeric_limits::max(); + unsigned best_partition = 0; - for (size_t i = 0; i < numsearchers; ++i) { - if (results[i].get_distances()[pos[i]] < best_distance) { - best_distance = results[i].get_distances()[pos[i]]; - best_partition = i; + for (size_t i = 0; i < numsearchers; ++i) + { + if (results[i].get_distances()[pos[i]] < best_distance) + { + best_distance = results[i].get_distances()[pos[i]]; + best_partition = i; + } + } + best_distances[k] = best_distance; + best_indices[k] = results[best_partition].get_indices()[pos[best_partition]]; + best_partitions[k] = best_partition; + if (results[best_partition].tags_enabled()) + best_tags[k] = results[best_partition].get_tags()[pos[best_partition]]; + std::cout << best_partition << " " << pos[best_partition] << std::endl; + pos[best_partition]++; } - } - best_distances[k] = best_distance; - best_indices[k] = - results[best_partition].get_indices()[pos[best_partition]]; - best_partitions[k] = best_partition; - if (results[best_partition].tags_enabled()) - best_tags[k] = results[best_partition].get_tags()[pos[best_partition]]; - std::cout << best_partition << " " << pos[best_partition] << std::endl; - pos[best_partition]++; - } - unsigned int total_time = 0; - for (size_t i = 0; i < numsearchers; ++i) - total_time += results[i].get_time(); - diskann::SearchResult result = - SearchResult(K, total_time, best_indices, best_distances, best_tags, - best_partitions); - - delete[] best_indices; - delete[] best_distances; - delete[] best_partitions; - delete[] best_tags; - - return result; - } else { - return results[0]; - } + unsigned int total_time = 0; + for (size_t i = 0; i < numsearchers; ++i) + total_time += results[i].get_time(); + diskann::SearchResult result = + SearchResult(K, total_time, best_indices, best_distances, best_tags, best_partitions); + + delete[] best_indices; + delete[] best_distances; + delete[] best_partitions; + delete[] best_tags; + + return result; + } + else + { + return results[0]; + } } -template -void Server::handle_post(web::http::http_request message) { - message.extract_string(true) - .then([=](utility::string_t body) { - int64_t queryId = -1; - unsigned int K = 0; - try { - T *queryVector = nullptr; - unsigned int dimensions = 0; - unsigned int Ls; - parseJson(body, K, queryId, queryVector, dimensions, Ls); - - auto startTime = std::chrono::high_resolution_clock::now(); - std::vector results; - - for (auto &searcher : _multi_searcher) - results.push_back( - searcher->search(queryVector, dimensions, (unsigned int)K, Ls)); - diskann::SearchResult result = aggregate_results(K, results); - diskann::aligned_free(queryVector); - web::json::value response = prepareResponse(queryId, K); - response[INDICES_KEY] = idsToJsonArray(result); - response[DISTANCES_KEY] = distancesToJsonArray(result); - if (result.tags_enabled()) - response[TAGS_KEY] = tagsToJsonArray(result); - if (result.partitions_enabled()) - response[PARTITION_KEY] = partitionsToJsonArray(result); - - response[TIME_TAKEN_KEY] = - std::chrono::duration_cast( - std::chrono::high_resolution_clock::now() - startTime) - .count(); - - std::cout << "Responding to: " << queryId << std::endl; - return std::make_pair(web::http::status_codes::OK, response); - } catch (const std::exception &ex) { - std::cerr << "Exception while processing query: " << queryId << ":" - << ex.what() << std::endl; - web::json::value response = prepareResponse(queryId, K); - response[ERROR_MESSAGE_KEY] = web::json::value::string(ex.what()); - return std::make_pair(web::http::status_codes::InternalError, - response); - } catch (...) { - std::cerr << "Uncaught exception while processing query: " << queryId; - web::json::value response = prepareResponse(queryId, K); - response[ERROR_MESSAGE_KEY] = web::json::value::string(UNKNOWN_ERROR); - return std::make_pair(web::http::status_codes::InternalError, - response); - } - }) - .then( - [=](std::pair response_status) { - try { - message.reply(response_status.first, response_status.second) - .wait(); - } catch (const std::exception &ex) { - std::cerr << "Exception while processing reply: " << ex.what() - << std::endl; +template void Server::handle_post(web::http::http_request message) +{ + message.extract_string(true) + .then([=](utility::string_t body) { + int64_t queryId = -1; + unsigned int K = 0; + try + { + T *queryVector = nullptr; + unsigned int dimensions = 0; + unsigned int Ls; + parseJson(body, K, queryId, queryVector, dimensions, Ls); + + auto startTime = std::chrono::high_resolution_clock::now(); + std::vector results; + + for (auto &searcher : _multi_searcher) + results.push_back(searcher->search(queryVector, dimensions, (unsigned int)K, Ls)); + diskann::SearchResult result = aggregate_results(K, results); + diskann::aligned_free(queryVector); + web::json::value response = prepareResponse(queryId, K); + response[INDICES_KEY] = idsToJsonArray(result); + response[DISTANCES_KEY] = distancesToJsonArray(result); + if (result.tags_enabled()) + response[TAGS_KEY] = tagsToJsonArray(result); + if (result.partitions_enabled()) + response[PARTITION_KEY] = partitionsToJsonArray(result); + + response[TIME_TAKEN_KEY] = std::chrono::duration_cast( + std::chrono::high_resolution_clock::now() - startTime) + .count(); + + std::cout << "Responding to: " << queryId << std::endl; + return std::make_pair(web::http::status_codes::OK, response); + } + catch (const std::exception &ex) + { + std::cerr << "Exception while processing query: " << queryId << ":" << ex.what() << std::endl; + web::json::value response = prepareResponse(queryId, K); + response[ERROR_MESSAGE_KEY] = web::json::value::string(ex.what()); + return std::make_pair(web::http::status_codes::InternalError, response); + } + catch (...) + { + std::cerr << "Uncaught exception while processing query: " << queryId; + web::json::value response = prepareResponse(queryId, K); + response[ERROR_MESSAGE_KEY] = web::json::value::string(UNKNOWN_ERROR); + return std::make_pair(web::http::status_codes::InternalError, response); + } + }) + .then([=](std::pair response_status) { + try + { + message.reply(response_status.first, response_status.second).wait(); + } + catch (const std::exception &ex) + { + std::cerr << "Exception while processing reply: " << ex.what() << std::endl; }; - }); + }); } -web::json::value Server::prepareResponse(const int64_t &queryId, const int k) { - web::json::value response = web::json::value::object(); - response[QUERY_ID_KEY] = queryId; - response[K_KEY] = k; +web::json::value Server::prepareResponse(const int64_t &queryId, const int k) +{ + web::json::value response = web::json::value::object(); + response[QUERY_ID_KEY] = queryId; + response[K_KEY] = k; - return response; + return response; } template -void Server::parseJson(const utility::string_t &body, unsigned int &k, - int64_t &queryId, T *&queryVector, - unsigned int &dimensions, unsigned &Ls) { - std::cout << body << std::endl; - web::json::value val = web::json::value::parse(body); - web::json::array queryArr = val.at(VECTOR_KEY).as_array(); - queryId = val.has_field(QUERY_ID_KEY) - ? val.at(QUERY_ID_KEY).as_number().to_int64() - : -1; - Ls = val.has_field(L_KEY) ? val.at(L_KEY).as_number().to_uint32() : DEFAULT_L; - k = val.at(K_KEY).as_integer(); - - if (k <= 0 || k > Ls) { - throw new std::invalid_argument( - "Num of expected NN (k) must be greater than zero and less than or " - "equal to Ls."); - } - if (queryArr.size() == 0) { - throw new std::invalid_argument("Query vector has zero elements."); - } - - dimensions = static_cast(queryArr.size()); - unsigned new_dim = ROUND_UP(dimensions, 8); - diskann::alloc_aligned((void **)&queryVector, new_dim * sizeof(T), - 8 * sizeof(T)); - memset(queryVector, 0, new_dim * sizeof(float)); - for (size_t i = 0; i < queryArr.size(); i++) { - queryVector[i] = (float)queryArr[i].as_double(); - } +void Server::parseJson(const utility::string_t &body, unsigned int &k, int64_t &queryId, T *&queryVector, + unsigned int &dimensions, unsigned &Ls) +{ + std::cout << body << std::endl; + web::json::value val = web::json::value::parse(body); + web::json::array queryArr = val.at(VECTOR_KEY).as_array(); + queryId = val.has_field(QUERY_ID_KEY) ? val.at(QUERY_ID_KEY).as_number().to_int64() : -1; + Ls = val.has_field(L_KEY) ? val.at(L_KEY).as_number().to_uint32() : DEFAULT_L; + k = val.at(K_KEY).as_integer(); + + if (k <= 0 || k > Ls) + { + throw new std::invalid_argument("Num of expected NN (k) must be greater than zero and less than or " + "equal to Ls."); + } + if (queryArr.size() == 0) + { + throw new std::invalid_argument("Query vector has zero elements."); + } + + dimensions = static_cast(queryArr.size()); + unsigned new_dim = ROUND_UP(dimensions, 8); + diskann::alloc_aligned((void **)&queryVector, new_dim * sizeof(T), 8 * sizeof(T)); + memset(queryVector, 0, new_dim * sizeof(float)); + for (size_t i = 0; i < queryArr.size(); i++) + { + queryVector[i] = (float)queryArr[i].as_double(); + } } template -web::json::value Server::toJsonArray( - const std::vector &v, - std::function valConverter) { - web::json::value rslts = web::json::value::array(); - for (size_t i = 0; i < v.size(); i++) { - auto jsonVal = valConverter(v[i]); - rslts[i] = jsonVal; - } - return rslts; +web::json::value Server::toJsonArray(const std::vector &v, std::function valConverter) +{ + web::json::value rslts = web::json::value::array(); + for (size_t i = 0; i < v.size(); i++) + { + auto jsonVal = valConverter(v[i]); + rslts[i] = jsonVal; + } + return rslts; } -web::json::value Server::idsToJsonArray(const diskann::SearchResult &result) { - web::json::value idArray = web::json::value::array(); - auto ids = result.get_indices(); - for (size_t i = 0; i < ids.size(); i++) { - auto idVal = web::json::value::number(ids[i]); - idArray[i] = idVal; - } - std::cout << "Vector size: " << ids.size() << std::endl; - return idArray; +web::json::value Server::idsToJsonArray(const diskann::SearchResult &result) +{ + web::json::value idArray = web::json::value::array(); + auto ids = result.get_indices(); + for (size_t i = 0; i < ids.size(); i++) + { + auto idVal = web::json::value::number(ids[i]); + idArray[i] = idVal; + } + std::cout << "Vector size: " << ids.size() << std::endl; + return idArray; } -web::json::value Server::distancesToJsonArray( - const diskann::SearchResult &result) { - web::json::value distArray = web::json::value::array(); - auto distances = result.get_distances(); - for (size_t i = 0; i < distances.size(); i++) { - distArray[i] = web::json::value::number(distances[i]); - } - return distArray; +web::json::value Server::distancesToJsonArray(const diskann::SearchResult &result) +{ + web::json::value distArray = web::json::value::array(); + auto distances = result.get_distances(); + for (size_t i = 0; i < distances.size(); i++) + { + distArray[i] = web::json::value::number(distances[i]); + } + return distArray; } -web::json::value Server::tagsToJsonArray(const diskann::SearchResult &result) { - web::json::value tagArray = web::json::value::array(); - auto tags = result.get_tags(); - for (size_t i = 0; i < tags.size(); i++) { - tagArray[i] = web::json::value::string(tags[i]); - } - return tagArray; +web::json::value Server::tagsToJsonArray(const diskann::SearchResult &result) +{ + web::json::value tagArray = web::json::value::array(); + auto tags = result.get_tags(); + for (size_t i = 0; i < tags.size(); i++) + { + tagArray[i] = web::json::value::string(tags[i]); + } + return tagArray; } -web::json::value Server::partitionsToJsonArray( - const diskann::SearchResult &result) { - web::json::value partitionArray = web::json::value::array(); - auto partitions = result.get_partitions(); - for (size_t i = 0; i < partitions.size(); i++) { - partitionArray[i] = web::json::value::number(partitions[i]); - } - return partitionArray; +web::json::value Server::partitionsToJsonArray(const diskann::SearchResult &result) +{ + web::json::value partitionArray = web::json::value::array(); + auto partitions = result.get_partitions(); + for (size_t i = 0; i < partitions.size(); i++) + { + partitionArray[i] = web::json::value::number(partitions[i]); + } + return partitionArray; } -}; // namespace diskann \ No newline at end of file +}; // namespace diskann \ No newline at end of file diff --git a/src/scratch.cpp b/src/scratch.cpp index 02cd6d975..02e594560 100644 --- a/src/scratch.cpp +++ b/src/scratch.cpp @@ -6,133 +6,130 @@ #include "scratch.h" -namespace diskann { +namespace diskann +{ // // Functions to manage scratch space for in-memory index based search // template -InMemQueryScratch::InMemQueryScratch(uint32_t search_l, uint32_t indexing_l, - uint32_t r, uint32_t maxc, size_t dim, - size_t aligned_dim, - size_t alignment_factor, - bool init_pq_scratch) - : _L(0), _R(r), _maxc(maxc) { - if (search_l == 0 || indexing_l == 0 || r == 0 || dim == 0) { - std::stringstream ss; - ss << "In InMemQueryScratch, one of search_l = " << search_l - << ", indexing_l = " << indexing_l << ", dim = " << dim - << " or r = " << r << " is zero." << std::endl; - throw diskann::ANNException(ss.str(), -1); - } - - // REFACTOR - // auto aligned_dim = ROUND_UP(dim, 8); - // alloc_aligned(((void **)&_aligned_query), aligned_dim * sizeof(T), 8 * - // sizeof(T)); memset(_aligned_query, 0, aligned_dim * sizeof(T)); - - alloc_aligned(((void **)&_aligned_query), aligned_dim * sizeof(T), - alignment_factor * sizeof(T)); - memset(_aligned_query, 0, aligned_dim * sizeof(T)); - - if (init_pq_scratch) - _pq_scratch = new PQScratch(MAX_GRAPH_DEGREE, aligned_dim); - else - _pq_scratch = nullptr; - - _occlude_factor.reserve(maxc); - _inserted_into_pool_bs = new boost::dynamic_bitset<>(); - _id_scratch.reserve(std::ceil(1.5 * GRAPH_SLACK_FACTOR * _R)); - _dist_scratch.reserve(std::ceil(1.5 * GRAPH_SLACK_FACTOR * _R)); - - resize_for_new_L(std::max(search_l, indexing_l)); +InMemQueryScratch::InMemQueryScratch(uint32_t search_l, uint32_t indexing_l, uint32_t r, uint32_t maxc, size_t dim, + size_t aligned_dim, size_t alignment_factor, bool init_pq_scratch) + : _L(0), _R(r), _maxc(maxc) +{ + if (search_l == 0 || indexing_l == 0 || r == 0 || dim == 0) + { + std::stringstream ss; + ss << "In InMemQueryScratch, one of search_l = " << search_l << ", indexing_l = " << indexing_l + << ", dim = " << dim << " or r = " << r << " is zero." << std::endl; + throw diskann::ANNException(ss.str(), -1); + } + + // REFACTOR + // auto aligned_dim = ROUND_UP(dim, 8); + // alloc_aligned(((void **)&_aligned_query), aligned_dim * sizeof(T), 8 * + // sizeof(T)); memset(_aligned_query, 0, aligned_dim * sizeof(T)); + + alloc_aligned(((void **)&_aligned_query), aligned_dim * sizeof(T), alignment_factor * sizeof(T)); + memset(_aligned_query, 0, aligned_dim * sizeof(T)); + + if (init_pq_scratch) + _pq_scratch = new PQScratch(MAX_GRAPH_DEGREE, aligned_dim); + else + _pq_scratch = nullptr; + + _occlude_factor.reserve(maxc); + _inserted_into_pool_bs = new boost::dynamic_bitset<>(); + _id_scratch.reserve(std::ceil(1.5 * GRAPH_SLACK_FACTOR * _R)); + _dist_scratch.reserve(std::ceil(1.5 * GRAPH_SLACK_FACTOR * _R)); + + resize_for_new_L(std::max(search_l, indexing_l)); } -template -void InMemQueryScratch::clear() { - _pool.clear(); - _best_l_nodes.clear(); - _occlude_factor.clear(); +template void InMemQueryScratch::clear() +{ + _pool.clear(); + _best_l_nodes.clear(); + _occlude_factor.clear(); - _inserted_into_pool_rs.clear(); - _inserted_into_pool_bs->reset(); + _inserted_into_pool_rs.clear(); + _inserted_into_pool_bs->reset(); - _id_scratch.clear(); - _dist_scratch.clear(); + _id_scratch.clear(); + _dist_scratch.clear(); - _expanded_nodes_set.clear(); - _expanded_nghrs_vec.clear(); - _occlude_list_output.clear(); + _expanded_nodes_set.clear(); + _expanded_nghrs_vec.clear(); + _occlude_list_output.clear(); } -template -void InMemQueryScratch::resize_for_new_L(uint32_t new_l) { - if (new_l > _L) { - _L = new_l; - _pool.reserve(3 * _L + _R); - _best_l_nodes.reserve(_L); - - _inserted_into_pool_rs.reserve(20 * _L); - } +template void InMemQueryScratch::resize_for_new_L(uint32_t new_l) +{ + if (new_l > _L) + { + _L = new_l; + _pool.reserve(3 * _L + _R); + _best_l_nodes.reserve(_L); + + _inserted_into_pool_rs.reserve(20 * _L); + } } -template -InMemQueryScratch::~InMemQueryScratch() { - if (_aligned_query != nullptr) { - aligned_free(_aligned_query); - } +template InMemQueryScratch::~InMemQueryScratch() +{ + if (_aligned_query != nullptr) + { + aligned_free(_aligned_query); + } - delete _pq_scratch; - delete _inserted_into_pool_bs; + delete _pq_scratch; + delete _inserted_into_pool_bs; } // // Functions to manage scratch space for SSD based search // -template -void SSDQueryScratch::reset() { - coord_idx = 0; - sector_idx = 0; - visited.clear(); - retset.clear(); - full_retset.clear(); +template void SSDQueryScratch::reset() +{ + coord_idx = 0; + sector_idx = 0; + visited.clear(); + retset.clear(); + full_retset.clear(); } -template -SSDQueryScratch::SSDQueryScratch(size_t aligned_dim, - size_t visited_reserve) { - size_t coord_alloc_size = ROUND_UP(MAX_N_CMPS * aligned_dim, 256); +template SSDQueryScratch::SSDQueryScratch(size_t aligned_dim, size_t visited_reserve) +{ + size_t coord_alloc_size = ROUND_UP(MAX_N_CMPS * aligned_dim, 256); - diskann::alloc_aligned((void **)&coord_scratch, coord_alloc_size, 256); - diskann::alloc_aligned((void **)§or_scratch, - (size_t)MAX_N_SECTOR_READS * (size_t)SECTOR_LEN, - SECTOR_LEN); - diskann::alloc_aligned((void **)&aligned_query_T, aligned_dim * sizeof(T), - 8 * sizeof(T)); + diskann::alloc_aligned((void **)&coord_scratch, coord_alloc_size, 256); + diskann::alloc_aligned((void **)§or_scratch, (size_t)MAX_N_SECTOR_READS * (size_t)SECTOR_LEN, SECTOR_LEN); + diskann::alloc_aligned((void **)&aligned_query_T, aligned_dim * sizeof(T), 8 * sizeof(T)); - _pq_scratch = new PQScratch(MAX_GRAPH_DEGREE, aligned_dim); + _pq_scratch = new PQScratch(MAX_GRAPH_DEGREE, aligned_dim); - memset(coord_scratch, 0, MAX_N_CMPS * aligned_dim); - memset(aligned_query_T, 0, aligned_dim * sizeof(T)); + memset(coord_scratch, 0, MAX_N_CMPS * aligned_dim); + memset(aligned_query_T, 0, aligned_dim * sizeof(T)); - visited.reserve(visited_reserve); - full_retset.reserve(visited_reserve); + visited.reserve(visited_reserve); + full_retset.reserve(visited_reserve); } -template -SSDQueryScratch::~SSDQueryScratch() { - diskann::aligned_free((void *)coord_scratch); - diskann::aligned_free((void *)sector_scratch); +template SSDQueryScratch::~SSDQueryScratch() +{ + diskann::aligned_free((void *)coord_scratch); + diskann::aligned_free((void *)sector_scratch); - delete[] _pq_scratch; + delete[] _pq_scratch; } template -SSDThreadData::SSDThreadData(size_t aligned_dim, size_t visited_reserve) - : scratch(aligned_dim, visited_reserve) {} +SSDThreadData::SSDThreadData(size_t aligned_dim, size_t visited_reserve) : scratch(aligned_dim, visited_reserve) +{ +} -template -void SSDThreadData::clear() { - scratch.reset(); +template void SSDThreadData::clear() +{ + scratch.reset(); } template DISKANN_DLLEXPORT class InMemQueryScratch; @@ -146,4 +143,4 @@ template DISKANN_DLLEXPORT class SSDQueryScratch; template DISKANN_DLLEXPORT class SSDThreadData; template DISKANN_DLLEXPORT class SSDThreadData; template DISKANN_DLLEXPORT class SSDThreadData; -} // namespace diskann \ No newline at end of file +} // namespace diskann \ No newline at end of file diff --git a/src/utils.cpp b/src/utils.cpp index 5f724964b..b675e656d 100644 --- a/src/utils.cpp +++ b/src/utils.cpp @@ -9,7 +9,7 @@ #include "aligned_file_reader.h" #endif -const uint32_t MAX_REQUEST_SIZE = 1024 * 1024 * 1024; // 64MB +const uint32_t MAX_REQUEST_SIZE = 1024 * 1024 * 1024; // 64MB const uint32_t MAX_SIMULTANEOUS_READ_REQUESTS = 128; #ifdef _WINDOWS @@ -17,44 +17,48 @@ const uint32_t MAX_SIMULTANEOUS_READ_REQUESTS = 128; // Taken from: // https://insufficientlycomplicated.wordpress.com/2011/11/07/detecting-intel-advanced-vector-extensions-avx-in-visual-studio/ -bool cpuHasAvxSupport() { - bool avxSupported = false; - - // Checking for AVX requires 3 things: - // 1) CPUID indicates that the OS uses XSAVE and XRSTORE - // instructions (allowing saving YMM registers on context - // switch) - // 2) CPUID indicates support for AVX - // 3) XGETBV indicates the AVX registers will be saved and - // restored on context switch - // - // Note that XGETBV is only available on 686 or later CPUs, so - // the instruction needs to be conditionally run. - int cpuInfo[4]; - __cpuid(cpuInfo, 1); - - bool osUsesXSAVE_XRSTORE = cpuInfo[2] & (1 << 27) || false; - bool cpuAVXSuport = cpuInfo[2] & (1 << 28) || false; - - if (osUsesXSAVE_XRSTORE && cpuAVXSuport) { - // Check if the OS will save the YMM registers - unsigned long long xcrFeatureMask = _xgetbv(_XCR_XFEATURE_ENABLED_MASK); - avxSupported = (xcrFeatureMask & 0x6) || false; - } - - return avxSupported; +bool cpuHasAvxSupport() +{ + bool avxSupported = false; + + // Checking for AVX requires 3 things: + // 1) CPUID indicates that the OS uses XSAVE and XRSTORE + // instructions (allowing saving YMM registers on context + // switch) + // 2) CPUID indicates support for AVX + // 3) XGETBV indicates the AVX registers will be saved and + // restored on context switch + // + // Note that XGETBV is only available on 686 or later CPUs, so + // the instruction needs to be conditionally run. + int cpuInfo[4]; + __cpuid(cpuInfo, 1); + + bool osUsesXSAVE_XRSTORE = cpuInfo[2] & (1 << 27) || false; + bool cpuAVXSuport = cpuInfo[2] & (1 << 28) || false; + + if (osUsesXSAVE_XRSTORE && cpuAVXSuport) + { + // Check if the OS will save the YMM registers + unsigned long long xcrFeatureMask = _xgetbv(_XCR_XFEATURE_ENABLED_MASK); + avxSupported = (xcrFeatureMask & 0x6) || false; + } + + return avxSupported; } -bool cpuHasAvx2Support() { - int cpuInfo[4]; - __cpuid(cpuInfo, 0); - int n = cpuInfo[0]; - if (n >= 7) { - __cpuidex(cpuInfo, 7, 0); - static int avx2Mask = 0x20; - return (cpuInfo[1] & avx2Mask) > 0; - } - return false; +bool cpuHasAvx2Support() +{ + int cpuInfo[4]; + __cpuid(cpuInfo, 0); + int n = cpuInfo[0]; + if (n >= 7) + { + __cpuidex(cpuInfo, 7, 0); + static int avx2Mask = 0x20; + return (cpuInfo[1] & avx2Mask) > 0; + } + return false; } bool AvxSupportedCPU = cpuHasAvxSupport(); @@ -66,423 +70,407 @@ bool Avx2SupportedCPU = true; bool AvxSupportedCPU = false; #endif -namespace diskann { +namespace diskann +{ -void block_convert(std::ofstream &writr, std::ifstream &readr, float *read_buf, - size_t npts, size_t ndims) { - readr.read((char *)read_buf, npts * ndims * sizeof(float)); - uint32_t ndims_u32 = (uint32_t)ndims; +void block_convert(std::ofstream &writr, std::ifstream &readr, float *read_buf, size_t npts, size_t ndims) +{ + readr.read((char *)read_buf, npts * ndims * sizeof(float)); + uint32_t ndims_u32 = (uint32_t)ndims; #pragma omp parallel for - for (int64_t i = 0; i < (int64_t)npts; i++) { - float norm_pt = std::numeric_limits::epsilon(); - for (uint32_t dim = 0; dim < ndims_u32; dim++) { - norm_pt += *(read_buf + i * ndims + dim) * *(read_buf + i * ndims + dim); + for (int64_t i = 0; i < (int64_t)npts; i++) + { + float norm_pt = std::numeric_limits::epsilon(); + for (uint32_t 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 (uint32_t dim = 0; dim < ndims_u32; dim++) + { + *(read_buf + i * ndims + dim) = *(read_buf + i * ndims + dim) / norm_pt; + } } - norm_pt = std::sqrt(norm_pt); - for (uint32_t 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)); + 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(int32_t)); - readr.read((char *)&ndims_s32, sizeof(int32_t)); - - writr.write((char *)&npts_s32, sizeof(int32_t)); - writr.write((char *)&ndims_s32, sizeof(int32_t)); - - size_t npts = (size_t)npts_s32; - size_t ndims = (size_t)ndims_s32; - diskann::cout << "Normalizing FLOAT vectors in file: " << inFileName - << std::endl; - diskann::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims - << std::endl; - - size_t blk_size = 131072; - size_t nblks = ROUND_UP(npts, blk_size) / blk_size; - diskann::cout << "# blks: " << nblks << std::endl; - - float *read_buf = new float[npts * ndims]; - for (size_t i = 0; i < nblks; i++) { - size_t 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; -} +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); -double calculate_recall(uint32_t num_queries, uint32_t *gold_std, - float *gs_dist, uint32_t dim_gs, uint32_t *our_results, - uint32_t dim_or, uint32_t recall_at) { - double total_recall = 0; - std::set gt, res; - - for (size_t i = 0; i < num_queries; i++) { - gt.clear(); - res.clear(); - uint32_t *gt_vec = gold_std + dim_gs * i; - uint32_t *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; - while (tie_breaker < dim_gs && - gt_dist_vec[tie_breaker] == gt_dist_vec[recall_at - 1]) - tie_breaker++; - } + int npts_s32, ndims_s32; + readr.read((char *)&npts_s32, sizeof(int32_t)); + readr.read((char *)&ndims_s32, sizeof(int32_t)); - gt.insert(gt_vec, gt_vec + tie_breaker); - res.insert(res_vec, - res_vec + recall_at); // change to recall_at for recall k@k - // or dim_or for k@dim_or - uint32_t cur_recall = 0; - for (auto &v : gt) { - if (res.find(v) != res.end()) { - cur_recall++; - } + writr.write((char *)&npts_s32, sizeof(int32_t)); + writr.write((char *)&ndims_s32, sizeof(int32_t)); + + size_t npts = (size_t)npts_s32; + size_t ndims = (size_t)ndims_s32; + diskann::cout << "Normalizing FLOAT vectors in file: " << inFileName << std::endl; + diskann::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims << std::endl; + + size_t blk_size = 131072; + size_t nblks = ROUND_UP(npts, blk_size) / blk_size; + diskann::cout << "# blks: " << nblks << std::endl; + + float *read_buf = new float[npts * ndims]; + for (size_t i = 0; i < nblks; i++) + { + size_t cblk_size = std::min(npts - i * blk_size, blk_size); + block_convert(writr, readr, read_buf, cblk_size, ndims); } - total_recall += cur_recall; - } - return total_recall / (num_queries) * (100.0 / recall_at); + delete[] read_buf; + + diskann::cout << "Wrote normalized points to file: " << outFileName << std::endl; } -double calculate_recall(uint32_t num_queries, uint32_t *gold_std, - float *gs_dist, uint32_t dim_gs, uint32_t *our_results, - uint32_t dim_or, uint32_t 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(); - uint32_t *gt_vec = gold_std + dim_gs * i; - uint32_t *res_vec = our_results + dim_or * i; - size_t tie_breaker = recall_at; - uint32_t active_points_count = 0; - uint32_t 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++; +double calculate_recall(uint32_t num_queries, uint32_t *gold_std, float *gs_dist, uint32_t dim_gs, + uint32_t *our_results, uint32_t dim_or, uint32_t recall_at) +{ + double total_recall = 0; + std::set gt, res; + + for (size_t i = 0; i < num_queries; i++) + { + gt.clear(); + res.clear(); + uint32_t *gt_vec = gold_std + dim_gs * i; + uint32_t *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; + while (tie_breaker < dim_gs && gt_dist_vec[tie_breaker] == gt_dist_vec[recall_at - 1]) + tie_breaker++; + } + + gt.insert(gt_vec, gt_vec + tie_breaker); + res.insert(res_vec, + res_vec + recall_at); // change to recall_at for recall k@k + // or dim_or for k@dim_or + uint32_t cur_recall = 0; + for (auto &v : gt) + { + if (res.find(v) != res.end()) + { + cur_recall++; + } + } + total_recall += cur_recall; } + return total_recall / (num_queries) * (100.0 / recall_at); +} - gt.insert(gt_vec, gt_vec + tie_breaker); - res.insert(res_vec, res_vec + recall_at); - uint32_t cur_recall = 0; - for (auto &v : res) { - if (gt.find(v) != gt.end()) { - cur_recall++; - } +double calculate_recall(uint32_t num_queries, uint32_t *gold_std, float *gs_dist, uint32_t dim_gs, + uint32_t *our_results, uint32_t dim_or, uint32_t 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(); + uint32_t *gt_vec = gold_std + dim_gs * i; + uint32_t *res_vec = our_results + dim_or * i; + size_t tie_breaker = recall_at; + uint32_t active_points_count = 0; + uint32_t 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); + uint32_t cur_recall = 0; + for (auto &v : res) + { + if (gt.find(v) != gt.end()) + { + cur_recall++; + } + } + total_recall += cur_recall; } - total_recall += cur_recall; - } - return ((double)(total_recall / (num_queries))) * - ((double)(100.0 / recall_at)); + return ((double)(total_recall / (num_queries))) * ((double)(100.0 / recall_at)); } -double calculate_range_search_recall( - uint32_t num_queries, std::vector> &groundtruth, - std::vector> &our_results) { - double total_recall = 0; - std::set gt, res; - - for (size_t i = 0; i < num_queries; i++) { - gt.clear(); - res.clear(); - - gt.insert(groundtruth[i].begin(), groundtruth[i].end()); - res.insert(our_results[i].begin(), our_results[i].end()); - uint32_t cur_recall = 0; - for (auto &v : gt) { - if (res.find(v) != res.end()) { - cur_recall++; - } +double calculate_range_search_recall(uint32_t num_queries, std::vector> &groundtruth, + std::vector> &our_results) +{ + double total_recall = 0; + std::set gt, res; + + for (size_t i = 0; i < num_queries; i++) + { + gt.clear(); + res.clear(); + + gt.insert(groundtruth[i].begin(), groundtruth[i].end()); + res.insert(our_results[i].begin(), our_results[i].end()); + uint32_t cur_recall = 0; + for (auto &v : gt) + { + if (res.find(v) != res.end()) + { + cur_recall++; + } + } + if (gt.size() != 0) + total_recall += ((100.0 * cur_recall) / gt.size()); + else + total_recall += 100; } - if (gt.size() != 0) - total_recall += ((100.0 * cur_recall) / gt.size()); - else - total_recall += 100; - } - return total_recall / (num_queries); + return total_recall / (num_queries); } #ifdef EXEC_ENV_OLS -void get_bin_metadata(AlignedFileReader &reader, size_t &npts, size_t &ndim, - size_t offset) { - std::vector readReqs; - AlignedRead readReq; - uint32_t buf[2]; // npts/ndim are uint32_ts. - - readReq.buf = buf; - readReq.offset = offset; - readReq.len = 2 * sizeof(uint32_t); - readReqs.push_back(readReq); - - IOContext &ctx = reader.get_ctx(); - reader.read(readReqs, ctx); // synchronous - if ((*(ctx.m_pRequestsStatus))[0] == IOContext::READ_SUCCESS) { - npts = buf[0]; - ndim = buf[1]; - diskann::cout << "File has: " << npts << " points, " << ndim - << " dimensions at offset: " << offset << std::endl; - } else { - std::stringstream str; - str << "Could not read binary metadata from index file at offset: " - << offset << std::endl; - throw diskann::ANNException(str.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } -} - -template -void load_bin(AlignedFileReader &reader, T *&data, size_t &npts, size_t &ndim, - size_t offset) { - // Code assumes that the reader is already setup correctly. - get_bin_metadata(reader, npts, ndim, offset); - data = new T[npts * ndim]; - - size_t data_size = npts * ndim * sizeof(T); - size_t write_offset = 0; - size_t read_start = offset + 2 * sizeof(uint32_t); - - // BingAlignedFileReader can only read uint32_t bytes of data. So, - // we limit ourselves even more to reading 1GB at a time. - std::vector readReqs; - while (data_size > 0) { +void get_bin_metadata(AlignedFileReader &reader, size_t &npts, size_t &ndim, size_t offset) +{ + std::vector readReqs; AlignedRead readReq; - readReq.buf = data + write_offset; - readReq.offset = read_start + write_offset; - readReq.len = data_size > MAX_REQUEST_SIZE ? MAX_REQUEST_SIZE : data_size; + uint32_t buf[2]; // npts/ndim are uint32_ts. + + readReq.buf = buf; + readReq.offset = offset; + readReq.len = 2 * sizeof(uint32_t); readReqs.push_back(readReq); - // in the corner case, the loop will not execute - data_size -= readReq.len; - write_offset += readReq.len; - } - IOContext &ctx = reader.get_ctx(); - reader.read(readReqs, ctx); - for (int i = 0; i < readReqs.size(); i++) { - // Since we are making sync calls, no request will be in the - // READ_WAIT state. - if ((*(ctx.m_pRequestsStatus))[i] != IOContext::READ_SUCCESS) { - std::stringstream str; - str << "Could not read binary data from index file at offset: " - << readReqs[i].offset << std::endl; - throw diskann::ANNException(str.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); + + IOContext &ctx = reader.get_ctx(); + reader.read(readReqs, ctx); // synchronous + if ((*(ctx.m_pRequestsStatus))[0] == IOContext::READ_SUCCESS) + { + npts = buf[0]; + ndim = buf[1]; + diskann::cout << "File has: " << npts << " points, " << ndim << " dimensions at offset: " << offset + << std::endl; + } + else + { + std::stringstream str; + str << "Could not read binary metadata from index file at offset: " << offset << std::endl; + throw diskann::ANNException(str.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } - } -} -template -void load_bin(AlignedFileReader &reader, std::unique_ptr &data, - size_t &npts, size_t &ndim, size_t offset) { - T *ptr = nullptr; - load_bin(reader, ptr, npts, ndim, offset); - data.reset(ptr); } -template -void copy_aligned_data_from_file(AlignedFileReader &reader, T *&data, - size_t &npts, size_t &ndim, - const size_t &rounded_dim, size_t offset) { - if (data == nullptr) { - diskann::cerr << "Memory was not allocated for " << data - << " before calling the load function. Exiting..." - << std::endl; - throw diskann::ANNException( - "Null pointer passed to copy_aligned_data_from_file()", -1, __FUNCSIG__, - __FILE__, __LINE__); - } - - size_t pts, dim; - get_bin_metadata(reader, pts, dim, offset); - - if (ndim != dim || npts != pts) { - std::stringstream ss; - ss << "Either file dimension: " << dim - << " is != passed dimension: " << ndim << " or file #pts: " << pts - << " is != passed #pts: " << npts << std::endl; - throw diskann::ANNException(ss.str(), -1, __FUNCSIG__, __FILE__, __LINE__); - } - - // Instead of reading one point of ndim size and setting (rounded_dim - dim) - // values to zero We'll set everything to zero and read in chunks of data at - // the appropriate locations. - size_t read_offset = offset + 2 * sizeof(uint32_t); - memset(data, 0, npts * rounded_dim * sizeof(T)); - int i = 0; - std::vector read_requests; - - while (i < npts) { - int j = 0; - read_requests.clear(); - while (j < MAX_SIMULTANEOUS_READ_REQUESTS && i < npts) { - AlignedRead read_req; - read_req.buf = data + i * rounded_dim; - read_req.len = dim * sizeof(T); - read_req.offset = read_offset + i * dim * sizeof(T); - read_requests.push_back(read_req); - i++; - j++; +template void load_bin(AlignedFileReader &reader, T *&data, size_t &npts, size_t &ndim, size_t offset) +{ + // Code assumes that the reader is already setup correctly. + get_bin_metadata(reader, npts, ndim, offset); + data = new T[npts * ndim]; + + size_t data_size = npts * ndim * sizeof(T); + size_t write_offset = 0; + size_t read_start = offset + 2 * sizeof(uint32_t); + + // BingAlignedFileReader can only read uint32_t bytes of data. So, + // we limit ourselves even more to reading 1GB at a time. + std::vector readReqs; + while (data_size > 0) + { + AlignedRead readReq; + readReq.buf = data + write_offset; + readReq.offset = read_start + write_offset; + readReq.len = data_size > MAX_REQUEST_SIZE ? MAX_REQUEST_SIZE : data_size; + readReqs.push_back(readReq); + // in the corner case, the loop will not execute + data_size -= readReq.len; + write_offset += readReq.len; } IOContext &ctx = reader.get_ctx(); - reader.read(read_requests, ctx); - for (int k = 0; k < read_requests.size(); k++) { - if ((*ctx.m_pRequestsStatus)[k] != IOContext::READ_SUCCESS) { - throw diskann::ANNException( - "Load data from file using AlignedReader failed.", -1, __FUNCSIG__, - __FILE__, __LINE__); - } + reader.read(readReqs, ctx); + for (int i = 0; i < readReqs.size(); i++) + { + // Since we are making sync calls, no request will be in the + // READ_WAIT state. + if ((*(ctx.m_pRequestsStatus))[i] != IOContext::READ_SUCCESS) + { + std::stringstream str; + str << "Could not read binary data from index file at offset: " << readReqs[i].offset << std::endl; + throw diskann::ANNException(str.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } } - } +} +template +void load_bin(AlignedFileReader &reader, std::unique_ptr &data, size_t &npts, size_t &ndim, size_t offset) +{ + T *ptr = nullptr; + load_bin(reader, ptr, npts, ndim, offset); + data.reset(ptr); } -// Unlike load_bin, assumes that data is already allocated 'size' entries template -void read_array(AlignedFileReader &reader, T *data, size_t size, - size_t offset) { - if (data == nullptr) { - throw diskann::ANNException("read_array requires an allocated buffer.", -1); - if (size * sizeof(T) > MAX_REQUEST_SIZE) { - std::stringstream ss; - ss << "Cannot read more than " << MAX_REQUEST_SIZE - << " bytes. Current request size: " << std::to_string(size) - << " sizeof(T): " << sizeof(T) << std::endl; - throw diskann::ANNException(ss.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); +void copy_aligned_data_from_file(AlignedFileReader &reader, T *&data, size_t &npts, size_t &ndim, + const size_t &rounded_dim, size_t offset) +{ + if (data == nullptr) + { + diskann::cerr << "Memory was not allocated for " << data << " before calling the load function. Exiting..." + << std::endl; + throw diskann::ANNException("Null pointer passed to copy_aligned_data_from_file()", -1, __FUNCSIG__, __FILE__, + __LINE__); + } + + size_t pts, dim; + get_bin_metadata(reader, pts, dim, offset); + + if (ndim != dim || npts != pts) + { + std::stringstream ss; + ss << "Either file dimension: " << dim << " is != passed dimension: " << ndim << " or file #pts: " << pts + << " is != passed #pts: " << npts << std::endl; + throw diskann::ANNException(ss.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } + + // Instead of reading one point of ndim size and setting (rounded_dim - dim) + // values to zero We'll set everything to zero and read in chunks of data at + // the appropriate locations. + size_t read_offset = offset + 2 * sizeof(uint32_t); + memset(data, 0, npts * rounded_dim * sizeof(T)); + int i = 0; std::vector read_requests; - AlignedRead read_req; - read_req.buf = data; - read_req.len = size * sizeof(T); - read_req.offset = offset; - read_requests.push_back(read_req); - IOContext &ctx = reader.get_ctx(); - reader.read(read_requests, ctx); - - if ((*(ctx.m_pRequestsStatus))[0] != IOContext::READ_SUCCESS) { - std::stringstream ss; - ss << "Failed to read_array() of size: " << size * sizeof(T) - << " at offset: " << offset << " from reader. " << std::endl; - throw diskann::ANNException(ss.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); + + while (i < npts) + { + int j = 0; + read_requests.clear(); + while (j < MAX_SIMULTANEOUS_READ_REQUESTS && i < npts) + { + AlignedRead read_req; + read_req.buf = data + i * rounded_dim; + read_req.len = dim * sizeof(T); + read_req.offset = read_offset + i * dim * sizeof(T); + read_requests.push_back(read_req); + i++; + j++; + } + IOContext &ctx = reader.get_ctx(); + reader.read(read_requests, ctx); + for (int k = 0; k < read_requests.size(); k++) + { + if ((*ctx.m_pRequestsStatus)[k] != IOContext::READ_SUCCESS) + { + throw diskann::ANNException("Load data from file using AlignedReader failed.", -1, __FUNCSIG__, + __FILE__, __LINE__); + } + } } - } } -template -void read_value(AlignedFileReader &reader, T &value, size_t offset) { - read_array(reader, &value, 1, offset); +// Unlike load_bin, assumes that data is already allocated 'size' entries +template void read_array(AlignedFileReader &reader, T *data, size_t size, size_t offset) +{ + if (data == nullptr) + { + throw diskann::ANNException("read_array requires an allocated buffer.", -1); + if (size * sizeof(T) > MAX_REQUEST_SIZE) + { + std::stringstream ss; + ss << "Cannot read more than " << MAX_REQUEST_SIZE + << " bytes. Current request size: " << std::to_string(size) << " sizeof(T): " << sizeof(T) << std::endl; + throw diskann::ANNException(ss.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } + std::vector read_requests; + AlignedRead read_req; + read_req.buf = data; + read_req.len = size * sizeof(T); + read_req.offset = offset; + read_requests.push_back(read_req); + IOContext &ctx = reader.get_ctx(); + reader.read(read_requests, ctx); + + if ((*(ctx.m_pRequestsStatus))[0] != IOContext::READ_SUCCESS) + { + std::stringstream ss; + ss << "Failed to read_array() of size: " << size * sizeof(T) << " at offset: " << offset << " from reader. " + << std::endl; + throw diskann::ANNException(ss.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } + } } -template DISKANN_DLLEXPORT void load_bin( - AlignedFileReader &reader, std::unique_ptr &data, size_t &npts, - size_t &ndim, size_t offset); -template DISKANN_DLLEXPORT void load_bin( - AlignedFileReader &reader, std::unique_ptr &data, size_t &npts, - size_t &ndim, size_t offset); -template DISKANN_DLLEXPORT void load_bin( - AlignedFileReader &reader, std::unique_ptr &data, size_t &npts, - size_t &ndim, size_t offset); -template DISKANN_DLLEXPORT void load_bin( - AlignedFileReader &reader, std::unique_ptr &data, size_t &npts, - size_t &ndim, size_t offset); -template DISKANN_DLLEXPORT void load_bin( - AlignedFileReader &reader, std::unique_ptr &data, size_t &npts, - size_t &ndim, size_t offset); -template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, - std::unique_ptr &data, - size_t &npts, size_t &ndim, - size_t offset); - -template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, - uint8_t *&data, size_t &npts, - size_t &ndim, size_t offset); -template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, - int64_t *&data, size_t &npts, - size_t &ndim, size_t offset); -template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, - uint64_t *&data, - size_t &npts, size_t &ndim, - size_t offset); -template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, - uint32_t *&data, - size_t &npts, size_t &ndim, - size_t offset); -template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, - int32_t *&data, size_t &npts, - size_t &ndim, size_t offset); - -template DISKANN_DLLEXPORT void copy_aligned_data_from_file( - AlignedFileReader &reader, uint8_t *&data, size_t &npts, size_t &dim, - const size_t &rounded_dim, size_t offset); -template DISKANN_DLLEXPORT void copy_aligned_data_from_file( - AlignedFileReader &reader, int8_t *&data, size_t &npts, size_t &dim, - const size_t &rounded_dim, size_t offset); -template DISKANN_DLLEXPORT void copy_aligned_data_from_file( - AlignedFileReader &reader, float *&data, size_t &npts, size_t &dim, - const size_t &rounded_dim, size_t offset); - -template DISKANN_DLLEXPORT void read_array(AlignedFileReader &reader, - char *data, size_t size, - size_t offset); - -template DISKANN_DLLEXPORT void read_array(AlignedFileReader &reader, - uint8_t *data, size_t size, - size_t offset); -template DISKANN_DLLEXPORT void read_array(AlignedFileReader &reader, - int8_t *data, size_t size, - size_t offset); -template DISKANN_DLLEXPORT void read_array(AlignedFileReader &reader, - uint32_t *data, - size_t size, - size_t offset); -template DISKANN_DLLEXPORT void read_array(AlignedFileReader &reader, - float *data, size_t size, +template void read_value(AlignedFileReader &reader, T &value, size_t offset) +{ + read_array(reader, &value, 1, offset); +} + +template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, std::unique_ptr &data, + size_t &npts, size_t &ndim, size_t offset); +template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, std::unique_ptr &data, + size_t &npts, size_t &ndim, size_t offset); +template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, std::unique_ptr &data, + size_t &npts, size_t &ndim, size_t offset); +template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, std::unique_ptr &data, + size_t &npts, size_t &ndim, size_t offset); +template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, std::unique_ptr &data, + size_t &npts, size_t &ndim, size_t offset); +template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, std::unique_ptr &data, size_t &npts, + size_t &ndim, size_t offset); + +template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, uint8_t *&data, size_t &npts, size_t &ndim, + size_t offset); +template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, int64_t *&data, size_t &npts, size_t &ndim, size_t offset); +template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, uint64_t *&data, size_t &npts, + size_t &ndim, size_t offset); +template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, uint32_t *&data, size_t &npts, + size_t &ndim, size_t offset); +template DISKANN_DLLEXPORT void load_bin(AlignedFileReader &reader, int32_t *&data, size_t &npts, size_t &ndim, + size_t offset); + +template DISKANN_DLLEXPORT void copy_aligned_data_from_file(AlignedFileReader &reader, uint8_t *&data, + size_t &npts, size_t &dim, + const size_t &rounded_dim, size_t offset); +template DISKANN_DLLEXPORT void copy_aligned_data_from_file(AlignedFileReader &reader, int8_t *&data, + size_t &npts, size_t &dim, + const size_t &rounded_dim, size_t offset); +template DISKANN_DLLEXPORT void copy_aligned_data_from_file(AlignedFileReader &reader, float *&data, + size_t &npts, size_t &dim, const size_t &rounded_dim, + size_t offset); + +template DISKANN_DLLEXPORT void read_array(AlignedFileReader &reader, char *data, size_t size, size_t offset); -template DISKANN_DLLEXPORT void read_value(AlignedFileReader &reader, - uint8_t &value, +template DISKANN_DLLEXPORT void read_array(AlignedFileReader &reader, uint8_t *data, size_t size, size_t offset); -template DISKANN_DLLEXPORT void read_value(AlignedFileReader &reader, - int8_t &value, - size_t offset); -template DISKANN_DLLEXPORT void read_value(AlignedFileReader &reader, - float &value, size_t offset); -template DISKANN_DLLEXPORT void read_value(AlignedFileReader &reader, - uint32_t &value, - size_t offset); -template DISKANN_DLLEXPORT void read_value(AlignedFileReader &reader, - uint64_t &value, +template DISKANN_DLLEXPORT void read_array(AlignedFileReader &reader, int8_t *data, size_t size, size_t offset); +template DISKANN_DLLEXPORT void read_array(AlignedFileReader &reader, uint32_t *data, size_t size, size_t offset); +template DISKANN_DLLEXPORT void read_array(AlignedFileReader &reader, float *data, size_t size, size_t offset); + +template DISKANN_DLLEXPORT void read_value(AlignedFileReader &reader, uint8_t &value, size_t offset); +template DISKANN_DLLEXPORT void read_value(AlignedFileReader &reader, int8_t &value, size_t offset); +template DISKANN_DLLEXPORT void read_value(AlignedFileReader &reader, float &value, size_t offset); +template DISKANN_DLLEXPORT void read_value(AlignedFileReader &reader, uint32_t &value, size_t offset); +template DISKANN_DLLEXPORT void read_value(AlignedFileReader &reader, uint64_t &value, size_t offset); #endif -} // namespace diskann +} // namespace diskann diff --git a/src/windows_aligned_file_reader.cpp b/src/windows_aligned_file_reader.cpp index a1911efd4..e99d5de34 100644 --- a/src/windows_aligned_file_reader.cpp +++ b/src/windows_aligned_file_reader.cpp @@ -9,158 +9,171 @@ #define SECTOR_LEN 4096 -void WindowsAlignedFileReader::open(const std::string &fname) { +void WindowsAlignedFileReader::open(const std::string &fname) +{ #ifdef UNICODE - m_filename = std::wstring(fname.begin(), fname.end()); + m_filename = std::wstring(fname.begin(), fname.end()); #else - m_filename = fname; + m_filename = fname; #endif - this->register_thread(); + this->register_thread(); } -void WindowsAlignedFileReader::close() { - for (auto &k_v : ctx_map) { - IOContext ctx = ctx_map[k_v.first]; - CloseHandle(ctx.fhandle); - } +void WindowsAlignedFileReader::close() +{ + for (auto &k_v : ctx_map) + { + IOContext ctx = ctx_map[k_v.first]; + CloseHandle(ctx.fhandle); + } } -void WindowsAlignedFileReader::register_thread() { - std::unique_lock lk(this->ctx_mut); - if (this->ctx_map.find(std::this_thread::get_id()) != ctx_map.end()) { - diskann::cout << "Warning:: Duplicate registration for thread_id : " - << std::this_thread::get_id() << std::endl; - } - - IOContext ctx; - ctx.fhandle = CreateFile(m_filename.c_str(), GENERIC_READ, FILE_SHARE_READ, - NULL, OPEN_EXISTING, - FILE_ATTRIBUTE_READONLY | FILE_FLAG_NO_BUFFERING | - FILE_FLAG_OVERLAPPED | FILE_FLAG_RANDOM_ACCESS, - NULL); - if (ctx.fhandle == INVALID_HANDLE_VALUE) { - diskann::cout << "Error opening " - << std::string(m_filename.begin(), m_filename.end()) - << " -- error=" << GetLastError() << std::endl; - } - - // create IOCompletionPort - ctx.iocp = CreateIoCompletionPort(ctx.fhandle, ctx.iocp, 0, 0); - - // create MAX_DEPTH # of reqs - for (uint64_t i = 0; i < MAX_IO_DEPTH; i++) { - OVERLAPPED os; - memset(&os, 0, sizeof(OVERLAPPED)); - // os.hEvent = CreateEventA(NULL, TRUE, FALSE, NULL); - ctx.reqs.push_back(os); - } - this->ctx_map.insert(std::make_pair(std::this_thread::get_id(), ctx)); +void WindowsAlignedFileReader::register_thread() +{ + std::unique_lock lk(this->ctx_mut); + if (this->ctx_map.find(std::this_thread::get_id()) != ctx_map.end()) + { + diskann::cout << "Warning:: Duplicate registration for thread_id : " << std::this_thread::get_id() << std::endl; + } + + IOContext ctx; + ctx.fhandle = CreateFile( + m_filename.c_str(), GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, + FILE_ATTRIBUTE_READONLY | FILE_FLAG_NO_BUFFERING | FILE_FLAG_OVERLAPPED | FILE_FLAG_RANDOM_ACCESS, NULL); + if (ctx.fhandle == INVALID_HANDLE_VALUE) + { + diskann::cout << "Error opening " << std::string(m_filename.begin(), m_filename.end()) + << " -- error=" << GetLastError() << std::endl; + } + + // create IOCompletionPort + ctx.iocp = CreateIoCompletionPort(ctx.fhandle, ctx.iocp, 0, 0); + + // create MAX_DEPTH # of reqs + for (uint64_t i = 0; i < MAX_IO_DEPTH; i++) + { + OVERLAPPED os; + memset(&os, 0, sizeof(OVERLAPPED)); + // os.hEvent = CreateEventA(NULL, TRUE, FALSE, NULL); + ctx.reqs.push_back(os); + } + this->ctx_map.insert(std::make_pair(std::this_thread::get_id(), ctx)); } -IOContext &WindowsAlignedFileReader::get_ctx() { - std::unique_lock lk(this->ctx_mut); - if (ctx_map.find(std::this_thread::get_id()) == ctx_map.end()) { - std::stringstream stream; - stream << "unable to find IOContext for thread_id : " - << std::this_thread::get_id() << "\n"; - throw diskann::ANNException(stream.str(), -2, __FUNCSIG__, __FILE__, - __LINE__); - } - IOContext &ctx = ctx_map[std::this_thread::get_id()]; - lk.unlock(); - return ctx; +IOContext &WindowsAlignedFileReader::get_ctx() +{ + std::unique_lock lk(this->ctx_mut); + if (ctx_map.find(std::this_thread::get_id()) == ctx_map.end()) + { + std::stringstream stream; + stream << "unable to find IOContext for thread_id : " << std::this_thread::get_id() << "\n"; + throw diskann::ANNException(stream.str(), -2, __FUNCSIG__, __FILE__, __LINE__); + } + IOContext &ctx = ctx_map[std::this_thread::get_id()]; + lk.unlock(); + return ctx; } -void WindowsAlignedFileReader::read(std::vector &read_reqs, - IOContext &ctx, bool async) { - using namespace std::chrono_literals; - // execute each request sequentially - size_t n_reqs = read_reqs.size(); - uint64_t n_batches = ROUND_UP(n_reqs, MAX_IO_DEPTH) / MAX_IO_DEPTH; - for (uint64_t i = 0; i < n_batches; i++) { - // reset all OVERLAPPED objects - for (auto &os : ctx.reqs) { - // HANDLE evt = os.hEvent; - memset(&os, 0, sizeof(os)); - // os.hEvent = evt; - - /* - if (ResetEvent(os.hEvent) == 0) { - diskann::cerr << "ResetEvent failed" << std::endl; - exit(-3); +void WindowsAlignedFileReader::read(std::vector &read_reqs, IOContext &ctx, bool async) +{ + using namespace std::chrono_literals; + // execute each request sequentially + size_t n_reqs = read_reqs.size(); + uint64_t n_batches = ROUND_UP(n_reqs, MAX_IO_DEPTH) / MAX_IO_DEPTH; + for (uint64_t i = 0; i < n_batches; i++) + { + // reset all OVERLAPPED objects + for (auto &os : ctx.reqs) + { + // HANDLE evt = os.hEvent; + memset(&os, 0, sizeof(os)); + // os.hEvent = evt; + + /* + if (ResetEvent(os.hEvent) == 0) { + diskann::cerr << "ResetEvent failed" << std::endl; + exit(-3); + } + */ } - */ - } - // batch start/end - uint64_t batch_start = MAX_IO_DEPTH * i; - uint64_t batch_size = - std::min((uint64_t)(n_reqs - batch_start), (uint64_t)MAX_IO_DEPTH); - - // fill OVERLAPPED and issue them - for (uint64_t j = 0; j < batch_size; j++) { - AlignedRead &req = read_reqs[batch_start + j]; - OVERLAPPED &os = ctx.reqs[j]; - - uint64_t offset = req.offset; - uint64_t nbytes = req.len; - char *read_buf = (char *)req.buf; - assert(IS_ALIGNED(read_buf, SECTOR_LEN)); - assert(IS_ALIGNED(offset, SECTOR_LEN)); - assert(IS_ALIGNED(nbytes, SECTOR_LEN)); - - // fill in OVERLAPPED struct - os.Offset = offset & 0xffffffff; - os.OffsetHigh = (offset >> 32); - - BOOL ret = ReadFile(ctx.fhandle, read_buf, nbytes, NULL, &os); - if (ret == FALSE) { - auto error = GetLastError(); - if (error != ERROR_IO_PENDING) { - diskann::cerr << "Error queuing IO -- " << error << "\n"; + // batch start/end + uint64_t batch_start = MAX_IO_DEPTH * i; + uint64_t batch_size = std::min((uint64_t)(n_reqs - batch_start), (uint64_t)MAX_IO_DEPTH); + + // fill OVERLAPPED and issue them + for (uint64_t j = 0; j < batch_size; j++) + { + AlignedRead &req = read_reqs[batch_start + j]; + OVERLAPPED &os = ctx.reqs[j]; + + uint64_t offset = req.offset; + uint64_t nbytes = req.len; + char *read_buf = (char *)req.buf; + assert(IS_ALIGNED(read_buf, SECTOR_LEN)); + assert(IS_ALIGNED(offset, SECTOR_LEN)); + assert(IS_ALIGNED(nbytes, SECTOR_LEN)); + + // fill in OVERLAPPED struct + os.Offset = offset & 0xffffffff; + os.OffsetHigh = (offset >> 32); + + BOOL ret = ReadFile(ctx.fhandle, read_buf, nbytes, NULL, &os); + if (ret == FALSE) + { + auto error = GetLastError(); + if (error != ERROR_IO_PENDING) + { + diskann::cerr << "Error queuing IO -- " << error << "\n"; + } + } + else + { + diskann::cerr << "Error queueing IO -- ReadFile returned TRUE" << std::endl; + } } - } else { - diskann::cerr << "Error queueing IO -- ReadFile returned TRUE" - << std::endl; - } - } - DWORD n_read = 0; - uint64_t n_complete = 0; - ULONG_PTR completion_key = 0; - OVERLAPPED *lp_os; - while (n_complete < batch_size) { - if (GetQueuedCompletionStatus(ctx.iocp, &n_read, &completion_key, &lp_os, - INFINITE) != 0) { - // successfully dequeued a completed I/O - n_complete++; - } else { - // failed to dequeue OR dequeued failed I/O - if (lp_os == NULL) { - DWORD error = GetLastError(); - if (error != WAIT_TIMEOUT) { - diskann::cerr << "GetQueuedCompletionStatus() failed " - "with error = " - << error << std::endl; - throw diskann::ANNException( - "GetQueuedCompletionStatus failed with error: ", error, - __FUNCSIG__, __FILE__, __LINE__); - } - // no completion packet dequeued ==> sleep for 5us and try - // again - std::this_thread::sleep_for(5us); - } else { - // completion packet for failed IO dequeued - auto op_idx = lp_os - ctx.reqs.data(); - std::stringstream stream; - stream << "I/O failed , offset: " << read_reqs[op_idx].offset - << "with error code: " << GetLastError() << std::endl; - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); + DWORD n_read = 0; + uint64_t n_complete = 0; + ULONG_PTR completion_key = 0; + OVERLAPPED *lp_os; + while (n_complete < batch_size) + { + if (GetQueuedCompletionStatus(ctx.iocp, &n_read, &completion_key, &lp_os, INFINITE) != 0) + { + // successfully dequeued a completed I/O + n_complete++; + } + else + { + // failed to dequeue OR dequeued failed I/O + if (lp_os == NULL) + { + DWORD error = GetLastError(); + if (error != WAIT_TIMEOUT) + { + diskann::cerr << "GetQueuedCompletionStatus() failed " + "with error = " + << error << std::endl; + throw diskann::ANNException("GetQueuedCompletionStatus failed with error: ", error, __FUNCSIG__, + __FILE__, __LINE__); + } + // no completion packet dequeued ==> sleep for 5us and try + // again + std::this_thread::sleep_for(5us); + } + else + { + // completion packet for failed IO dequeued + auto op_idx = lp_os - ctx.reqs.data(); + std::stringstream stream; + stream << "I/O failed , offset: " << read_reqs[op_idx].offset + << "with error code: " << GetLastError() << std::endl; + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } + } } - } } - } } #endif #endif diff --git a/tests/build_disk_index.cpp b/tests/build_disk_index.cpp index b60c71e8b..3d097458b 100644 --- a/tests/build_disk_index.cpp +++ b/tests/build_disk_index.cpp @@ -12,193 +12,178 @@ namespace po = boost::program_options; -int main(int argc, char **argv) { - std::string data_type, dist_fn, data_path, index_path_prefix, codebook_prefix, - label_file, universal_label, label_type; - uint32_t num_threads, R, L, disk_PQ, build_PQ, QD, Lf, filter_threshold; - float B, M; - bool append_reorder_data = false; - bool use_opq = false; +int main(int argc, char **argv) +{ + std::string data_type, dist_fn, data_path, index_path_prefix, codebook_prefix, label_file, universal_label, + label_type; + uint32_t num_threads, R, L, disk_PQ, build_PQ, QD, Lf, filter_threshold; + float B, M; + bool append_reorder_data = false; + bool use_opq = false; - po::options_description desc{"Arguments"}; - try { - desc.add_options()("help,h", "Print information on arguments"); - desc.add_options()("data_type", - po::value(&data_type)->required(), - "data type "); - desc.add_options()("dist_fn", po::value(&dist_fn)->required(), - "distance function "); - desc.add_options()("data_path", - po::value(&data_path)->required(), - "Input data file in bin format"); - desc.add_options()("index_path_prefix", - po::value(&index_path_prefix)->required(), - "Path prefix for saving index file components"); - desc.add_options()("max_degree,R", - po::value(&R)->default_value(64), - "Maximum graph degree"); - desc.add_options()( - "Lbuild,L", po::value(&L)->default_value(100), - "Build complexity, higher value results in better graphs"); - desc.add_options()("search_DRAM_budget,B", po::value(&B)->required(), - "DRAM budget in GB for searching the index to set the " - "compressed level for data while search happens"); - desc.add_options()("build_DRAM_budget,M", po::value(&M)->required(), - "DRAM budget in GB for building the index"); - desc.add_options()( - "num_threads,T", - po::value(&num_threads)->default_value(omp_get_num_procs()), - "Number of threads used for building index (defaults to " - "omp_get_num_procs())"); - desc.add_options()("QD", po::value(&QD)->default_value(0), - " Quantized Dimension for compression"); - desc.add_options()( - "codebook_prefix", - po::value(&codebook_prefix)->default_value(""), - "Path prefix for pre-trained codebook"); - desc.add_options()("PQ_disk_bytes", - po::value(&disk_PQ)->default_value(0), - "Number of bytes to which vectors should be compressed " - "on SSD; 0 for no compression"); - desc.add_options()("append_reorder_data", - po::bool_switch()->default_value(false), - "Include full precision data in the index. Use only in " - "conjuction with compressed data on SSD."); - desc.add_options()("build_PQ_bytes", - po::value(&build_PQ)->default_value(0), - "Number of PQ bytes to build the index; 0 for full " - "precision build"); - desc.add_options()("use_opq", po::bool_switch()->default_value(false), - "Use Optimized Product Quantization (OPQ)."); - desc.add_options()( - "label_file", po::value(&label_file)->default_value(""), - "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"); - desc.add_options()( - "universal_label", - po::value(&universal_label)->default_value(""), - "Universal label, Use only in conjuction 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 point instead of comma " - "separated filters for that point"); - desc.add_options()("FilteredLbuild,Lf", - po::value(&Lf)->default_value(0), - "Build complexity for filtered points, higher value " - "results in better graphs"); - desc.add_options()( - "filter_threshold,F", - po::value(&filter_threshold)->default_value(0), - "Threshold to break up the existing nodes to generate new graph " - "internally where each node has a maximum F labels."); - desc.add_options()( - "label_type", - po::value(&label_type)->default_value("uint"), - "Storage type of Labels , default value is uint which " - "will consume memory 4 bytes per filter"); + po::options_description desc{"Arguments"}; + try + { + desc.add_options()("help,h", "Print information on arguments"); + desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); + desc.add_options()("dist_fn", po::value(&dist_fn)->required(), "distance function "); + desc.add_options()("data_path", po::value(&data_path)->required(), + "Input data file in bin format"); + desc.add_options()("index_path_prefix", po::value(&index_path_prefix)->required(), + "Path prefix for saving index file components"); + desc.add_options()("max_degree,R", po::value(&R)->default_value(64), "Maximum graph degree"); + desc.add_options()("Lbuild,L", po::value(&L)->default_value(100), + "Build complexity, higher value results in better graphs"); + desc.add_options()("search_DRAM_budget,B", po::value(&B)->required(), + "DRAM budget in GB for searching the index to set the " + "compressed level for data while search happens"); + desc.add_options()("build_DRAM_budget,M", po::value(&M)->required(), + "DRAM budget in GB for building the index"); + desc.add_options()("num_threads,T", po::value(&num_threads)->default_value(omp_get_num_procs()), + "Number of threads used for building index (defaults to " + "omp_get_num_procs())"); + desc.add_options()("QD", po::value(&QD)->default_value(0), " Quantized Dimension for compression"); + desc.add_options()("codebook_prefix", po::value(&codebook_prefix)->default_value(""), + "Path prefix for pre-trained codebook"); + desc.add_options()("PQ_disk_bytes", po::value(&disk_PQ)->default_value(0), + "Number of bytes to which vectors should be compressed " + "on SSD; 0 for no compression"); + desc.add_options()("append_reorder_data", po::bool_switch()->default_value(false), + "Include full precision data in the index. Use only in " + "conjuction with compressed data on SSD."); + desc.add_options()("build_PQ_bytes", po::value(&build_PQ)->default_value(0), + "Number of PQ bytes to build the index; 0 for full " + "precision build"); + desc.add_options()("use_opq", po::bool_switch()->default_value(false), + "Use Optimized Product Quantization (OPQ)."); + desc.add_options()("label_file", po::value(&label_file)->default_value(""), + "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"); + desc.add_options()("universal_label", po::value(&universal_label)->default_value(""), + "Universal label, Use only in conjuction 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 point instead of comma " + "separated filters for that point"); + desc.add_options()("FilteredLbuild,Lf", po::value(&Lf)->default_value(0), + "Build complexity for filtered points, higher value " + "results in better graphs"); + desc.add_options()("filter_threshold,F", po::value(&filter_threshold)->default_value(0), + "Threshold to break up the existing nodes to generate new graph " + "internally where each node has a maximum F labels."); + desc.add_options()("label_type", po::value(&label_type)->default_value("uint"), + "Storage type of Labels , default value is uint which " + "will consume memory 4 bytes per filter"); - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) { - std::cout << desc; - return 0; + po::variables_map vm; + po::store(po::parse_command_line(argc, argv, desc), vm); + if (vm.count("help")) + { + std::cout << desc; + return 0; + } + po::notify(vm); + if (vm["append_reorder_data"].as()) + append_reorder_data = true; + if (vm["use_opq"].as()) + use_opq = true; + } + catch (const std::exception &ex) + { + std::cerr << ex.what() << '\n'; + return -1; } - po::notify(vm); - if (vm["append_reorder_data"].as()) append_reorder_data = true; - if (vm["use_opq"].as()) use_opq = true; - } catch (const std::exception &ex) { - std::cerr << ex.what() << '\n'; - return -1; - } - - bool use_filters = false; - if (label_file != "") { - use_filters = true; - } - diskann::Metric metric; - if (dist_fn == std::string("l2")) - metric = diskann::Metric::L2; - else if (dist_fn == std::string("mips")) - metric = diskann::Metric::INNER_PRODUCT; - else { - std::cout << "Error. Only l2 and mips distance functions are supported" - << std::endl; - return -1; - } + bool use_filters = false; + if (label_file != "") + { + use_filters = true; + } - if (append_reorder_data) { - if (disk_PQ == 0) { - std::cout << "Error: It is not necessary to append data for reordering " - "when vectors are not compressed on disk." - << std::endl; - return -1; + diskann::Metric metric; + if (dist_fn == std::string("l2")) + metric = diskann::Metric::L2; + else if (dist_fn == std::string("mips")) + metric = diskann::Metric::INNER_PRODUCT; + else + { + std::cout << "Error. Only l2 and mips distance functions are supported" << std::endl; + return -1; } - if (data_type != std::string("float")) { - std::cout << "Error: Appending data for reordering currently only " - "supported for float data type." - << std::endl; - return -1; + + if (append_reorder_data) + { + if (disk_PQ == 0) + { + std::cout << "Error: It is not necessary to append data for reordering " + "when vectors are not compressed on disk." + << std::endl; + return -1; + } + if (data_type != std::string("float")) + { + std::cout << "Error: Appending data for reordering currently only " + "supported for float data type." + << std::endl; + return -1; + } } - } - std::string params = std::string(std::to_string(R)) + " " + - std::string(std::to_string(L)) + " " + - std::string(std::to_string(B)) + " " + - std::string(std::to_string(M)) + " " + - std::string(std::to_string(num_threads)) + " " + - std::string(std::to_string(disk_PQ)) + " " + - std::string(std::to_string(append_reorder_data)) + " " + - std::string(std::to_string(build_PQ)) + " " + - std::string(std::to_string(QD)); + std::string params = std::string(std::to_string(R)) + " " + std::string(std::to_string(L)) + " " + + std::string(std::to_string(B)) + " " + std::string(std::to_string(M)) + " " + + std::string(std::to_string(num_threads)) + " " + std::string(std::to_string(disk_PQ)) + " " + + std::string(std::to_string(append_reorder_data)) + " " + + std::string(std::to_string(build_PQ)) + " " + std::string(std::to_string(QD)); - try { - if (label_file != "" && label_type == "ushort") { - if (data_type == std::string("int8")) - return diskann::build_disk_index( - data_path.c_str(), index_path_prefix.c_str(), params.c_str(), - metric, use_opq, codebook_prefix, use_filters, label_file, - universal_label, filter_threshold, Lf); - else if (data_type == std::string("uint8")) - return diskann::build_disk_index( - data_path.c_str(), index_path_prefix.c_str(), params.c_str(), - metric, use_opq, codebook_prefix, use_filters, label_file, - universal_label, filter_threshold, Lf); - else if (data_type == std::string("float")) - return diskann::build_disk_index( - data_path.c_str(), index_path_prefix.c_str(), params.c_str(), - metric, use_opq, codebook_prefix, use_filters, label_file, - universal_label, filter_threshold, Lf); - else { - diskann::cerr << "Error. Unsupported data type" << std::endl; - return -1; - } - } else { - if (data_type == std::string("int8")) - return diskann::build_disk_index( - data_path.c_str(), index_path_prefix.c_str(), params.c_str(), - metric, use_opq, codebook_prefix, use_filters, label_file, - universal_label, filter_threshold, Lf); - else if (data_type == std::string("uint8")) - return diskann::build_disk_index( - data_path.c_str(), index_path_prefix.c_str(), params.c_str(), - metric, use_opq, codebook_prefix, use_filters, label_file, - universal_label, filter_threshold, Lf); - else if (data_type == std::string("float")) - return diskann::build_disk_index( - data_path.c_str(), index_path_prefix.c_str(), params.c_str(), - metric, use_opq, codebook_prefix, use_filters, label_file, - universal_label, filter_threshold, Lf); - else { - diskann::cerr << "Error. Unsupported data type" << std::endl; + try + { + if (label_file != "" && label_type == "ushort") + { + if (data_type == std::string("int8")) + return diskann::build_disk_index(data_path.c_str(), index_path_prefix.c_str(), params.c_str(), + metric, use_opq, codebook_prefix, use_filters, label_file, + universal_label, filter_threshold, Lf); + else if (data_type == std::string("uint8")) + return diskann::build_disk_index( + data_path.c_str(), index_path_prefix.c_str(), params.c_str(), metric, use_opq, codebook_prefix, + use_filters, label_file, universal_label, filter_threshold, Lf); + else if (data_type == std::string("float")) + return diskann::build_disk_index( + data_path.c_str(), index_path_prefix.c_str(), params.c_str(), metric, use_opq, codebook_prefix, + use_filters, label_file, universal_label, filter_threshold, Lf); + else + { + diskann::cerr << "Error. Unsupported data type" << std::endl; + return -1; + } + } + else + { + if (data_type == std::string("int8")) + return diskann::build_disk_index(data_path.c_str(), index_path_prefix.c_str(), params.c_str(), + metric, use_opq, codebook_prefix, use_filters, label_file, + universal_label, filter_threshold, Lf); + else if (data_type == std::string("uint8")) + return diskann::build_disk_index(data_path.c_str(), index_path_prefix.c_str(), params.c_str(), + metric, use_opq, codebook_prefix, use_filters, label_file, + universal_label, filter_threshold, Lf); + else if (data_type == std::string("float")) + return diskann::build_disk_index(data_path.c_str(), index_path_prefix.c_str(), params.c_str(), + metric, use_opq, codebook_prefix, use_filters, label_file, + universal_label, filter_threshold, Lf); + else + { + diskann::cerr << "Error. Unsupported data type" << std::endl; + return -1; + } + } + } + catch (const std::exception &e) + { + std::cout << std::string(e.what()) << std::endl; + diskann::cerr << "Index build failed." << std::endl; return -1; - } } - } catch (const std::exception &e) { - std::cout << std::string(e.what()) << std::endl; - diskann::cerr << "Index build failed." << std::endl; - return -1; - } } diff --git a/tests/build_memory_index.cpp b/tests/build_memory_index.cpp index 5146afa60..3712350c3 100644 --- a/tests/build_memory_index.cpp +++ b/tests/build_memory_index.cpp @@ -21,196 +21,185 @@ namespace po = boost::program_options; template -int build_in_memory_index(const diskann::Metric &metric, - const std::string &data_path, const uint32_t R, - const uint32_t L, const float alpha, - const std::string &save_path, - const uint32_t num_threads, const bool use_pq_build, - const size_t num_pq_bytes, const bool use_opq, - const std::string &label_file, - const std::string &universal_label, - const uint32_t Lf) { - diskann::IndexWriteParameters paras = - diskann::IndexWriteParametersBuilder(L, R) - .with_filter_list_size(Lf) - .with_alpha(alpha) - .with_saturate_graph(false) - .with_num_threads(num_threads) - .build(); - std::string labels_file_to_use = save_path + "_label_formatted.txt"; - std::string mem_labels_int_map_file = save_path + "_labels_map.txt"; +int build_in_memory_index(const diskann::Metric &metric, const std::string &data_path, const uint32_t R, + const uint32_t L, const float alpha, const std::string &save_path, const uint32_t num_threads, + const bool use_pq_build, const size_t num_pq_bytes, const bool use_opq, + const std::string &label_file, const std::string &universal_label, const uint32_t Lf) +{ + diskann::IndexWriteParameters paras = diskann::IndexWriteParametersBuilder(L, R) + .with_filter_list_size(Lf) + .with_alpha(alpha) + .with_saturate_graph(false) + .with_num_threads(num_threads) + .build(); + std::string labels_file_to_use = save_path + "_label_formatted.txt"; + std::string mem_labels_int_map_file = save_path + "_labels_map.txt"; - size_t data_num, data_dim; - diskann::get_bin_metadata(data_path, data_num, data_dim); + size_t data_num, data_dim; + diskann::get_bin_metadata(data_path, data_num, data_dim); - diskann::Index index(metric, data_dim, data_num, false, - false, false, use_pq_build, - num_pq_bytes, use_opq); - auto s = std::chrono::high_resolution_clock::now(); - if (label_file == "") { - index.build(data_path.c_str(), data_num, paras); - } else { - convert_labels_string_to_int(label_file, labels_file_to_use, - mem_labels_int_map_file, universal_label); - if (universal_label != "") { - LabelT unv_label_as_num = 0; - index.set_universal_label(unv_label_as_num); + diskann::Index index(metric, data_dim, data_num, false, false, false, use_pq_build, num_pq_bytes, + use_opq); + auto s = std::chrono::high_resolution_clock::now(); + if (label_file == "") + { + index.build(data_path.c_str(), data_num, paras); } - index.build_filtered_index(data_path.c_str(), labels_file_to_use, data_num, - paras); - } - std::chrono::duration diff = - std::chrono::high_resolution_clock::now() - s; + else + { + convert_labels_string_to_int(label_file, labels_file_to_use, mem_labels_int_map_file, universal_label); + if (universal_label != "") + { + LabelT unv_label_as_num = 0; + index.set_universal_label(unv_label_as_num); + } + index.build_filtered_index(data_path.c_str(), labels_file_to_use, data_num, paras); + } + std::chrono::duration diff = std::chrono::high_resolution_clock::now() - s; - std::cout << "Indexing time: " << diff.count() << "\n"; - index.save(save_path.c_str()); - if (label_file != "") std::remove(labels_file_to_use.c_str()); - return 0; + std::cout << "Indexing time: " << diff.count() << "\n"; + index.save(save_path.c_str()); + if (label_file != "") + std::remove(labels_file_to_use.c_str()); + return 0; } -int main(int argc, char **argv) { - std::string data_type, dist_fn, data_path, index_path_prefix, label_file, - universal_label, label_type; - uint32_t num_threads, R, L, Lf, build_PQ_bytes; - float alpha; - bool use_pq_build, use_opq; +int main(int argc, char **argv) +{ + std::string data_type, dist_fn, data_path, index_path_prefix, label_file, universal_label, label_type; + uint32_t num_threads, R, L, Lf, build_PQ_bytes; + float alpha; + bool use_pq_build, use_opq; - po::options_description desc{"Arguments"}; - try { - desc.add_options()("help,h", "Print information on arguments"); - desc.add_options()("data_type", - po::value(&data_type)->required(), - "data type "); - desc.add_options()("dist_fn", po::value(&dist_fn)->required(), - "distance function "); - desc.add_options()("data_path", - po::value(&data_path)->required(), - "Input data file in bin format"); - desc.add_options()("index_path_prefix", - po::value(&index_path_prefix)->required(), - "Path prefix for saving index file components"); - desc.add_options()("max_degree,R", - po::value(&R)->default_value(64), - "Maximum graph degree"); - desc.add_options()( - "Lbuild,L", po::value(&L)->default_value(100), - "Build complexity, higher value results in better graphs"); - desc.add_options()("alpha", po::value(&alpha)->default_value(1.2f), - "alpha controls density and diameter of graph, set " - "1 for sparse graph, " - "1.2 or 1.4 for denser graphs with lower diameter"); - desc.add_options()( - "num_threads,T", - po::value(&num_threads)->default_value(omp_get_num_procs()), - "Number of threads used for building index (defaults to " - "omp_get_num_procs())"); - desc.add_options()( - "build_PQ_bytes", - po::value(&build_PQ_bytes)->default_value(0), - "Number of PQ bytes to build the index; 0 for full precision " - "build"); - desc.add_options()("use_opq", po::bool_switch()->default_value(false), - "Set true for OPQ compression while using PQ " - "distance comparisons for " - "building the index, and false for PQ compression"); - desc.add_options()( - "label_file", po::value(&label_file)->default_value(""), - "Input label file in txt format for Filtered Index search. " - "The file should contain comma separated filters for each node " - "with each line corresponding to a graph node"); - desc.add_options()( - "universal_label", - po::value(&universal_label)->default_value(""), - "Universal label, if using it, only in conjunction with " - "labels_file"); - desc.add_options()("FilteredLbuild,Lf", - po::value(&Lf)->default_value(0), - "Build complexity for filtered points, higher value " - "results in better graphs"); - desc.add_options()( - "label_type", - po::value(&label_type)->default_value("uint"), - "Storage type of Labels , default value is uint which " - "will consume memory 4 bytes per filter"); + po::options_description desc{"Arguments"}; + try + { + desc.add_options()("help,h", "Print information on arguments"); + desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); + desc.add_options()("dist_fn", po::value(&dist_fn)->required(), + "distance function "); + desc.add_options()("data_path", po::value(&data_path)->required(), + "Input data file in bin format"); + desc.add_options()("index_path_prefix", po::value(&index_path_prefix)->required(), + "Path prefix for saving index file components"); + desc.add_options()("max_degree,R", po::value(&R)->default_value(64), "Maximum graph degree"); + desc.add_options()("Lbuild,L", po::value(&L)->default_value(100), + "Build complexity, higher value results in better graphs"); + desc.add_options()("alpha", po::value(&alpha)->default_value(1.2f), + "alpha controls density and diameter of graph, set " + "1 for sparse graph, " + "1.2 or 1.4 for denser graphs with lower diameter"); + desc.add_options()("num_threads,T", po::value(&num_threads)->default_value(omp_get_num_procs()), + "Number of threads used for building index (defaults to " + "omp_get_num_procs())"); + desc.add_options()("build_PQ_bytes", po::value(&build_PQ_bytes)->default_value(0), + "Number of PQ bytes to build the index; 0 for full precision " + "build"); + desc.add_options()("use_opq", po::bool_switch()->default_value(false), + "Set true for OPQ compression while using PQ " + "distance comparisons for " + "building the index, and false for PQ compression"); + desc.add_options()("label_file", po::value(&label_file)->default_value(""), + "Input label file in txt format for Filtered Index search. " + "The file should contain comma separated filters for each node " + "with each line corresponding to a graph node"); + desc.add_options()("universal_label", po::value(&universal_label)->default_value(""), + "Universal label, if using it, only in conjunction with " + "labels_file"); + desc.add_options()("FilteredLbuild,Lf", po::value(&Lf)->default_value(0), + "Build complexity for filtered points, higher value " + "results in better graphs"); + desc.add_options()("label_type", po::value(&label_type)->default_value("uint"), + "Storage type of Labels , default value is uint which " + "will consume memory 4 bytes per filter"); - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) { - std::cout << desc; - return 0; + po::variables_map vm; + po::store(po::parse_command_line(argc, argv, desc), vm); + if (vm.count("help")) + { + std::cout << desc; + return 0; + } + po::notify(vm); + use_pq_build = (build_PQ_bytes > 0); + use_opq = vm["use_opq"].as(); + } + catch (const std::exception &ex) + { + std::cerr << ex.what() << '\n'; + return -1; } - po::notify(vm); - use_pq_build = (build_PQ_bytes > 0); - use_opq = vm["use_opq"].as(); - } catch (const std::exception &ex) { - std::cerr << ex.what() << '\n'; - return -1; - } - - diskann::Metric metric; - if (dist_fn == std::string("mips")) { - metric = diskann::Metric::INNER_PRODUCT; - } else if (dist_fn == std::string("l2")) { - metric = diskann::Metric::L2; - } else if (dist_fn == std::string("cosine")) { - metric = diskann::Metric::COSINE; - } else { - std::cout << "Unsupported distance function. Currently only L2/ Inner " - "Product/Cosine are supported." - << std::endl; - return -1; - } - try { - diskann::cout << "Starting index build with R: " << R << " Lbuild: " << L - << " alpha: " << alpha << " #threads: " << num_threads - << std::endl; - if (label_file != "" && label_type == "ushort") { - if (data_type == std::string("int8")) - return build_in_memory_index( - metric, data_path, R, L, alpha, index_path_prefix, num_threads, - use_pq_build, build_PQ_bytes, use_opq, label_file, universal_label, - Lf); - else if (data_type == std::string("uint8")) - return build_in_memory_index( - metric, data_path, R, L, alpha, index_path_prefix, num_threads, - use_pq_build, build_PQ_bytes, use_opq, label_file, universal_label, - Lf); - else if (data_type == std::string("float")) - return build_in_memory_index( - metric, data_path, R, L, alpha, index_path_prefix, num_threads, - use_pq_build, build_PQ_bytes, use_opq, label_file, universal_label, - Lf); - else { - std::cout << "Unsupported type. Use one of int8, uint8 or float." + diskann::Metric metric; + if (dist_fn == std::string("mips")) + { + metric = diskann::Metric::INNER_PRODUCT; + } + else if (dist_fn == std::string("l2")) + { + metric = diskann::Metric::L2; + } + else if (dist_fn == std::string("cosine")) + { + metric = diskann::Metric::COSINE; + } + else + { + std::cout << "Unsupported distance function. Currently only L2/ Inner " + "Product/Cosine are supported." << std::endl; return -1; - } - } else { - if (data_type == std::string("int8")) - return build_in_memory_index( - metric, data_path, R, L, alpha, index_path_prefix, num_threads, - use_pq_build, build_PQ_bytes, use_opq, label_file, universal_label, - Lf); - else if (data_type == std::string("uint8")) - return build_in_memory_index( - metric, data_path, R, L, alpha, index_path_prefix, num_threads, - use_pq_build, build_PQ_bytes, use_opq, label_file, universal_label, - Lf); - else if (data_type == std::string("float")) - return build_in_memory_index( - metric, data_path, R, L, alpha, index_path_prefix, num_threads, - use_pq_build, build_PQ_bytes, use_opq, label_file, universal_label, - Lf); - else { - std::cout << "Unsupported type. Use one of int8, uint8 or float." - << std::endl; + } + + try + { + diskann::cout << "Starting index build with R: " << R << " Lbuild: " << L << " alpha: " << alpha + << " #threads: " << num_threads << std::endl; + if (label_file != "" && label_type == "ushort") + { + if (data_type == std::string("int8")) + return build_in_memory_index( + metric, data_path, R, L, alpha, index_path_prefix, num_threads, use_pq_build, build_PQ_bytes, + use_opq, label_file, universal_label, Lf); + else if (data_type == std::string("uint8")) + return build_in_memory_index( + metric, data_path, R, L, alpha, index_path_prefix, num_threads, use_pq_build, build_PQ_bytes, + use_opq, label_file, universal_label, Lf); + else if (data_type == std::string("float")) + return build_in_memory_index( + metric, data_path, R, L, alpha, index_path_prefix, num_threads, use_pq_build, build_PQ_bytes, + use_opq, label_file, universal_label, Lf); + else + { + std::cout << "Unsupported type. Use one of int8, uint8 or float." << std::endl; + return -1; + } + } + else + { + if (data_type == std::string("int8")) + return build_in_memory_index(metric, data_path, R, L, alpha, index_path_prefix, num_threads, + use_pq_build, build_PQ_bytes, use_opq, label_file, universal_label, + Lf); + else if (data_type == std::string("uint8")) + return build_in_memory_index(metric, data_path, R, L, alpha, index_path_prefix, num_threads, + use_pq_build, build_PQ_bytes, use_opq, label_file, + universal_label, Lf); + else if (data_type == std::string("float")) + return build_in_memory_index(metric, data_path, R, L, alpha, index_path_prefix, num_threads, + use_pq_build, build_PQ_bytes, use_opq, label_file, universal_label, + Lf); + else + { + std::cout << "Unsupported type. Use one of int8, uint8 or float." << std::endl; + return -1; + } + } + } + catch (const std::exception &e) + { + std::cout << std::string(e.what()) << std::endl; + diskann::cerr << "Index build failed." << std::endl; return -1; - } } - } catch (const std::exception &e) { - std::cout << std::string(e.what()) << std::endl; - diskann::cerr << "Index build failed." << std::endl; - return -1; - } } diff --git a/tests/build_stitched_index.cpp b/tests/build_stitched_index.cpp index c19ea6da1..770a11f83 100644 --- a/tests/build_stitched_index.cpp +++ b/tests/build_stitched_index.cpp @@ -20,28 +20,29 @@ #include "utils.h" namespace po = boost::program_options; -typedef std::tuple>, uint64_t> - stitch_indices_return_values; +typedef std::tuple>, uint64_t> stitch_indices_return_values; /* * Inline function to display progress bar. */ -inline void print_progress(double percentage) { - int val = (int)(percentage * 100); - int lpad = (int)(percentage * PBWIDTH); - int rpad = PBWIDTH - lpad; - printf("\r%3d%% [%.*s%*s]", val, lpad, PBSTR, rpad, ""); - fflush(stdout); +inline void print_progress(double percentage) +{ + int val = (int)(percentage * 100); + int lpad = (int)(percentage * PBWIDTH); + int rpad = PBWIDTH - lpad; + printf("\r%3d%% [%.*s%*s]", val, lpad, PBSTR, rpad, ""); + fflush(stdout); } /* * Inline function to generate a random integer in a range. */ -inline size_t random(size_t range_from, size_t range_to) { - std::random_device rand_dev; - std::mt19937 generator(rand_dev()); - std::uniform_int_distribution distr(range_from, range_to); - return distr(generator); +inline size_t random(size_t range_from, size_t range_to) +{ + std::random_device rand_dev; + std::mt19937 generator(rand_dev()); + std::uniform_int_distribution distr(range_from, range_to); + return distr(generator); } /* @@ -49,63 +50,52 @@ inline size_t random(size_t range_from, size_t range_to) { * * Arguments are merely the inputs from the command line. */ -void handle_args(int argc, char **argv, std::string &data_type, - path &input_data_path, path &final_index_path_prefix, - path &label_data_path, std::string &universal_label, - uint32_t &num_threads, uint32_t &R, uint32_t &L, - uint32_t &stitched_R, float &alpha) { - po::options_description desc{"Arguments"}; - try { - desc.add_options()("help,h", "Print information on arguments"); - desc.add_options()("data_type", - po::value(&data_type)->required(), - "data type "); - desc.add_options()("data_path", - po::value(&input_data_path)->required(), - "Input data file in bin format"); - desc.add_options()("index_path_prefix", - po::value(&final_index_path_prefix)->required(), - "Path prefix for saving index file components"); - desc.add_options()("max_degree,R", - po::value(&R)->default_value(64), - "Maximum graph degree"); - desc.add_options()( - "Lbuild,L", po::value(&L)->default_value(100), - "Build complexity, higher value results in better graphs"); - desc.add_options()("stitched_R", - po::value(&stitched_R)->default_value(100), - "Degree to prune final graph down to"); - desc.add_options()("alpha", po::value(&alpha)->default_value(1.2f), - "alpha controls density and diameter of graph, set " - "1 for sparse graph, " - "1.2 or 1.4 for denser graphs with lower diameter"); - desc.add_options()( - "num_threads,T", - po::value(&num_threads)->default_value(omp_get_num_procs()), - "Number of threads used for building index (defaults to " - "omp_get_num_procs())"); - desc.add_options()("label_file", - po::value(&label_data_path)->default_value(""), - "Input label file in txt format if present"); - desc.add_options()( - "universal_label", - po::value(&universal_label)->default_value(""), - "If a point comes with the specified universal label (and only the " - "univ. " - "label), then the point is considered to have every possible " - "label"); - - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) { - std::cout << desc; - exit(0); +void handle_args(int argc, char **argv, std::string &data_type, path &input_data_path, path &final_index_path_prefix, + path &label_data_path, std::string &universal_label, uint32_t &num_threads, uint32_t &R, uint32_t &L, + uint32_t &stitched_R, float &alpha) +{ + po::options_description desc{"Arguments"}; + try + { + desc.add_options()("help,h", "Print information on arguments"); + desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); + desc.add_options()("data_path", po::value(&input_data_path)->required(), "Input data file in bin format"); + desc.add_options()("index_path_prefix", po::value(&final_index_path_prefix)->required(), + "Path prefix for saving index file components"); + desc.add_options()("max_degree,R", po::value(&R)->default_value(64), "Maximum graph degree"); + desc.add_options()("Lbuild,L", po::value(&L)->default_value(100), + "Build complexity, higher value results in better graphs"); + desc.add_options()("stitched_R", po::value(&stitched_R)->default_value(100), + "Degree to prune final graph down to"); + desc.add_options()("alpha", po::value(&alpha)->default_value(1.2f), + "alpha controls density and diameter of graph, set " + "1 for sparse graph, " + "1.2 or 1.4 for denser graphs with lower diameter"); + desc.add_options()("num_threads,T", po::value(&num_threads)->default_value(omp_get_num_procs()), + "Number of threads used for building index (defaults to " + "omp_get_num_procs())"); + desc.add_options()("label_file", po::value(&label_data_path)->default_value(""), + "Input label file in txt format if present"); + desc.add_options()("universal_label", po::value(&universal_label)->default_value(""), + "If a point comes with the specified universal label (and only the " + "univ. " + "label), then the point is considered to have every possible " + "label"); + + po::variables_map vm; + po::store(po::parse_command_line(argc, argv, desc), vm); + if (vm.count("help")) + { + std::cout << desc; + exit(0); + } + po::notify(vm); + } + catch (const std::exception &ex) + { + std::cerr << ex.what() << '\n'; + throw; } - po::notify(vm); - } catch (const std::exception &ex) { - std::cerr << ex.what() << '\n'; - throw; - } } /* @@ -116,109 +106,98 @@ void handle_args(int argc, char **argv, std::string &data_type, * 3. data (redundant for static indices) * 4. labels (redundant for static indices) */ -void save_full_index(path final_index_path_prefix, path input_data_path, - uint64_t final_index_size, +void save_full_index(path final_index_path_prefix, path input_data_path, uint64_t final_index_size, std::vector> stitched_graph, - tsl::robin_map entry_points, - std::string universal_label, path label_data_path) { - // aux. file 1 - auto saving_index_timer = std::chrono::high_resolution_clock::now(); - std::ifstream original_label_data_stream; - original_label_data_stream.exceptions(std::ios::badbit | std::ios::failbit); - original_label_data_stream.open(label_data_path, std::ios::binary); - std::ofstream new_label_data_stream; - new_label_data_stream.exceptions(std::ios::badbit | std::ios::failbit); - new_label_data_stream.open(final_index_path_prefix + "_labels.txt", - std::ios::binary); - new_label_data_stream << original_label_data_stream.rdbuf(); - original_label_data_stream.close(); - new_label_data_stream.close(); - - // aux. file 2 - std::ifstream original_input_data_stream; - original_input_data_stream.exceptions(std::ios::badbit | std::ios::failbit); - original_input_data_stream.open(input_data_path, std::ios::binary); - std::ofstream new_input_data_stream; - new_input_data_stream.exceptions(std::ios::badbit | std::ios::failbit); - new_input_data_stream.open(final_index_path_prefix + ".data", - std::ios::binary); - new_input_data_stream << original_input_data_stream.rdbuf(); - original_input_data_stream.close(); - new_input_data_stream.close(); - - // aux. file 3 - std::ofstream labels_to_medoids_writer; - labels_to_medoids_writer.exceptions(std::ios::badbit | std::ios::failbit); - labels_to_medoids_writer.open(final_index_path_prefix + - "_labels_to_medoids.txt"); - for (auto iter : entry_points) - labels_to_medoids_writer << iter.first << ", " << iter.second << std::endl; - labels_to_medoids_writer.close(); - - // aux. file 4 (only if we're using a universal label) - if (universal_label != "") { - std::ofstream universal_label_writer; - universal_label_writer.exceptions(std::ios::badbit | std::ios::failbit); - universal_label_writer.open(final_index_path_prefix + - "_universal_label.txt"); - universal_label_writer << universal_label << std::endl; - universal_label_writer.close(); - } - - // main index - uint64_t index_num_frozen_points = 0, index_num_edges = 0; - uint32_t index_max_observed_degree = 0, index_entry_point = 0; - const size_t METADATA = 2 * sizeof(uint64_t) + 2 * sizeof(uint32_t); - for (auto &point_neighbors : stitched_graph) { - index_max_observed_degree = - std::max(index_max_observed_degree, (uint32_t)point_neighbors.size()); - } - - std::ofstream stitched_graph_writer; - stitched_graph_writer.exceptions(std::ios::badbit | std::ios::failbit); - stitched_graph_writer.open(final_index_path_prefix, std::ios_base::binary); - - stitched_graph_writer.write((char *)&final_index_size, sizeof(uint64_t)); - stitched_graph_writer.write((char *)&index_max_observed_degree, - sizeof(uint32_t)); - stitched_graph_writer.write((char *)&index_entry_point, sizeof(uint32_t)); - stitched_graph_writer.write((char *)&index_num_frozen_points, - sizeof(uint64_t)); - - size_t bytes_written = METADATA; - for (uint32_t node_point = 0; node_point < stitched_graph.size(); - node_point++) { - uint32_t current_node_num_neighbors = stitched_graph[node_point].size(); - std::vector current_node_neighbors = stitched_graph[node_point]; - stitched_graph_writer.write((char *)¤t_node_num_neighbors, - sizeof(uint32_t)); - bytes_written += sizeof(uint32_t); - for (const auto ¤t_node_neighbor : current_node_neighbors) { - stitched_graph_writer.write((char *)¤t_node_neighbor, - sizeof(uint32_t)); - bytes_written += sizeof(uint32_t); + tsl::robin_map entry_points, std::string universal_label, + path label_data_path) +{ + // aux. file 1 + auto saving_index_timer = std::chrono::high_resolution_clock::now(); + std::ifstream original_label_data_stream; + original_label_data_stream.exceptions(std::ios::badbit | std::ios::failbit); + original_label_data_stream.open(label_data_path, std::ios::binary); + std::ofstream new_label_data_stream; + new_label_data_stream.exceptions(std::ios::badbit | std::ios::failbit); + new_label_data_stream.open(final_index_path_prefix + "_labels.txt", std::ios::binary); + new_label_data_stream << original_label_data_stream.rdbuf(); + original_label_data_stream.close(); + new_label_data_stream.close(); + + // aux. file 2 + std::ifstream original_input_data_stream; + original_input_data_stream.exceptions(std::ios::badbit | std::ios::failbit); + original_input_data_stream.open(input_data_path, std::ios::binary); + std::ofstream new_input_data_stream; + new_input_data_stream.exceptions(std::ios::badbit | std::ios::failbit); + new_input_data_stream.open(final_index_path_prefix + ".data", std::ios::binary); + new_input_data_stream << original_input_data_stream.rdbuf(); + original_input_data_stream.close(); + new_input_data_stream.close(); + + // aux. file 3 + std::ofstream labels_to_medoids_writer; + labels_to_medoids_writer.exceptions(std::ios::badbit | std::ios::failbit); + labels_to_medoids_writer.open(final_index_path_prefix + "_labels_to_medoids.txt"); + for (auto iter : entry_points) + labels_to_medoids_writer << iter.first << ", " << iter.second << std::endl; + labels_to_medoids_writer.close(); + + // aux. file 4 (only if we're using a universal label) + if (universal_label != "") + { + std::ofstream universal_label_writer; + universal_label_writer.exceptions(std::ios::badbit | std::ios::failbit); + universal_label_writer.open(final_index_path_prefix + "_universal_label.txt"); + universal_label_writer << universal_label << std::endl; + universal_label_writer.close(); + } + + // main index + uint64_t index_num_frozen_points = 0, index_num_edges = 0; + uint32_t index_max_observed_degree = 0, index_entry_point = 0; + const size_t METADATA = 2 * sizeof(uint64_t) + 2 * sizeof(uint32_t); + for (auto &point_neighbors : stitched_graph) + { + index_max_observed_degree = std::max(index_max_observed_degree, (uint32_t)point_neighbors.size()); } - index_num_edges += current_node_num_neighbors; - } - if (bytes_written != final_index_size) { - std::cerr << "Error: written bytes does not match allocated space" + std::ofstream stitched_graph_writer; + stitched_graph_writer.exceptions(std::ios::badbit | std::ios::failbit); + stitched_graph_writer.open(final_index_path_prefix, std::ios_base::binary); + + stitched_graph_writer.write((char *)&final_index_size, sizeof(uint64_t)); + stitched_graph_writer.write((char *)&index_max_observed_degree, sizeof(uint32_t)); + stitched_graph_writer.write((char *)&index_entry_point, sizeof(uint32_t)); + stitched_graph_writer.write((char *)&index_num_frozen_points, sizeof(uint64_t)); + + size_t bytes_written = METADATA; + for (uint32_t node_point = 0; node_point < stitched_graph.size(); node_point++) + { + uint32_t current_node_num_neighbors = stitched_graph[node_point].size(); + std::vector current_node_neighbors = stitched_graph[node_point]; + stitched_graph_writer.write((char *)¤t_node_num_neighbors, sizeof(uint32_t)); + bytes_written += sizeof(uint32_t); + for (const auto ¤t_node_neighbor : current_node_neighbors) + { + stitched_graph_writer.write((char *)¤t_node_neighbor, sizeof(uint32_t)); + bytes_written += sizeof(uint32_t); + } + index_num_edges += current_node_num_neighbors; + } + + if (bytes_written != final_index_size) + { + std::cerr << "Error: written bytes does not match allocated space" << std::endl; + throw; + } + + stitched_graph_writer.close(); + + std::chrono::duration saving_index_time = std::chrono::high_resolution_clock::now() - saving_index_timer; + std::cout << "Stitched graph written in " << saving_index_time.count() << " seconds" << std::endl; + std::cout << "Stitched graph average degree: " << ((float)index_num_edges) / ((float)(stitched_graph.size())) << std::endl; - throw; - } - - stitched_graph_writer.close(); - - std::chrono::duration saving_index_time = - std::chrono::high_resolution_clock::now() - saving_index_timer; - std::cout << "Stitched graph written in " << saving_index_time.count() - << " seconds" << std::endl; - std::cout << "Stitched graph average degree: " - << ((float)index_num_edges) / ((float)(stitched_graph.size())) - << std::endl; - std::cout << "Stitched graph max degree: " << index_max_observed_degree - << std::endl - << std::endl; + std::cout << "Stitched graph max degree: " << index_max_observed_degree << std::endl << std::endl; } /* @@ -229,55 +208,52 @@ void save_full_index(path final_index_path_prefix, path input_data_path, */ template stitch_indices_return_values stitch_label_indices( - path final_index_path_prefix, uint32_t total_number_of_points, - label_set all_labels, + path final_index_path_prefix, uint32_t total_number_of_points, label_set all_labels, tsl::robin_map labels_to_number_of_points, tsl::robin_map &label_entry_points, - tsl::robin_map> - label_id_to_orig_id_map) { - size_t final_index_size = 0; - std::vector> stitched_graph(total_number_of_points); - - auto stitching_index_timer = std::chrono::high_resolution_clock::now(); - for (const auto &lbl : all_labels) { - path curr_label_index_path(final_index_path_prefix + "_" + lbl); - std::vector> curr_label_index; - uint64_t curr_label_index_size; - uint32_t curr_label_entry_point; - - std::tie(curr_label_index, curr_label_index_size) = - diskann::load_label_index(curr_label_index_path, - labels_to_number_of_points[lbl]); - curr_label_entry_point = random(0, curr_label_index.size()); - label_entry_points[lbl] = - label_id_to_orig_id_map[lbl][curr_label_entry_point]; - - for (uint32_t node_point = 0; node_point < curr_label_index.size(); - node_point++) { - uint32_t original_point_id = label_id_to_orig_id_map[lbl][node_point]; - for (auto &node_neighbor : curr_label_index[node_point]) { - uint32_t original_neighbor_id = - label_id_to_orig_id_map[lbl][node_neighbor]; - std::vector curr_point_neighbors = - stitched_graph[original_point_id]; - if (std::find(curr_point_neighbors.begin(), curr_point_neighbors.end(), - original_neighbor_id) == curr_point_neighbors.end()) { - stitched_graph[original_point_id].push_back(original_neighbor_id); - final_index_size += sizeof(uint32_t); + tsl::robin_map> label_id_to_orig_id_map) +{ + size_t final_index_size = 0; + std::vector> stitched_graph(total_number_of_points); + + auto stitching_index_timer = std::chrono::high_resolution_clock::now(); + for (const auto &lbl : all_labels) + { + path curr_label_index_path(final_index_path_prefix + "_" + lbl); + std::vector> curr_label_index; + uint64_t curr_label_index_size; + uint32_t curr_label_entry_point; + + std::tie(curr_label_index, curr_label_index_size) = + diskann::load_label_index(curr_label_index_path, labels_to_number_of_points[lbl]); + curr_label_entry_point = random(0, curr_label_index.size()); + label_entry_points[lbl] = label_id_to_orig_id_map[lbl][curr_label_entry_point]; + + for (uint32_t node_point = 0; node_point < curr_label_index.size(); node_point++) + { + uint32_t original_point_id = label_id_to_orig_id_map[lbl][node_point]; + for (auto &node_neighbor : curr_label_index[node_point]) + { + uint32_t original_neighbor_id = label_id_to_orig_id_map[lbl][node_neighbor]; + std::vector curr_point_neighbors = stitched_graph[original_point_id]; + if (std::find(curr_point_neighbors.begin(), curr_point_neighbors.end(), original_neighbor_id) == + curr_point_neighbors.end()) + { + stitched_graph[original_point_id].push_back(original_neighbor_id); + final_index_size += sizeof(uint32_t); + } + } } - } } - } - const size_t METADATA = 2 * sizeof(uint64_t) + 2 * sizeof(uint32_t); - final_index_size += (total_number_of_points * sizeof(uint32_t) + METADATA); + const size_t METADATA = 2 * sizeof(uint64_t) + 2 * sizeof(uint32_t); + final_index_size += (total_number_of_points * sizeof(uint32_t) + METADATA); - std::chrono::duration stitching_index_time = - std::chrono::high_resolution_clock::now() - stitching_index_timer; - std::cout << "stitched graph generated in memory in " - << stitching_index_time.count() << " seconds" << std::endl; + std::chrono::duration stitching_index_time = + std::chrono::high_resolution_clock::now() - stitching_index_timer; + std::cout << "stitched graph generated in memory in " << stitching_index_time.count() << " seconds" << std::endl; - return std::make_tuple(stitched_graph, final_index_size); + return std::make_tuple(stitched_graph, final_index_size); } /* @@ -288,37 +264,31 @@ stitch_indices_return_values stitch_label_indices( * and pruned graph. */ template -void prune_and_save(path final_index_path_prefix, path full_index_path_prefix, - path input_data_path, - std::vector> stitched_graph, - uint32_t stitched_R, - tsl::robin_map label_entry_points, - std::string universal_label, path label_data_path, - uint32_t num_threads) { - size_t dimension, number_of_label_points; - auto diskann_cout_buffer = diskann::cout.rdbuf(nullptr); - auto std_cout_buffer = std::cout.rdbuf(nullptr); - auto pruning_index_timer = std::chrono::high_resolution_clock::now(); - - diskann::get_bin_metadata(input_data_path, number_of_label_points, dimension); - diskann::Index index(diskann::Metric::L2, dimension, - number_of_label_points, false, false); - - // not searching this index, set search_l to 0 - index.load(full_index_path_prefix.c_str(), num_threads, 1); - - std::cout << "parsing labels" << std::endl; - - index.prune_all_neighbors(stitched_R, 750, 1.2); - index.save((final_index_path_prefix).c_str()); - - diskann::cout.rdbuf(diskann_cout_buffer); - std::cout.rdbuf(std_cout_buffer); - std::chrono::duration pruning_index_time = - std::chrono::high_resolution_clock::now() - pruning_index_timer; - std::cout << "pruning performed in " << pruning_index_time.count() - << " seconds\n" - << std::endl; +void prune_and_save(path final_index_path_prefix, path full_index_path_prefix, path input_data_path, + std::vector> stitched_graph, uint32_t stitched_R, + tsl::robin_map label_entry_points, std::string universal_label, + path label_data_path, uint32_t num_threads) +{ + size_t dimension, number_of_label_points; + auto diskann_cout_buffer = diskann::cout.rdbuf(nullptr); + auto std_cout_buffer = std::cout.rdbuf(nullptr); + auto pruning_index_timer = std::chrono::high_resolution_clock::now(); + + diskann::get_bin_metadata(input_data_path, number_of_label_points, dimension); + diskann::Index index(diskann::Metric::L2, dimension, number_of_label_points, false, false); + + // not searching this index, set search_l to 0 + index.load(full_index_path_prefix.c_str(), num_threads, 1); + + std::cout << "parsing labels" << std::endl; + + index.prune_all_neighbors(stitched_R, 750, 1.2); + index.save((final_index_path_prefix).c_str()); + + diskann::cout.rdbuf(diskann_cout_buffer); + std::cout.rdbuf(std_cout_buffer); + std::chrono::duration pruning_index_time = std::chrono::high_resolution_clock::now() - pruning_index_timer; + std::cout << "pruning performed in " << pruning_index_time.count() << " seconds\n" << std::endl; } /* @@ -329,157 +299,131 @@ void prune_and_save(path final_index_path_prefix, path full_index_path_prefix, * 2. the separate diskANN indices built for each label * 3. the '.data' file created while generating the indices */ -void clean_up_artifacts(path input_data_path, path final_index_path_prefix, - label_set all_labels) { - for (const auto &lbl : all_labels) { - path curr_label_input_data_path(input_data_path + "_" + lbl); - path curr_label_index_path(final_index_path_prefix + "_" + lbl); - path curr_label_index_path_data(curr_label_index_path + ".data"); - - if (std::remove(curr_label_index_path.c_str()) != 0) throw; - if (std::remove(curr_label_input_data_path.c_str()) != 0) throw; - if (std::remove(curr_label_index_path_data.c_str()) != 0) throw; - } +void clean_up_artifacts(path input_data_path, path final_index_path_prefix, label_set all_labels) +{ + for (const auto &lbl : all_labels) + { + path curr_label_input_data_path(input_data_path + "_" + lbl); + path curr_label_index_path(final_index_path_prefix + "_" + lbl); + path curr_label_index_path_data(curr_label_index_path + ".data"); + + if (std::remove(curr_label_index_path.c_str()) != 0) + throw; + if (std::remove(curr_label_input_data_path.c_str()) != 0) + throw; + if (std::remove(curr_label_index_path_data.c_str()) != 0) + throw; + } } -int main(int argc, char **argv) { - // 1. handle cmdline inputs - std::string data_type; - path input_data_path, final_index_path_prefix, label_data_path; - std::string universal_label; - uint32_t num_threads, R, L, stitched_R; - float alpha; +int main(int argc, char **argv) +{ + // 1. handle cmdline inputs + std::string data_type; + path input_data_path, final_index_path_prefix, label_data_path; + std::string universal_label; + uint32_t num_threads, R, L, stitched_R; + float alpha; - auto index_timer = std::chrono::high_resolution_clock::now(); - handle_args(argc, argv, data_type, input_data_path, final_index_path_prefix, - label_data_path, universal_label, num_threads, R, L, stitched_R, - alpha); + auto index_timer = std::chrono::high_resolution_clock::now(); + handle_args(argc, argv, data_type, input_data_path, final_index_path_prefix, label_data_path, universal_label, + num_threads, R, L, stitched_R, alpha); - path labels_file_to_use = final_index_path_prefix + "_label_formatted.txt"; - path labels_map_file = final_index_path_prefix + "_labels_map.txt"; + path labels_file_to_use = final_index_path_prefix + "_label_formatted.txt"; + path labels_map_file = final_index_path_prefix + "_labels_map.txt"; - convert_labels_string_to_int(label_data_path, labels_file_to_use, - labels_map_file, universal_label); + convert_labels_string_to_int(label_data_path, labels_file_to_use, labels_map_file, universal_label); - // 2. parse label file and create necessary data structures - std::vector point_ids_to_labels; - tsl::robin_map labels_to_number_of_points; - label_set all_labels; + // 2. parse label file and create necessary data structures + std::vector point_ids_to_labels; + tsl::robin_map labels_to_number_of_points; + label_set all_labels; - std::tie(point_ids_to_labels, labels_to_number_of_points, all_labels) = - diskann::parse_label_file(labels_file_to_use, universal_label); + std::tie(point_ids_to_labels, labels_to_number_of_points, all_labels) = + diskann::parse_label_file(labels_file_to_use, universal_label); - // 3. for each label, make a separate data file - tsl::robin_map> label_id_to_orig_id_map; - uint32_t total_number_of_points = point_ids_to_labels.size(); + // 3. for each label, make a separate data file + tsl::robin_map> label_id_to_orig_id_map; + uint32_t total_number_of_points = point_ids_to_labels.size(); #ifndef _WINDOWS - if (data_type == "uint8") - label_id_to_orig_id_map = - diskann::generate_label_specific_vector_files( - input_data_path, labels_to_number_of_points, point_ids_to_labels, - all_labels); - else if (data_type == "int8") - label_id_to_orig_id_map = - diskann::generate_label_specific_vector_files( - input_data_path, labels_to_number_of_points, point_ids_to_labels, - all_labels); - else if (data_type == "float") - label_id_to_orig_id_map = - diskann::generate_label_specific_vector_files( - input_data_path, labels_to_number_of_points, point_ids_to_labels, - all_labels); - else - throw; + if (data_type == "uint8") + label_id_to_orig_id_map = diskann::generate_label_specific_vector_files( + input_data_path, labels_to_number_of_points, point_ids_to_labels, all_labels); + else if (data_type == "int8") + label_id_to_orig_id_map = diskann::generate_label_specific_vector_files( + input_data_path, labels_to_number_of_points, point_ids_to_labels, all_labels); + else if (data_type == "float") + label_id_to_orig_id_map = diskann::generate_label_specific_vector_files( + input_data_path, labels_to_number_of_points, point_ids_to_labels, all_labels); + else + throw; #else - if (data_type == "uint8") - label_id_to_orig_id_map = - diskann::generate_label_specific_vector_files_compat( - input_data_path, labels_to_number_of_points, point_ids_to_labels, - all_labels); - else if (data_type == "int8") - label_id_to_orig_id_map = - diskann::generate_label_specific_vector_files_compat( - input_data_path, labels_to_number_of_points, point_ids_to_labels, - all_labels); - else if (data_type == "float") - label_id_to_orig_id_map = - diskann::generate_label_specific_vector_files_compat( - input_data_path, labels_to_number_of_points, point_ids_to_labels, - all_labels); - else - throw; + if (data_type == "uint8") + label_id_to_orig_id_map = diskann::generate_label_specific_vector_files_compat( + input_data_path, labels_to_number_of_points, point_ids_to_labels, all_labels); + else if (data_type == "int8") + label_id_to_orig_id_map = diskann::generate_label_specific_vector_files_compat( + input_data_path, labels_to_number_of_points, point_ids_to_labels, all_labels); + else if (data_type == "float") + label_id_to_orig_id_map = diskann::generate_label_specific_vector_files_compat( + input_data_path, labels_to_number_of_points, point_ids_to_labels, all_labels); + else + throw; #endif - // 4. for each created data file, create a vanilla diskANN index - if (data_type == "uint8") - diskann::generate_label_indices( - input_data_path, final_index_path_prefix, all_labels, R, L, alpha, - num_threads); - else if (data_type == "int8") - diskann::generate_label_indices(input_data_path, - final_index_path_prefix, all_labels, - R, L, alpha, num_threads); - else if (data_type == "float") - diskann::generate_label_indices(input_data_path, - final_index_path_prefix, all_labels, - R, L, alpha, num_threads); - else - throw; - - // 5. "stitch" the indices together - std::vector> stitched_graph; - tsl::robin_map label_entry_points; - uint64_t stitched_graph_size; - - if (data_type == "uint8") - std::tie(stitched_graph, stitched_graph_size) = - stitch_label_indices( - final_index_path_prefix, total_number_of_points, all_labels, - labels_to_number_of_points, label_entry_points, - label_id_to_orig_id_map); - else if (data_type == "int8") - std::tie(stitched_graph, stitched_graph_size) = - stitch_label_indices( - final_index_path_prefix, total_number_of_points, all_labels, - labels_to_number_of_points, label_entry_points, - label_id_to_orig_id_map); - else if (data_type == "float") - std::tie(stitched_graph, stitched_graph_size) = stitch_label_indices( - final_index_path_prefix, total_number_of_points, all_labels, - labels_to_number_of_points, label_entry_points, - label_id_to_orig_id_map); - else - throw; - path full_index_path_prefix = final_index_path_prefix + "_full"; - // 5a. save the stitched graph to disk - save_full_index(full_index_path_prefix, input_data_path, stitched_graph_size, - stitched_graph, label_entry_points, universal_label, - labels_file_to_use); - - // 6. run a prune on the stitched index, and save to disk - if (data_type == "uint8") - prune_and_save(final_index_path_prefix, full_index_path_prefix, - input_data_path, stitched_graph, stitched_R, - label_entry_points, universal_label, - labels_file_to_use, num_threads); - else if (data_type == "int8") - prune_and_save(final_index_path_prefix, full_index_path_prefix, - input_data_path, stitched_graph, stitched_R, - label_entry_points, universal_label, - labels_file_to_use, num_threads); - else if (data_type == "float") - prune_and_save(final_index_path_prefix, full_index_path_prefix, - input_data_path, stitched_graph, stitched_R, - label_entry_points, universal_label, - labels_file_to_use, num_threads); - else - throw; - - std::chrono::duration index_time = - std::chrono::high_resolution_clock::now() - index_timer; - std::cout << "pruned/stitched graph generated in " << index_time.count() - << " seconds" << std::endl; - - clean_up_artifacts(input_data_path, final_index_path_prefix, all_labels); + // 4. for each created data file, create a vanilla diskANN index + if (data_type == "uint8") + diskann::generate_label_indices(input_data_path, final_index_path_prefix, all_labels, R, L, alpha, + num_threads); + else if (data_type == "int8") + diskann::generate_label_indices(input_data_path, final_index_path_prefix, all_labels, R, L, alpha, + num_threads); + else if (data_type == "float") + diskann::generate_label_indices(input_data_path, final_index_path_prefix, all_labels, R, L, alpha, + num_threads); + else + throw; + + // 5. "stitch" the indices together + std::vector> stitched_graph; + tsl::robin_map label_entry_points; + uint64_t stitched_graph_size; + + if (data_type == "uint8") + std::tie(stitched_graph, stitched_graph_size) = + stitch_label_indices(final_index_path_prefix, total_number_of_points, all_labels, + labels_to_number_of_points, label_entry_points, label_id_to_orig_id_map); + else if (data_type == "int8") + std::tie(stitched_graph, stitched_graph_size) = + stitch_label_indices(final_index_path_prefix, total_number_of_points, all_labels, + labels_to_number_of_points, label_entry_points, label_id_to_orig_id_map); + else if (data_type == "float") + std::tie(stitched_graph, stitched_graph_size) = + stitch_label_indices(final_index_path_prefix, total_number_of_points, all_labels, + labels_to_number_of_points, label_entry_points, label_id_to_orig_id_map); + else + throw; + path full_index_path_prefix = final_index_path_prefix + "_full"; + // 5a. save the stitched graph to disk + save_full_index(full_index_path_prefix, input_data_path, stitched_graph_size, stitched_graph, label_entry_points, + universal_label, labels_file_to_use); + + // 6. run a prune on the stitched index, and save to disk + if (data_type == "uint8") + prune_and_save(final_index_path_prefix, full_index_path_prefix, input_data_path, stitched_graph, + stitched_R, label_entry_points, universal_label, labels_file_to_use, num_threads); + else if (data_type == "int8") + prune_and_save(final_index_path_prefix, full_index_path_prefix, input_data_path, stitched_graph, + stitched_R, label_entry_points, universal_label, labels_file_to_use, num_threads); + else if (data_type == "float") + prune_and_save(final_index_path_prefix, full_index_path_prefix, input_data_path, stitched_graph, + stitched_R, label_entry_points, universal_label, labels_file_to_use, num_threads); + else + throw; + + std::chrono::duration index_time = std::chrono::high_resolution_clock::now() - index_timer; + std::cout << "pruned/stitched graph generated in " << index_time.count() << " seconds" << std::endl; + + clean_up_artifacts(input_data_path, final_index_path_prefix, all_labels); } diff --git a/tests/range_search_disk_index.cpp b/tests/range_search_disk_index.cpp index 671925d9c..a67dac378 100644 --- a/tests/range_search_disk_index.cpp +++ b/tests/range_search_disk_index.cpp @@ -33,334 +33,332 @@ namespace po = boost::program_options; #define WARMUP false -void print_stats(std::string category, std::vector percentiles, - std::vector results) { - diskann::cout << std::setw(20) << category << ": " << std::flush; - for (uint32_t s = 0; s < percentiles.size(); s++) { - diskann::cout << std::setw(8) << percentiles[s] << "%"; - } - diskann::cout << std::endl; - diskann::cout << std::setw(22) << " " << std::flush; - for (uint32_t s = 0; s < percentiles.size(); s++) { - diskann::cout << std::setw(9) << results[s]; - } - diskann::cout << std::endl; +void print_stats(std::string category, std::vector percentiles, std::vector results) +{ + diskann::cout << std::setw(20) << category << ": " << std::flush; + for (uint32_t s = 0; s < percentiles.size(); s++) + { + diskann::cout << std::setw(8) << percentiles[s] << "%"; + } + diskann::cout << std::endl; + diskann::cout << std::setw(22) << " " << std::flush; + for (uint32_t s = 0; s < percentiles.size(); s++) + { + diskann::cout << std::setw(9) << results[s]; + } + diskann::cout << std::endl; } template -int search_disk_index(diskann::Metric &metric, - const std::string &index_path_prefix, - const std::string &query_file, std::string >_file, - const uint32_t num_threads, const float search_range, - const uint32_t beamwidth, - const uint32_t num_nodes_to_cache, - const std::vector &Lvec) { - std::string pq_prefix = index_path_prefix + "_pq"; - std::string disk_index_file = index_path_prefix + "_disk.index"; - std::string warmup_query_file = index_path_prefix + "_sample_data.bin"; - - diskann::cout << "Search parameters: #threads: " << num_threads << ", "; - if (beamwidth <= 0) - diskann::cout << "beamwidth to be optimized for each L value" << std::endl; - else - diskann::cout << " beamwidth: " << beamwidth << std::endl; - - // load query bin - T *query = nullptr; - std::vector> groundtruth_ids; - size_t query_num, query_dim, query_aligned_dim, gt_num; - diskann::load_aligned_bin(query_file, query, query_num, query_dim, - query_aligned_dim); - - bool calc_recall_flag = false; - if (gt_file != std::string("null") && file_exists(gt_file)) { - diskann::load_range_truthset( - gt_file, groundtruth_ids, - gt_num); // use for range search type of truthset - // diskann::prune_truthset_for_range(gt_file, search_range, - // groundtruth_ids, gt_num); // use for traditional truthset - if (gt_num != query_num) { - diskann::cout - << "Error. Mismatch in number of queries and ground truth data" - << std::endl; - return -1; +int search_disk_index(diskann::Metric &metric, const std::string &index_path_prefix, const std::string &query_file, + std::string >_file, const uint32_t num_threads, const float search_range, + const uint32_t beamwidth, const uint32_t num_nodes_to_cache, const std::vector &Lvec) +{ + std::string pq_prefix = index_path_prefix + "_pq"; + std::string disk_index_file = index_path_prefix + "_disk.index"; + std::string warmup_query_file = index_path_prefix + "_sample_data.bin"; + + diskann::cout << "Search parameters: #threads: " << num_threads << ", "; + if (beamwidth <= 0) + diskann::cout << "beamwidth to be optimized for each L value" << std::endl; + else + diskann::cout << " beamwidth: " << beamwidth << std::endl; + + // load query bin + T *query = nullptr; + std::vector> groundtruth_ids; + size_t query_num, query_dim, query_aligned_dim, gt_num; + diskann::load_aligned_bin(query_file, query, query_num, query_dim, query_aligned_dim); + + bool calc_recall_flag = false; + if (gt_file != std::string("null") && file_exists(gt_file)) + { + diskann::load_range_truthset(gt_file, groundtruth_ids, + gt_num); // use for range search type of truthset + // diskann::prune_truthset_for_range(gt_file, search_range, + // groundtruth_ids, gt_num); // use for traditional truthset + if (gt_num != query_num) + { + diskann::cout << "Error. Mismatch in number of queries and ground truth data" << std::endl; + return -1; + } + calc_recall_flag = true; } - calc_recall_flag = true; - } - std::shared_ptr reader = nullptr; + std::shared_ptr reader = nullptr; #ifdef _WINDOWS #ifndef USE_BING_INFRA - reader.reset(new WindowsAlignedFileReader()); + reader.reset(new WindowsAlignedFileReader()); #else - reader.reset(new diskann::BingAlignedFileReader()); + reader.reset(new diskann::BingAlignedFileReader()); #endif #else - reader.reset(new LinuxAlignedFileReader()); + reader.reset(new LinuxAlignedFileReader()); #endif - std::unique_ptr> _pFlashIndex( - new diskann::PQFlashIndex(reader, metric)); - - int res = _pFlashIndex->load(num_threads, index_path_prefix.c_str()); - - if (res != 0) { - return res; - } - // 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->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_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, - warmup_dim, warmup_aligned_dim); - } else { - warmup_num = (std::min)((uint32_t)150000, (uint32_t)15000 * num_threads); - warmup_dim = query_dim; - warmup_aligned_dim = query_aligned_dim; - diskann::alloc_aligned(((void **)&warmup), - warmup_num * warmup_aligned_dim * sizeof(T), - 8 * sizeof(T)); - std::memset(warmup, 0, warmup_num * warmup_aligned_dim * sizeof(T)); - std::random_device rd; - std::mt19937 gen(rd()); - std::uniform_int_distribution<> dis(-128, 127); - for (uint32_t i = 0; i < warmup_num; i++) { - for (uint32_t d = 0; d < warmup_dim; d++) { - warmup[i * warmup_aligned_dim + d] = (T)dis(gen); - } - } + std::unique_ptr> _pFlashIndex( + new diskann::PQFlashIndex(reader, metric)); + + int res = _pFlashIndex->load(num_threads, index_path_prefix.c_str()); + + if (res != 0) + { + return res; } - diskann::cout << "Warming up index... " << std::flush; - std::vector warmup_result_ids_64(warmup_num, 0); - std::vector warmup_result_dists(warmup_num, 0); + // 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->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_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, warmup_dim, warmup_aligned_dim); + } + else + { + warmup_num = (std::min)((uint32_t)150000, (uint32_t)15000 * num_threads); + warmup_dim = query_dim; + warmup_aligned_dim = query_aligned_dim; + diskann::alloc_aligned(((void **)&warmup), warmup_num * warmup_aligned_dim * sizeof(T), 8 * sizeof(T)); + std::memset(warmup, 0, warmup_num * warmup_aligned_dim * sizeof(T)); + std::random_device rd; + std::mt19937 gen(rd()); + std::uniform_int_distribution<> dis(-128, 127); + for (uint32_t i = 0; i < warmup_num; i++) + { + for (uint32_t d = 0; d < warmup_dim; d++) + { + warmup[i * warmup_aligned_dim + d] = (T)dis(gen); + } + } + } + diskann::cout << "Warming up index... " << std::flush; + std::vector warmup_result_ids_64(warmup_num, 0); + std::vector warmup_result_dists(warmup_num, 0); #pragma omp parallel for schedule(dynamic, 1) - for (int64_t 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); + for (int64_t 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); + } + diskann::cout << "..done" << std::endl; } - diskann::cout << "..done" << std::endl; - } - - diskann::cout.setf(std::ios_base::fixed, std::ios_base::floatfield); - diskann::cout.precision(2); - - std::string recall_string = "Recall@rng=" + std::to_string(search_range); - diskann::cout << std::setw(6) << "L" << std::setw(12) << "Beamwidth" - << std::setw(16) << "QPS" << std::setw(16) << "Mean Latency" - << std::setw(16) << "99.9 Latency" << std::setw(16) - << "Mean IOs" << std::setw(16) << "CPU (s)"; - if (calc_recall_flag) { - diskann::cout << std::setw(16) << recall_string << std::endl; - } else - diskann::cout << std::endl; - diskann::cout - << "===============================================================" - "===========================================" - << std::endl; - std::vector>> query_result_ids(Lvec.size()); + diskann::cout.setf(std::ios_base::fixed, std::ios_base::floatfield); + diskann::cout.precision(2); + + std::string recall_string = "Recall@rng=" + std::to_string(search_range); + diskann::cout << std::setw(6) << "L" << std::setw(12) << "Beamwidth" << std::setw(16) << "QPS" << std::setw(16) + << "Mean Latency" << std::setw(16) << "99.9 Latency" << std::setw(16) << "Mean IOs" << std::setw(16) + << "CPU (s)"; + if (calc_recall_flag) + { + diskann::cout << std::setw(16) << recall_string << std::endl; + } + else + diskann::cout << std::endl; + diskann::cout << "===============================================================" + "===========================================" + << std::endl; - uint32_t optimized_beamwidth = 2; - uint32_t max_list_size = 10000; + std::vector>> query_result_ids(Lvec.size()); - for (uint32_t test_id = 0; test_id < Lvec.size(); test_id++) { - uint64_t L = Lvec[test_id]; + uint32_t optimized_beamwidth = 2; + uint32_t max_list_size = 10000; - if (beamwidth <= 0) { - optimized_beamwidth = - optimize_beamwidth(_pFlashIndex, warmup, warmup_num, - warmup_aligned_dim, L, optimized_beamwidth); - } else - optimized_beamwidth = beamwidth; + for (uint32_t test_id = 0; test_id < Lvec.size(); test_id++) + { + uint64_t L = Lvec[test_id]; + + if (beamwidth <= 0) + { + optimized_beamwidth = + optimize_beamwidth(_pFlashIndex, warmup, warmup_num, warmup_aligned_dim, L, optimized_beamwidth); + } + else + optimized_beamwidth = beamwidth; - query_result_ids[test_id].clear(); - query_result_ids[test_id].resize(query_num); + query_result_ids[test_id].clear(); + query_result_ids[test_id].resize(query_num); - diskann::QueryStats *stats = new diskann::QueryStats[query_num]; + diskann::QueryStats *stats = new diskann::QueryStats[query_num]; - auto s = std::chrono::high_resolution_clock::now(); + auto s = std::chrono::high_resolution_clock::now(); #pragma omp parallel for schedule(dynamic, 1) - for (int64_t i = 0; i < (int64_t)query_num; i++) { - std::vector indices; - std::vector distances; - uint32_t res_count = _pFlashIndex->range_search( - query + (i * query_aligned_dim), search_range, L, max_list_size, - indices, distances, optimized_beamwidth, stats + i); - query_result_ids[test_id][i].reserve(res_count); - query_result_ids[test_id][i].resize(res_count); - for (uint32_t idx = 0; idx < res_count; idx++) - query_result_ids[test_id][i][idx] = indices[idx]; - } - auto e = std::chrono::high_resolution_clock::now(); - std::chrono::duration diff = e - s; - auto qps = (1.0 * query_num) / (1.0 * diff.count()); - - auto mean_latency = diskann::get_mean_stats( - stats, query_num, - [](const diskann::QueryStats &stats) { return stats.total_us; }); - - auto latency_999 = diskann::get_percentile_stats( - stats, query_num, 0.999, - [](const diskann::QueryStats &stats) { return stats.total_us; }); - - auto mean_ios = diskann::get_mean_stats( - stats, query_num, - [](const diskann::QueryStats &stats) { return stats.n_ios; }); - - float mean_cpuus = diskann::get_mean_stats( - stats, query_num, - [](const diskann::QueryStats &stats) { return stats.cpu_us; }); - - float recall = 0; - float ratio_of_sums = 0; - if (calc_recall_flag) { - recall = diskann::calculate_range_search_recall( - query_num, groundtruth_ids, query_result_ids[test_id]); - - uint32_t total_true_positive = 0; - uint32_t total_positive = 0; - for (uint32_t i = 0; i < query_num; i++) { - total_true_positive += query_result_ids[test_id][i].size(); - total_positive += groundtruth_ids[i].size(); - } - - ratio_of_sums = (1.0 * total_true_positive) / (1.0 * total_positive); + for (int64_t i = 0; i < (int64_t)query_num; i++) + { + std::vector indices; + std::vector distances; + uint32_t res_count = + _pFlashIndex->range_search(query + (i * query_aligned_dim), search_range, L, max_list_size, indices, + distances, optimized_beamwidth, stats + i); + query_result_ids[test_id][i].reserve(res_count); + query_result_ids[test_id][i].resize(res_count); + for (uint32_t idx = 0; idx < res_count; idx++) + query_result_ids[test_id][i][idx] = indices[idx]; + } + auto e = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = e - s; + auto qps = (1.0 * query_num) / (1.0 * diff.count()); + + auto mean_latency = diskann::get_mean_stats( + stats, query_num, [](const diskann::QueryStats &stats) { return stats.total_us; }); + + auto latency_999 = diskann::get_percentile_stats( + stats, query_num, 0.999, [](const diskann::QueryStats &stats) { return stats.total_us; }); + + auto mean_ios = diskann::get_mean_stats(stats, query_num, + [](const diskann::QueryStats &stats) { return stats.n_ios; }); + + float mean_cpuus = diskann::get_mean_stats( + stats, query_num, [](const diskann::QueryStats &stats) { return stats.cpu_us; }); + + float recall = 0; + float ratio_of_sums = 0; + if (calc_recall_flag) + { + recall = diskann::calculate_range_search_recall(query_num, groundtruth_ids, query_result_ids[test_id]); + + uint32_t total_true_positive = 0; + uint32_t total_positive = 0; + for (uint32_t i = 0; i < query_num; i++) + { + total_true_positive += query_result_ids[test_id][i].size(); + total_positive += groundtruth_ids[i].size(); + } + + ratio_of_sums = (1.0 * total_true_positive) / (1.0 * total_positive); + } + + diskann::cout << std::setw(6) << L << std::setw(12) << optimized_beamwidth << std::setw(16) << qps + << std::setw(16) << mean_latency << std::setw(16) << latency_999 << std::setw(16) << mean_ios + << std::setw(16) << mean_cpuus; + if (calc_recall_flag) + { + diskann::cout << std::setw(16) << recall << "," << ratio_of_sums << std::endl; + } + else + diskann::cout << std::endl; } - diskann::cout << std::setw(6) << L << std::setw(12) << optimized_beamwidth - << std::setw(16) << qps << std::setw(16) << mean_latency - << std::setw(16) << latency_999 << std::setw(16) << mean_ios - << std::setw(16) << mean_cpuus; - if (calc_recall_flag) { - diskann::cout << std::setw(16) << recall << "," << ratio_of_sums - << std::endl; - } else - diskann::cout << std::endl; - } - - diskann::cout << "Done searching. " << std::endl; - - diskann::aligned_free(query); - if (warmup != nullptr) diskann::aligned_free(warmup); - return 0; + diskann::cout << "Done searching. " << std::endl; + + diskann::aligned_free(query); + if (warmup != nullptr) + diskann::aligned_free(warmup); + return 0; } -int main(int argc, char **argv) { - std::string data_type, dist_fn, index_path_prefix, result_path_prefix, - query_file, gt_file; - uint32_t num_threads, W, num_nodes_to_cache; - std::vector Lvec; - float range; - - po::options_description desc{"Arguments"}; - try { - desc.add_options()("help,h", "Print information on arguments"); - desc.add_options()("data_type", - po::value(&data_type)->required(), - "data type "); - desc.add_options()("dist_fn", po::value(&dist_fn)->required(), - "distance function "); - desc.add_options()("index_path_prefix", - po::value(&index_path_prefix)->required(), - "Path prefix to the index"); - desc.add_options()("query_file", - po::value(&query_file)->required(), - "Query file in binary format"); - desc.add_options()( - "gt_file", - po::value(>_file)->default_value(std::string("null")), - "ground truth file for the queryset"); - desc.add_options()("range_threshold,K", - po::value(&range)->required(), - "Number of neighbors to be returned"); - desc.add_options()("search_list,L", - po::value>(&Lvec)->multitoken(), - "List of L values of search"); - desc.add_options()("beamwidth,W", po::value(&W)->default_value(2), - "Beamwidth for search"); - desc.add_options()( - "num_nodes_to_cache", - po::value(&num_nodes_to_cache)->default_value(100000), - "Beamwidth for search"); - desc.add_options()( - "num_threads,T", - po::value(&num_threads)->default_value(omp_get_num_procs()), - "Number of threads used for building index (defaults to " - "omp_get_num_procs())"); - - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) { - std::cout << desc; - return 0; +int main(int argc, char **argv) +{ + std::string data_type, dist_fn, index_path_prefix, result_path_prefix, query_file, gt_file; + uint32_t num_threads, W, num_nodes_to_cache; + std::vector Lvec; + float range; + + po::options_description desc{"Arguments"}; + try + { + desc.add_options()("help,h", "Print information on arguments"); + desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); + desc.add_options()("dist_fn", po::value(&dist_fn)->required(), + "distance function "); + desc.add_options()("index_path_prefix", po::value(&index_path_prefix)->required(), + "Path prefix to the index"); + desc.add_options()("query_file", po::value(&query_file)->required(), + "Query file in binary format"); + desc.add_options()("gt_file", po::value(>_file)->default_value(std::string("null")), + "ground truth file for the queryset"); + desc.add_options()("range_threshold,K", po::value(&range)->required(), + "Number of neighbors to be returned"); + desc.add_options()("search_list,L", po::value>(&Lvec)->multitoken(), + "List of L values of search"); + desc.add_options()("beamwidth,W", po::value(&W)->default_value(2), "Beamwidth for search"); + desc.add_options()("num_nodes_to_cache", po::value(&num_nodes_to_cache)->default_value(100000), + "Beamwidth for search"); + desc.add_options()("num_threads,T", po::value(&num_threads)->default_value(omp_get_num_procs()), + "Number of threads used for building index (defaults to " + "omp_get_num_procs())"); + + po::variables_map vm; + po::store(po::parse_command_line(argc, argv, desc), vm); + if (vm.count("help")) + { + std::cout << desc; + return 0; + } + po::notify(vm); + } + catch (const std::exception &ex) + { + std::cerr << ex.what() << '\n'; + return -1; + } + + diskann::Metric metric; + if (dist_fn == std::string("mips")) + { + metric = diskann::Metric::INNER_PRODUCT; + } + else if (dist_fn == std::string("l2")) + { + metric = diskann::Metric::L2; + } + else if (dist_fn == std::string("cosine")) + { + metric = diskann::Metric::COSINE; + } + else + { + std::cout << "Unsupported distance function. Currently only L2/ Inner " + "Product/Cosine are supported." + << std::endl; + return -1; + } + + if ((data_type != std::string("float")) && (metric == diskann::Metric::INNER_PRODUCT)) + { + std::cout << "Currently support only floating point data for Inner Product." << std::endl; + return -1; + } + + try + { + if (data_type == std::string("float")) + return search_disk_index(metric, index_path_prefix, query_file, gt_file, num_threads, range, W, + num_nodes_to_cache, Lvec); + else if (data_type == std::string("int8")) + return search_disk_index(metric, index_path_prefix, query_file, gt_file, num_threads, range, W, + num_nodes_to_cache, Lvec); + else if (data_type == std::string("uint8")) + return search_disk_index(metric, index_path_prefix, query_file, gt_file, num_threads, range, W, + num_nodes_to_cache, Lvec); + else + { + std::cerr << "Unsupported data type. Use float or int8 or uint8" << std::endl; + return -1; + } } - po::notify(vm); - } catch (const std::exception &ex) { - std::cerr << ex.what() << '\n'; - return -1; - } - - diskann::Metric metric; - if (dist_fn == std::string("mips")) { - metric = diskann::Metric::INNER_PRODUCT; - } else if (dist_fn == std::string("l2")) { - metric = diskann::Metric::L2; - } else if (dist_fn == std::string("cosine")) { - metric = diskann::Metric::COSINE; - } else { - std::cout << "Unsupported distance function. Currently only L2/ Inner " - "Product/Cosine are supported." - << std::endl; - return -1; - } - - if ((data_type != std::string("float")) && - (metric == diskann::Metric::INNER_PRODUCT)) { - std::cout << "Currently support only floating point data for Inner Product." - << std::endl; - return -1; - } - - try { - if (data_type == std::string("float")) - return search_disk_index(metric, index_path_prefix, query_file, - gt_file, num_threads, range, W, - num_nodes_to_cache, Lvec); - else if (data_type == std::string("int8")) - return search_disk_index(metric, index_path_prefix, query_file, - gt_file, num_threads, range, W, - num_nodes_to_cache, Lvec); - else if (data_type == std::string("uint8")) - return search_disk_index(metric, index_path_prefix, query_file, - gt_file, num_threads, range, W, - num_nodes_to_cache, Lvec); - else { - std::cerr << "Unsupported data type. Use float or int8 or uint8" - << std::endl; - return -1; + catch (const std::exception &e) + { + std::cout << std::string(e.what()) << std::endl; + diskann::cerr << "Index search failed." << std::endl; + return -1; } - } catch (const std::exception &e) { - std::cout << std::string(e.what()) << std::endl; - diskann::cerr << "Index search failed." << std::endl; - return -1; - } } diff --git a/tests/restapi/client.cpp b/tests/restapi/client.cpp index 40d83ac41..fdf4414dd 100644 --- a/tests/restapi/client.cpp +++ b/tests/restapi/client.cpp @@ -23,92 +23,102 @@ using namespace diskann; namespace po = boost::program_options; template -void query_loop(const std::string &ip_addr_port, const std::string &query_file, - const unsigned nq, const unsigned Ls, const unsigned k_value) { - web::http::client::http_client client(U(ip_addr_port)); +void query_loop(const std::string &ip_addr_port, const std::string &query_file, const unsigned nq, const unsigned Ls, + const unsigned k_value) +{ + web::http::client::http_client client(U(ip_addr_port)); - T *data; - size_t npts = 1, ndims = 128, rounded_dim = 128; - diskann::load_aligned_bin(query_file, data, npts, ndims, rounded_dim); + T *data; + size_t npts = 1, ndims = 128, rounded_dim = 128; + diskann::load_aligned_bin(query_file, data, npts, ndims, rounded_dim); - for (unsigned i = 0; i < nq; ++i) { - T *vec = data + i * rounded_dim; - web::http::http_request http_query(methods::POST); - web::json::value queryJson = web::json::value::object(); - queryJson[QUERY_ID_KEY] = i; - queryJson[K_KEY] = k_value; - queryJson[L_KEY] = Ls; - for (size_t i = 0; i < ndims; ++i) { - queryJson[VECTOR_KEY][i] = web::json::value::number(vec[i]); - } - http_query.set_body(queryJson); + for (unsigned i = 0; i < nq; ++i) + { + T *vec = data + i * rounded_dim; + web::http::http_request http_query(methods::POST); + web::json::value queryJson = web::json::value::object(); + queryJson[QUERY_ID_KEY] = i; + queryJson[K_KEY] = k_value; + queryJson[L_KEY] = Ls; + for (size_t i = 0; i < ndims; ++i) + { + queryJson[VECTOR_KEY][i] = web::json::value::number(vec[i]); + } + http_query.set_body(queryJson); - client.request(http_query) - .then([](web::http::http_response response) - -> pplx::task { - if (response.status_code() == status_codes::OK) { - return response.extract_string(); - } - std::cerr << "Query failed" << std::endl; - return pplx::task_from_result(utility::string_t()); - }) - .then([](pplx::task previousTask) { - try { - std::cout << previousTask.get() << std::endl; - } catch (http_exception const &e) { - std::wcout << e.what() << std::endl; - } - }) - .wait(); - } + client.request(http_query) + .then([](web::http::http_response response) -> pplx::task { + if (response.status_code() == status_codes::OK) + { + return response.extract_string(); + } + std::cerr << "Query failed" << std::endl; + return pplx::task_from_result(utility::string_t()); + }) + .then([](pplx::task previousTask) { + try + { + std::cout << previousTask.get() << std::endl; + } + catch (http_exception const &e) + { + std::wcout << e.what() << std::endl; + } + }) + .wait(); + } } -int main(int argc, char *argv[]) { - std::string data_type, query_file, address; - uint32_t num_queries; - uint32_t l_search, k_value; +int main(int argc, char *argv[]) +{ + std::string data_type, query_file, address; + uint32_t num_queries; + uint32_t l_search, k_value; - po::options_description desc{"Arguments"}; - try { - desc.add_options()("help,h", "Print information on arguments"); - desc.add_options()("data_type", - po::value(&data_type)->required(), - "data type "); - desc.add_options()("address", po::value(&address)->required(), - "Web server address"); - desc.add_options()("query_file", - po::value(&query_file)->required(), - "File containing the queries to search"); - desc.add_options()("num_queries,Q", - po::value(&num_queries)->required(), - "Number of queries to search"); - desc.add_options()("l_search", po::value(&l_search)->required(), - "Value of L"); - desc.add_options()("k_value,K", - po::value(&k_value)->default_value(10), - "Value of K (default 10)"); - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) { - std::cout << desc; - return 0; + po::options_description desc{"Arguments"}; + try + { + desc.add_options()("help,h", "Print information on arguments"); + desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); + desc.add_options()("address", po::value(&address)->required(), "Web server address"); + desc.add_options()("query_file", po::value(&query_file)->required(), + "File containing the queries to search"); + desc.add_options()("num_queries,Q", po::value(&num_queries)->required(), + "Number of queries to search"); + desc.add_options()("l_search", po::value(&l_search)->required(), "Value of L"); + desc.add_options()("k_value,K", po::value(&k_value)->default_value(10), "Value of K (default 10)"); + po::variables_map vm; + po::store(po::parse_command_line(argc, argv, desc), vm); + if (vm.count("help")) + { + std::cout << desc; + return 0; + } + po::notify(vm); + } + catch (const std::exception &ex) + { + std::cerr << ex.what() << std::endl; + return -1; } - po::notify(vm); - } catch (const std::exception &ex) { - std::cerr << ex.what() << std::endl; - return -1; - } - if (data_type == std::string("float")) { - query_loop(address, query_file, num_queries, l_search, k_value); - } else if (data_type == std::string("int8")) { - query_loop(address, query_file, num_queries, l_search, k_value); - } else if (data_type == std::string("uint8")) { - query_loop(address, query_file, num_queries, l_search, k_value); - } else { - std::cerr << "Unsupported type " << argv[2] << std::endl; - return -1; - } + if (data_type == std::string("float")) + { + query_loop(address, query_file, num_queries, l_search, k_value); + } + else if (data_type == std::string("int8")) + { + query_loop(address, query_file, num_queries, l_search, k_value); + } + else if (data_type == std::string("uint8")) + { + query_loop(address, query_file, num_queries, l_search, k_value); + } + else + { + std::cerr << "Unsupported type " << argv[2] << std::endl; + return -1; + } - return 0; + return 0; } \ No newline at end of file diff --git a/tests/restapi/inmem_server.cpp b/tests/restapi/inmem_server.cpp index fb0fb6223..11da541ff 100644 --- a/tests/restapi/inmem_server.cpp +++ b/tests/restapi/inmem_server.cpp @@ -17,115 +17,122 @@ namespace po = boost::program_options; std::unique_ptr g_httpServer(nullptr); std::vector> g_inMemorySearch; -void setup(const utility::string_t &address, const std::string &typestring) { - web::http::uri_builder uriBldr(address); - auto uri = uriBldr.to_uri(); +void setup(const utility::string_t &address, const std::string &typestring) +{ + web::http::uri_builder uriBldr(address); + auto uri = uriBldr.to_uri(); - std::cout << "Attempting to start server on " << uri.to_string() << std::endl; + std::cout << "Attempting to start server on " << uri.to_string() << std::endl; - g_httpServer = - std::unique_ptr(new Server(uri, g_inMemorySearch, typestring)); - std::cout << "Created a server object" << std::endl; + g_httpServer = std::unique_ptr(new Server(uri, g_inMemorySearch, typestring)); + std::cout << "Created a server object" << std::endl; - g_httpServer->open().wait(); - ucout << U"Listening for requests on: " << address << std::endl; + g_httpServer->open().wait(); + ucout << U"Listening for requests on: " << address << std::endl; } -void teardown(const utility::string_t &address) { - g_httpServer->close().wait(); +void teardown(const utility::string_t &address) +{ + g_httpServer->close().wait(); } -int main(int argc, char *argv[]) { - std::string data_type, index_file, data_file, address, dist_fn, tags_file; - uint32_t num_threads; - uint32_t l_search; +int main(int argc, char *argv[]) +{ + std::string data_type, index_file, data_file, address, dist_fn, tags_file; + uint32_t num_threads; + uint32_t l_search; - po::options_description desc{"Arguments"}; - try { - desc.add_options()("help,h", "Print information on arguments"); - desc.add_options()("data_type", - po::value(&data_type)->required(), - "data type "); - desc.add_options()("address", po::value(&address)->required(), - "Web server address"); - desc.add_options()("data_file", - po::value(&data_file)->required(), - "File containing the data found in the index"); - desc.add_options()("index_path_prefix", - po::value(&index_file)->required(), - "Path prefix for saving index file components"); - desc.add_options()("num_threads,T", - po::value(&num_threads)->required(), - "Number of threads used for building index"); - desc.add_options()("l_search", po::value(&l_search)->required(), - "Value of L"); - desc.add_options()("dist_fn", - po::value(&dist_fn)->default_value("l2"), - "distance function "); - desc.add_options()( - "tags_file", - po::value(&tags_file)->default_value(std::string()), - "Tags file location"); - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) { - std::cout << desc; - return 0; + po::options_description desc{"Arguments"}; + try + { + desc.add_options()("help,h", "Print information on arguments"); + desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); + desc.add_options()("address", po::value(&address)->required(), "Web server address"); + desc.add_options()("data_file", po::value(&data_file)->required(), + "File containing the data found in the index"); + desc.add_options()("index_path_prefix", po::value(&index_file)->required(), + "Path prefix for saving index file components"); + desc.add_options()("num_threads,T", po::value(&num_threads)->required(), + "Number of threads used for building index"); + desc.add_options()("l_search", po::value(&l_search)->required(), "Value of L"); + desc.add_options()("dist_fn", po::value(&dist_fn)->default_value("l2"), + "distance function "); + desc.add_options()("tags_file", po::value(&tags_file)->default_value(std::string()), + "Tags file location"); + po::variables_map vm; + po::store(po::parse_command_line(argc, argv, desc), vm); + if (vm.count("help")) + { + std::cout << desc; + return 0; + } + po::notify(vm); + } + catch (const std::exception &ex) + { + std::cerr << ex.what() << std::endl; + return -1; + } + diskann::Metric metric; + if (dist_fn == std::string("l2")) + metric = diskann::Metric::L2; + else if (dist_fn == std::string("mips")) + metric = diskann::Metric::INNER_PRODUCT; + else + { + std::cout << "Error. Only l2 and mips distance functions are supported" << std::endl; + return -1; } - po::notify(vm); - } catch (const std::exception &ex) { - std::cerr << ex.what() << std::endl; - return -1; - } - diskann::Metric metric; - if (dist_fn == std::string("l2")) - metric = diskann::Metric::L2; - else if (dist_fn == std::string("mips")) - metric = diskann::Metric::INNER_PRODUCT; - else { - std::cout << "Error. Only l2 and mips distance functions are supported" - << std::endl; - return -1; - } - if (data_type == std::string("float")) { - auto searcher = - std::unique_ptr(new diskann::InMemorySearch( - data_file, index_file, tags_file, metric, num_threads, l_search)); - g_inMemorySearch.push_back(std::move(searcher)); - } else if (data_type == std::string("int8")) { - auto searcher = std::unique_ptr( - new diskann::InMemorySearch(data_file, index_file, tags_file, - metric, num_threads, l_search)); - g_inMemorySearch.push_back(std::move(searcher)); - } else if (data_type == std::string("uint8")) { - auto searcher = std::unique_ptr( - new diskann::InMemorySearch(data_file, index_file, tags_file, - metric, num_threads, l_search)); - g_inMemorySearch.push_back(std::move(searcher)); - } else { - std::cerr << "Unsupported data type " << argv[2] << std::endl; - } + if (data_type == std::string("float")) + { + auto searcher = std::unique_ptr( + new diskann::InMemorySearch(data_file, index_file, tags_file, metric, num_threads, l_search)); + g_inMemorySearch.push_back(std::move(searcher)); + } + else if (data_type == std::string("int8")) + { + auto searcher = std::unique_ptr( + new diskann::InMemorySearch(data_file, index_file, tags_file, metric, num_threads, l_search)); + g_inMemorySearch.push_back(std::move(searcher)); + } + else if (data_type == std::string("uint8")) + { + auto searcher = std::unique_ptr( + new diskann::InMemorySearch(data_file, index_file, tags_file, metric, num_threads, l_search)); + g_inMemorySearch.push_back(std::move(searcher)); + } + else + { + std::cerr << "Unsupported data type " << argv[2] << std::endl; + } - while (1) { - try { - setup(address, data_type); - std::cout << "Type 'exit' (case-sensitive) to exit" << std::endl; - std::string line; - std::getline(std::cin, line); - if (line == "exit") { - teardown(address); - g_httpServer->close().wait(); - exit(0); - } - } catch (const std::exception &ex) { - std::cerr << "Exception occurred: " << ex.what() << std::endl; - std::cerr << "Restarting HTTP server"; - teardown(address); - } catch (...) { - std::cerr << "Unknown exception occurreed" << std::endl; - std::cerr << "Restarting HTTP server"; - teardown(address); + while (1) + { + try + { + setup(address, data_type); + std::cout << "Type 'exit' (case-sensitive) to exit" << std::endl; + std::string line; + std::getline(std::cin, line); + if (line == "exit") + { + teardown(address); + g_httpServer->close().wait(); + exit(0); + } + } + catch (const std::exception &ex) + { + std::cerr << "Exception occurred: " << ex.what() << std::endl; + std::cerr << "Restarting HTTP server"; + teardown(address); + } + catch (...) + { + std::cerr << "Unknown exception occurreed" << std::endl; + std::cerr << "Restarting HTTP server"; + teardown(address); + } } - } } diff --git a/tests/restapi/main.cpp b/tests/restapi/main.cpp index f1c5a78c7..cb48d6787 100644 --- a/tests/restapi/main.cpp +++ b/tests/restapi/main.cpp @@ -9,65 +9,75 @@ std::unique_ptr g_httpServer(nullptr); std::unique_ptr g_inMemorySearch(nullptr); -void setup(const utility::string_t &address) { - web::http::uri_builder uriBldr(address); - auto uri = uriBldr.to_uri(); +void setup(const utility::string_t &address) +{ + web::http::uri_builder uriBldr(address); + auto uri = uriBldr.to_uri(); - std::wcout << L"Attempting to start server on " << uri.to_string() - << std::endl; + std::wcout << L"Attempting to start server on " << uri.to_string() << std::endl; - g_httpServer = std::unique_ptr(new Server(uri, g_inMemorySearch)); - g_httpServer->open().wait(); + g_httpServer = std::unique_ptr(new Server(uri, g_inMemorySearch)); + g_httpServer->open().wait(); - ucout << U"Listening for requests on: " << address << std::endl; + ucout << U"Listening for requests on: " << address << std::endl; } -void teardown(const utility::string_t &address) { - g_httpServer->close().wait(); +void teardown(const utility::string_t &address) +{ + g_httpServer->close().wait(); } -void loadIndex(const char *indexFile, const char *baseFile, - const char *idsFile) { - auto nsgSearch = - new diskann::InMemorySearch(baseFile, indexFile, idsFile, diskann::L2); - g_inMemorySearch = std::unique_ptr(nsgSearch); +void loadIndex(const char *indexFile, const char *baseFile, const char *idsFile) +{ + auto nsgSearch = new diskann::InMemorySearch(baseFile, indexFile, idsFile, diskann::L2); + g_inMemorySearch = std::unique_ptr(nsgSearch); } -std::wstring getHostingAddress(const char *hostNameAndPort) { - wchar_t buffer[4096]; - mbstowcs_s(nullptr, buffer, sizeof(buffer) / sizeof(buffer[0]), - hostNameAndPort, sizeof(buffer) / sizeof(buffer[0])); - return std::wstring(buffer); +std::wstring getHostingAddress(const char *hostNameAndPort) +{ + wchar_t buffer[4096]; + mbstowcs_s(nullptr, buffer, sizeof(buffer) / sizeof(buffer[0]), hostNameAndPort, + sizeof(buffer) / sizeof(buffer[0])); + return std::wstring(buffer); } -int main(int argc, char *argv[]) { - if (argc != 5) { - std::cout << "Usage: nsg_server " - " " - << std::endl; - exit(1); - } +int main(int argc, char *argv[]) +{ + if (argc != 5) + { + std::cout << "Usage: nsg_server " + " " + << std::endl; + exit(1); + } - auto address = getHostingAddress(argv[1]); - loadIndex(argv[2], argv[3], argv[4]); - while (1) { - try { - setup(address); - std::cout << "Type 'exit' (case-sensitive) to exit" << std::endl; - std::string line; - std::getline(std::cin, line); - if (line == "exit") { - teardown(address); - exit(0); - } - } catch (const std::exception &ex) { - std::cerr << "Exception occurred: " << ex.what() << std::endl; - std::cerr << "Restarting HTTP server"; - teardown(address); - } catch (...) { - std::cerr << "Unknown exception occurreed" << std::endl; - std::cerr << "Restarting HTTP server"; - teardown(address); + auto address = getHostingAddress(argv[1]); + loadIndex(argv[2], argv[3], argv[4]); + while (1) + { + try + { + setup(address); + std::cout << "Type 'exit' (case-sensitive) to exit" << std::endl; + std::string line; + std::getline(std::cin, line); + if (line == "exit") + { + teardown(address); + exit(0); + } + } + catch (const std::exception &ex) + { + std::cerr << "Exception occurred: " << ex.what() << std::endl; + std::cerr << "Restarting HTTP server"; + teardown(address); + } + catch (...) + { + std::cerr << "Unknown exception occurreed" << std::endl; + std::cerr << "Restarting HTTP server"; + teardown(address); + } } - } } diff --git a/tests/restapi/multiple_ssdindex_server.cpp b/tests/restapi/multiple_ssdindex_server.cpp index c550cc0ca..89cb06fcb 100644 --- a/tests/restapi/multiple_ssdindex_server.cpp +++ b/tests/restapi/multiple_ssdindex_server.cpp @@ -18,153 +18,165 @@ namespace po = boost::program_options; std::unique_ptr g_httpServer(nullptr); std::vector> g_ssdSearch; -void setup(const utility::string_t &address, const std::string &typestring) { - web::http::uri_builder uriBldr(address); - auto uri = uriBldr.to_uri(); +void setup(const utility::string_t &address, const std::string &typestring) +{ + web::http::uri_builder uriBldr(address); + auto uri = uriBldr.to_uri(); - std::cout << "Attempting to start server on " << uri.to_string() << std::endl; + std::cout << "Attempting to start server on " << uri.to_string() << std::endl; - g_httpServer = - std::unique_ptr(new Server(uri, g_ssdSearch, typestring)); - std::cout << "Created a server object" << std::endl; + g_httpServer = std::unique_ptr(new Server(uri, g_ssdSearch, typestring)); + std::cout << "Created a server object" << std::endl; - g_httpServer->open().wait(); - ucout << U"Listening for requests on: " << address << std::endl; + g_httpServer->open().wait(); + ucout << U"Listening for requests on: " << address << std::endl; } -void teardown(const utility::string_t &address) { - g_httpServer->close().wait(); +void teardown(const utility::string_t &address) +{ + g_httpServer->close().wait(); } -int main(int argc, char *argv[]) { - std::string data_type, index_prefix_paths, address, dist_fn, tags_file; - uint32_t num_nodes_to_cache; - uint32_t num_threads; - - po::options_description desc{"Arguments"}; - try { - desc.add_options()("help,h", "Print information on arguments"); - desc.add_options()("address", po::value(&address)->required(), - "Web server address"); - desc.add_options()("data_type", - po::value(&data_type)->required(), - "data type "); - desc.add_options()("index_prefix_paths", - po::value(&index_prefix_paths)->required(), - "Path prefix for loading index file components"); - desc.add_options()( - "num_nodes_to_cache", - po::value(&num_nodes_to_cache)->default_value(0), - "Number of nodes to cache during search"); - desc.add_options()( - "num_threads,T", - po::value(&num_threads)->default_value(omp_get_num_procs()), - "Number of threads used for building index (defaults to " - "omp_get_num_procs())"); - desc.add_options()("dist_fn", - po::value(&dist_fn)->default_value("l2"), - "distance function "); - desc.add_options()( - "tags_file", - po::value(&tags_file)->default_value(std::string()), - "Tags file location"); - - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) { - std::cout << desc; - return 0; +int main(int argc, char *argv[]) +{ + std::string data_type, index_prefix_paths, address, dist_fn, tags_file; + uint32_t num_nodes_to_cache; + uint32_t num_threads; + + po::options_description desc{"Arguments"}; + try + { + desc.add_options()("help,h", "Print information on arguments"); + desc.add_options()("address", po::value(&address)->required(), "Web server address"); + desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); + desc.add_options()("index_prefix_paths", po::value(&index_prefix_paths)->required(), + "Path prefix for loading index file components"); + desc.add_options()("num_nodes_to_cache", po::value(&num_nodes_to_cache)->default_value(0), + "Number of nodes to cache during search"); + desc.add_options()("num_threads,T", po::value(&num_threads)->default_value(omp_get_num_procs()), + "Number of threads used for building index (defaults to " + "omp_get_num_procs())"); + desc.add_options()("dist_fn", po::value(&dist_fn)->default_value("l2"), + "distance function "); + desc.add_options()("tags_file", po::value(&tags_file)->default_value(std::string()), + "Tags file location"); + + po::variables_map vm; + po::store(po::parse_command_line(argc, argv, desc), vm); + if (vm.count("help")) + { + std::cout << desc; + return 0; + } + po::notify(vm); } - po::notify(vm); - } catch (const std::exception &ex) { - std::cerr << ex.what() << std::endl; - return -1; - } - - diskann::Metric metric; - if (dist_fn == std::string("l2")) - metric = diskann::Metric::L2; - else if (dist_fn == std::string("mips")) - metric = diskann::Metric::INNER_PRODUCT; - else { - std::cout << "Error. Only l2 and mips distance functions are supported" - << std::endl; - return -1; - } - - std::vector> index_tag_paths; - std::ifstream index_in(index_prefix_paths); - if (!index_in.is_open()) { - std::cerr << "Could not open " << index_prefix_paths << std::endl; - exit(-1); - } - std::ifstream tags_in(tags_file); - if (!tags_in.is_open()) { - std::cerr << "Could not open " << tags_file << std::endl; - exit(-1); - } - std::string prefix, tagfile; - while (std::getline(index_in, prefix)) { - if (std::getline(tags_in, tagfile)) { - index_tag_paths.push_back(std::make_pair(prefix, tagfile)); - } else { - std::cerr << "The number of tags specified does not match the number of " - "indices specified" - << std::endl; - exit(-1); + catch (const std::exception &ex) + { + std::cerr << ex.what() << std::endl; + return -1; } - } - index_in.close(); - tags_in.close(); - - if (data_type == std::string("float")) { - for (auto &index_tag : index_tag_paths) { - auto searcher = std::unique_ptr( - new diskann::PQFlashSearch(index_tag.first.c_str(), - num_nodes_to_cache, num_threads, - index_tag.second.c_str(), metric)); - g_ssdSearch.push_back(std::move(searcher)); + + diskann::Metric metric; + if (dist_fn == std::string("l2")) + metric = diskann::Metric::L2; + else if (dist_fn == std::string("mips")) + metric = diskann::Metric::INNER_PRODUCT; + else + { + std::cout << "Error. Only l2 and mips distance functions are supported" << std::endl; + return -1; + } + + std::vector> index_tag_paths; + std::ifstream index_in(index_prefix_paths); + if (!index_in.is_open()) + { + std::cerr << "Could not open " << index_prefix_paths << std::endl; + exit(-1); + } + std::ifstream tags_in(tags_file); + if (!tags_in.is_open()) + { + std::cerr << "Could not open " << tags_file << std::endl; + exit(-1); } - } else if (data_type == std::string("int8")) { - for (auto &index_tag : index_tag_paths) { - auto searcher = std::unique_ptr( - new diskann::PQFlashSearch(index_tag.first.c_str(), - num_nodes_to_cache, num_threads, - index_tag.second.c_str(), metric)); - g_ssdSearch.push_back(std::move(searcher)); + std::string prefix, tagfile; + while (std::getline(index_in, prefix)) + { + if (std::getline(tags_in, tagfile)) + { + index_tag_paths.push_back(std::make_pair(prefix, tagfile)); + } + else + { + std::cerr << "The number of tags specified does not match the number of " + "indices specified" + << std::endl; + exit(-1); + } } - } else if (data_type == std::string("uint8")) { - for (auto &index_tag : index_tag_paths) { - auto searcher = std::unique_ptr( - new diskann::PQFlashSearch( - index_tag.first.c_str(), num_nodes_to_cache, num_threads, - index_tag.second.c_str(), metric)); - g_ssdSearch.push_back(std::move(searcher)); + index_in.close(); + tags_in.close(); + + if (data_type == std::string("float")) + { + for (auto &index_tag : index_tag_paths) + { + auto searcher = std::unique_ptr(new diskann::PQFlashSearch( + index_tag.first.c_str(), num_nodes_to_cache, num_threads, index_tag.second.c_str(), metric)); + g_ssdSearch.push_back(std::move(searcher)); + } } - } else { - std::cerr << "Unsupported data type " << data_type << std::endl; - exit(-1); - } - - while (1) { - try { - setup(address, data_type); - std::cout << "Type 'exit' (case-sensitive) to exit" << std::endl; - std::string line; - std::getline(std::cin, line); - if (line == "exit") { - teardown(address); - g_httpServer->close().wait(); - exit(0); - } - } catch (const std::exception &ex) { - std::cerr << "Exception occurred: " << ex.what() << std::endl; - std::cerr << "Restarting HTTP server"; - teardown(address); - } catch (...) { - std::cerr << "Unknown exception occurreed" << std::endl; - std::cerr << "Restarting HTTP server"; - teardown(address); + else if (data_type == std::string("int8")) + { + for (auto &index_tag : index_tag_paths) + { + auto searcher = std::unique_ptr(new diskann::PQFlashSearch( + index_tag.first.c_str(), num_nodes_to_cache, num_threads, index_tag.second.c_str(), metric)); + g_ssdSearch.push_back(std::move(searcher)); + } + } + else if (data_type == std::string("uint8")) + { + for (auto &index_tag : index_tag_paths) + { + auto searcher = std::unique_ptr(new diskann::PQFlashSearch( + index_tag.first.c_str(), num_nodes_to_cache, num_threads, index_tag.second.c_str(), metric)); + g_ssdSearch.push_back(std::move(searcher)); + } + } + else + { + std::cerr << "Unsupported data type " << data_type << std::endl; + exit(-1); + } + + while (1) + { + try + { + setup(address, data_type); + std::cout << "Type 'exit' (case-sensitive) to exit" << std::endl; + std::string line; + std::getline(std::cin, line); + if (line == "exit") + { + teardown(address); + g_httpServer->close().wait(); + exit(0); + } + } + catch (const std::exception &ex) + { + std::cerr << "Exception occurred: " << ex.what() << std::endl; + std::cerr << "Restarting HTTP server"; + teardown(address); + } + catch (...) + { + std::cerr << "Unknown exception occurreed" << std::endl; + std::cerr << "Restarting HTTP server"; + teardown(address); + } } - } } diff --git a/tests/restapi/ssd_server.cpp b/tests/restapi/ssd_server.cpp index 9364e122e..d17997374 100644 --- a/tests/restapi/ssd_server.cpp +++ b/tests/restapi/ssd_server.cpp @@ -18,120 +18,124 @@ namespace po = boost::program_options; std::unique_ptr g_httpServer(nullptr); std::vector> g_ssdSearch; -void setup(const utility::string_t &address, const std::string &typestring) { - web::http::uri_builder uriBldr(address); - auto uri = uriBldr.to_uri(); +void setup(const utility::string_t &address, const std::string &typestring) +{ + web::http::uri_builder uriBldr(address); + auto uri = uriBldr.to_uri(); - std::cout << "Attempting to start server on " << uri.to_string() << std::endl; + std::cout << "Attempting to start server on " << uri.to_string() << std::endl; - g_httpServer = - std::unique_ptr(new Server(uri, g_ssdSearch, typestring)); - std::cout << "Created a server object" << std::endl; + g_httpServer = std::unique_ptr(new Server(uri, g_ssdSearch, typestring)); + std::cout << "Created a server object" << std::endl; - g_httpServer->open().wait(); - ucout << U"Listening for requests on: " << address << std::endl; + g_httpServer->open().wait(); + ucout << U"Listening for requests on: " << address << std::endl; } -void teardown(const utility::string_t &address) { - g_httpServer->close().wait(); +void teardown(const utility::string_t &address) +{ + g_httpServer->close().wait(); } -int main(int argc, char *argv[]) { - std::string data_type, index_path_prefix, address, dist_fn, tags_file; - uint32_t num_nodes_to_cache; - uint32_t num_threads; +int main(int argc, char *argv[]) +{ + std::string data_type, index_path_prefix, address, dist_fn, tags_file; + uint32_t num_nodes_to_cache; + uint32_t num_threads; - po::options_description desc{"Arguments"}; - try { - desc.add_options()("help,h", "Print information on arguments"); - desc.add_options()("data_type", - po::value(&data_type)->required(), - "data type "); - desc.add_options()("address", po::value(&address)->required(), - "Web server address"); - desc.add_options()("index_path_prefix", - po::value(&index_path_prefix)->required(), - "Path prefix for loading index file components"); - desc.add_options()( - "num_nodes_to_cache", - po::value(&num_nodes_to_cache)->default_value(0), - "Number of nodes to cache during search"); - desc.add_options()( - "num_threads,T", - po::value(&num_threads)->default_value(omp_get_num_procs()), - "Number of threads used for building index (defaults to " - "omp_get_num_procs())"); - desc.add_options()("dist_fn", - po::value(&dist_fn)->default_value("l2"), - "distance function "); - desc.add_options()( - "tags_file", - po::value(&tags_file)->default_value(std::string()), - "Tags file location"); - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) { - std::cout << desc; - return 0; + po::options_description desc{"Arguments"}; + try + { + desc.add_options()("help,h", "Print information on arguments"); + desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); + desc.add_options()("address", po::value(&address)->required(), "Web server address"); + desc.add_options()("index_path_prefix", po::value(&index_path_prefix)->required(), + "Path prefix for loading index file components"); + desc.add_options()("num_nodes_to_cache", po::value(&num_nodes_to_cache)->default_value(0), + "Number of nodes to cache during search"); + desc.add_options()("num_threads,T", po::value(&num_threads)->default_value(omp_get_num_procs()), + "Number of threads used for building index (defaults to " + "omp_get_num_procs())"); + desc.add_options()("dist_fn", po::value(&dist_fn)->default_value("l2"), + "distance function "); + desc.add_options()("tags_file", po::value(&tags_file)->default_value(std::string()), + "Tags file location"); + po::variables_map vm; + po::store(po::parse_command_line(argc, argv, desc), vm); + if (vm.count("help")) + { + std::cout << desc; + return 0; + } + po::notify(vm); + } + catch (const std::exception &ex) + { + std::cerr << ex.what() << std::endl; + return -1; } - po::notify(vm); - } catch (const std::exception &ex) { - std::cerr << ex.what() << std::endl; - return -1; - } - diskann::Metric metric; - if (dist_fn == std::string("l2")) - metric = diskann::Metric::L2; - else if (dist_fn == std::string("mips")) - metric = diskann::Metric::INNER_PRODUCT; - else { - std::cout << "Error. Only l2 and mips distance functions are supported" - << std::endl; - return -1; - } + diskann::Metric metric; + if (dist_fn == std::string("l2")) + metric = diskann::Metric::L2; + else if (dist_fn == std::string("mips")) + metric = diskann::Metric::INNER_PRODUCT; + else + { + std::cout << "Error. Only l2 and mips distance functions are supported" << std::endl; + return -1; + } - if (data_type == std::string("float")) { - auto searcher = std::unique_ptr( - new diskann::PQFlashSearch(index_path_prefix, num_nodes_to_cache, - num_threads, tags_file, metric)); - g_ssdSearch.push_back(std::move(searcher)); - } else if (data_type == std::string("int8")) { - auto searcher = - std::unique_ptr(new diskann::PQFlashSearch( - index_path_prefix, num_nodes_to_cache, num_threads, tags_file, - metric)); - g_ssdSearch.push_back(std::move(searcher)); - } else if (data_type == std::string("uint8")) { - auto searcher = std::unique_ptr( - new diskann::PQFlashSearch(index_path_prefix, - num_nodes_to_cache, num_threads, - tags_file, metric)); - g_ssdSearch.push_back(std::move(searcher)); - } else { - std::cerr << "Unsupported data type " << argv[2] << std::endl; - exit(-1); - } + if (data_type == std::string("float")) + { + auto searcher = std::unique_ptr( + new diskann::PQFlashSearch(index_path_prefix, num_nodes_to_cache, num_threads, tags_file, metric)); + g_ssdSearch.push_back(std::move(searcher)); + } + else if (data_type == std::string("int8")) + { + auto searcher = std::unique_ptr( + new diskann::PQFlashSearch(index_path_prefix, num_nodes_to_cache, num_threads, tags_file, metric)); + g_ssdSearch.push_back(std::move(searcher)); + } + else if (data_type == std::string("uint8")) + { + auto searcher = std::unique_ptr( + new diskann::PQFlashSearch(index_path_prefix, num_nodes_to_cache, num_threads, tags_file, metric)); + g_ssdSearch.push_back(std::move(searcher)); + } + else + { + std::cerr << "Unsupported data type " << argv[2] << std::endl; + exit(-1); + } - while (1) { - try { - setup(address, data_type); - std::cout << "Type 'exit' (case-sensitive) to exit" << std::endl; - std::string line; - std::getline(std::cin, line); - if (line == "exit") { - teardown(address); - g_httpServer->close().wait(); - exit(0); - } - } catch (const std::exception &ex) { - std::cerr << "Exception occurred: " << ex.what() << std::endl; - std::cerr << "Restarting HTTP server"; - teardown(address); - } catch (...) { - std::cerr << "Unknown exception occurreed" << std::endl; - std::cerr << "Restarting HTTP server"; - teardown(address); + while (1) + { + try + { + setup(address, data_type); + std::cout << "Type 'exit' (case-sensitive) to exit" << std::endl; + std::string line; + std::getline(std::cin, line); + if (line == "exit") + { + teardown(address); + g_httpServer->close().wait(); + exit(0); + } + } + catch (const std::exception &ex) + { + std::cerr << "Exception occurred: " << ex.what() << std::endl; + std::cerr << "Restarting HTTP server"; + teardown(address); + } + catch (...) + { + std::cerr << "Unknown exception occurreed" << std::endl; + std::cerr << "Restarting HTTP server"; + teardown(address); + } } - } } diff --git a/tests/search_disk_index.cpp b/tests/search_disk_index.cpp index e4a87ac7d..02d734c74 100644 --- a/tests/search_disk_index.cpp +++ b/tests/search_disk_index.cpp @@ -30,460 +30,452 @@ namespace po = boost::program_options; -void print_stats(std::string category, std::vector percentiles, - std::vector results) { - diskann::cout << std::setw(20) << category << ": " << std::flush; - for (uint32_t s = 0; s < percentiles.size(); s++) { - diskann::cout << std::setw(8) << percentiles[s] << "%"; - } - diskann::cout << std::endl; - diskann::cout << std::setw(22) << " " << std::flush; - for (uint32_t s = 0; s < percentiles.size(); s++) { - diskann::cout << std::setw(9) << results[s]; - } - diskann::cout << std::endl; +void print_stats(std::string category, std::vector percentiles, std::vector results) +{ + diskann::cout << std::setw(20) << category << ": " << std::flush; + for (uint32_t s = 0; s < percentiles.size(); s++) + { + diskann::cout << std::setw(8) << percentiles[s] << "%"; + } + diskann::cout << std::endl; + diskann::cout << std::setw(22) << " " << std::flush; + for (uint32_t s = 0; s < percentiles.size(); s++) + { + diskann::cout << std::setw(9) << results[s]; + } + diskann::cout << std::endl; } template -int search_disk_index( - diskann::Metric &metric, const std::string &index_path_prefix, - const std::string &result_output_prefix, const std::string &query_file, - std::string >_file, const uint32_t num_threads, const uint32_t recall_at, - const uint32_t beamwidth, const uint32_t num_nodes_to_cache, - const uint32_t search_io_limit, const std::vector &Lvec, - const float fail_if_recall_below, - const std::vector &query_filters, - const bool use_reorder_data = false) { - diskann::cout << "Search parameters: #threads: " << num_threads << ", "; - if (beamwidth <= 0) - diskann::cout << "beamwidth to be optimized for each L value" << std::flush; - else - diskann::cout << " beamwidth: " << beamwidth << std::flush; - if (search_io_limit == std::numeric_limits::max()) - diskann::cout << "." << std::endl; - else - diskann::cout << ", io_limit: " << search_io_limit << "." << std::endl; - - std::string warmup_query_file = index_path_prefix + "_sample_data.bin"; - - // load query bin - T *query = nullptr; - uint32_t *gt_ids = nullptr; - float *gt_dists = nullptr; - size_t query_num, query_dim, query_aligned_dim, gt_num, gt_dim; - diskann::load_aligned_bin(query_file, query, query_num, query_dim, - query_aligned_dim); - - bool filtered_search = false; - if (!query_filters.empty()) { - filtered_search = true; - if (query_filters.size() != 1 && query_filters.size() != query_num) { - std::cout << "Error. Mismatch in number of queries and size of query " - "filters file" - << std::endl; - return -1; // To return -1 or some other error handling? +int search_disk_index(diskann::Metric &metric, const std::string &index_path_prefix, + const std::string &result_output_prefix, const std::string &query_file, std::string >_file, + const uint32_t num_threads, const uint32_t recall_at, const uint32_t beamwidth, + const uint32_t num_nodes_to_cache, const uint32_t search_io_limit, + const std::vector &Lvec, const float fail_if_recall_below, + const std::vector &query_filters, const bool use_reorder_data = false) +{ + diskann::cout << "Search parameters: #threads: " << num_threads << ", "; + if (beamwidth <= 0) + diskann::cout << "beamwidth to be optimized for each L value" << std::flush; + else + diskann::cout << " beamwidth: " << beamwidth << std::flush; + if (search_io_limit == std::numeric_limits::max()) + diskann::cout << "." << std::endl; + else + diskann::cout << ", io_limit: " << search_io_limit << "." << std::endl; + + std::string warmup_query_file = index_path_prefix + "_sample_data.bin"; + + // load query bin + T *query = nullptr; + uint32_t *gt_ids = nullptr; + float *gt_dists = nullptr; + size_t query_num, query_dim, query_aligned_dim, gt_num, gt_dim; + diskann::load_aligned_bin(query_file, query, query_num, query_dim, query_aligned_dim); + + bool filtered_search = false; + if (!query_filters.empty()) + { + filtered_search = true; + if (query_filters.size() != 1 && query_filters.size() != query_num) + { + std::cout << "Error. Mismatch in number of queries and size of query " + "filters file" + << std::endl; + return -1; // To return -1 or some other error handling? + } } - } - - bool calc_recall_flag = false; - if (gt_file != std::string("null") && gt_file != std::string("NULL") && - file_exists(gt_file)) { - diskann::load_truthset(gt_file, gt_ids, gt_dists, gt_num, gt_dim); - if (gt_num != query_num) { - diskann::cout - << "Error. Mismatch in number of queries and ground truth data" - << std::endl; + + bool calc_recall_flag = false; + if (gt_file != std::string("null") && gt_file != std::string("NULL") && file_exists(gt_file)) + { + diskann::load_truthset(gt_file, gt_ids, gt_dists, gt_num, gt_dim); + if (gt_num != query_num) + { + diskann::cout << "Error. Mismatch in number of queries and ground truth data" << std::endl; + } + calc_recall_flag = true; } - calc_recall_flag = true; - } - std::shared_ptr reader = nullptr; + std::shared_ptr reader = nullptr; #ifdef _WINDOWS #ifndef USE_BING_INFRA - reader.reset(new WindowsAlignedFileReader()); + reader.reset(new WindowsAlignedFileReader()); #else - reader.reset(new diskann::BingAlignedFileReader()); + reader.reset(new diskann::BingAlignedFileReader()); #endif #else - reader.reset(new LinuxAlignedFileReader()); + reader.reset(new LinuxAlignedFileReader()); #endif - std::unique_ptr> _pFlashIndex( - new diskann::PQFlashIndex(reader, metric)); - - int res = _pFlashIndex->load(num_threads, index_path_prefix.c_str()); - - if (res != 0) { - return res; - } - // 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); - if (num_nodes_to_cache > 0) - _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_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, - warmup_dim, warmup_aligned_dim); - } else { - warmup_num = (std::min)((uint32_t)150000, (uint32_t)15000 * num_threads); - warmup_dim = query_dim; - warmup_aligned_dim = query_aligned_dim; - diskann::alloc_aligned(((void **)&warmup), - warmup_num * warmup_aligned_dim * sizeof(T), - 8 * sizeof(T)); - std::memset(warmup, 0, warmup_num * warmup_aligned_dim * sizeof(T)); - std::random_device rd; - std::mt19937 gen(rd()); - std::uniform_int_distribution<> dis(-128, 127); - for (uint32_t i = 0; i < warmup_num; i++) { - for (uint32_t d = 0; d < warmup_dim; d++) { - warmup[i * warmup_aligned_dim + d] = (T)dis(gen); - } - } + std::unique_ptr> _pFlashIndex( + new diskann::PQFlashIndex(reader, metric)); + + int res = _pFlashIndex->load(num_threads, index_path_prefix.c_str()); + + if (res != 0) + { + return res; } - diskann::cout << "Warming up index... " << std::flush; - std::vector warmup_result_ids_64(warmup_num, 0); - std::vector warmup_result_dists(warmup_num, 0); + // 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); + if (num_nodes_to_cache > 0) + _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_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, warmup_dim, warmup_aligned_dim); + } + else + { + warmup_num = (std::min)((uint32_t)150000, (uint32_t)15000 * num_threads); + warmup_dim = query_dim; + warmup_aligned_dim = query_aligned_dim; + diskann::alloc_aligned(((void **)&warmup), warmup_num * warmup_aligned_dim * sizeof(T), 8 * sizeof(T)); + std::memset(warmup, 0, warmup_num * warmup_aligned_dim * sizeof(T)); + std::random_device rd; + std::mt19937 gen(rd()); + std::uniform_int_distribution<> dis(-128, 127); + for (uint32_t i = 0; i < warmup_num; i++) + { + for (uint32_t d = 0; d < warmup_dim; d++) + { + warmup[i * warmup_aligned_dim + d] = (T)dis(gen); + } + } + } + diskann::cout << "Warming up index... " << std::flush; + std::vector warmup_result_ids_64(warmup_num, 0); + std::vector warmup_result_dists(warmup_num, 0); #pragma omp parallel for schedule(dynamic, 1) - for (int64_t 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); + for (int64_t 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); + } + diskann::cout << "..done" << std::endl; } - diskann::cout << "..done" << std::endl; - } - - diskann::cout.setf(std::ios_base::fixed, std::ios_base::floatfield); - diskann::cout.precision(2); - - std::string recall_string = "Recall@" + std::to_string(recall_at); - diskann::cout << std::setw(6) << "L" << std::setw(12) << "Beamwidth" - << std::setw(16) << "QPS" << std::setw(16) << "Mean Latency" - << std::setw(16) << "99.9 Latency" << std::setw(16) - << "Mean IOs" << std::setw(16) << "CPU (s)"; - if (calc_recall_flag) { - diskann::cout << std::setw(16) << recall_string << std::endl; - } else - diskann::cout << std::endl; - diskann::cout - << "===============================================================" - "=======================================================" - << std::endl; - std::vector> query_result_ids(Lvec.size()); - std::vector> query_result_dists(Lvec.size()); + diskann::cout.setf(std::ios_base::fixed, std::ios_base::floatfield); + diskann::cout.precision(2); - uint32_t optimized_beamwidth = 2; + std::string recall_string = "Recall@" + std::to_string(recall_at); + diskann::cout << std::setw(6) << "L" << std::setw(12) << "Beamwidth" << std::setw(16) << "QPS" << std::setw(16) + << "Mean Latency" << std::setw(16) << "99.9 Latency" << std::setw(16) << "Mean IOs" << std::setw(16) + << "CPU (s)"; + if (calc_recall_flag) + { + diskann::cout << std::setw(16) << recall_string << std::endl; + } + else + diskann::cout << std::endl; + diskann::cout << "===============================================================" + "=======================================================" + << std::endl; - float best_recall = 0.0; + std::vector> query_result_ids(Lvec.size()); + std::vector> query_result_dists(Lvec.size()); - for (uint32_t test_id = 0; test_id < Lvec.size(); test_id++) { - uint64_t L = Lvec[test_id]; + uint32_t optimized_beamwidth = 2; - if (L < recall_at) { - diskann::cout << "Ignoring search with L:" << L - << " since it's smaller than K:" << recall_at << std::endl; - continue; - } + float best_recall = 0.0; - if (beamwidth <= 0) { - diskann::cout << "Tuning beamwidth.." << std::endl; - optimized_beamwidth = - optimize_beamwidth(_pFlashIndex, warmup, warmup_num, - warmup_aligned_dim, L, optimized_beamwidth); - } else - optimized_beamwidth = beamwidth; + for (uint32_t test_id = 0; test_id < Lvec.size(); test_id++) + { + uint64_t L = Lvec[test_id]; - query_result_ids[test_id].resize(recall_at * query_num); - query_result_dists[test_id].resize(recall_at * query_num); + if (L < recall_at) + { + diskann::cout << "Ignoring search with L:" << L << " since it's smaller than K:" << recall_at << std::endl; + continue; + } - auto stats = new diskann::QueryStats[query_num]; + if (beamwidth <= 0) + { + diskann::cout << "Tuning beamwidth.." << std::endl; + optimized_beamwidth = + optimize_beamwidth(_pFlashIndex, warmup, warmup_num, warmup_aligned_dim, L, optimized_beamwidth); + } + else + optimized_beamwidth = beamwidth; - std::vector query_result_ids_64(recall_at * query_num); - auto s = std::chrono::high_resolution_clock::now(); + query_result_ids[test_id].resize(recall_at * query_num); + query_result_dists[test_id].resize(recall_at * query_num); + + auto stats = new diskann::QueryStats[query_num]; + + std::vector query_result_ids_64(recall_at * query_num); + auto s = std::chrono::high_resolution_clock::now(); #pragma omp parallel for schedule(dynamic, 1) - for (int64_t i = 0; i < (int64_t)query_num; i++) { - if (!filtered_search) { - _pFlashIndex->cached_beam_search( - query + (i * query_aligned_dim), recall_at, L, - query_result_ids_64.data() + (i * recall_at), - query_result_dists[test_id].data() + (i * recall_at), - optimized_beamwidth, use_reorder_data, stats + i); - } else { - LabelT label_for_search; - if (query_filters.size() == 1) { // one label for all queries - label_for_search = - _pFlashIndex->get_converted_label(query_filters[0]); - } else { // one label for each query - label_for_search = - _pFlashIndex->get_converted_label(query_filters[i]); + for (int64_t i = 0; i < (int64_t)query_num; i++) + { + if (!filtered_search) + { + _pFlashIndex->cached_beam_search(query + (i * query_aligned_dim), recall_at, L, + query_result_ids_64.data() + (i * recall_at), + query_result_dists[test_id].data() + (i * recall_at), + optimized_beamwidth, use_reorder_data, stats + i); + } + else + { + LabelT label_for_search; + if (query_filters.size() == 1) + { // one label for all queries + label_for_search = _pFlashIndex->get_converted_label(query_filters[0]); + } + else + { // one label for each query + label_for_search = _pFlashIndex->get_converted_label(query_filters[i]); + } + _pFlashIndex->cached_beam_search( + query + (i * query_aligned_dim), recall_at, L, query_result_ids_64.data() + (i * recall_at), + query_result_dists[test_id].data() + (i * recall_at), optimized_beamwidth, true, label_for_search, + use_reorder_data, 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()); + + diskann::convert_types(query_result_ids_64.data(), query_result_ids[test_id].data(), + query_num, recall_at); + + auto mean_latency = diskann::get_mean_stats( + stats, query_num, [](const diskann::QueryStats &stats) { return stats.total_us; }); + + auto latency_999 = diskann::get_percentile_stats( + stats, query_num, 0.999, [](const diskann::QueryStats &stats) { return stats.total_us; }); + + auto mean_ios = diskann::get_mean_stats(stats, query_num, + [](const diskann::QueryStats &stats) { return stats.n_ios; }); + + auto mean_cpuus = diskann::get_mean_stats(stats, query_num, + [](const diskann::QueryStats &stats) { return stats.cpu_us; }); + + 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); + best_recall = std::max(recall, best_recall); } - _pFlashIndex->cached_beam_search( - query + (i * query_aligned_dim), recall_at, L, - query_result_ids_64.data() + (i * recall_at), - query_result_dists[test_id].data() + (i * recall_at), - optimized_beamwidth, true, label_for_search, use_reorder_data, - stats + i); - } + + diskann::cout << std::setw(6) << L << std::setw(12) << optimized_beamwidth << std::setw(16) << qps + << std::setw(16) << mean_latency << std::setw(16) << latency_999 << std::setw(16) << mean_ios + << std::setw(16) << mean_cpuus; + if (calc_recall_flag) + { + diskann::cout << std::setw(16) << recall << std::endl; + } + else + diskann::cout << std::endl; + delete[] stats; } - auto e = std::chrono::high_resolution_clock::now(); - std::chrono::duration diff = e - s; - float qps = (1.0 * query_num) / (1.0 * diff.count()); - - diskann::convert_types(query_result_ids_64.data(), - query_result_ids[test_id].data(), - query_num, recall_at); - - auto mean_latency = diskann::get_mean_stats( - stats, query_num, - [](const diskann::QueryStats &stats) { return stats.total_us; }); - - auto latency_999 = diskann::get_percentile_stats( - stats, query_num, 0.999, - [](const diskann::QueryStats &stats) { return stats.total_us; }); - - auto mean_ios = diskann::get_mean_stats( - stats, query_num, - [](const diskann::QueryStats &stats) { return stats.n_ios; }); - - auto mean_cpuus = diskann::get_mean_stats( - stats, query_num, - [](const diskann::QueryStats &stats) { return stats.cpu_us; }); - - 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); - best_recall = std::max(recall, best_recall); + + diskann::cout << "Done searching. Now saving results " << std::endl; + uint64_t test_id = 0; + for (auto L : Lvec) + { + if (L < recall_at) + continue; + + std::string cur_result_path = result_output_prefix + "_" + std::to_string(L) + "_idx_uint32.bin"; + diskann::save_bin(cur_result_path, query_result_ids[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::cout << std::setw(6) << L << std::setw(12) << optimized_beamwidth - << std::setw(16) << qps << std::setw(16) << mean_latency - << std::setw(16) << latency_999 << std::setw(16) << mean_ios - << std::setw(16) << mean_cpuus; - if (calc_recall_flag) { - diskann::cout << std::setw(16) << recall << std::endl; - } else - diskann::cout << std::endl; - delete[] stats; - } - - diskann::cout << "Done searching. Now saving results " << std::endl; - uint64_t test_id = 0; - for (auto L : Lvec) { - if (L < recall_at) continue; - - std::string cur_result_path = - result_output_prefix + "_" + std::to_string(L) + "_idx_uint32.bin"; - diskann::save_bin(cur_result_path, - query_result_ids[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); - return best_recall >= fail_if_recall_below ? 0 : -1; + diskann::aligned_free(query); + if (warmup != nullptr) + diskann::aligned_free(warmup); + return best_recall >= fail_if_recall_below ? 0 : -1; } -int main(int argc, char **argv) { - std::string data_type, dist_fn, index_path_prefix, result_path_prefix, - query_file, gt_file, filter_label, label_type, query_filters_file; - uint32_t num_threads, K, W, num_nodes_to_cache, search_io_limit; - std::vector Lvec; - bool use_reorder_data = false; - float fail_if_recall_below = 0.0f; - - po::options_description desc{"Arguments"}; - try { - desc.add_options()("help,h", "Print information on arguments"); - desc.add_options()("data_type", - po::value(&data_type)->required(), - "data type "); - desc.add_options()("dist_fn", po::value(&dist_fn)->required(), - "distance function "); - desc.add_options()("index_path_prefix", - po::value(&index_path_prefix)->required(), - "Path prefix to the index"); - desc.add_options()("result_path", - po::value(&result_path_prefix)->required(), - "Path prefix for saving results of the queries"); - desc.add_options()("query_file", - po::value(&query_file)->required(), - "Query file in binary format"); - desc.add_options()( - "gt_file", - po::value(>_file)->default_value(std::string("null")), - "ground truth file for the queryset"); - desc.add_options()("recall_at,K", po::value(&K)->required(), - "Number of neighbors to be returned"); - desc.add_options()("search_list,L", - po::value>(&Lvec)->multitoken(), - "List of L values of search"); - desc.add_options()("beamwidth,W", po::value(&W)->default_value(2), - "Beamwidth for search. Set 0 to optimize internally."); - desc.add_options()( - "num_nodes_to_cache", - po::value(&num_nodes_to_cache)->default_value(0), - "Beamwidth for search"); - desc.add_options()( - "search_io_limit", - po::value(&search_io_limit) - ->default_value(std::numeric_limits::max()), - "Max #IOs for search"); - desc.add_options()( - "num_threads,T", - po::value(&num_threads)->default_value(omp_get_num_procs()), - "Number of threads used for building index (defaults to " - "omp_get_num_procs())"); - desc.add_options()("use_reorder_data", - po::bool_switch()->default_value(false), - "Include full precision data in the index. Use only in " - "conjuction with compressed data on SSD."); - desc.add_options()( - "filter_label", - po::value(&filter_label)->default_value(std::string("")), - "Filter Label for Filtered Search"); - desc.add_options()("query_filters_file", - po::value(&query_filters_file) - ->default_value(std::string("")), - "Filter file for Queries for Filtered Search "); - desc.add_options()( - "label_type", - po::value(&label_type)->default_value("uint"), - "Storage type of Labels , default value is uint which " - "will consume memory 4 bytes per filter"); - desc.add_options()( - "fail_if_recall_below", - po::value(&fail_if_recall_below)->default_value(0.0f), - "If set to a value >0 and <100%, program returns -1 if best recall " - "found is below this threshold. "); - - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) { - std::cout << desc; - return 0; +int main(int argc, char **argv) +{ + std::string data_type, dist_fn, index_path_prefix, result_path_prefix, query_file, gt_file, filter_label, + label_type, query_filters_file; + uint32_t num_threads, K, W, num_nodes_to_cache, search_io_limit; + std::vector Lvec; + bool use_reorder_data = false; + float fail_if_recall_below = 0.0f; + + po::options_description desc{"Arguments"}; + try + { + desc.add_options()("help,h", "Print information on arguments"); + desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); + desc.add_options()("dist_fn", po::value(&dist_fn)->required(), + "distance function "); + desc.add_options()("index_path_prefix", po::value(&index_path_prefix)->required(), + "Path prefix to the index"); + desc.add_options()("result_path", po::value(&result_path_prefix)->required(), + "Path prefix for saving results of the queries"); + desc.add_options()("query_file", po::value(&query_file)->required(), + "Query file in binary format"); + desc.add_options()("gt_file", po::value(>_file)->default_value(std::string("null")), + "ground truth file for the queryset"); + desc.add_options()("recall_at,K", po::value(&K)->required(), "Number of neighbors to be returned"); + desc.add_options()("search_list,L", po::value>(&Lvec)->multitoken(), + "List of L values of search"); + desc.add_options()("beamwidth,W", po::value(&W)->default_value(2), + "Beamwidth for search. Set 0 to optimize internally."); + desc.add_options()("num_nodes_to_cache", po::value(&num_nodes_to_cache)->default_value(0), + "Beamwidth for search"); + desc.add_options()("search_io_limit", + po::value(&search_io_limit)->default_value(std::numeric_limits::max()), + "Max #IOs for search"); + desc.add_options()("num_threads,T", po::value(&num_threads)->default_value(omp_get_num_procs()), + "Number of threads used for building index (defaults to " + "omp_get_num_procs())"); + desc.add_options()("use_reorder_data", po::bool_switch()->default_value(false), + "Include full precision data in the index. Use only in " + "conjuction with compressed data on SSD."); + desc.add_options()("filter_label", po::value(&filter_label)->default_value(std::string("")), + "Filter Label for Filtered Search"); + desc.add_options()("query_filters_file", + po::value(&query_filters_file)->default_value(std::string("")), + "Filter file for Queries for Filtered Search "); + desc.add_options()("label_type", po::value(&label_type)->default_value("uint"), + "Storage type of Labels , default value is uint which " + "will consume memory 4 bytes per filter"); + desc.add_options()("fail_if_recall_below", po::value(&fail_if_recall_below)->default_value(0.0f), + "If set to a value >0 and <100%, program returns -1 if best recall " + "found is below this threshold. "); + + po::variables_map vm; + po::store(po::parse_command_line(argc, argv, desc), vm); + if (vm.count("help")) + { + std::cout << desc; + return 0; + } + po::notify(vm); + if (vm["use_reorder_data"].as()) + use_reorder_data = true; + } + catch (const std::exception &ex) + { + std::cerr << ex.what() << '\n'; + return -1; + } + + diskann::Metric metric; + if (dist_fn == std::string("mips")) + { + metric = diskann::Metric::INNER_PRODUCT; + } + else if (dist_fn == std::string("l2")) + { + metric = diskann::Metric::L2; + } + else if (dist_fn == std::string("cosine")) + { + metric = diskann::Metric::COSINE; } - po::notify(vm); - if (vm["use_reorder_data"].as()) use_reorder_data = true; - } catch (const std::exception &ex) { - std::cerr << ex.what() << '\n'; - return -1; - } - - diskann::Metric metric; - if (dist_fn == std::string("mips")) { - metric = diskann::Metric::INNER_PRODUCT; - } else if (dist_fn == std::string("l2")) { - metric = diskann::Metric::L2; - } else if (dist_fn == std::string("cosine")) { - metric = diskann::Metric::COSINE; - } else { - std::cout << "Unsupported distance function. Currently only L2/ Inner " - "Product/Cosine are supported." - << std::endl; - return -1; - } - - if ((data_type != std::string("float")) && - (metric == diskann::Metric::INNER_PRODUCT)) { - std::cout << "Currently support only floating point data for Inner Product." - << std::endl; - return -1; - } - - if (use_reorder_data && data_type != std::string("float")) { - std::cout << "Error: Reorder data for reordering currently only " - "supported for float data type." - << std::endl; - return -1; - } - - if (filter_label != "" && query_filters_file != "") { - std::cerr - << "Only one of filter_label and query_filters_file should be provided" - << std::endl; - return -1; - } - - std::vector query_filters; - if (filter_label != "") { - query_filters.push_back(filter_label); - } else if (query_filters_file != "") { - query_filters = read_file_to_vector_of_strings(query_filters_file); - } - - try { - if (!query_filters.empty() && label_type == "ushort") { - if (data_type == std::string("float")) - return search_disk_index( - metric, index_path_prefix, result_path_prefix, query_file, gt_file, - num_threads, K, W, num_nodes_to_cache, search_io_limit, Lvec, - fail_if_recall_below, query_filters, use_reorder_data); - else if (data_type == std::string("int8")) - return search_disk_index( - metric, index_path_prefix, result_path_prefix, query_file, gt_file, - num_threads, K, W, num_nodes_to_cache, search_io_limit, Lvec, - fail_if_recall_below, query_filters, use_reorder_data); - else if (data_type == std::string("uint8")) - return search_disk_index( - metric, index_path_prefix, result_path_prefix, query_file, gt_file, - num_threads, K, W, num_nodes_to_cache, search_io_limit, Lvec, - fail_if_recall_below, query_filters, use_reorder_data); - else { - std::cerr << "Unsupported data type. Use float or int8 or uint8" + else + { + std::cout << "Unsupported distance function. Currently only L2/ Inner " + "Product/Cosine are supported." << std::endl; return -1; - } - } else { - if (data_type == std::string("float")) - return search_disk_index( - metric, index_path_prefix, result_path_prefix, query_file, gt_file, - num_threads, K, W, num_nodes_to_cache, search_io_limit, Lvec, - fail_if_recall_below, query_filters, use_reorder_data); - else if (data_type == std::string("int8")) - return search_disk_index( - metric, index_path_prefix, result_path_prefix, query_file, gt_file, - num_threads, K, W, num_nodes_to_cache, search_io_limit, Lvec, - fail_if_recall_below, query_filters, use_reorder_data); - else if (data_type == std::string("uint8")) - return search_disk_index( - metric, index_path_prefix, result_path_prefix, query_file, gt_file, - num_threads, K, W, num_nodes_to_cache, search_io_limit, Lvec, - fail_if_recall_below, query_filters, use_reorder_data); - else { - std::cerr << "Unsupported data type. Use float or int8 or uint8" + } + + if ((data_type != std::string("float")) && (metric == diskann::Metric::INNER_PRODUCT)) + { + std::cout << "Currently support only floating point data for Inner Product." << std::endl; + return -1; + } + + if (use_reorder_data && data_type != std::string("float")) + { + std::cout << "Error: Reorder data for reordering currently only " + "supported for float data type." << std::endl; return -1; - } } - } catch (const std::exception &e) { - std::cout << std::string(e.what()) << std::endl; - diskann::cerr << "Index search failed." << std::endl; - return -1; - } + + if (filter_label != "" && query_filters_file != "") + { + std::cerr << "Only one of filter_label and query_filters_file should be provided" << std::endl; + return -1; + } + + std::vector query_filters; + if (filter_label != "") + { + query_filters.push_back(filter_label); + } + else if (query_filters_file != "") + { + query_filters = read_file_to_vector_of_strings(query_filters_file); + } + + try + { + if (!query_filters.empty() && label_type == "ushort") + { + if (data_type == std::string("float")) + return search_disk_index( + metric, index_path_prefix, result_path_prefix, query_file, gt_file, num_threads, K, W, + num_nodes_to_cache, search_io_limit, Lvec, fail_if_recall_below, query_filters, use_reorder_data); + else if (data_type == std::string("int8")) + return search_disk_index( + metric, index_path_prefix, result_path_prefix, query_file, gt_file, num_threads, K, W, + num_nodes_to_cache, search_io_limit, Lvec, fail_if_recall_below, query_filters, use_reorder_data); + else if (data_type == std::string("uint8")) + return search_disk_index( + metric, index_path_prefix, result_path_prefix, query_file, gt_file, num_threads, K, W, + num_nodes_to_cache, search_io_limit, Lvec, fail_if_recall_below, query_filters, use_reorder_data); + else + { + std::cerr << "Unsupported data type. Use float or int8 or uint8" << std::endl; + return -1; + } + } + else + { + if (data_type == std::string("float")) + return search_disk_index(metric, index_path_prefix, result_path_prefix, query_file, gt_file, + num_threads, K, W, num_nodes_to_cache, search_io_limit, Lvec, + fail_if_recall_below, query_filters, use_reorder_data); + else if (data_type == std::string("int8")) + return search_disk_index(metric, index_path_prefix, result_path_prefix, query_file, gt_file, + num_threads, K, W, num_nodes_to_cache, search_io_limit, Lvec, + fail_if_recall_below, query_filters, use_reorder_data); + else if (data_type == std::string("uint8")) + return search_disk_index(metric, index_path_prefix, result_path_prefix, query_file, gt_file, + num_threads, K, W, num_nodes_to_cache, search_io_limit, Lvec, + fail_if_recall_below, query_filters, use_reorder_data); + else + { + std::cerr << "Unsupported data type. Use float or int8 or uint8" << std::endl; + return -1; + } + } + } + catch (const std::exception &e) + { + std::cout << std::string(e.what()) << std::endl; + diskann::cerr << "Index search failed." << std::endl; + return -1; + } } \ No newline at end of file diff --git a/tests/search_memory_index.cpp b/tests/search_memory_index.cpp index e9ba47c3e..02c96db24 100644 --- a/tests/search_memory_index.cpp +++ b/tests/search_memory_index.cpp @@ -24,395 +24,411 @@ 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 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) { - // Load the query file - T *query = nullptr; - uint32_t *gt_ids = nullptr; - float *gt_dists = nullptr; - size_t query_num, query_dim, query_aligned_dim, gt_num, gt_dim; - diskann::load_aligned_bin(query_file, query, query_num, query_dim, - query_aligned_dim); - - // Check for ground truth - bool calc_recall_flag = false; - if (truthset_file != std::string("null") && file_exists(truthset_file)) { - diskann::load_truthset(truthset_file, 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; +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 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) +{ + // Load the query file + T *query = nullptr; + uint32_t *gt_ids = nullptr; + float *gt_dists = nullptr; + size_t query_num, query_dim, query_aligned_dim, gt_num, gt_dim; + diskann::load_aligned_bin(query_file, query, query_num, query_dim, query_aligned_dim); + + // Check for ground truth + bool calc_recall_flag = false; + if (truthset_file != std::string("null") && file_exists(truthset_file)) + { + diskann::load_truthset(truthset_file, 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; + } + calc_recall_flag = true; + } + else + { + diskann::cout << " Truthset file " << truthset_file << " not found. Not computing recall." << std::endl; + } + + bool filtered_search = false; + if (!query_filters.empty()) + { + filtered_search = true; + if (query_filters.size() != 1 && query_filters.size() != query_num) + { + std::cout << "Error. Mismatch in number of queries and size of query " + "filters file" + << std::endl; + return -1; // To return -1 or some other error handling? + } + } + + using TagT = uint32_t; + const bool concurrent = false, pq_dist_build = false, use_opq = false; + const size_t num_pq_chunks = 0; + using IndexType = diskann::Index; + const size_t num_frozen_pts = IndexType::get_graph_num_frozen_points(index_path); + IndexType index(metric, query_dim, 0, dynamic, tags, concurrent, pq_dist_build, num_pq_chunks, use_opq, + num_frozen_pts); + std::cout << "Index class instantiated" << std::endl; + index.load(index_path.c_str(), num_threads, *(std::max_element(Lvec.begin(), Lvec.end()))); + std::cout << "Index loaded" << std::endl; + if (metric == diskann::FAST_L2) + index.optimize_index_layout(); + + std::cout << "Using " << num_threads << " threads to search" << std::endl; + std::cout.setf(std::ios_base::fixed, std::ios_base::floatfield); + std::cout.precision(2); + const std::string qps_title = show_qps_per_thread ? "QPS/thread" : "QPS"; + uint32_t table_width = 0; + if (tags) + { + std::cout << std::setw(4) << "Ls" << std::setw(12) << qps_title << std::setw(20) << "Mean Latency (mus)" + << std::setw(15) << "99.9 Latency"; + table_width += 4 + 12 + 20 + 15; + } + else + { + std::cout << std::setw(4) << "Ls" << std::setw(12) << qps_title << std::setw(18) << "Avg dist cmps" + << std::setw(20) << "Mean Latency (mus)" << std::setw(15) << "99.9 Latency"; + table_width += 4 + 12 + 18 + 20 + 15; } - calc_recall_flag = true; - } else { - diskann::cout << " Truthset file " << truthset_file - << " not found. Not computing recall." << std::endl; - } - - bool filtered_search = false; - if (!query_filters.empty()) { - filtered_search = true; - if (query_filters.size() != 1 && query_filters.size() != query_num) { - std::cout << "Error. Mismatch in number of queries and size of query " - "filters file" - << std::endl; - return -1; // To return -1 or some other error handling? + uint32_t recalls_to_print = 0; + const uint32_t first_recall = print_all_recalls ? 1 : recall_at; + if (calc_recall_flag) + { + for (uint32_t curr_recall = first_recall; curr_recall <= recall_at; curr_recall++) + { + std::cout << std::setw(12) << ("Recall@" + std::to_string(curr_recall)); + } + recalls_to_print = recall_at + 1 - first_recall; + table_width += recalls_to_print * 12; } - } - - using TagT = uint32_t; - const bool concurrent = false, pq_dist_build = false, use_opq = false; - const size_t num_pq_chunks = 0; - using IndexType = diskann::Index; - const size_t num_frozen_pts = - IndexType::get_graph_num_frozen_points(index_path); - IndexType index(metric, query_dim, 0, dynamic, tags, concurrent, - pq_dist_build, num_pq_chunks, use_opq, num_frozen_pts); - std::cout << "Index class instantiated" << std::endl; - index.load(index_path.c_str(), num_threads, - *(std::max_element(Lvec.begin(), Lvec.end()))); - std::cout << "Index loaded" << std::endl; - if (metric == diskann::FAST_L2) index.optimize_index_layout(); - - std::cout << "Using " << num_threads << " threads to search" << std::endl; - std::cout.setf(std::ios_base::fixed, std::ios_base::floatfield); - std::cout.precision(2); - const std::string qps_title = show_qps_per_thread ? "QPS/thread" : "QPS"; - uint32_t table_width = 0; - if (tags) { - std::cout << std::setw(4) << "Ls" << std::setw(12) << qps_title - << std::setw(20) << "Mean Latency (mus)" << std::setw(15) - << "99.9 Latency"; - table_width += 4 + 12 + 20 + 15; - } else { - std::cout << std::setw(4) << "Ls" << std::setw(12) << qps_title - << std::setw(18) << "Avg dist cmps" << std::setw(20) - << "Mean Latency (mus)" << std::setw(15) << "99.9 Latency"; - table_width += 4 + 12 + 18 + 20 + 15; - } - uint32_t recalls_to_print = 0; - const uint32_t first_recall = print_all_recalls ? 1 : recall_at; - if (calc_recall_flag) { - for (uint32_t curr_recall = first_recall; curr_recall <= recall_at; - curr_recall++) { - std::cout << std::setw(12) << ("Recall@" + std::to_string(curr_recall)); + std::cout << std::endl; + std::cout << std::string(table_width, '=') << std::endl; + + std::vector> query_result_ids(Lvec.size()); + std::vector> query_result_dists(Lvec.size()); + std::vector latency_stats(query_num, 0); + std::vector cmp_stats; + if (not tags) + { + cmp_stats = std::vector(query_num, 0); } - recalls_to_print = recall_at + 1 - first_recall; - table_width += recalls_to_print * 12; - } - std::cout << std::endl; - std::cout << std::string(table_width, '=') << std::endl; - - std::vector> query_result_ids(Lvec.size()); - std::vector> query_result_dists(Lvec.size()); - std::vector latency_stats(query_num, 0); - std::vector cmp_stats; - if (not tags) { - cmp_stats = std::vector(query_num, 0); - } - - std::vector query_result_tags; - if (tags) { - query_result_tags.resize(recall_at * query_num); - } - - float best_recall = 0.0; - - for (uint32_t test_id = 0; test_id < Lvec.size(); test_id++) { - uint64_t L = Lvec[test_id]; - if (L < recall_at) { - diskann::cout << "Ignoring search with L:" << L - << " since it's smaller than K:" << recall_at << std::endl; - continue; + + std::vector query_result_tags; + if (tags) + { + query_result_tags.resize(recall_at * query_num); } - query_result_ids[test_id].resize(recall_at * query_num); - query_result_dists[test_id].resize(recall_at * query_num); - std::vector res = std::vector(); + float best_recall = 0.0; + + for (uint32_t test_id = 0; test_id < Lvec.size(); test_id++) + { + uint64_t L = Lvec[test_id]; + if (L < recall_at) + { + diskann::cout << "Ignoring search with L:" << L << " since it's smaller than K:" << recall_at << std::endl; + continue; + } + + query_result_ids[test_id].resize(recall_at * query_num); + query_result_dists[test_id].resize(recall_at * query_num); + std::vector res = std::vector(); - auto s = std::chrono::high_resolution_clock::now(); - omp_set_num_threads(num_threads); + auto s = std::chrono::high_resolution_clock::now(); + omp_set_num_threads(num_threads); #pragma omp parallel for schedule(dynamic, 1) - for (int64_t i = 0; i < (int64_t)query_num; i++) { - auto qs = std::chrono::high_resolution_clock::now(); - if (filtered_search) { - LabelT filter_label_as_num; - if (query_filters.size() == 1) { - filter_label_as_num = index.get_converted_label(query_filters[0]); - } else { - filter_label_as_num = index.get_converted_label(query_filters[i]); + for (int64_t i = 0; i < (int64_t)query_num; i++) + { + auto qs = std::chrono::high_resolution_clock::now(); + if (filtered_search) + { + LabelT filter_label_as_num; + if (query_filters.size() == 1) + { + filter_label_as_num = index.get_converted_label(query_filters[0]); + } + else + { + filter_label_as_num = index.get_converted_label(query_filters[i]); + } + auto retval = index.search_with_filters(query + i * query_aligned_dim, filter_label_as_num, recall_at, + L, query_result_ids[test_id].data() + i * recall_at, + query_result_dists[test_id].data() + i * recall_at); + cmp_stats[i] = retval.second; + } + else if (metric == diskann::FAST_L2) + { + index.search_with_optimized_layout(query + i * query_aligned_dim, recall_at, L, + query_result_ids[test_id].data() + i * recall_at); + } + else if (tags) + { + index.search_with_tags(query + i * query_aligned_dim, recall_at, L, + query_result_tags.data() + i * recall_at, nullptr, res); + for (int64_t r = 0; r < (int64_t)recall_at; r++) + { + query_result_ids[test_id][recall_at * i + r] = query_result_tags[recall_at * i + r]; + } + } + else + { + cmp_stats[i] = index + .search(query + i * query_aligned_dim, recall_at, L, + query_result_ids[test_id].data() + i * recall_at) + .second; + } + auto qe = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = qe - qs; + latency_stats[i] = diff.count() * 1000000; } - auto retval = index.search_with_filters( - query + i * query_aligned_dim, filter_label_as_num, recall_at, L, - query_result_ids[test_id].data() + i * recall_at, - query_result_dists[test_id].data() + i * recall_at); - cmp_stats[i] = retval.second; - } else if (metric == diskann::FAST_L2) { - index.search_with_optimized_layout( - query + i * query_aligned_dim, recall_at, L, - query_result_ids[test_id].data() + i * recall_at); - } else if (tags) { - index.search_with_tags(query + i * query_aligned_dim, recall_at, L, - query_result_tags.data() + i * recall_at, - nullptr, res); - for (int64_t r = 0; r < (int64_t)recall_at; r++) { - query_result_ids[test_id][recall_at * i + r] = - query_result_tags[recall_at * i + r]; + std::chrono::duration diff = std::chrono::high_resolution_clock::now() - s; + + float displayed_qps = static_cast(query_num) / diff.count(); + + if (show_qps_per_thread) + displayed_qps /= num_threads; + + std::vector recalls; + if (calc_recall_flag) + { + recalls.reserve(recalls_to_print); + for (uint32_t curr_recall = first_recall; curr_recall <= recall_at; curr_recall++) + { + recalls.push_back(diskann::calculate_recall(query_num, gt_ids, gt_dists, gt_dim, + query_result_ids[test_id].data(), recall_at, curr_recall)); + } } - } else { - cmp_stats[i] = - index - .search(query + i * query_aligned_dim, recall_at, L, - query_result_ids[test_id].data() + i * recall_at) - .second; - } - auto qe = std::chrono::high_resolution_clock::now(); - std::chrono::duration diff = qe - qs; - latency_stats[i] = diff.count() * 1000000; - } - std::chrono::duration diff = - std::chrono::high_resolution_clock::now() - s; - - float displayed_qps = static_cast(query_num) / diff.count(); - - if (show_qps_per_thread) displayed_qps /= num_threads; - - std::vector recalls; - if (calc_recall_flag) { - recalls.reserve(recalls_to_print); - for (uint32_t curr_recall = first_recall; curr_recall <= recall_at; - curr_recall++) { - recalls.push_back(diskann::calculate_recall( - query_num, gt_ids, gt_dists, gt_dim, - query_result_ids[test_id].data(), recall_at, curr_recall)); - } + + std::sort(latency_stats.begin(), latency_stats.end()); + float mean_latency = + std::accumulate(latency_stats.begin(), latency_stats.end(), 0.0) / static_cast(query_num); + + float avg_cmps = (float)std::accumulate(cmp_stats.begin(), cmp_stats.end(), 0) / (float)query_num; + + if (tags) + { + std::cout << std::setw(4) << L << std::setw(12) << displayed_qps << std::setw(20) << (float)mean_latency + << std::setw(15) << (float)latency_stats[(uint64_t)(0.999 * query_num)]; + } + else + { + std::cout << std::setw(4) << L << std::setw(12) << displayed_qps << std::setw(18) << avg_cmps + << std::setw(20) << (float)mean_latency << std::setw(15) + << (float)latency_stats[(uint64_t)(0.999 * query_num)]; + } + for (float recall : recalls) + { + std::cout << std::setw(12) << recall; + best_recall = std::max(recall, best_recall); + } + std::cout << std::endl; } - std::sort(latency_stats.begin(), latency_stats.end()); - float mean_latency = - std::accumulate(latency_stats.begin(), latency_stats.end(), 0.0) / - static_cast(query_num); - - float avg_cmps = - (float)std::accumulate(cmp_stats.begin(), cmp_stats.end(), 0) / - (float)query_num; - - if (tags) { - std::cout << std::setw(4) << L << std::setw(12) << displayed_qps - << std::setw(20) << (float)mean_latency << std::setw(15) - << (float)latency_stats[(uint64_t)(0.999 * query_num)]; - } else { - std::cout << std::setw(4) << L << std::setw(12) << displayed_qps - << std::setw(18) << avg_cmps << std::setw(20) - << (float)mean_latency << std::setw(15) - << (float)latency_stats[(uint64_t)(0.999 * query_num)]; + std::cout << "Done searching. Now saving results " << std::endl; + uint64_t test_id = 0; + for (auto L : Lvec) + { + if (L < recall_at) + { + diskann::cout << "Ignoring search with L:" << L << " since it's smaller than K:" << recall_at << std::endl; + continue; + } + std::string cur_result_path = result_path_prefix + "_" + std::to_string(L) + "_idx_uint32.bin"; + diskann::save_bin(cur_result_path, query_result_ids[test_id].data(), query_num, recall_at); + test_id++; } - for (float recall : recalls) { - std::cout << std::setw(12) << recall; - best_recall = std::max(recall, best_recall); + + diskann::aligned_free(query); + + return best_recall >= fail_if_recall_below ? 0 : -1; +} + +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; + std::vector Lvec; + bool print_all_recalls, dynamic, tags, show_qps_per_thread; + float fail_if_recall_below = 0.0f; + + po::options_description desc{"Arguments"}; + try + { + desc.add_options()("help,h", "Print information on arguments"); + desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); + desc.add_options()("dist_fn", po::value(&dist_fn)->required(), + "distance function "); + desc.add_options()("index_path_prefix", po::value(&index_path_prefix)->required(), + "Path prefix to the index"); + desc.add_options()("result_path", po::value(&result_path)->required(), + "Path prefix for saving results of the queries"); + desc.add_options()("query_file", po::value(&query_file)->required(), + "Query file in binary format"); + desc.add_options()("filter_label", po::value(&filter_label)->default_value(std::string("")), + "Filter Label for Filtered Search"); + desc.add_options()("query_filters_file", + po::value(&query_filters_file)->default_value(std::string("")), + "Filter file for Queries for Filtered Search "); + desc.add_options()("label_type", po::value(&label_type)->default_value("uint"), + "Storage type of Labels , default value is uint which " + "will consume memory 4 bytes per filter"); + desc.add_options()("gt_file", po::value(>_file)->default_value(std::string("null")), + "ground truth file for the queryset"); + desc.add_options()("recall_at,K", po::value(&K)->required(), "Number of neighbors to be returned"); + desc.add_options()("print_all_recalls", po::bool_switch(&print_all_recalls), + "Print recalls at all positions, from 1 up to specified " + "recall_at value"); + desc.add_options()("search_list,L", po::value>(&Lvec)->multitoken(), + "List of L values of search"); + desc.add_options()("num_threads,T", po::value(&num_threads)->default_value(omp_get_num_procs()), + "Number of threads used for building index (defaults to " + "omp_get_num_procs())"); + desc.add_options()("dynamic", po::value(&dynamic)->default_value(false), + "Whether the index is dynamic. Default false."); + desc.add_options()("tags", po::value(&tags)->default_value(false), + "Whether to search with tags. Default false."); + desc.add_options()("qps_per_thread", po::bool_switch(&show_qps_per_thread), + "Print overall QPS divided by the number of threads in " + "the output table"); + desc.add_options()("fail_if_recall_below", po::value(&fail_if_recall_below)->default_value(0.0f), + "If set to a value >0 and <100%, program returns -1 if best recall " + "found is below this threshold. "); + + po::variables_map vm; + po::store(po::parse_command_line(argc, argv, desc), vm); + if (vm.count("help")) + { + std::cout << desc; + return 0; + } + po::notify(vm); } - std::cout << std::endl; - } - - std::cout << "Done searching. Now saving results " << std::endl; - uint64_t test_id = 0; - for (auto L : Lvec) { - if (L < recall_at) { - diskann::cout << "Ignoring search with L:" << L - << " since it's smaller than K:" << recall_at << std::endl; - continue; + catch (const std::exception &ex) + { + std::cerr << ex.what() << '\n'; + return -1; } - std::string cur_result_path = - result_path_prefix + "_" + std::to_string(L) + "_idx_uint32.bin"; - diskann::save_bin(cur_result_path, - query_result_ids[test_id].data(), query_num, - recall_at); - test_id++; - } - diskann::aligned_free(query); + diskann::Metric metric; + if ((dist_fn == std::string("mips")) && (data_type == std::string("float"))) + { + metric = diskann::Metric::INNER_PRODUCT; + } + else if (dist_fn == std::string("l2")) + { + metric = diskann::Metric::L2; + } + else if (dist_fn == std::string("cosine")) + { + metric = diskann::Metric::COSINE; + } + else if ((dist_fn == std::string("fast_l2")) && (data_type == std::string("float"))) + { + metric = diskann::Metric::FAST_L2; + } + else + { + std::cout << "Unsupported distance function. Currently only l2/ cosine are " + "supported in general, and mips/fast_l2 only for floating " + "point data." + << std::endl; + return -1; + } - return best_recall >= fail_if_recall_below ? 0 : -1; -} + if (dynamic && not tags) + { + std::cerr << "Tags must be enabled while searching dynamically built indices" << std::endl; + return -1; + } -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; - std::vector Lvec; - bool print_all_recalls, dynamic, tags, show_qps_per_thread; - float fail_if_recall_below = 0.0f; - - po::options_description desc{"Arguments"}; - try { - desc.add_options()("help,h", "Print information on arguments"); - desc.add_options()("data_type", - po::value(&data_type)->required(), - "data type "); - desc.add_options()("dist_fn", po::value(&dist_fn)->required(), - "distance function "); - desc.add_options()("index_path_prefix", - po::value(&index_path_prefix)->required(), - "Path prefix to the index"); - desc.add_options()("result_path", - po::value(&result_path)->required(), - "Path prefix for saving results of the queries"); - desc.add_options()("query_file", - po::value(&query_file)->required(), - "Query file in binary format"); - desc.add_options()( - "filter_label", - po::value(&filter_label)->default_value(std::string("")), - "Filter Label for Filtered Search"); - desc.add_options()("query_filters_file", - po::value(&query_filters_file) - ->default_value(std::string("")), - "Filter file for Queries for Filtered Search "); - desc.add_options()( - "label_type", - po::value(&label_type)->default_value("uint"), - "Storage type of Labels , default value is uint which " - "will consume memory 4 bytes per filter"); - desc.add_options()( - "gt_file", - po::value(>_file)->default_value(std::string("null")), - "ground truth file for the queryset"); - desc.add_options()("recall_at,K", po::value(&K)->required(), - "Number of neighbors to be returned"); - desc.add_options()("print_all_recalls", po::bool_switch(&print_all_recalls), - "Print recalls at all positions, from 1 up to specified " - "recall_at value"); - desc.add_options()("search_list,L", - po::value>(&Lvec)->multitoken(), - "List of L values of search"); - desc.add_options()( - "num_threads,T", - po::value(&num_threads)->default_value(omp_get_num_procs()), - "Number of threads used for building index (defaults to " - "omp_get_num_procs())"); - desc.add_options()("dynamic", - po::value(&dynamic)->default_value(false), - "Whether the index is dynamic. Default false."); - desc.add_options()("tags", po::value(&tags)->default_value(false), - "Whether to search with tags. Default false."); - desc.add_options()("qps_per_thread", po::bool_switch(&show_qps_per_thread), - "Print overall QPS divided by the number of threads in " - "the output table"); - desc.add_options()( - "fail_if_recall_below", - po::value(&fail_if_recall_below)->default_value(0.0f), - "If set to a value >0 and <100%, program returns -1 if best recall " - "found is below this threshold. "); - - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) { - std::cout << desc; - return 0; + if (fail_if_recall_below < 0.0 || fail_if_recall_below >= 100.0) + { + std::cerr << "fail_if_recall_below parameter must be between 0 and 100%" << std::endl; + return -1; } - po::notify(vm); - } catch (const std::exception &ex) { - std::cerr << ex.what() << '\n'; - return -1; - } - - diskann::Metric metric; - if ((dist_fn == std::string("mips")) && (data_type == std::string("float"))) { - metric = diskann::Metric::INNER_PRODUCT; - } else if (dist_fn == std::string("l2")) { - metric = diskann::Metric::L2; - } else if (dist_fn == std::string("cosine")) { - metric = diskann::Metric::COSINE; - } else if ((dist_fn == std::string("fast_l2")) && - (data_type == std::string("float"))) { - metric = diskann::Metric::FAST_L2; - } else { - std::cout << "Unsupported distance function. Currently only l2/ cosine are " - "supported in general, and mips/fast_l2 only for floating " - "point data." - << std::endl; - return -1; - } - - if (dynamic && not tags) { - std::cerr - << "Tags must be enabled while searching dynamically built indices" - << std::endl; - return -1; - } - - if (fail_if_recall_below < 0.0 || fail_if_recall_below >= 100.0) { - std::cerr << "fail_if_recall_below parameter must be between 0 and 100%" - << std::endl; - return -1; - } - - if (filter_label != "" && query_filters_file != "") { - std::cerr - << "Only one of filter_label and query_filters_file should be provided" - << std::endl; - return -1; - } - - std::vector query_filters; - if (filter_label != "") { - query_filters.push_back(filter_label); - } else if (query_filters_file != "") { - query_filters = read_file_to_vector_of_strings(query_filters_file); - } - - try { - if (!query_filters.empty() && label_type == "ushort") { - 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, 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, 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, - num_threads, K, print_all_recalls, Lvec, dynamic, tags, - show_qps_per_thread, query_filters, fail_if_recall_below); - } else { - std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; + + if (filter_label != "" && query_filters_file != "") + { + std::cerr << "Only one of filter_label and query_filters_file should be provided" << std::endl; return -1; - } - } else { - 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, 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, 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, - num_threads, K, print_all_recalls, Lvec, dynamic, tags, - show_qps_per_thread, query_filters, fail_if_recall_below); - } else { - std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; + } + + std::vector query_filters; + if (filter_label != "") + { + query_filters.push_back(filter_label); + } + else if (query_filters_file != "") + { + query_filters = read_file_to_vector_of_strings(query_filters_file); + } + + try + { + if (!query_filters.empty() && label_type == "ushort") + { + 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, + 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, + 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, + num_threads, K, print_all_recalls, Lvec, dynamic, tags, + show_qps_per_thread, query_filters, fail_if_recall_below); + } + else + { + std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; + return -1; + } + } + else + { + 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, 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, 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, + num_threads, K, print_all_recalls, Lvec, dynamic, tags, + show_qps_per_thread, query_filters, fail_if_recall_below); + } + else + { + std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; + return -1; + } + } + } + catch (std::exception &e) + { + std::cout << std::string(e.what()) << std::endl; + diskann::cerr << "Index search failed." << std::endl; return -1; - } } - } catch (std::exception &e) { - std::cout << std::string(e.what()) << std::endl; - diskann::cerr << "Index search failed." << std::endl; - return -1; - } } diff --git a/tests/test_insert_deletes_consolidate.cpp b/tests/test_insert_deletes_consolidate.cpp index 75579fb1f..c42972402 100644 --- a/tests/test_insert_deletes_consolidate.cpp +++ b/tests/test_insert_deletes_consolidate.cpp @@ -25,452 +25,406 @@ namespace po = boost::program_options; // load_aligned_bin modified to read pieces of the file, but using ifstream // instead of cached_ifstream. template -inline void load_aligned_bin_part(const std::string &bin_file, T *data, - size_t offset_points, size_t points_to_read) { - diskann::Timer timer; - std::ifstream reader; - reader.exceptions(std::ios::failbit | std::ios::badbit); - reader.open(bin_file, std::ios::binary | std::ios::ate); - size_t actual_file_size = reader.tellg(); - reader.seekg(0, std::ios::beg); - - int npts_i32, dim_i32; - reader.read((char *)&npts_i32, sizeof(int)); - reader.read((char *)&dim_i32, sizeof(int)); - size_t npts = (uint32_t)npts_i32; - size_t dim = (uint32_t)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; - std::cout << stream.str(); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); - } - - if (offset_points + points_to_read > npts) { - std::stringstream stream; - stream << "Error. Not enough points in file. Requested " << offset_points - << " offset and " << points_to_read << " points, but have only " - << npts << " points" << std::endl; - std::cout << stream.str(); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); - } - - reader.seekg(2 * sizeof(uint32_t) + offset_points * dim * sizeof(T)); - - const size_t rounded_dim = ROUND_UP(dim, 8); - - for (size_t i = 0; i < points_to_read; i++) { - reader.read((char *)(data + i * rounded_dim), dim * sizeof(T)); - memset(data + i * rounded_dim + dim, 0, (rounded_dim - dim) * sizeof(T)); - } - reader.close(); - - const double elapsedSeconds = timer.elapsed() / 1000000.0; - std::cout << "Read " << points_to_read << " points using non-cached reads in " - << elapsedSeconds << std::endl; +inline void load_aligned_bin_part(const std::string &bin_file, T *data, size_t offset_points, size_t points_to_read) +{ + diskann::Timer timer; + std::ifstream reader; + reader.exceptions(std::ios::failbit | std::ios::badbit); + reader.open(bin_file, std::ios::binary | std::ios::ate); + size_t actual_file_size = reader.tellg(); + reader.seekg(0, std::ios::beg); + + int npts_i32, dim_i32; + reader.read((char *)&npts_i32, sizeof(int)); + reader.read((char *)&dim_i32, sizeof(int)); + size_t npts = (uint32_t)npts_i32; + size_t dim = (uint32_t)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; + std::cout << stream.str(); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } + + if (offset_points + points_to_read > npts) + { + std::stringstream stream; + stream << "Error. Not enough points in file. Requested " << offset_points << " offset and " << points_to_read + << " points, but have only " << npts << " points" << std::endl; + std::cout << stream.str(); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } + + reader.seekg(2 * sizeof(uint32_t) + offset_points * dim * sizeof(T)); + + const size_t rounded_dim = ROUND_UP(dim, 8); + + for (size_t i = 0; i < points_to_read; i++) + { + reader.read((char *)(data + i * rounded_dim), dim * sizeof(T)); + memset(data + i * rounded_dim + dim, 0, (rounded_dim - dim) * sizeof(T)); + } + reader.close(); + + const double elapsedSeconds = timer.elapsed() / 1000000.0; + std::cout << "Read " << points_to_read << " points using non-cached reads in " << elapsedSeconds << std::endl; } -std::string get_save_filename(const std::string &save_path, - size_t points_to_skip, size_t points_deleted, - size_t last_point_threshold) { - std::string final_path = save_path; - if (points_to_skip > 0) { - final_path += "skip" + std::to_string(points_to_skip) + "-"; - } - - final_path += "del" + std::to_string(points_deleted) + "-"; - final_path += std::to_string(last_point_threshold); - return final_path; +std::string get_save_filename(const std::string &save_path, size_t points_to_skip, size_t points_deleted, + size_t last_point_threshold) +{ + std::string final_path = save_path; + if (points_to_skip > 0) + { + final_path += "skip" + std::to_string(points_to_skip) + "-"; + } + + final_path += "del" + std::to_string(points_deleted) + "-"; + final_path += std::to_string(last_point_threshold); + return final_path; } template -void insert_till_next_checkpoint(diskann::Index &index, size_t start, - size_t end, size_t thread_count, T *data, - size_t aligned_dim) { - diskann::Timer insert_timer; +void insert_till_next_checkpoint(diskann::Index &index, size_t start, size_t end, size_t thread_count, T *data, + size_t aligned_dim) +{ + diskann::Timer insert_timer; #pragma omp parallel for num_threads(thread_count) schedule(dynamic) - for (int64_t j = start; j < (int64_t)end; j++) { - index.insert_point(&data[(j - start) * aligned_dim], - 1 + static_cast(j)); - } - const double elapsedSeconds = insert_timer.elapsed() / 1000000.0; - std::cout << "Insertion time " << elapsedSeconds << " seconds (" - << (end - start) / elapsedSeconds << " points/second overall, " - << (end - start) / elapsedSeconds / thread_count - << " per thread)\n "; + for (int64_t j = start; j < (int64_t)end; j++) + { + index.insert_point(&data[(j - start) * aligned_dim], 1 + static_cast(j)); + } + const double elapsedSeconds = insert_timer.elapsed() / 1000000.0; + std::cout << "Insertion time " << elapsedSeconds << " seconds (" << (end - start) / elapsedSeconds + << " points/second overall, " << (end - start) / elapsedSeconds / thread_count << " per thread)\n "; } template -void delete_from_beginning(diskann::Index &index, - diskann::IndexWriteParameters &delete_params, - size_t points_to_skip, - size_t points_to_delete_from_beginning) { - try { - std::cout << std::endl - << "Lazy deleting points " << points_to_skip << " to " - << points_to_skip + points_to_delete_from_beginning << "... "; - for (size_t i = points_to_skip; - i < points_to_skip + points_to_delete_from_beginning; ++i) - index.lazy_delete(i + 1); // Since tags are data location + 1 - std::cout << "done." << std::endl; - - auto report = index.consolidate_deletes(delete_params); - std::cout << "#active points: " << report._active_points << std::endl - << "max points: " << report._max_points << std::endl - << "empty slots: " << report._empty_slots << std::endl - << "deletes processed: " << report._slots_released << std::endl - << "latest delete size: " << report._delete_set_size << std::endl - << "rate: (" << points_to_delete_from_beginning / report._time - << " points/second overall, " - << points_to_delete_from_beginning / report._time / - delete_params.num_threads - << " per thread)" << std::endl; - } catch (std::system_error &e) { - std::cout << "Exception caught in deletion thread: " << e.what() - << std::endl; - } +void delete_from_beginning(diskann::Index &index, diskann::IndexWriteParameters &delete_params, + size_t points_to_skip, size_t points_to_delete_from_beginning) +{ + try + { + std::cout << std::endl + << "Lazy deleting points " << points_to_skip << " to " + << points_to_skip + points_to_delete_from_beginning << "... "; + for (size_t i = points_to_skip; i < points_to_skip + points_to_delete_from_beginning; ++i) + index.lazy_delete(i + 1); // Since tags are data location + 1 + std::cout << "done." << std::endl; + + auto report = index.consolidate_deletes(delete_params); + std::cout << "#active points: " << report._active_points << std::endl + << "max points: " << report._max_points << std::endl + << "empty slots: " << report._empty_slots << std::endl + << "deletes processed: " << report._slots_released << std::endl + << "latest delete size: " << report._delete_set_size << std::endl + << "rate: (" << points_to_delete_from_beginning / report._time << " points/second overall, " + << points_to_delete_from_beginning / report._time / delete_params.num_threads << " per thread)" + << std::endl; + } + catch (std::system_error &e) + { + std::cout << "Exception caught in deletion thread: " << e.what() << std::endl; + } } template -void build_incremental_index( - const std::string &data_path, const uint32_t L, const uint32_t R, - const float alpha, const uint32_t thread_count, size_t points_to_skip, - size_t max_points_to_insert, size_t beginning_index_size, - float start_point_norm, 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) { - diskann::IndexWriteParameters params = - diskann::IndexWriteParametersBuilder(L, R) - .with_max_occlusion_size(500) // C = 500 - .with_alpha(alpha) - .with_num_rounds(1) - .with_num_threads(thread_count) - .with_num_frozen_points(num_start_pts) - .build(); - - size_t dim, aligned_dim; - size_t num_points; - - diskann::get_bin_metadata(data_path, num_points, dim); - aligned_dim = ROUND_UP(dim, 8); - - if (points_to_skip > num_points) { - throw diskann::ANNException("Asked to skip more points than in data file", - -1, __FUNCSIG__, __FILE__, __LINE__); - } - - if (max_points_to_insert == 0) { - max_points_to_insert = num_points; - } - - if (points_to_skip + max_points_to_insert > num_points) { - max_points_to_insert = num_points - points_to_skip; - std::cerr << "WARNING: Reducing max_points_to_insert to " - << max_points_to_insert - << " points since the data file has only that many" << std::endl; - } - - using TagT = uint32_t; - const bool enable_tags = true; - - diskann::Index index(diskann::L2, dim, max_points_to_insert, true, - params, L, thread_count, enable_tags, - concurrent); - - size_t current_point_offset = points_to_skip; - const size_t last_point_threshold = points_to_skip + max_points_to_insert; - - if (beginning_index_size > max_points_to_insert) { - beginning_index_size = max_points_to_insert; - std::cerr << "WARNING: Reducing beginning index size to " - << beginning_index_size - << " points since the data file has only that many" << std::endl; - } - if (checkpoints_per_snapshot > 0 && - beginning_index_size > points_per_checkpoint) { - beginning_index_size = points_per_checkpoint; - std::cerr << "WARNING: Reducing beginning index size to " - << beginning_index_size << std::endl; - } - - T *data = nullptr; - diskann::alloc_aligned((void **)&data, - std::max(points_per_checkpoint, beginning_index_size) * - aligned_dim * sizeof(T), - 8 * sizeof(T)); - - std::vector tags(beginning_index_size); - std::iota(tags.begin(), tags.end(), - 1 + static_cast(current_point_offset)); - - load_aligned_bin_part(data_path, data, current_point_offset, - beginning_index_size); - std::cout << "load aligned bin succeeded" << std::endl; - diskann::Timer timer; - - if (beginning_index_size > 0) { - index.build(data, beginning_index_size, params, tags); - index.enable_delete(); - } else { - index.set_start_points_at_random(static_cast(start_point_norm)); - index.enable_delete(); - } - - const double elapsedSeconds = timer.elapsed() / 1000000.0; - std::cout << "Initial non-incremental index build time for " - << beginning_index_size << " points took " << elapsedSeconds - << " seconds (" << beginning_index_size / elapsedSeconds - << " points/second)\n "; - - current_point_offset = beginning_index_size; - - if (points_to_delete_from_beginning > max_points_to_insert) { - points_to_delete_from_beginning = - static_cast(max_points_to_insert); - std::cerr << "WARNING: Reducing points to delete from beginning to " - << points_to_delete_from_beginning - << " points since the data file has only that many" << std::endl; - } - - if (concurrent) { - int sub_threads = (thread_count + 1) / 2; - bool delete_launched = false; - std::future delete_task; +void build_incremental_index(const std::string &data_path, const uint32_t L, const uint32_t R, const float alpha, + const uint32_t thread_count, size_t points_to_skip, size_t max_points_to_insert, + size_t beginning_index_size, float start_point_norm, 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) +{ + diskann::IndexWriteParameters params = diskann::IndexWriteParametersBuilder(L, R) + .with_max_occlusion_size(500) // C = 500 + .with_alpha(alpha) + .with_num_rounds(1) + .with_num_threads(thread_count) + .with_num_frozen_points(num_start_pts) + .build(); + + size_t dim, aligned_dim; + size_t num_points; + + diskann::get_bin_metadata(data_path, num_points, dim); + aligned_dim = ROUND_UP(dim, 8); + + if (points_to_skip > num_points) + { + throw diskann::ANNException("Asked to skip more points than in data file", -1, __FUNCSIG__, __FILE__, __LINE__); + } + + if (max_points_to_insert == 0) + { + max_points_to_insert = num_points; + } + + if (points_to_skip + max_points_to_insert > num_points) + { + max_points_to_insert = num_points - points_to_skip; + std::cerr << "WARNING: Reducing max_points_to_insert to " << max_points_to_insert + << " points since the data file has only that many" << std::endl; + } + + using TagT = uint32_t; + const bool enable_tags = true; + diskann::Index index(diskann::L2, dim, max_points_to_insert, true, params, L, thread_count, enable_tags, + concurrent); + + size_t current_point_offset = points_to_skip; + const size_t last_point_threshold = points_to_skip + max_points_to_insert; + + if (beginning_index_size > max_points_to_insert) + { + beginning_index_size = max_points_to_insert; + std::cerr << "WARNING: Reducing beginning index size to " << beginning_index_size + << " points since the data file has only that many" << std::endl; + } + if (checkpoints_per_snapshot > 0 && beginning_index_size > points_per_checkpoint) + { + beginning_index_size = points_per_checkpoint; + std::cerr << "WARNING: Reducing beginning index size to " << beginning_index_size << std::endl; + } + + T *data = nullptr; + diskann::alloc_aligned( + (void **)&data, std::max(points_per_checkpoint, beginning_index_size) * aligned_dim * sizeof(T), 8 * sizeof(T)); + + std::vector tags(beginning_index_size); + std::iota(tags.begin(), tags.end(), 1 + static_cast(current_point_offset)); + + load_aligned_bin_part(data_path, data, current_point_offset, beginning_index_size); + std::cout << "load aligned bin succeeded" << std::endl; diskann::Timer timer; - for (size_t start = current_point_offset; start < last_point_threshold; - start += points_per_checkpoint, - current_point_offset += points_per_checkpoint) { - const size_t end = - std::min(start + points_per_checkpoint, last_point_threshold); - std::cout << std::endl - << "Inserting from " << start << " to " << end << std::endl; - - auto insert_task = std::async(std::launch::async, [&]() { - load_aligned_bin_part(data_path, data, start, end - start); - insert_till_next_checkpoint(index, start, end, sub_threads, data, - aligned_dim); - }); - insert_task.wait(); - - if (!delete_launched && end >= start_deletes_after && - end >= points_to_skip + points_to_delete_from_beginning) { - delete_launched = true; - diskann::IndexWriteParameters delete_params = - diskann::IndexWriteParametersBuilder(params) - .with_num_threads(sub_threads) - .build(); - - delete_task = std::async(std::launch::async, [&]() { - delete_from_beginning(index, delete_params, points_to_skip, - points_to_delete_from_beginning); - }); - } + if (beginning_index_size > 0) + { + index.build(data, beginning_index_size, params, tags); + index.enable_delete(); } - delete_task.wait(); - - std::cout << "Time Elapsed " << timer.elapsed() / 1000 << "ms\n"; - const auto save_path_inc = get_save_filename( - save_path + ".after-concurrent-delete-", points_to_skip, - points_to_delete_from_beginning, last_point_threshold); - index.save(save_path_inc.c_str(), true); - } else { - size_t last_snapshot_points_threshold = 0; - size_t num_checkpoints_till_snapshot = checkpoints_per_snapshot; - - for (size_t start = current_point_offset; start < last_point_threshold; - start += points_per_checkpoint, - current_point_offset += points_per_checkpoint) { - const size_t end = - std::min(start + points_per_checkpoint, last_point_threshold); - std::cout << std::endl - << "Inserting from " << start << " to " << end << std::endl; - - load_aligned_bin_part(data_path, data, start, end - start); - insert_till_next_checkpoint(index, start, end, thread_count, data, - aligned_dim); - - if (checkpoints_per_snapshot > 0 && - --num_checkpoints_till_snapshot == 0) { - diskann::Timer save_timer; - - const auto save_path_inc = - get_save_filename(save_path + ".inc-", points_to_skip, - points_to_delete_from_beginning, end); - index.save(save_path_inc.c_str(), false); - const double elapsedSeconds = save_timer.elapsed() / 1000000.0; - const size_t points_saved = end - points_to_skip; - - std::cout << "Saved " << points_saved << " points in " << elapsedSeconds - << " seconds (" << points_saved / elapsedSeconds - << " points/second)\n"; - - num_checkpoints_till_snapshot = checkpoints_per_snapshot; - last_snapshot_points_threshold = end; - } - - std::cout << "Number of points in the index post insertion " << end - << std::endl; + else + { + index.set_start_points_at_random(static_cast(start_point_norm)); + index.enable_delete(); } - if (checkpoints_per_snapshot > 0 && - last_snapshot_points_threshold != last_point_threshold) { - const auto save_path_inc = get_save_filename( - save_path + ".inc-", points_to_skip, points_to_delete_from_beginning, - last_point_threshold); - // index.save(save_path_inc.c_str(), false); + const double elapsedSeconds = timer.elapsed() / 1000000.0; + std::cout << "Initial non-incremental index build time for " << beginning_index_size << " points took " + << elapsedSeconds << " seconds (" << beginning_index_size / elapsedSeconds << " points/second)\n "; + + current_point_offset = beginning_index_size; + + if (points_to_delete_from_beginning > max_points_to_insert) + { + points_to_delete_from_beginning = static_cast(max_points_to_insert); + std::cerr << "WARNING: Reducing points to delete from beginning to " << points_to_delete_from_beginning + << " points since the data file has only that many" << std::endl; } - if (points_to_delete_from_beginning > 0) { - delete_from_beginning(index, params, points_to_skip, - points_to_delete_from_beginning); + if (concurrent) + { + int sub_threads = (thread_count + 1) / 2; + bool delete_launched = false; + std::future delete_task; + + diskann::Timer timer; + + for (size_t start = current_point_offset; start < last_point_threshold; + start += points_per_checkpoint, current_point_offset += points_per_checkpoint) + { + const size_t end = std::min(start + points_per_checkpoint, last_point_threshold); + std::cout << std::endl << "Inserting from " << start << " to " << end << std::endl; + + auto insert_task = std::async(std::launch::async, [&]() { + load_aligned_bin_part(data_path, data, start, end - start); + insert_till_next_checkpoint(index, start, end, sub_threads, data, aligned_dim); + }); + insert_task.wait(); + + if (!delete_launched && end >= start_deletes_after && + end >= points_to_skip + points_to_delete_from_beginning) + { + delete_launched = true; + diskann::IndexWriteParameters delete_params = + diskann::IndexWriteParametersBuilder(params).with_num_threads(sub_threads).build(); + + delete_task = std::async(std::launch::async, [&]() { + delete_from_beginning(index, delete_params, points_to_skip, points_to_delete_from_beginning); + }); + } + } + delete_task.wait(); + + std::cout << "Time Elapsed " << timer.elapsed() / 1000 << "ms\n"; + const auto save_path_inc = get_save_filename(save_path + ".after-concurrent-delete-", points_to_skip, + points_to_delete_from_beginning, last_point_threshold); + index.save(save_path_inc.c_str(), true); + } + else + { + size_t last_snapshot_points_threshold = 0; + size_t num_checkpoints_till_snapshot = checkpoints_per_snapshot; + + for (size_t start = current_point_offset; start < last_point_threshold; + start += points_per_checkpoint, current_point_offset += points_per_checkpoint) + { + const size_t end = std::min(start + points_per_checkpoint, last_point_threshold); + std::cout << std::endl << "Inserting from " << start << " to " << end << std::endl; + + load_aligned_bin_part(data_path, data, start, end - start); + insert_till_next_checkpoint(index, start, end, thread_count, data, aligned_dim); + + if (checkpoints_per_snapshot > 0 && --num_checkpoints_till_snapshot == 0) + { + diskann::Timer save_timer; + + const auto save_path_inc = + get_save_filename(save_path + ".inc-", points_to_skip, points_to_delete_from_beginning, end); + index.save(save_path_inc.c_str(), false); + const double elapsedSeconds = save_timer.elapsed() / 1000000.0; + const size_t points_saved = end - points_to_skip; + + std::cout << "Saved " << points_saved << " points in " << elapsedSeconds << " seconds (" + << points_saved / elapsedSeconds << " points/second)\n"; + + num_checkpoints_till_snapshot = checkpoints_per_snapshot; + last_snapshot_points_threshold = end; + } + + std::cout << "Number of points in the index post insertion " << end << std::endl; + } + + if (checkpoints_per_snapshot > 0 && last_snapshot_points_threshold != last_point_threshold) + { + const auto save_path_inc = get_save_filename(save_path + ".inc-", points_to_skip, + points_to_delete_from_beginning, last_point_threshold); + // index.save(save_path_inc.c_str(), false); + } + + if (points_to_delete_from_beginning > 0) + { + delete_from_beginning(index, params, points_to_skip, points_to_delete_from_beginning); + } + const auto save_path_inc = get_save_filename(save_path + ".after-delete-", points_to_skip, + points_to_delete_from_beginning, last_point_threshold); + index.save(save_path_inc.c_str(), true); } - const auto save_path_inc = get_save_filename( - save_path + ".after-delete-", points_to_skip, - points_to_delete_from_beginning, last_point_threshold); - index.save(save_path_inc.c_str(), true); - } - diskann::aligned_free(data); + diskann::aligned_free(data); } -int main(int argc, char **argv) { - std::string data_type, dist_fn, data_path, index_path_prefix; - 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; - bool concurrent; - - po::options_description desc{"Arguments"}; - try { - desc.add_options()("help,h", "Print information on arguments"); - desc.add_options()("data_type", - po::value(&data_type)->required(), - "data type "); - desc.add_options()("dist_fn", po::value(&dist_fn)->required(), - "distance function "); - desc.add_options()("data_path", - po::value(&data_path)->required(), - "Input data file in bin format"); - desc.add_options()("index_path_prefix", - po::value(&index_path_prefix)->required(), - "Path prefix for saving index file components"); - desc.add_options()("max_degree,R", - po::value(&R)->default_value(64), - "Maximum graph degree"); - desc.add_options()( - "Lbuild,L", po::value(&L)->default_value(100), - "Build complexity, higher value results in better graphs"); - desc.add_options()("alpha", po::value(&alpha)->default_value(1.2f), - "alpha controls density and diameter of graph, set " - "1 for sparse graph, " - "1.2 or 1.4 for denser graphs with lower diameter"); - desc.add_options()( - "num_threads,T", - po::value(&num_threads)->default_value(omp_get_num_procs()), - "Number of threads used for building index (defaults to " - "omp_get_num_procs())"); - desc.add_options()("points_to_skip", - po::value(&points_to_skip)->required(), - "Skip these first set of points from file"); - desc.add_options()( - "max_points_to_insert", - po::value(&max_points_to_insert)->default_value(0), - "These number of points from the file are inserted after " - "points_to_skip"); - desc.add_options()("beginning_index_size", - po::value(&beginning_index_size)->required(), - "Batch build will be called on these set of points"); - desc.add_options()( - "points_per_checkpoint", - po::value(&points_per_checkpoint)->required(), - "Insertions are done in batches of points_per_checkpoint"); - desc.add_options()( - "checkpoints_per_snapshot", - po::value(&checkpoints_per_snapshot)->required(), - "Save the index to disk every few checkpoints"); - desc.add_options()( - "points_to_delete_from_beginning", - po::value(&points_to_delete_from_beginning)->required(), ""); - desc.add_options()("do_concurrent", - po::value(&concurrent)->default_value(false), ""); - desc.add_options()( - "start_deletes_after", - po::value(&start_deletes_after)->default_value(0), ""); - desc.add_options()( - "start_point_norm", - po::value(&start_point_norm)->default_value(0), - "Set the start point to a random point on a sphere of this radius"); - desc.add_options()( - "num_start_points", - po::value(&num_start_pts)->default_value(0), - "Set the number of random start (frozen) points to use when " - "inserting and searching"); - - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) { - std::cout << desc; - return 0; +int main(int argc, char **argv) +{ + std::string data_type, dist_fn, data_path, index_path_prefix; + 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; + bool concurrent; + + po::options_description desc{"Arguments"}; + try + { + desc.add_options()("help,h", "Print information on arguments"); + desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); + desc.add_options()("dist_fn", po::value(&dist_fn)->required(), "distance function "); + desc.add_options()("data_path", po::value(&data_path)->required(), + "Input data file in bin format"); + desc.add_options()("index_path_prefix", po::value(&index_path_prefix)->required(), + "Path prefix for saving index file components"); + desc.add_options()("max_degree,R", po::value(&R)->default_value(64), "Maximum graph degree"); + desc.add_options()("Lbuild,L", po::value(&L)->default_value(100), + "Build complexity, higher value results in better graphs"); + desc.add_options()("alpha", po::value(&alpha)->default_value(1.2f), + "alpha controls density and diameter of graph, set " + "1 for sparse graph, " + "1.2 or 1.4 for denser graphs with lower diameter"); + desc.add_options()("num_threads,T", po::value(&num_threads)->default_value(omp_get_num_procs()), + "Number of threads used for building index (defaults to " + "omp_get_num_procs())"); + desc.add_options()("points_to_skip", po::value(&points_to_skip)->required(), + "Skip these first set of points from file"); + desc.add_options()("max_points_to_insert", po::value(&max_points_to_insert)->default_value(0), + "These number of points from the file are inserted after " + "points_to_skip"); + desc.add_options()("beginning_index_size", po::value(&beginning_index_size)->required(), + "Batch build will be called on these set of points"); + desc.add_options()("points_per_checkpoint", po::value(&points_per_checkpoint)->required(), + "Insertions are done in batches of points_per_checkpoint"); + desc.add_options()("checkpoints_per_snapshot", po::value(&checkpoints_per_snapshot)->required(), + "Save the index to disk every few checkpoints"); + desc.add_options()("points_to_delete_from_beginning", + po::value(&points_to_delete_from_beginning)->required(), ""); + desc.add_options()("do_concurrent", po::value(&concurrent)->default_value(false), ""); + desc.add_options()("start_deletes_after", po::value(&start_deletes_after)->default_value(0), ""); + desc.add_options()("start_point_norm", po::value(&start_point_norm)->default_value(0), + "Set the start point to a random point on a sphere of this radius"); + desc.add_options()("num_start_points", po::value(&num_start_pts)->default_value(0), + "Set the number of random start (frozen) points to use when " + "inserting and searching"); + + po::variables_map vm; + po::store(po::parse_command_line(argc, argv, desc), vm); + if (vm.count("help")) + { + std::cout << desc; + return 0; + } + po::notify(vm); + if (beginning_index_size == 0) + if (start_point_norm == 0) + { + std::cout << "When beginning_index_size is 0, use a start " + "point with " + "appropriate norm" + << std::endl; + return -1; + } } - po::notify(vm); - if (beginning_index_size == 0) - if (start_point_norm == 0) { - std::cout << "When beginning_index_size is 0, use a start " - "point with " - "appropriate norm" - << std::endl; + catch (const std::exception &ex) + { + std::cerr << ex.what() << '\n'; return -1; - } - } catch (const std::exception &ex) { - std::cerr << ex.what() << '\n'; - return -1; - } - - try { - if (data_type == std::string("int8")) - build_incremental_index( - data_path, L, R, alpha, num_threads, 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); - else if (data_type == std::string("uint8")) - build_incremental_index( - data_path, L, R, alpha, num_threads, 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); - else if (data_type == std::string("float")) - build_incremental_index( - data_path, L, R, alpha, num_threads, 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); - else - std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; - } catch (const std::exception &e) { - std::cerr << "Caught exception: " << e.what() << std::endl; - exit(-1); - } catch (...) { - std::cerr << "Caught unknown exception" << std::endl; - exit(-1); - } - - return 0; + } + + try + { + if (data_type == std::string("int8")) + build_incremental_index(data_path, L, R, alpha, num_threads, 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); + else if (data_type == std::string("uint8")) + build_incremental_index(data_path, L, R, alpha, num_threads, 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); + else if (data_type == std::string("float")) + build_incremental_index(data_path, L, R, alpha, num_threads, 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); + else + std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; + } + catch (const std::exception &e) + { + std::cerr << "Caught exception: " << e.what() << std::endl; + exit(-1); + } + catch (...) + { + std::cerr << "Caught unknown exception" << std::endl; + exit(-1); + } + + return 0; } diff --git a/tests/test_streaming_scenario.cpp b/tests/test_streaming_scenario.cpp index ab61bb140..e1fe80c83 100644 --- a/tests/test_streaming_scenario.cpp +++ b/tests/test_streaming_scenario.cpp @@ -25,378 +25,354 @@ namespace po = boost::program_options; // load_aligned_bin modified to read pieces of the file, but using ifstream // instead of cached_ifstream. template -inline void load_aligned_bin_part(const std::string &bin_file, T *data, - size_t offset_points, size_t points_to_read) { - std::ifstream reader; - reader.exceptions(std::ios::failbit | std::ios::badbit); - reader.open(bin_file, std::ios::binary | std::ios::ate); - size_t actual_file_size = reader.tellg(); - reader.seekg(0, std::ios::beg); - - int npts_i32, dim_i32; - reader.read((char *)&npts_i32, sizeof(int)); - reader.read((char *)&dim_i32, sizeof(int)); - size_t npts = (uint32_t)npts_i32; - size_t dim = (uint32_t)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; - std::cout << stream.str(); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); - } - - if (offset_points + points_to_read > npts) { - std::stringstream stream; - stream << "Error. Not enough points in file. Requested " << offset_points - << " offset and " << points_to_read << " points, but have only " - << npts << " points" << std::endl; - std::cout << stream.str(); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); - } - - reader.seekg(2 * sizeof(uint32_t) + offset_points * dim * sizeof(T)); - - const size_t rounded_dim = ROUND_UP(dim, 8); - - for (size_t i = 0; i < points_to_read; i++) { - reader.read((char *)(data + i * rounded_dim), dim * sizeof(T)); - memset(data + i * rounded_dim + dim, 0, (rounded_dim - dim) * sizeof(T)); - } - reader.close(); +inline void load_aligned_bin_part(const std::string &bin_file, T *data, size_t offset_points, size_t points_to_read) +{ + std::ifstream reader; + reader.exceptions(std::ios::failbit | std::ios::badbit); + reader.open(bin_file, std::ios::binary | std::ios::ate); + size_t actual_file_size = reader.tellg(); + reader.seekg(0, std::ios::beg); + + int npts_i32, dim_i32; + reader.read((char *)&npts_i32, sizeof(int)); + reader.read((char *)&dim_i32, sizeof(int)); + size_t npts = (uint32_t)npts_i32; + size_t dim = (uint32_t)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; + std::cout << stream.str(); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } + + if (offset_points + points_to_read > npts) + { + std::stringstream stream; + stream << "Error. Not enough points in file. Requested " << offset_points << " offset and " << points_to_read + << " points, but have only " << npts << " points" << std::endl; + std::cout << stream.str(); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } + + reader.seekg(2 * sizeof(uint32_t) + offset_points * dim * sizeof(T)); + + const size_t rounded_dim = ROUND_UP(dim, 8); + + for (size_t i = 0; i < points_to_read; i++) + { + reader.read((char *)(data + i * rounded_dim), dim * sizeof(T)); + memset(data + i * rounded_dim + dim, 0, (rounded_dim - dim) * sizeof(T)); + } + reader.close(); } -std::string get_save_filename(const std::string &save_path, - size_t active_window, size_t consolidate_interval, - size_t max_points_to_insert) { - std::string final_path = save_path; - final_path += "act" + std::to_string(active_window) + "-"; - final_path += "cons" + std::to_string(consolidate_interval) + "-"; - final_path += "max" + std::to_string(max_points_to_insert); - return final_path; +std::string get_save_filename(const std::string &save_path, size_t active_window, size_t consolidate_interval, + size_t max_points_to_insert) +{ + std::string final_path = save_path; + final_path += "act" + std::to_string(active_window) + "-"; + final_path += "cons" + std::to_string(consolidate_interval) + "-"; + final_path += "max" + std::to_string(max_points_to_insert); + return final_path; } template -void insert_next_batch(diskann::Index &index, size_t start, - size_t end, size_t insert_threads, T *data, - size_t aligned_dim) { - try { - diskann::Timer insert_timer; - std::cout << std::endl - << "Inserting from " << start << " to " << end << std::endl; - - size_t num_failed = 0; +void insert_next_batch(diskann::Index &index, size_t start, size_t end, size_t insert_threads, T *data, + size_t aligned_dim) +{ + try + { + diskann::Timer insert_timer; + std::cout << std::endl << "Inserting from " << start << " to " << end << std::endl; + + size_t num_failed = 0; #pragma omp parallel for num_threads(insert_threads) schedule(dynamic) reduction(+ : num_failed) - for (int64_t j = start; j < (int64_t)end; j++) { - if (index.insert_point(&data[(j - start) * aligned_dim], - 1 + static_cast(j)) != 0) { - std::cerr << "Insert failed " << j << std::endl; - num_failed++; - } + for (int64_t j = start; j < (int64_t)end; j++) + { + if (index.insert_point(&data[(j - start) * aligned_dim], 1 + static_cast(j)) != 0) + { + std::cerr << "Insert failed " << j << std::endl; + num_failed++; + } + } + const double elapsedSeconds = insert_timer.elapsed() / 1000000.0; + std::cout << "Insertion time " << elapsedSeconds << " seconds (" << (end - start) / elapsedSeconds + << " points/second overall, " << (end - start) / elapsedSeconds / insert_threads << " per thread)" + << std::endl; + if (num_failed > 0) + std::cout << num_failed << " of " << end - start << "inserts failed" << std::endl; + } + catch (std::system_error &e) + { + std::cout << "Exiting after catching exception in insertion task: " << e.what() << std::endl; } - const double elapsedSeconds = insert_timer.elapsed() / 1000000.0; - std::cout << "Insertion time " << elapsedSeconds << " seconds (" - << (end - start) / elapsedSeconds << " points/second overall, " - << (end - start) / elapsedSeconds / insert_threads - << " per thread)" << std::endl; - if (num_failed > 0) - std::cout << num_failed << " of " << end - start << "inserts failed" - << std::endl; - } catch (std::system_error &e) { - std::cout << "Exiting after catching exception in insertion task: " - << e.what() << std::endl; - } } template -void delete_and_consolidate(diskann::Index &index, - diskann::IndexWriteParameters &delete_params, - size_t start, size_t end) { - try { - std::cout << std::endl - << "Lazy deleting points " << start << " to " << end << "... "; - for (size_t i = start; i < end; ++i) index.lazy_delete(1 + i); - std::cout << "lazy delete done." << std::endl; - - auto report = index.consolidate_deletes(delete_params); - while (report._status != - diskann::consolidation_report::status_code::SUCCESS) { - int wait_time = 5; - if (report._status == - diskann::consolidation_report::status_code::LOCK_FAIL) { - diskann::cerr << "Unable to acquire consolidate delete lock after " - << "deleting points " << start << " to " << end - << ". Will retry in " << wait_time << "seconds." - << std::endl; - } else if (report._status == diskann::consolidation_report::status_code:: - INCONSISTENT_COUNT_ERROR) { - diskann::cerr << "Inconsistent counts in data structure. " - << "Will retry in " << wait_time << "seconds." - << std::endl; - } else { - std::cerr << "Exiting after unknown error in consolidate delete" - << std::endl; +void delete_and_consolidate(diskann::Index &index, diskann::IndexWriteParameters &delete_params, + size_t start, size_t end) +{ + try + { + std::cout << std::endl << "Lazy deleting points " << start << " to " << end << "... "; + for (size_t i = start; i < end; ++i) + index.lazy_delete(1 + i); + std::cout << "lazy delete done." << std::endl; + + auto report = index.consolidate_deletes(delete_params); + while (report._status != diskann::consolidation_report::status_code::SUCCESS) + { + int wait_time = 5; + if (report._status == diskann::consolidation_report::status_code::LOCK_FAIL) + { + diskann::cerr << "Unable to acquire consolidate delete lock after " + << "deleting points " << start << " to " << end << ". Will retry in " << wait_time + << "seconds." << std::endl; + } + else if (report._status == diskann::consolidation_report::status_code::INCONSISTENT_COUNT_ERROR) + { + diskann::cerr << "Inconsistent counts in data structure. " + << "Will retry in " << wait_time << "seconds." << std::endl; + } + else + { + std::cerr << "Exiting after unknown error in consolidate delete" << std::endl; + exit(-1); + } + std::this_thread::sleep_for(std::chrono::seconds(wait_time)); + report = index.consolidate_deletes(delete_params); + } + auto points_processed = report._active_points + report._slots_released; + auto deletion_rate = points_processed / report._time; + std::cout << "#active points: " << report._active_points << std::endl + << "max points: " << report._max_points << std::endl + << "empty slots: " << report._empty_slots << std::endl + << "deletes processed: " << report._slots_released << std::endl + << "latest delete size: " << report._delete_set_size << std::endl + << "Deletion rate: " << deletion_rate << "/sec " + << "Deletion rate: " << deletion_rate / delete_params.num_threads << "/thread/sec " << std::endl; + } + catch (std::system_error &e) + { + std::cerr << "Exiting after catching exception in deletion task: " << e.what() << std::endl; exit(-1); - } - std::this_thread::sleep_for(std::chrono::seconds(wait_time)); - report = index.consolidate_deletes(delete_params); } - auto points_processed = report._active_points + report._slots_released; - auto deletion_rate = points_processed / report._time; - std::cout << "#active points: " << report._active_points << std::endl - << "max points: " << report._max_points << std::endl - << "empty slots: " << report._empty_slots << std::endl - << "deletes processed: " << report._slots_released << std::endl - << "latest delete size: " << report._delete_set_size << std::endl - << "Deletion rate: " << deletion_rate << "/sec " - << "Deletion rate: " << deletion_rate / delete_params.num_threads - << "/thread/sec " << std::endl; - } catch (std::system_error &e) { - std::cerr << "Exiting after catching exception in deletion task: " - << e.what() << std::endl; - exit(-1); - } } template -void build_incremental_index(const std::string &data_path, const uint32_t L, - const uint32_t R, const float alpha, - const uint32_t insert_threads, - const uint32_t consolidate_threads, - size_t max_points_to_insert, size_t active_window, - size_t consolidate_interval, - const float start_point_norm, - uint32_t num_start_pts, - const std::string &save_path) { - const uint32_t C = 500; - const bool saturate_graph = false; - - diskann::IndexWriteParameters params = - diskann::IndexWriteParametersBuilder(L, R) - .with_max_occlusion_size(C) - .with_alpha(alpha) - .with_saturate_graph(saturate_graph) - .with_num_rounds(1) - .with_num_threads(insert_threads) - .with_num_frozen_points(num_start_pts) - .build(); - - diskann::IndexWriteParameters delete_params = - diskann::IndexWriteParametersBuilder(L, R) - .with_max_occlusion_size(C) - .with_alpha(alpha) - .with_saturate_graph(saturate_graph) - .with_num_rounds(1) - .with_num_threads(consolidate_threads) - .build(); - - size_t dim, aligned_dim; - size_t num_points; - - diskann::get_bin_metadata(data_path, num_points, dim); - diskann::cout << "metadata: file " << data_path << " has " << num_points - << " points in " << dim << " dims" << std::endl; - aligned_dim = ROUND_UP(dim, 8); - - if (max_points_to_insert == 0) { - max_points_to_insert = num_points; - } - - if (num_points < max_points_to_insert) - throw diskann::ANNException(std::string("num_points(") + - std::to_string(num_points) + - ") < max_points_to_insert(" + - std::to_string(max_points_to_insert) + ")", - -1, __FUNCSIG__, __FILE__, __LINE__); - - if (max_points_to_insert < active_window + consolidate_interval) - throw diskann::ANNException( - "ERROR: max_points_to_insert < " - "active_window + consolidate_interval", - -1, __FUNCSIG__, __FILE__, __LINE__); - - if (consolidate_interval < max_points_to_insert / 1000) - throw diskann::ANNException("ERROR: consolidate_interval is too small", -1, - __FUNCSIG__, __FILE__, __LINE__); - - using TagT = uint32_t; - using LabelT = uint32_t; - const bool enable_tags = true; - - diskann::Index index( - diskann::L2, dim, active_window + 4 * consolidate_interval, true, params, - L, insert_threads, enable_tags, true); - index.set_start_points_at_random(static_cast(start_point_norm)); - index.enable_delete(); - - T *data = nullptr; - diskann::alloc_aligned( - (void **)&data, - std::max(consolidate_interval, active_window) * aligned_dim * sizeof(T), - 8 * sizeof(T)); - - std::vector tags(max_points_to_insert); - std::iota(tags.begin(), tags.end(), static_cast(0)); - - diskann::Timer timer; - - std::vector> delete_tasks; - - auto insert_task = std::async(std::launch::async, [&]() { - load_aligned_bin_part(data_path, data, 0, active_window); - insert_next_batch(index, 0, active_window, insert_threads, data, - aligned_dim); - }); - insert_task.wait(); - - for (size_t start = active_window; - start + consolidate_interval <= max_points_to_insert; - start += consolidate_interval) { - auto end = std::min(start + consolidate_interval, max_points_to_insert); +void build_incremental_index(const std::string &data_path, const uint32_t L, const uint32_t R, const float alpha, + const uint32_t insert_threads, const uint32_t consolidate_threads, + size_t max_points_to_insert, size_t active_window, size_t consolidate_interval, + const float start_point_norm, uint32_t num_start_pts, const std::string &save_path) +{ + const uint32_t C = 500; + const bool saturate_graph = false; + + diskann::IndexWriteParameters params = diskann::IndexWriteParametersBuilder(L, R) + .with_max_occlusion_size(C) + .with_alpha(alpha) + .with_saturate_graph(saturate_graph) + .with_num_rounds(1) + .with_num_threads(insert_threads) + .with_num_frozen_points(num_start_pts) + .build(); + + diskann::IndexWriteParameters delete_params = diskann::IndexWriteParametersBuilder(L, R) + .with_max_occlusion_size(C) + .with_alpha(alpha) + .with_saturate_graph(saturate_graph) + .with_num_rounds(1) + .with_num_threads(consolidate_threads) + .build(); + + size_t dim, aligned_dim; + size_t num_points; + + diskann::get_bin_metadata(data_path, num_points, dim); + diskann::cout << "metadata: file " << data_path << " has " << num_points << " points in " << dim << " dims" + << std::endl; + aligned_dim = ROUND_UP(dim, 8); + + if (max_points_to_insert == 0) + { + max_points_to_insert = num_points; + } + + if (num_points < max_points_to_insert) + throw diskann::ANNException(std::string("num_points(") + std::to_string(num_points) + + ") < max_points_to_insert(" + std::to_string(max_points_to_insert) + ")", + -1, __FUNCSIG__, __FILE__, __LINE__); + + if (max_points_to_insert < active_window + consolidate_interval) + throw diskann::ANNException("ERROR: max_points_to_insert < " + "active_window + consolidate_interval", + -1, __FUNCSIG__, __FILE__, __LINE__); + + if (consolidate_interval < max_points_to_insert / 1000) + throw diskann::ANNException("ERROR: consolidate_interval is too small", -1, __FUNCSIG__, __FILE__, __LINE__); + + using TagT = uint32_t; + using LabelT = uint32_t; + const bool enable_tags = true; + + diskann::Index index(diskann::L2, dim, active_window + 4 * consolidate_interval, true, params, L, + insert_threads, enable_tags, true); + index.set_start_points_at_random(static_cast(start_point_norm)); + index.enable_delete(); + + T *data = nullptr; + diskann::alloc_aligned((void **)&data, std::max(consolidate_interval, active_window) * aligned_dim * sizeof(T), + 8 * sizeof(T)); + + std::vector tags(max_points_to_insert); + std::iota(tags.begin(), tags.end(), static_cast(0)); + + diskann::Timer timer; + + std::vector> delete_tasks; + auto insert_task = std::async(std::launch::async, [&]() { - load_aligned_bin_part(data_path, data, start, end - start); - insert_next_batch(index, start, end, insert_threads, data, aligned_dim); + load_aligned_bin_part(data_path, data, 0, active_window); + insert_next_batch(index, 0, active_window, insert_threads, data, aligned_dim); }); insert_task.wait(); - if (delete_tasks.size() > 0) delete_tasks[delete_tasks.size() - 1].wait(); - if (start >= active_window + consolidate_interval) { - auto start_del = start - active_window - consolidate_interval; - auto end_del = start - active_window; - - delete_tasks.emplace_back(std::async(std::launch::async, [&]() { - delete_and_consolidate(index, delete_params, start_del, end_del); - })); + for (size_t start = active_window; start + consolidate_interval <= max_points_to_insert; + start += consolidate_interval) + { + auto end = std::min(start + consolidate_interval, max_points_to_insert); + auto insert_task = std::async(std::launch::async, [&]() { + load_aligned_bin_part(data_path, data, start, end - start); + insert_next_batch(index, start, end, insert_threads, data, aligned_dim); + }); + insert_task.wait(); + + if (delete_tasks.size() > 0) + delete_tasks[delete_tasks.size() - 1].wait(); + if (start >= active_window + consolidate_interval) + { + auto start_del = start - active_window - consolidate_interval; + auto end_del = start - active_window; + + delete_tasks.emplace_back(std::async( + std::launch::async, [&]() { delete_and_consolidate(index, delete_params, start_del, end_del); })); + } } - } - if (delete_tasks.size() > 0) delete_tasks[delete_tasks.size() - 1].wait(); + if (delete_tasks.size() > 0) + delete_tasks[delete_tasks.size() - 1].wait(); - std::cout << "Time Elapsed " << timer.elapsed() / 1000 << "ms\n"; - const auto save_path_inc = - get_save_filename(save_path + ".after-streaming-", active_window, - consolidate_interval, max_points_to_insert); - index.save(save_path_inc.c_str(), true); + std::cout << "Time Elapsed " << timer.elapsed() / 1000 << "ms\n"; + const auto save_path_inc = + get_save_filename(save_path + ".after-streaming-", active_window, consolidate_interval, max_points_to_insert); + index.save(save_path_inc.c_str(), true); - diskann::aligned_free(data); + diskann::aligned_free(data); } -int main(int argc, char **argv) { - std::string data_type, dist_fn, data_path, index_path_prefix; - uint32_t insert_threads, consolidate_threads; - uint32_t R, L, num_start_pts; - float alpha, start_point_norm; - size_t max_points_to_insert, active_window, consolidate_interval; - - po::options_description desc{"Arguments"}; - try { - desc.add_options()("help,h", "Print information on arguments"); - desc.add_options()("data_type", - po::value(&data_type)->required(), - "data type "); - desc.add_options()("dist_fn", po::value(&dist_fn)->required(), - "distance function "); - desc.add_options()("data_path", - po::value(&data_path)->required(), - "Input data file in bin format"); - desc.add_options()("index_path_prefix", - po::value(&index_path_prefix)->required(), - "Path prefix for saving index file components"); - desc.add_options()("max_degree,R", - po::value(&R)->default_value(64), - "Maximum graph degree"); - desc.add_options()( - "Lbuild,L", po::value(&L)->default_value(100), - "Build complexity, higher value results in better graphs"); - desc.add_options()("alpha", po::value(&alpha)->default_value(1.2f), - "alpha controls density and diameter of graph, set " - "1 for sparse graph, " - "1.2 or 1.4 for denser graphs with lower diameter"); - desc.add_options()( - "insert_threads", - po::value(&insert_threads) - ->default_value(omp_get_num_procs() / 2), - "Number of threads used for inserting into the index (defaults to " - "omp_get_num_procs()/2)"); - desc.add_options()("consolidate_threads", - po::value(&consolidate_threads) - ->default_value(omp_get_num_procs() / 2), - "Number of threads used for consolidating deletes to " - "the index (defaults to omp_get_num_procs()/2)"); - - desc.add_options()( - "max_points_to_insert", - po::value(&max_points_to_insert)->default_value(0), - "The number of points from the file that the program streams " - "over "); - desc.add_options()("active_window", - po::value(&active_window)->required(), - "Program maintains an index over an active window of " - "this size that slides through the data"); - desc.add_options()( - "consolidate_interval", - po::value(&consolidate_interval)->required(), - "The program simultaneously adds this number of points to the " - "right of " - "the window while deleting the same number from the left"); - desc.add_options()( - "start_point_norm", po::value(&start_point_norm)->required(), - "Set the start point to a random point on a sphere of this radius"); - desc.add_options()( - "num_start_points", - po::value(&num_start_pts)->default_value(0), - "Set the number of random start (frozen) points to use when " - "inserting and searching"); - - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) { - std::cout << desc; - return 0; +int main(int argc, char **argv) +{ + std::string data_type, dist_fn, data_path, index_path_prefix; + uint32_t insert_threads, consolidate_threads; + uint32_t R, L, num_start_pts; + float alpha, start_point_norm; + size_t max_points_to_insert, active_window, consolidate_interval; + + po::options_description desc{"Arguments"}; + try + { + desc.add_options()("help,h", "Print information on arguments"); + desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); + desc.add_options()("dist_fn", po::value(&dist_fn)->required(), "distance function "); + desc.add_options()("data_path", po::value(&data_path)->required(), + "Input data file in bin format"); + desc.add_options()("index_path_prefix", po::value(&index_path_prefix)->required(), + "Path prefix for saving index file components"); + desc.add_options()("max_degree,R", po::value(&R)->default_value(64), "Maximum graph degree"); + desc.add_options()("Lbuild,L", po::value(&L)->default_value(100), + "Build complexity, higher value results in better graphs"); + desc.add_options()("alpha", po::value(&alpha)->default_value(1.2f), + "alpha controls density and diameter of graph, set " + "1 for sparse graph, " + "1.2 or 1.4 for denser graphs with lower diameter"); + desc.add_options()("insert_threads", + po::value(&insert_threads)->default_value(omp_get_num_procs() / 2), + "Number of threads used for inserting into the index (defaults to " + "omp_get_num_procs()/2)"); + desc.add_options()("consolidate_threads", + po::value(&consolidate_threads)->default_value(omp_get_num_procs() / 2), + "Number of threads used for consolidating deletes to " + "the index (defaults to omp_get_num_procs()/2)"); + + desc.add_options()("max_points_to_insert", po::value(&max_points_to_insert)->default_value(0), + "The number of points from the file that the program streams " + "over "); + desc.add_options()("active_window", po::value(&active_window)->required(), + "Program maintains an index over an active window of " + "this size that slides through the data"); + desc.add_options()("consolidate_interval", po::value(&consolidate_interval)->required(), + "The program simultaneously adds this number of points to the " + "right of " + "the window while deleting the same number from the left"); + desc.add_options()("start_point_norm", po::value(&start_point_norm)->required(), + "Set the start point to a random point on a sphere of this radius"); + desc.add_options()("num_start_points", po::value(&num_start_pts)->default_value(0), + "Set the number of random start (frozen) points to use when " + "inserting and searching"); + + po::variables_map vm; + po::store(po::parse_command_line(argc, argv, desc), vm); + if (vm.count("help")) + { + std::cout << desc; + return 0; + } + po::notify(vm); + if (start_point_norm == 0) + { + std::cout << "When beginning_index_size is 0, use a start point with " + "appropriate norm" + << std::endl; + return -1; + } + } + catch (const std::exception &ex) + { + std::cerr << ex.what() << '\n'; + return -1; + } + + try + { + if (data_type == std::string("int8")) + build_incremental_index(data_path, L, R, alpha, insert_threads, consolidate_threads, + max_points_to_insert, active_window, consolidate_interval, start_point_norm, + num_start_pts, index_path_prefix); + else if (data_type == std::string("uint8")) + build_incremental_index(data_path, L, R, alpha, insert_threads, consolidate_threads, + max_points_to_insert, active_window, consolidate_interval, + start_point_norm, num_start_pts, index_path_prefix); + else if (data_type == std::string("float")) + build_incremental_index(data_path, L, R, alpha, insert_threads, consolidate_threads, + max_points_to_insert, active_window, consolidate_interval, start_point_norm, + num_start_pts, index_path_prefix); + else + std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; } - po::notify(vm); - if (start_point_norm == 0) { - std::cout << "When beginning_index_size is 0, use a start point with " - "appropriate norm" - << std::endl; - return -1; + catch (const std::exception &e) + { + std::cerr << "Caught exception: " << e.what() << std::endl; + exit(-1); + } + catch (...) + { + std::cerr << "Caught unknown exception" << std::endl; + exit(-1); } - } catch (const std::exception &ex) { - std::cerr << ex.what() << '\n'; - return -1; - } - - try { - if (data_type == std::string("int8")) - build_incremental_index( - data_path, L, R, alpha, insert_threads, consolidate_threads, - max_points_to_insert, active_window, consolidate_interval, - start_point_norm, num_start_pts, index_path_prefix); - else if (data_type == std::string("uint8")) - build_incremental_index( - data_path, L, R, alpha, insert_threads, consolidate_threads, - max_points_to_insert, active_window, consolidate_interval, - start_point_norm, num_start_pts, index_path_prefix); - else if (data_type == std::string("float")) - build_incremental_index( - data_path, L, R, alpha, insert_threads, consolidate_threads, - max_points_to_insert, active_window, consolidate_interval, - start_point_norm, num_start_pts, index_path_prefix); - else - std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; - } catch (const std::exception &e) { - std::cerr << "Caught exception: " << e.what() << std::endl; - exit(-1); - } catch (...) { - std::cerr << "Caught unknown exception" << std::endl; - exit(-1); - } - - return 0; + + return 0; } diff --git a/tests/utils/bin_to_fvecs.cpp b/tests/utils/bin_to_fvecs.cpp index f35038409..e9a6a8ecc 100644 --- a/tests/utils/bin_to_fvecs.cpp +++ b/tests/utils/bin_to_fvecs.cpp @@ -4,58 +4,60 @@ #include #include "util.h" -void block_convert(std::ifstream &writr, std::ofstream &readr, float *read_buf, - float *write_buf, uint64_t npts, uint64_t ndims) { - writr.write((char *)read_buf, - npts * (ndims * sizeof(float) + sizeof(unsigned))); +void block_convert(std::ifstream &writr, std::ofstream &readr, float *read_buf, float *write_buf, uint64_t npts, + uint64_t ndims) +{ + writr.write((char *)read_buf, npts * (ndims * sizeof(float) + sizeof(unsigned))); #pragma omp parallel for - for (uint64_t i = 0; i < npts; i++) { - memcpy(write_buf + i * ndims, (read_buf + i * (ndims + 1)) + 1, - ndims * sizeof(float)); - } - readr.read((char *)write_buf, npts * ndims * sizeof(float)); + for (uint64_t i = 0; i < npts; i++) + { + memcpy(write_buf + i * ndims, (read_buf + i * (ndims + 1)) + 1, ndims * sizeof(float)); + } + readr.read((char *)write_buf, npts * ndims * sizeof(float)); } -int main(int argc, char **argv) { - if (argc != 3) { - std::cout << argv[0] << " input_bin output_fvecs" << std::endl; - exit(-1); - } - std::ifstream readr(argv[1], std::ios::binary); - int npts_s32; - int ndims_s32; - readr.read((char *)&npts_s32, sizeof(int32_t)); - readr.read((char *)&ndims_s32, sizeof(int32_t)); - size_t npts = npts_s32; - size_t ndims = ndims_s32; - uint32_t ndims_u32 = (uint32_t)ndims_s32; - // uint64_t fsize = writr.tellg(); - readr.seekg(0, std::ios::beg); - - unsigned ndims_u32; - writr.write((char *)&ndims_u32, sizeof(unsigned)); - writr.seekg(0, std::ios::beg); - uint64_t ndims = (uint64_t)ndims_u32; - uint64_t npts = fsize / ((ndims + 1) * sizeof(float)); - std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims - << std::endl; - - uint64_t blk_size = 131072; - uint64_t nblks = ROUND_UP(npts, blk_size) / blk_size; - std::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]; - for (uint64_t i = 0; i < nblks; i++) { - uint64_t 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; - } - - delete[] read_buf; - delete[] write_buf; - - writr.close(); - readr.close(); +int main(int argc, char **argv) +{ + if (argc != 3) + { + std::cout << argv[0] << " input_bin output_fvecs" << std::endl; + exit(-1); + } + std::ifstream readr(argv[1], std::ios::binary); + int npts_s32; + int ndims_s32; + readr.read((char *)&npts_s32, sizeof(int32_t)); + readr.read((char *)&ndims_s32, sizeof(int32_t)); + size_t npts = npts_s32; + size_t ndims = ndims_s32; + uint32_t ndims_u32 = (uint32_t)ndims_s32; + // uint64_t fsize = writr.tellg(); + readr.seekg(0, std::ios::beg); + + unsigned ndims_u32; + writr.write((char *)&ndims_u32, sizeof(unsigned)); + writr.seekg(0, std::ios::beg); + uint64_t ndims = (uint64_t)ndims_u32; + uint64_t npts = fsize / ((ndims + 1) * sizeof(float)); + std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims << std::endl; + + uint64_t blk_size = 131072; + uint64_t nblks = ROUND_UP(npts, blk_size) / blk_size; + std::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]; + for (uint64_t i = 0; i < nblks; i++) + { + uint64_t 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; + } + + 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 0fdb1a902..7851bef6d 100644 --- a/tests/utils/bin_to_tsv.cpp +++ b/tests/utils/bin_to_tsv.cpp @@ -5,64 +5,65 @@ #include "utils.h" template -void block_convert(std::ofstream &writer, std::ifstream &reader, T *read_buf, - size_t npts, size_t ndims) { - reader.read((char *)read_buf, npts * ndims * sizeof(float)); +void block_convert(std::ofstream &writer, std::ifstream &reader, T *read_buf, size_t npts, size_t ndims) +{ + reader.read((char *)read_buf, npts * ndims * sizeof(float)); - for (size_t i = 0; i < npts; i++) { - for (size_t d = 0; d < ndims; d++) { - writer << read_buf[d + i * ndims]; - if (d < ndims - 1) - writer << "\t"; - else - writer << "\n"; + for (size_t i = 0; i < npts; i++) + { + for (size_t d = 0; d < ndims; d++) + { + writer << read_buf[d + i * ndims]; + if (d < ndims - 1) + writer << "\t"; + else + writer << "\n"; + } } - } } -int main(int argc, char **argv) { - if (argc != 4) { - std::cout << argv[0] << " input_bin output_tsv" - << std::endl; - exit(-1); - } - std::string type_string(argv[1]); - if ((type_string != std::string("float")) && - (type_string != std::string("int8")) && - (type_string != std::string("uin8"))) { - std::cerr << "Error: type not supported. Use float/int8/uint8" << std::endl; - } +int main(int argc, char **argv) +{ + if (argc != 4) + { + std::cout << argv[0] << " input_bin output_tsv" << std::endl; + exit(-1); + } + std::string type_string(argv[1]); + if ((type_string != std::string("float")) && (type_string != std::string("int8")) && + (type_string != std::string("uin8"))) + { + std::cerr << "Error: type not supported. Use float/int8/uint8" << std::endl; + } - std::ifstream reader(argv[2], std::ios::binary); - uint32_t npts_u32; - uint32_t ndims_u32; - reader.read((char *)&npts_u32, sizeof(uint32_t)); - reader.read((char *)&ndims_u32, sizeof(uint32_t)); - size_t npts = npts_u32; - size_t ndims = ndims_u32; - std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims - << std::endl; + std::ifstream reader(argv[2], std::ios::binary); + uint32_t npts_u32; + uint32_t ndims_u32; + reader.read((char *)&npts_u32, sizeof(uint32_t)); + reader.read((char *)&ndims_u32, sizeof(uint32_t)); + size_t npts = npts_u32; + size_t ndims = ndims_u32; + std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims << std::endl; - size_t blk_size = 131072; - size_t nblks = ROUND_UP(npts, blk_size) / blk_size; + size_t blk_size = 131072; + size_t nblks = ROUND_UP(npts, blk_size) / blk_size; - std::ofstream writer(argv[3]); - char *read_buf = new char[blk_size * ndims * 4]; - for (size_t i = 0; i < nblks; i++) { - size_t 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); - else if (type_string == std::string("int8")) - 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); - std::cout << "Block #" << i << " written" << std::endl; - } + std::ofstream writer(argv[3]); + char *read_buf = new char[blk_size * ndims * 4]; + for (size_t i = 0; i < nblks; i++) + { + size_t 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); + else if (type_string == std::string("int8")) + 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); + std::cout << "Block #" << i << " written" << std::endl; + } - delete[] read_buf; + delete[] read_buf; - writer.close(); - reader.close(); + writer.close(); + reader.close(); } diff --git a/tests/utils/calculate_recall.cpp b/tests/utils/calculate_recall.cpp index 04d16c7e4..a45eb19d7 100644 --- a/tests/utils/calculate_recall.cpp +++ b/tests/utils/calculate_recall.cpp @@ -12,42 +12,43 @@ #include "utils.h" #include "disk_utils.h" -int main(int argc, char **argv) { - if (argc != 4) { - std::cout << argv[0] << " " - << std::endl; - return -1; - } - uint32_t *gold_std = NULL; - float *gs_dist = nullptr; - uint32_t *our_results = NULL; - float *or_dist = nullptr; - size_t points_num, points_num_gs, points_num_or; - size_t dim_gs; - size_t dim_or; - diskann::load_truthset(argv[1], gold_std, gs_dist, points_num_gs, dim_gs); - diskann::load_truthset(argv[2], our_results, or_dist, points_num_or, dim_or); +int main(int argc, char **argv) +{ + if (argc != 4) + { + std::cout << argv[0] << " " << std::endl; + return -1; + } + uint32_t *gold_std = NULL; + float *gs_dist = nullptr; + uint32_t *our_results = NULL; + float *or_dist = nullptr; + size_t points_num, points_num_gs, points_num_or; + size_t dim_gs; + size_t dim_or; + diskann::load_truthset(argv[1], gold_std, gs_dist, points_num_gs, dim_gs); + diskann::load_truthset(argv[2], our_results, or_dist, points_num_or, dim_or); - if (points_num_gs != points_num_or) { - std::cout << "Error. Number of queries mismatch in ground truth and " - "our results" - << std::endl; - return -1; - } - points_num = points_num_gs; + if (points_num_gs != points_num_or) + { + std::cout << "Error. Number of queries mismatch in ground truth and " + "our results" + << std::endl; + return -1; + } + points_num = points_num_gs; - uint32_t recall_at = std::atoi(argv[3]); + 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; - 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); + 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; + 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); - // double avg_recall = (recall*1.0)/(points_num*1.0); - std::cout << "Avg. recall@" << recall_at << " is " << recall_val << "\n"; + // double avg_recall = (recall*1.0)/(points_num*1.0); + std::cout << "Avg. recall@" << recall_at << " is " << recall_val << "\n"; } diff --git a/tests/utils/compute_groundtruth.cpp b/tests/utils/compute_groundtruth.cpp index 587cdd58a..991f29ff6 100644 --- a/tests/utils/compute_groundtruth.cpp +++ b/tests/utils/compute_groundtruth.cpp @@ -40,535 +40,534 @@ typedef std::string path; namespace po = boost::program_options; -template -T div_round_up(const T numerator, const T denominator) { - return (numerator % denominator == 0) ? (numerator / denominator) - : 1 + (numerator / denominator); +template T div_round_up(const T numerator, const T denominator) +{ + return (numerator % denominator == 0) ? (numerator / denominator) : 1 + (numerator / denominator); } using pairIF = std::pair; -struct cmpmaxstruct { - bool operator()(const pairIF &l, const pairIF &r) { - return l.second < r.second; - }; +struct cmpmaxstruct +{ + bool operator()(const pairIF &l, const pairIF &r) + { + return l.second < r.second; + }; }; -using maxPQIFCS = - std::priority_queue, cmpmaxstruct>; +using maxPQIFCS = std::priority_queue, cmpmaxstruct>; -template -T *aligned_malloc(const size_t n, const size_t alignment) { +template T *aligned_malloc(const size_t n, const size_t alignment) +{ #ifdef _WINDOWS - return (T *)_aligned_malloc(sizeof(T) * n, alignment); + return (T *)_aligned_malloc(sizeof(T) * n, alignment); #else - return static_cast(aligned_alloc(alignment, sizeof(T) * n)); + return static_cast(aligned_alloc(alignment, sizeof(T) * n)); #endif } -inline bool custom_dist(const std::pair &a, - const std::pair &b) { - return a.second < b.second; +inline bool custom_dist(const std::pair &a, const std::pair &b) +{ + return a.second < b.second; } -void compute_l2sq(float *const points_l2sq, const float *const matrix, - const int64_t num_points, const int dim) { - assert(points_l2sq != NULL); +void compute_l2sq(float *const points_l2sq, const float *const matrix, const int64_t num_points, const int dim) +{ + assert(points_l2sq != NULL); #pragma omp parallel for schedule(static, 65536) - for (int64_t d = 0; d < num_points; ++d) - points_l2sq[d] = cblas_sdot(dim, matrix + (ptrdiff_t)d * (ptrdiff_t)dim, 1, - matrix + (ptrdiff_t)d * (ptrdiff_t)dim, 1); + for (int64_t d = 0; d < num_points; ++d) + points_l2sq[d] = + cblas_sdot(dim, matrix + (ptrdiff_t)d * (ptrdiff_t)dim, 1, matrix + (ptrdiff_t)d * (ptrdiff_t)dim, 1); } -void distsq_to_points( - const size_t dim, - float *dist_matrix, // Col Major, cols are queries, rows are points - size_t npoints, const float *const points, - const float *const points_l2sq, // points in Col major - size_t nqueries, const float *const queries, - const float *const queries_l2sq, // queries in Col major - float *ones_vec = NULL) // Scratchspace of num_data size and init to 1.0 +void distsq_to_points(const size_t dim, + float *dist_matrix, // Col Major, cols are queries, rows are points + size_t npoints, const float *const points, + const float *const points_l2sq, // points in Col major + size_t nqueries, const float *const queries, + const float *const queries_l2sq, // queries in Col major + float *ones_vec = NULL) // Scratchspace of num_data size and init to 1.0 { - bool ones_vec_alloc = false; - if (ones_vec == NULL) { - ones_vec = new float[nqueries > npoints ? nqueries : npoints]; - std::fill_n(ones_vec, nqueries > npoints ? nqueries : npoints, (float)1.0); - ones_vec_alloc = true; - } - cblas_sgemm(CblasColMajor, CblasTrans, CblasNoTrans, npoints, nqueries, dim, - (float)-2.0, points, dim, queries, dim, (float)0.0, dist_matrix, - npoints); - cblas_sgemm(CblasColMajor, CblasNoTrans, CblasTrans, npoints, nqueries, 1, - (float)1.0, points_l2sq, npoints, ones_vec, nqueries, (float)1.0, - dist_matrix, npoints); - cblas_sgemm(CblasColMajor, CblasNoTrans, CblasTrans, npoints, nqueries, 1, - (float)1.0, ones_vec, npoints, queries_l2sq, nqueries, (float)1.0, - dist_matrix, npoints); - if (ones_vec_alloc) delete[] ones_vec; + bool ones_vec_alloc = false; + if (ones_vec == NULL) + { + ones_vec = new float[nqueries > npoints ? nqueries : npoints]; + std::fill_n(ones_vec, nqueries > npoints ? nqueries : npoints, (float)1.0); + ones_vec_alloc = true; + } + cblas_sgemm(CblasColMajor, CblasTrans, CblasNoTrans, npoints, nqueries, dim, (float)-2.0, points, dim, queries, dim, + (float)0.0, dist_matrix, npoints); + cblas_sgemm(CblasColMajor, CblasNoTrans, CblasTrans, npoints, nqueries, 1, (float)1.0, points_l2sq, npoints, + ones_vec, nqueries, (float)1.0, dist_matrix, npoints); + cblas_sgemm(CblasColMajor, CblasNoTrans, CblasTrans, npoints, nqueries, 1, (float)1.0, ones_vec, npoints, + queries_l2sq, nqueries, (float)1.0, dist_matrix, npoints); + if (ones_vec_alloc) + delete[] ones_vec; } -void inner_prod_to_points( - const size_t dim, - float *dist_matrix, // Col Major, cols are queries, rows are points - size_t npoints, const float *const points, size_t nqueries, - const float *const queries, - float *ones_vec = NULL) // Scratchspace of num_data size and init to 1.0 +void inner_prod_to_points(const size_t dim, + float *dist_matrix, // Col Major, cols are queries, rows are points + size_t npoints, const float *const points, size_t nqueries, const float *const queries, + float *ones_vec = NULL) // Scratchspace of num_data size and init to 1.0 { - bool ones_vec_alloc = false; - if (ones_vec == NULL) { - ones_vec = new float[nqueries > npoints ? nqueries : npoints]; - std::fill_n(ones_vec, nqueries > npoints ? nqueries : npoints, (float)1.0); - ones_vec_alloc = true; - } - cblas_sgemm(CblasColMajor, CblasTrans, CblasNoTrans, npoints, nqueries, dim, - (float)-1.0, points, dim, queries, dim, (float)0.0, dist_matrix, - npoints); - - if (ones_vec_alloc) delete[] ones_vec; + bool ones_vec_alloc = false; + if (ones_vec == NULL) + { + ones_vec = new float[nqueries > npoints ? nqueries : npoints]; + std::fill_n(ones_vec, nqueries > npoints ? nqueries : npoints, (float)1.0); + ones_vec_alloc = true; + } + cblas_sgemm(CblasColMajor, CblasTrans, CblasNoTrans, npoints, nqueries, dim, (float)-1.0, points, dim, queries, dim, + (float)0.0, dist_matrix, npoints); + + if (ones_vec_alloc) + delete[] ones_vec; } -void exact_knn( - const size_t dim, const size_t k, - int *const closest_points, // k * num_queries preallocated, col - // major, queries columns - float *const dist_closest_points, // k * num_queries - // preallocated, Dist to - // corresponding closes_points - size_t npoints, - float *points_in, // points in Col major - size_t nqueries, float *queries_in, - diskann::Metric metric = diskann::Metric::L2) // queries in Col major +void exact_knn(const size_t dim, const size_t k, + int *const closest_points, // k * num_queries preallocated, col + // major, queries columns + float *const dist_closest_points, // k * num_queries + // preallocated, Dist to + // corresponding closes_points + size_t npoints, + float *points_in, // points in Col major + size_t nqueries, float *queries_in, + diskann::Metric metric = diskann::Metric::L2) // queries in Col major { - float *points_l2sq = new float[npoints]; - float *queries_l2sq = new float[nqueries]; - compute_l2sq(points_l2sq, points_in, npoints, dim); - compute_l2sq(queries_l2sq, queries_in, nqueries, dim); - - float *points = points_in; - float *queries = queries_in; - - if (metric == diskann::Metric::COSINE) { // we convert cosine distance as - // normalized L2 distnace - points = new float[npoints * dim]; - queries = new float[nqueries * dim]; + float *points_l2sq = new float[npoints]; + float *queries_l2sq = new float[nqueries]; + compute_l2sq(points_l2sq, points_in, npoints, dim); + compute_l2sq(queries_l2sq, queries_in, nqueries, dim); + + float *points = points_in; + float *queries = queries_in; + + if (metric == diskann::Metric::COSINE) + { // we convert cosine distance as + // normalized L2 distnace + points = new float[npoints * dim]; + queries = new float[nqueries * dim]; #pragma omp parallel for schedule(static, 4096) - for (int64_t i = 0; i < (int64_t)npoints; i++) { - float norm = std::sqrt(points_l2sq[i]); - if (norm == 0) { - norm = std::numeric_limits::epsilon(); - } - for (uint32_t j = 0; j < dim; j++) { - points[i * dim + j] = points_in[i * dim + j] / norm; - } - } + for (int64_t i = 0; i < (int64_t)npoints; i++) + { + float norm = std::sqrt(points_l2sq[i]); + if (norm == 0) + { + norm = std::numeric_limits::epsilon(); + } + for (uint32_t j = 0; j < dim; j++) + { + points[i * dim + j] = points_in[i * dim + j] / norm; + } + } #pragma omp parallel for schedule(static, 4096) - for (int64_t i = 0; i < (int64_t)nqueries; i++) { - float norm = std::sqrt(queries_l2sq[i]); - if (norm == 0) { - norm = std::numeric_limits::epsilon(); - } - for (uint32_t j = 0; j < dim; j++) { - queries[i * dim + j] = queries_in[i * dim + j] / norm; - } + for (int64_t i = 0; i < (int64_t)nqueries; i++) + { + float norm = std::sqrt(queries_l2sq[i]); + if (norm == 0) + { + norm = std::numeric_limits::epsilon(); + } + for (uint32_t j = 0; j < dim; j++) + { + queries[i * dim + j] = queries_in[i * dim + j] / norm; + } + } + // recalculate norms after normalizing, they should all be one. + compute_l2sq(points_l2sq, points, npoints, dim); + compute_l2sq(queries_l2sq, queries, nqueries, dim); } - // recalculate norms after normalizing, they should all be one. - compute_l2sq(points_l2sq, points, npoints, dim); - compute_l2sq(queries_l2sq, queries, nqueries, dim); - } - - std::cout << "Going to compute " << k << " NNs for " << nqueries - << " queries over " << npoints << " points in " << dim - << " dimensions using"; - if (metric == diskann::Metric::INNER_PRODUCT) - std::cout << " MIPS "; - else if (metric == diskann::Metric::COSINE) - std::cout << " Cosine "; - else - std::cout << " L2 "; - std::cout << "distance fn. " << std::endl; - - size_t q_batch_size = (1 << 9); - float *dist_matrix = new float[(size_t)q_batch_size * (size_t)npoints]; - - for (size_t b = 0; b < div_round_up(nqueries, q_batch_size); ++b) { - int64_t q_b = b * q_batch_size; - int64_t q_e = - ((b + 1) * q_batch_size > nqueries) ? nqueries : (b + 1) * q_batch_size; - - if (metric == diskann::Metric::L2 || metric == diskann::Metric::COSINE) { - 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); - } else { - inner_prod_to_points(dim, dist_matrix, npoints, points, q_e - q_b, - queries + (ptrdiff_t)q_b * (ptrdiff_t)dim); - } - std::cout << "Computed distances for queries: [" << q_b << "," << q_e << ")" - << std::endl; + + std::cout << "Going to compute " << k << " NNs for " << nqueries << " queries over " << npoints << " points in " + << dim << " dimensions using"; + if (metric == diskann::Metric::INNER_PRODUCT) + std::cout << " MIPS "; + else if (metric == diskann::Metric::COSINE) + std::cout << " Cosine "; + else + std::cout << " L2 "; + std::cout << "distance fn. " << std::endl; + + size_t q_batch_size = (1 << 9); + float *dist_matrix = new float[(size_t)q_batch_size * (size_t)npoints]; + + for (size_t b = 0; b < div_round_up(nqueries, q_batch_size); ++b) + { + int64_t q_b = b * q_batch_size; + int64_t q_e = ((b + 1) * q_batch_size > nqueries) ? nqueries : (b + 1) * q_batch_size; + + if (metric == diskann::Metric::L2 || metric == diskann::Metric::COSINE) + { + 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); + } + else + { + inner_prod_to_points(dim, dist_matrix, npoints, points, q_e - q_b, + queries + (ptrdiff_t)q_b * (ptrdiff_t)dim); + } + std::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++) { - maxPQIFCS point_dist; - for (size_t p = 0; p < k; p++) - point_dist.emplace( - p, dist_matrix[(ptrdiff_t)p + - (ptrdiff_t)(q - q_b) * (ptrdiff_t)npoints]); - for (size_t p = k; p < npoints; p++) { - if (point_dist.top().second > - dist_matrix[(ptrdiff_t)p + - (ptrdiff_t)(q - q_b) * (ptrdiff_t)npoints]) - point_dist.emplace( - p, dist_matrix[(ptrdiff_t)p + - (ptrdiff_t)(q - q_b) * (ptrdiff_t)npoints]); - if (point_dist.size() > k) point_dist.pop(); - } - for (ptrdiff_t l = 0; l < (ptrdiff_t)k; ++l) { - closest_points[(ptrdiff_t)(k - 1 - l) + (ptrdiff_t)q * (ptrdiff_t)k] = - point_dist.top().first; - dist_closest_points[(ptrdiff_t)(k - 1 - l) + - (ptrdiff_t)q * (ptrdiff_t)k] = - point_dist.top().second; - point_dist.pop(); - } - assert(std::is_sorted( - dist_closest_points + (ptrdiff_t)q * (ptrdiff_t)k, - dist_closest_points + (ptrdiff_t)(q + 1) * (ptrdiff_t)k)); + for (long long q = q_b; q < q_e; q++) + { + maxPQIFCS point_dist; + for (size_t p = 0; p < k; p++) + point_dist.emplace(p, dist_matrix[(ptrdiff_t)p + (ptrdiff_t)(q - q_b) * (ptrdiff_t)npoints]); + for (size_t p = k; p < npoints; p++) + { + if (point_dist.top().second > dist_matrix[(ptrdiff_t)p + (ptrdiff_t)(q - q_b) * (ptrdiff_t)npoints]) + point_dist.emplace(p, dist_matrix[(ptrdiff_t)p + (ptrdiff_t)(q - q_b) * (ptrdiff_t)npoints]); + if (point_dist.size() > k) + point_dist.pop(); + } + for (ptrdiff_t l = 0; l < (ptrdiff_t)k; ++l) + { + closest_points[(ptrdiff_t)(k - 1 - l) + (ptrdiff_t)q * (ptrdiff_t)k] = point_dist.top().first; + dist_closest_points[(ptrdiff_t)(k - 1 - l) + (ptrdiff_t)q * (ptrdiff_t)k] = point_dist.top().second; + point_dist.pop(); + } + assert(std::is_sorted(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; } - std::cout << "Computed exact k-NN for queries: [" << q_b << "," << q_e - << ")" << std::endl; - } - delete[] dist_matrix; + delete[] dist_matrix; - delete[] points_l2sq; - delete[] queries_l2sq; + delete[] points_l2sq; + delete[] queries_l2sq; - if (metric == diskann::Metric::COSINE) { - delete[] points; - delete[] queries; - } + if (metric == diskann::Metric::COSINE) + { + delete[] points; + delete[] queries; + } } -template -inline int get_num_parts(const char *filename) { - std::ifstream reader; - reader.exceptions(std::ios::failbit | std::ios::badbit); - reader.open(filename, std::ios::binary); - std::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; - 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; - return num_parts; +template inline int get_num_parts(const char *filename) +{ + std::ifstream reader; + reader.exceptions(std::ios::failbit | std::ios::badbit); + reader.open(filename, std::ios::binary); + std::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; + 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; + return num_parts; } template -inline void load_bin_as_float(const char *filename, float *&data, size_t &npts, - size_t &ndims, int part_num) { - std::ifstream reader; - reader.exceptions(std::ios::failbit | std::ios::badbit); - reader.open(filename, std::ios::binary); - std::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)); - uint64_t start_id = part_num * PARTSIZE; - uint64_t end_id = (std::min)(start_id + PARTSIZE, (uint64_t)npts_i32); - npts = end_id - start_id; - ndims = (uint64_t)ndims_i32; - std::cout << "#pts in part = " << npts << ", #dims = " << ndims - << ", size = " << npts * ndims * sizeof(T) << "B" << std::endl; - - reader.seekg(start_id * ndims * sizeof(T) + 2 * sizeof(uint32_t), - std::ios::beg); - T *data_T = new T[npts * ndims]; - reader.read((char *)data_T, sizeof(T) * npts * ndims); - std::cout << "Finished reading part of the bin file." << std::endl; - reader.close(); - data = aligned_malloc(npts * ndims, ALIGNMENT); +inline void load_bin_as_float(const char *filename, float *&data, size_t &npts, size_t &ndims, int part_num) +{ + std::ifstream reader; + reader.exceptions(std::ios::failbit | std::ios::badbit); + reader.open(filename, std::ios::binary); + std::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)); + uint64_t start_id = part_num * PARTSIZE; + uint64_t end_id = (std::min)(start_id + PARTSIZE, (uint64_t)npts_i32); + npts = end_id - start_id; + ndims = (uint64_t)ndims_i32; + std::cout << "#pts in part = " << npts << ", #dims = " << ndims << ", size = " << npts * ndims * sizeof(T) << "B" + << std::endl; + + reader.seekg(start_id * ndims * sizeof(T) + 2 * sizeof(uint32_t), std::ios::beg); + T *data_T = new T[npts * ndims]; + reader.read((char *)data_T, sizeof(T) * npts * ndims); + std::cout << "Finished reading part of the bin file." << std::endl; + reader.close(); + data = aligned_malloc(npts * ndims, ALIGNMENT); #pragma omp parallel for schedule(dynamic, 32768) - for (int64_t i = 0; i < (int64_t)npts; i++) { - for (int64_t j = 0; j < (int64_t)ndims; j++) { - float cur_val_float = (float)data_T[i * ndims + j]; - std::memcpy((char *)(data + i * ndims + j), (char *)&cur_val_float, - sizeof(float)); + for (int64_t i = 0; i < (int64_t)npts; i++) + { + for (int64_t j = 0; j < (int64_t)ndims; j++) + { + float cur_val_float = (float)data_T[i * ndims + j]; + std::memcpy((char *)(data + i * ndims + j), (char *)&cur_val_float, sizeof(float)); + } } - } - delete[] data_T; - std::cout << "Finished converting part data to float." << std::endl; + delete[] data_T; + std::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; - writer.exceptions(std::ios::failbit | std::ios::badbit); - writer.open(filename, std::ios::binary | std::ios::out); - std::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; - - writer.write((char *)data, npts * ndims * sizeof(T)); - writer.close(); - std::cout << "Finished writing bin" << std::endl; +template inline void save_bin(const std::string filename, T *data, size_t npts, size_t ndims) +{ + std::ofstream writer; + writer.exceptions(std::ios::failbit | std::ios::badbit); + writer.open(filename, std::ios::binary | std::ios::out); + std::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; + + writer.write((char *)data, npts * ndims * sizeof(T)); + writer.close(); + std::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) { - 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(uint32_t) + 2 * sizeof(int) << "B" - << std::endl; - - writer.write((char *)data, npts * ndims * sizeof(uint32_t)); - writer.write((char *)distances, npts * ndims * sizeof(float)); - writer.close(); - std::cout << "Finished writing truthset" << std::endl; +inline void save_groundtruth_as_one_file(const std::string filename, int32_t *data, float *distances, size_t npts, + size_t ndims) +{ + 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(uint32_t) + 2 * sizeof(int) + << "B" << std::endl; + + writer.write((char *)data, npts * ndims * sizeof(uint32_t)); + writer.write((char *)distances, npts * ndims * sizeof(float)); + writer.close(); + std::cout << "Finished writing truthset" << std::endl; } template -std::vector>> processUnfilteredParts( - const std::string &base_file, size_t &nqueries, size_t &npoints, - size_t &dim, size_t &k, float *query_data, const diskann::Metric &metric, - std::vector &location_to_tag) { - float *base_data; - int num_parts = get_num_parts(base_file.c_str()); - std::vector>> res(nqueries); - for (int p = 0; p < num_parts; p++) { - size_t start_id = p * PARTSIZE; - load_bin_as_float(base_file.c_str(), base_data, npoints, dim, p); - - int *closest_points_part = new int[nqueries * k]; - float *dist_closest_points_part = new float[nqueries * k]; - - uint32_t part_k; - part_k = k < npoints ? k : npoints; - exact_knn(dim, part_k, closest_points_part, dist_closest_points_part, - npoints, base_data, nqueries, query_data, metric); - - for (size_t i = 0; i < nqueries; i++) { - for (size_t j = 0; j < part_k; j++) { - if (!location_to_tag.empty()) - if (location_to_tag[closest_points_part[i * k + j] + start_id] == 0) - continue; - - res[i].push_back(std::make_pair( - (uint32_t)(closest_points_part[i * part_k + j] + start_id), - dist_closest_points_part[i * part_k + j])); - } +std::vector>> processUnfilteredParts(const std::string &base_file, + size_t &nqueries, size_t &npoints, + size_t &dim, size_t &k, float *query_data, + const diskann::Metric &metric, + std::vector &location_to_tag) +{ + float *base_data; + int num_parts = get_num_parts(base_file.c_str()); + std::vector>> res(nqueries); + for (int p = 0; p < num_parts; p++) + { + size_t start_id = p * PARTSIZE; + load_bin_as_float(base_file.c_str(), base_data, npoints, dim, p); + + int *closest_points_part = new int[nqueries * k]; + float *dist_closest_points_part = new float[nqueries * k]; + + uint32_t part_k; + part_k = k < npoints ? k : npoints; + exact_knn(dim, part_k, closest_points_part, dist_closest_points_part, npoints, base_data, nqueries, query_data, + metric); + + for (size_t i = 0; i < nqueries; i++) + { + for (size_t j = 0; j < part_k; j++) + { + if (!location_to_tag.empty()) + if (location_to_tag[closest_points_part[i * k + j] + start_id] == 0) + continue; + + res[i].push_back(std::make_pair((uint32_t)(closest_points_part[i * part_k + j] + start_id), + dist_closest_points_part[i * part_k + j])); + } + } + + delete[] closest_points_part; + delete[] dist_closest_points_part; + + diskann::aligned_free(base_data); } - - delete[] closest_points_part; - delete[] dist_closest_points_part; - - diskann::aligned_free(base_data); - } - return res; + return res; }; template -int aux_main(const std::string &base_file, const std::string &query_file, - const std::string >_file, size_t k, - const diskann::Metric &metric, - const std::string &tags_file = std::string("")) { - size_t npoints, nqueries, dim; - - float *query_data; - - load_bin_as_float(query_file.c_str(), query_data, nqueries, dim, 0); - if (nqueries > PARTSIZE) - std::cerr << "WARNING: #Queries provided (" << nqueries - << ") is greater than " << PARTSIZE - << ". Computing GT only for the first " << PARTSIZE << " queries." - << std::endl; - - // load tags - const bool tags_enabled = tags_file.empty() ? false : true; - std::vector location_to_tag = - diskann::loadTags(tags_file, base_file); - - int *closest_points = new int[nqueries * k]; - float *dist_closest_points = new float[nqueries * k]; - - std::vector>> results = - processUnfilteredParts(base_file, nqueries, npoints, dim, k, - query_data, metric, location_to_tag); - - for (size_t i = 0; i < nqueries; i++) { - std::vector> &cur_res = results[i]; - std::sort(cur_res.begin(), cur_res.end(), custom_dist); - size_t j = 0; - for (auto iter : cur_res) { - if (j == k) break; - if (tags_enabled) { - std::uint32_t index_with_tag = location_to_tag[iter.first]; - closest_points[i * k + j] = (int32_t)index_with_tag; - } else { - closest_points[i * k + j] = (int32_t)iter.first; - } - - if (metric == diskann::Metric::INNER_PRODUCT) - dist_closest_points[i * k + j] = -iter.second; - else - dist_closest_points[i * k + j] = iter.second; - - ++j; +int aux_main(const std::string &base_file, const std::string &query_file, const std::string >_file, size_t k, + const diskann::Metric &metric, const std::string &tags_file = std::string("")) +{ + size_t npoints, nqueries, dim; + + float *query_data; + + load_bin_as_float(query_file.c_str(), query_data, nqueries, dim, 0); + if (nqueries > PARTSIZE) + std::cerr << "WARNING: #Queries provided (" << nqueries << ") is greater than " << PARTSIZE + << ". Computing GT only for the first " << PARTSIZE << " queries." << std::endl; + + // load tags + const bool tags_enabled = tags_file.empty() ? false : true; + std::vector location_to_tag = diskann::loadTags(tags_file, base_file); + + int *closest_points = new int[nqueries * k]; + float *dist_closest_points = new float[nqueries * k]; + + std::vector>> results = + processUnfilteredParts(base_file, nqueries, npoints, dim, k, query_data, metric, location_to_tag); + + for (size_t i = 0; i < nqueries; i++) + { + std::vector> &cur_res = results[i]; + std::sort(cur_res.begin(), cur_res.end(), custom_dist); + size_t j = 0; + for (auto iter : cur_res) + { + if (j == k) + break; + if (tags_enabled) + { + std::uint32_t index_with_tag = location_to_tag[iter.first]; + closest_points[i * k + j] = (int32_t)index_with_tag; + } + else + { + closest_points[i * k + j] = (int32_t)iter.first; + } + + if (metric == diskann::Metric::INNER_PRODUCT) + dist_closest_points[i * k + j] = -iter.second; + else + dist_closest_points[i * k + j] = iter.second; + + ++j; + } + if (j < k) + std::cout << "WARNING: found less than k GT entries for query " << i << std::endl; } - if (j < k) - std::cout << "WARNING: found less than k GT entries for query " << i - << std::endl; - } - - save_groundtruth_as_one_file(gt_file, closest_points, dist_closest_points, - nqueries, k); - delete[] closest_points; - delete[] dist_closest_points; - diskann::aligned_free(query_data); - - return 0; + + save_groundtruth_as_one_file(gt_file, closest_points, dist_closest_points, nqueries, k); + delete[] closest_points; + delete[] dist_closest_points; + diskann::aligned_free(query_data); + + return 0; } -void load_truthset(const std::string &bin_file, uint32_t *&ids, float *&dists, - size_t &npts, size_t &dim) { - size_t read_blk_size = 64 * 1024 * 1024; - cached_ifstream reader(bin_file, read_blk_size); - diskann::cout << "Reading truthset file " << bin_file.c_str() << " ..." - << std::endl; - size_t actual_file_size = reader.get_file_size(); - - int npts_i32, dim_i32; - reader.read((char *)&npts_i32, sizeof(int)); - reader.read((char *)&dim_i32, sizeof(int)); - npts = (uint32_t)npts_i32; - dim = (uint32_t)dim_i32; - - diskann::cout << "Metadata: #pts = " << npts << ", #dims = " << dim << "... " - << std::endl; - - int truthset_type = -1; // 1 means truthset has ids and distances, 2 means - // only ids, -1 is error - size_t expected_file_size_with_dists = - 2 * npts * dim * sizeof(uint32_t) + 2 * sizeof(uint32_t); - - if (actual_file_size == expected_file_size_with_dists) truthset_type = 1; - - size_t expected_file_size_just_ids = - npts * dim * sizeof(uint32_t) + 2 * sizeof(uint32_t); - - if (actual_file_size == expected_file_size_just_ids) truthset_type = 2; - - if (truthset_type == -1) { - std::stringstream stream; - stream << "Error. File size mismatch. File should have bin format, with " - "npts followed by ngt followed by npts*ngt ids and optionally " - "followed by npts*ngt distance values; actual size: " - << actual_file_size - << ", expected: " << expected_file_size_with_dists << " or " - << expected_file_size_just_ids; - diskann::cout << stream.str(); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); - } - - ids = new uint32_t[npts * dim]; - reader.read((char *)ids, npts * dim * sizeof(uint32_t)); - - if (truthset_type == 1) { - dists = new float[npts * dim]; - reader.read((char *)dists, npts * dim * sizeof(float)); - } +void load_truthset(const std::string &bin_file, uint32_t *&ids, float *&dists, size_t &npts, size_t &dim) +{ + size_t read_blk_size = 64 * 1024 * 1024; + cached_ifstream reader(bin_file, read_blk_size); + diskann::cout << "Reading truthset file " << bin_file.c_str() << " ..." << std::endl; + size_t actual_file_size = reader.get_file_size(); + + int npts_i32, dim_i32; + reader.read((char *)&npts_i32, sizeof(int)); + reader.read((char *)&dim_i32, sizeof(int)); + npts = (uint32_t)npts_i32; + dim = (uint32_t)dim_i32; + + diskann::cout << "Metadata: #pts = " << npts << ", #dims = " << dim << "... " << std::endl; + + int truthset_type = -1; // 1 means truthset has ids and distances, 2 means + // only ids, -1 is error + size_t expected_file_size_with_dists = 2 * npts * dim * sizeof(uint32_t) + 2 * sizeof(uint32_t); + + if (actual_file_size == expected_file_size_with_dists) + truthset_type = 1; + + size_t expected_file_size_just_ids = npts * dim * sizeof(uint32_t) + 2 * sizeof(uint32_t); + + if (actual_file_size == expected_file_size_just_ids) + truthset_type = 2; + + if (truthset_type == -1) + { + std::stringstream stream; + stream << "Error. File size mismatch. File should have bin format, with " + "npts followed by ngt followed by npts*ngt ids and optionally " + "followed by npts*ngt distance values; actual size: " + << actual_file_size << ", expected: " << expected_file_size_with_dists << " or " + << expected_file_size_just_ids; + diskann::cout << stream.str(); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } + + ids = new uint32_t[npts * dim]; + reader.read((char *)ids, npts * dim * sizeof(uint32_t)); + + if (truthset_type == 1) + { + dists = new float[npts * dim]; + reader.read((char *)dists, npts * dim * sizeof(float)); + } } -int main(int argc, char **argv) { - std::string data_type, dist_fn, base_file, query_file, gt_file, tags_file; - uint64_t K; - - try { - po::options_description desc{"Arguments"}; - - desc.add_options()("help,h", "Print information on arguments"); - - desc.add_options()("data_type", - po::value(&data_type)->required(), - "data type "); - desc.add_options()("dist_fn", po::value(&dist_fn)->required(), - "distance function "); - desc.add_options()("base_file", - po::value(&base_file)->required(), - "File containing the base vectors in binary format"); - desc.add_options()("query_file", - po::value(&query_file)->required(), - "File containing the query vectors in binary format"); - desc.add_options()("gt_file", po::value(>_file)->required(), - "File name for the writing ground truth in binary " - "format, please don' append .bin at end if " - "no filter_label or filter_label_file is provided it " - "will save the file with '.bin' at end." - "else it will save the file as filename_label.bin"); - desc.add_options()("K", po::value(&K)->required(), - "Number of ground truth nearest neighbors to compute"); - desc.add_options()( - "tags_file", - po::value(&tags_file)->default_value(std::string()), - "File containing the tags in binary format"); - - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) { - std::cout << desc; - return 0; +int main(int argc, char **argv) +{ + std::string data_type, dist_fn, base_file, query_file, gt_file, tags_file; + uint64_t K; + + try + { + po::options_description desc{"Arguments"}; + + desc.add_options()("help,h", "Print information on arguments"); + + desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); + desc.add_options()("dist_fn", po::value(&dist_fn)->required(), "distance function "); + desc.add_options()("base_file", po::value(&base_file)->required(), + "File containing the base vectors in binary format"); + desc.add_options()("query_file", po::value(&query_file)->required(), + "File containing the query vectors in binary format"); + desc.add_options()("gt_file", po::value(>_file)->required(), + "File name for the writing ground truth in binary " + "format, please don' append .bin at end if " + "no filter_label or filter_label_file is provided it " + "will save the file with '.bin' at end." + "else it will save the file as filename_label.bin"); + desc.add_options()("K", po::value(&K)->required(), + "Number of ground truth nearest neighbors to compute"); + desc.add_options()("tags_file", po::value(&tags_file)->default_value(std::string()), + "File containing the tags in binary format"); + + po::variables_map vm; + po::store(po::parse_command_line(argc, argv, desc), vm); + if (vm.count("help")) + { + std::cout << desc; + return 0; + } + po::notify(vm); + } + catch (const std::exception &ex) + { + std::cerr << ex.what() << '\n'; + return -1; + } + + if (data_type != std::string("float") && data_type != std::string("int8") && data_type != std::string("uint8")) + { + std::cout << "Unsupported type. float, int8 and uint8 types are supported." << std::endl; + return -1; + } + + diskann::Metric metric; + if (dist_fn == std::string("l2")) + { + metric = diskann::Metric::L2; + } + else if (dist_fn == std::string("mips")) + { + metric = diskann::Metric::INNER_PRODUCT; + } + else if (dist_fn == std::string("cosine")) + { + metric = diskann::Metric::COSINE; + } + else + { + std::cerr << "Unsupported distance function. Use l2/mips/cosine." << std::endl; + return -1; + } + + try + { + if (data_type == std::string("float")) + aux_main(base_file, query_file, gt_file, K, metric, tags_file); + if (data_type == std::string("int8")) + aux_main(base_file, query_file, gt_file, K, metric, tags_file); + if (data_type == std::string("uint8")) + aux_main(base_file, query_file, gt_file, K, metric, tags_file); + } + catch (const std::exception &e) + { + std::cout << std::string(e.what()) << std::endl; + diskann::cerr << "Compute GT failed." << std::endl; + return -1; } - po::notify(vm); - } catch (const std::exception &ex) { - std::cerr << ex.what() << '\n'; - return -1; - } - - if (data_type != std::string("float") && data_type != std::string("int8") && - data_type != std::string("uint8")) { - std::cout << "Unsupported type. float, int8 and uint8 types are supported." - << std::endl; - return -1; - } - - diskann::Metric metric; - if (dist_fn == std::string("l2")) { - metric = diskann::Metric::L2; - } else if (dist_fn == std::string("mips")) { - metric = diskann::Metric::INNER_PRODUCT; - } else if (dist_fn == std::string("cosine")) { - metric = diskann::Metric::COSINE; - } else { - std::cerr << "Unsupported distance function. Use l2/mips/cosine." - << std::endl; - return -1; - } - - try { - if (data_type == std::string("float")) - aux_main(base_file, query_file, gt_file, K, metric, tags_file); - if (data_type == std::string("int8")) - aux_main(base_file, query_file, gt_file, K, metric, tags_file); - if (data_type == std::string("uint8")) - aux_main(base_file, query_file, gt_file, K, metric, tags_file); - } catch (const std::exception &e) { - std::cout << std::string(e.what()) << std::endl; - diskann::cerr << "Compute GT failed." << std::endl; - return -1; - } } diff --git a/tests/utils/compute_groundtruth_for_filters.cpp b/tests/utils/compute_groundtruth_for_filters.cpp index 4b8745f4a..eb962257d 100644 --- a/tests/utils/compute_groundtruth_for_filters.cpp +++ b/tests/utils/compute_groundtruth_for_filters.cpp @@ -41,878 +41,887 @@ typedef std::string path; namespace po = boost::program_options; -template -T div_round_up(const T numerator, const T denominator) { - return (numerator % denominator == 0) ? (numerator / denominator) - : 1 + (numerator / denominator); +template T div_round_up(const T numerator, const T denominator) +{ + return (numerator % denominator == 0) ? (numerator / denominator) : 1 + (numerator / denominator); } using pairIF = std::pair; -struct cmpmaxstruct { - bool operator()(const pairIF &l, const pairIF &r) { - return l.second < r.second; - }; +struct cmpmaxstruct +{ + bool operator()(const pairIF &l, const pairIF &r) + { + return l.second < r.second; + }; }; -using maxPQIFCS = - std::priority_queue, cmpmaxstruct>; +using maxPQIFCS = std::priority_queue, cmpmaxstruct>; -template -T *aligned_malloc(const size_t n, const size_t alignment) { +template T *aligned_malloc(const size_t n, const size_t alignment) +{ #ifdef _WINDOWS - return (T *)_aligned_malloc(sizeof(T) * n, alignment); + return (T *)_aligned_malloc(sizeof(T) * n, alignment); #else - return static_cast(aligned_alloc(alignment, sizeof(T) * n)); + return static_cast(aligned_alloc(alignment, sizeof(T) * n)); #endif } -inline bool custom_dist(const std::pair &a, - const std::pair &b) { - return a.second < b.second; +inline bool custom_dist(const std::pair &a, const std::pair &b) +{ + return a.second < b.second; } -void compute_l2sq(float *const points_l2sq, const float *const matrix, - const int64_t num_points, const int dim) { - assert(points_l2sq != NULL); +void compute_l2sq(float *const points_l2sq, const float *const matrix, const int64_t num_points, const int dim) +{ + assert(points_l2sq != NULL); #pragma omp parallel for schedule(static, 65536) - for (int64_t d = 0; d < num_points; ++d) - points_l2sq[d] = cblas_sdot(dim, matrix + (ptrdiff_t)d * (ptrdiff_t)dim, 1, - matrix + (ptrdiff_t)d * (ptrdiff_t)dim, 1); + for (int64_t d = 0; d < num_points; ++d) + points_l2sq[d] = + cblas_sdot(dim, matrix + (ptrdiff_t)d * (ptrdiff_t)dim, 1, matrix + (ptrdiff_t)d * (ptrdiff_t)dim, 1); } -void distsq_to_points( - const size_t dim, - float *dist_matrix, // Col Major, cols are queries, rows are points - size_t npoints, const float *const points, - const float *const points_l2sq, // points in Col major - size_t nqueries, const float *const queries, - const float *const queries_l2sq, // queries in Col major - float *ones_vec = NULL) // Scratchspace of num_data size and init to 1.0 +void distsq_to_points(const size_t dim, + float *dist_matrix, // Col Major, cols are queries, rows are points + size_t npoints, const float *const points, + const float *const points_l2sq, // points in Col major + size_t nqueries, const float *const queries, + const float *const queries_l2sq, // queries in Col major + float *ones_vec = NULL) // Scratchspace of num_data size and init to 1.0 { - bool ones_vec_alloc = false; - if (ones_vec == NULL) { - ones_vec = new float[nqueries > npoints ? nqueries : npoints]; - std::fill_n(ones_vec, nqueries > npoints ? nqueries : npoints, (float)1.0); - ones_vec_alloc = true; - } - cblas_sgemm(CblasColMajor, CblasTrans, CblasNoTrans, npoints, nqueries, dim, - (float)-2.0, points, dim, queries, dim, (float)0.0, dist_matrix, - npoints); - cblas_sgemm(CblasColMajor, CblasNoTrans, CblasTrans, npoints, nqueries, 1, - (float)1.0, points_l2sq, npoints, ones_vec, nqueries, (float)1.0, - dist_matrix, npoints); - cblas_sgemm(CblasColMajor, CblasNoTrans, CblasTrans, npoints, nqueries, 1, - (float)1.0, ones_vec, npoints, queries_l2sq, nqueries, (float)1.0, - dist_matrix, npoints); - if (ones_vec_alloc) delete[] ones_vec; + bool ones_vec_alloc = false; + if (ones_vec == NULL) + { + ones_vec = new float[nqueries > npoints ? nqueries : npoints]; + std::fill_n(ones_vec, nqueries > npoints ? nqueries : npoints, (float)1.0); + ones_vec_alloc = true; + } + cblas_sgemm(CblasColMajor, CblasTrans, CblasNoTrans, npoints, nqueries, dim, (float)-2.0, points, dim, queries, dim, + (float)0.0, dist_matrix, npoints); + cblas_sgemm(CblasColMajor, CblasNoTrans, CblasTrans, npoints, nqueries, 1, (float)1.0, points_l2sq, npoints, + ones_vec, nqueries, (float)1.0, dist_matrix, npoints); + cblas_sgemm(CblasColMajor, CblasNoTrans, CblasTrans, npoints, nqueries, 1, (float)1.0, ones_vec, npoints, + queries_l2sq, nqueries, (float)1.0, dist_matrix, npoints); + if (ones_vec_alloc) + delete[] ones_vec; } -void inner_prod_to_points( - const size_t dim, - float *dist_matrix, // Col Major, cols are queries, rows are points - size_t npoints, const float *const points, size_t nqueries, - const float *const queries, - float *ones_vec = NULL) // Scratchspace of num_data size and init to 1.0 +void inner_prod_to_points(const size_t dim, + float *dist_matrix, // Col Major, cols are queries, rows are points + size_t npoints, const float *const points, size_t nqueries, const float *const queries, + float *ones_vec = NULL) // Scratchspace of num_data size and init to 1.0 { - bool ones_vec_alloc = false; - if (ones_vec == NULL) { - ones_vec = new float[nqueries > npoints ? nqueries : npoints]; - std::fill_n(ones_vec, nqueries > npoints ? nqueries : npoints, (float)1.0); - ones_vec_alloc = true; - } - cblas_sgemm(CblasColMajor, CblasTrans, CblasNoTrans, npoints, nqueries, dim, - (float)-1.0, points, dim, queries, dim, (float)0.0, dist_matrix, - npoints); - - if (ones_vec_alloc) delete[] ones_vec; + bool ones_vec_alloc = false; + if (ones_vec == NULL) + { + ones_vec = new float[nqueries > npoints ? nqueries : npoints]; + std::fill_n(ones_vec, nqueries > npoints ? nqueries : npoints, (float)1.0); + ones_vec_alloc = true; + } + cblas_sgemm(CblasColMajor, CblasTrans, CblasNoTrans, npoints, nqueries, dim, (float)-1.0, points, dim, queries, dim, + (float)0.0, dist_matrix, npoints); + + if (ones_vec_alloc) + delete[] ones_vec; } -void exact_knn( - const size_t dim, const size_t k, - int *const closest_points, // k * num_queries preallocated, col - // major, queries columns - float *const dist_closest_points, // k * num_queries - // preallocated, Dist to - // corresponding closes_points - size_t npoints, - float *points_in, // points in Col major - size_t nqueries, float *queries_in, - diskann::Metric metric = diskann::Metric::L2) // queries in Col major +void exact_knn(const size_t dim, const size_t k, + int *const closest_points, // k * num_queries preallocated, col + // major, queries columns + float *const dist_closest_points, // k * num_queries + // preallocated, Dist to + // corresponding closes_points + size_t npoints, + float *points_in, // points in Col major + size_t nqueries, float *queries_in, + diskann::Metric metric = diskann::Metric::L2) // queries in Col major { - float *points_l2sq = new float[npoints]; - float *queries_l2sq = new float[nqueries]; - compute_l2sq(points_l2sq, points_in, npoints, dim); - compute_l2sq(queries_l2sq, queries_in, nqueries, dim); - - float *points = points_in; - float *queries = queries_in; - - if (metric == diskann::Metric::COSINE) { // we convert cosine distance as - // normalized L2 distnace - points = new float[npoints * dim]; - queries = new float[nqueries * dim]; + float *points_l2sq = new float[npoints]; + float *queries_l2sq = new float[nqueries]; + compute_l2sq(points_l2sq, points_in, npoints, dim); + compute_l2sq(queries_l2sq, queries_in, nqueries, dim); + + float *points = points_in; + float *queries = queries_in; + + if (metric == diskann::Metric::COSINE) + { // we convert cosine distance as + // normalized L2 distnace + points = new float[npoints * dim]; + queries = new float[nqueries * dim]; #pragma omp parallel for schedule(static, 4096) - for (int64_t i = 0; i < (int64_t)npoints; i++) { - float norm = std::sqrt(points_l2sq[i]); - if (norm == 0) { - norm = std::numeric_limits::epsilon(); - } - for (uint32_t j = 0; j < dim; j++) { - points[i * dim + j] = points_in[i * dim + j] / norm; - } - } + for (int64_t i = 0; i < (int64_t)npoints; i++) + { + float norm = std::sqrt(points_l2sq[i]); + if (norm == 0) + { + norm = std::numeric_limits::epsilon(); + } + for (uint32_t j = 0; j < dim; j++) + { + points[i * dim + j] = points_in[i * dim + j] / norm; + } + } #pragma omp parallel for schedule(static, 4096) - for (int64_t i = 0; i < (int64_t)nqueries; i++) { - float norm = std::sqrt(queries_l2sq[i]); - if (norm == 0) { - norm = std::numeric_limits::epsilon(); - } - for (uint32_t j = 0; j < dim; j++) { - queries[i * dim + j] = queries_in[i * dim + j] / norm; - } - } - // recalculate norms after normalizing, they should all be one. - compute_l2sq(points_l2sq, points, npoints, dim); - compute_l2sq(queries_l2sq, queries, nqueries, dim); - } - - std::cout << "Going to compute " << k << " NNs for " << nqueries - << " queries over " << npoints << " points in " << dim - << " dimensions using"; - if (metric == diskann::Metric::INNER_PRODUCT) - std::cout << " MIPS "; - else if (metric == diskann::Metric::COSINE) - std::cout << " Cosine "; - else - std::cout << " L2 "; - std::cout << "distance fn. " << std::endl; - - size_t q_batch_size = (1 << 9); - float *dist_matrix = new float[(size_t)q_batch_size * (size_t)npoints]; - - for (uint64_t b = 0; b < div_round_up(nqueries, q_batch_size); ++b) { - int64_t q_b = b * q_batch_size; - int64_t q_e = - ((b + 1) * q_batch_size > nqueries) ? nqueries : (b + 1) * q_batch_size; - - if (metric == diskann::Metric::L2 || metric == diskann::Metric::COSINE) { - 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); - } else { - inner_prod_to_points(dim, dist_matrix, npoints, points, q_e - q_b, - queries + (ptrdiff_t)q_b * (ptrdiff_t)dim); + for (int64_t i = 0; i < (int64_t)nqueries; i++) + { + float norm = std::sqrt(queries_l2sq[i]); + if (norm == 0) + { + norm = std::numeric_limits::epsilon(); + } + for (uint32_t j = 0; j < dim; j++) + { + queries[i * dim + j] = queries_in[i * dim + j] / norm; + } + } + // recalculate norms after normalizing, they should all be one. + compute_l2sq(points_l2sq, points, npoints, dim); + compute_l2sq(queries_l2sq, queries, nqueries, dim); } - std::cout << "Computed distances for queries: [" << q_b << "," << q_e << ")" - << std::endl; + + std::cout << "Going to compute " << k << " NNs for " << nqueries << " queries over " << npoints << " points in " + << dim << " dimensions using"; + if (metric == diskann::Metric::INNER_PRODUCT) + std::cout << " MIPS "; + else if (metric == diskann::Metric::COSINE) + std::cout << " Cosine "; + else + std::cout << " L2 "; + std::cout << "distance fn. " << std::endl; + + size_t q_batch_size = (1 << 9); + float *dist_matrix = new float[(size_t)q_batch_size * (size_t)npoints]; + + for (uint64_t b = 0; b < div_round_up(nqueries, q_batch_size); ++b) + { + int64_t q_b = b * q_batch_size; + int64_t q_e = ((b + 1) * q_batch_size > nqueries) ? nqueries : (b + 1) * q_batch_size; + + if (metric == diskann::Metric::L2 || metric == diskann::Metric::COSINE) + { + 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); + } + else + { + inner_prod_to_points(dim, dist_matrix, npoints, points, q_e - q_b, + queries + (ptrdiff_t)q_b * (ptrdiff_t)dim); + } + std::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++) { - maxPQIFCS point_dist; - for (uint64_t p = 0; p < k; p++) - point_dist.emplace( - p, dist_matrix[(ptrdiff_t)p + - (ptrdiff_t)(q - q_b) * (ptrdiff_t)npoints]); - for (size_t p = k; p < npoints; p++) { - if (point_dist.top().second > - dist_matrix[(ptrdiff_t)p + - (ptrdiff_t)(q - q_b) * (ptrdiff_t)npoints]) - point_dist.emplace( - p, dist_matrix[(ptrdiff_t)p + - (ptrdiff_t)(q - q_b) * (ptrdiff_t)npoints]); - if (point_dist.size() > k) point_dist.pop(); - } - for (ptrdiff_t l = 0; l < (ptrdiff_t)k; ++l) { - closest_points[(ptrdiff_t)(k - 1 - l) + (ptrdiff_t)q * (ptrdiff_t)k] = - point_dist.top().first; - dist_closest_points[(ptrdiff_t)(k - 1 - l) + - (ptrdiff_t)q * (ptrdiff_t)k] = - point_dist.top().second; - point_dist.pop(); - } - assert(std::is_sorted( - dist_closest_points + (ptrdiff_t)q * (ptrdiff_t)k, - dist_closest_points + (ptrdiff_t)(q + 1) * (ptrdiff_t)k)); + for (long long q = q_b; q < q_e; q++) + { + maxPQIFCS point_dist; + for (uint64_t p = 0; p < k; p++) + point_dist.emplace(p, dist_matrix[(ptrdiff_t)p + (ptrdiff_t)(q - q_b) * (ptrdiff_t)npoints]); + for (size_t p = k; p < npoints; p++) + { + if (point_dist.top().second > dist_matrix[(ptrdiff_t)p + (ptrdiff_t)(q - q_b) * (ptrdiff_t)npoints]) + point_dist.emplace(p, dist_matrix[(ptrdiff_t)p + (ptrdiff_t)(q - q_b) * (ptrdiff_t)npoints]); + if (point_dist.size() > k) + point_dist.pop(); + } + for (ptrdiff_t l = 0; l < (ptrdiff_t)k; ++l) + { + closest_points[(ptrdiff_t)(k - 1 - l) + (ptrdiff_t)q * (ptrdiff_t)k] = point_dist.top().first; + dist_closest_points[(ptrdiff_t)(k - 1 - l) + (ptrdiff_t)q * (ptrdiff_t)k] = point_dist.top().second; + point_dist.pop(); + } + assert(std::is_sorted(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; } - std::cout << "Computed exact k-NN for queries: [" << q_b << "," << q_e - << ")" << std::endl; - } - delete[] dist_matrix; + delete[] dist_matrix; - delete[] points_l2sq; - delete[] queries_l2sq; + delete[] points_l2sq; + delete[] queries_l2sq; - if (metric == diskann::Metric::COSINE) { - delete[] points; - delete[] queries; - } + if (metric == diskann::Metric::COSINE) + { + delete[] points; + delete[] queries; + } } -template -inline int get_num_parts(const char *filename) { - std::ifstream reader; - reader.exceptions(std::ios::failbit | std::ios::badbit); - reader.open(filename, std::ios::binary); - std::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; - 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; - return num_parts; +template inline int get_num_parts(const char *filename) +{ + std::ifstream reader; + reader.exceptions(std::ios::failbit | std::ios::badbit); + reader.open(filename, std::ios::binary); + std::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; + 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; + return num_parts; } template -inline void load_bin_as_float(const char *filename, float *&data, - size_t &npts_u64, size_t &ndims_u64, - int part_num) { - std::ifstream reader; - reader.exceptions(std::ios::failbit | std::ios::badbit); - reader.open(filename, std::ios::binary); - std::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)); - uint64_t start_id = part_num * PARTSIZE; - uint64_t end_id = (std::min)(start_id + PARTSIZE, (uint64_t)npts_i32); - npts_u64 = end_id - start_id; - ndims_u64 = (uint64_t)ndims_i32; - std::cout << "#pts in part = " << npts_u64 << ", #dims = " << ndims_u64 - << ", size = " << npts_u64 * ndims_u64 * sizeof(T) << "B" - << std::endl; - - reader.seekg(start_id * ndims_u64 * sizeof(T) + 2 * sizeof(uint32_t), - std::ios::beg); - T *data_T = new T[npts_u64 * ndims_u64]; - reader.read((char *)data_T, sizeof(T) * npts_u64 * ndims_u64); - std::cout << "Finished reading part of the bin file." << std::endl; - reader.close(); - data = aligned_malloc(npts_u64 * ndims_u64, ALIGNMENT); +inline void load_bin_as_float(const char *filename, float *&data, size_t &npts_u64, size_t &ndims_u64, int part_num) +{ + std::ifstream reader; + reader.exceptions(std::ios::failbit | std::ios::badbit); + reader.open(filename, std::ios::binary); + std::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)); + uint64_t start_id = part_num * PARTSIZE; + uint64_t end_id = (std::min)(start_id + PARTSIZE, (uint64_t)npts_i32); + npts_u64 = end_id - start_id; + ndims_u64 = (uint64_t)ndims_i32; + std::cout << "#pts in part = " << npts_u64 << ", #dims = " << ndims_u64 + << ", size = " << npts_u64 * ndims_u64 * sizeof(T) << "B" << std::endl; + + reader.seekg(start_id * ndims_u64 * sizeof(T) + 2 * sizeof(uint32_t), std::ios::beg); + T *data_T = new T[npts_u64 * ndims_u64]; + reader.read((char *)data_T, sizeof(T) * npts_u64 * ndims_u64); + std::cout << "Finished reading part of the bin file." << std::endl; + reader.close(); + data = aligned_malloc(npts_u64 * ndims_u64, ALIGNMENT); #pragma omp parallel for schedule(dynamic, 32768) - for (int64_t i = 0; i < (int64_t)npts_u64; i++) { - for (int64_t j = 0; j < (int64_t)ndims_u64; j++) { - float cur_val_float = (float)data_T[i * ndims_u64 + j]; - std::memcpy((char *)(data + i * ndims_u64 + j), (char *)&cur_val_float, - sizeof(float)); + for (int64_t i = 0; i < (int64_t)npts_u64; i++) + { + for (int64_t j = 0; j < (int64_t)ndims_u64; j++) + { + float cur_val_float = (float)data_T[i * ndims_u64 + j]; + std::memcpy((char *)(data + i * ndims_u64 + j), (char *)&cur_val_float, sizeof(float)); + } } - } - delete[] data_T; - std::cout << "Finished converting part data to float." << std::endl; + delete[] data_T; + std::cout << "Finished converting part data to float." << std::endl; } template -inline std::vector load_filtered_bin_as_float( - const char *filename, float *&data, size_t &npts, size_t &ndims, - int part_num, const char *label_file, const std::string &filter_label, - const std::string &universal_label, size_t &npoints_filt, - std::vector> &pts_to_labels) { - std::ifstream reader(filename, std::ios::binary); - if (reader.fail()) { - throw diskann::ANNException(std::string("Failed to open file ") + filename, - -1); - } - - std::cout << "Reading bin file " << filename << " ...\n"; - int npts_i32, ndims_i32; - std::vector rev_map; - reader.read((char *)&npts_i32, sizeof(int)); - reader.read((char *)&ndims_i32, sizeof(int)); - uint64_t start_id = part_num * PARTSIZE; - uint64_t end_id = (std::min)(start_id + PARTSIZE, (uint64_t)npts_i32); - npts = end_id - start_id; - ndims = (uint32_t)ndims_i32; - uint64_t nptsuint64_t = (uint64_t)npts; - uint64_t ndimsuint64_t = (uint64_t)ndims; - npoints_filt = 0; - std::cout << "#pts in part = " << npts << ", #dims = " << ndims - << ", size = " << nptsuint64_t * ndimsuint64_t * sizeof(T) << "B" - << std::endl; - std::cout << "start and end ids: " << start_id << ", " << end_id << std::endl; - reader.seekg(start_id * ndims * sizeof(T) + 2 * sizeof(uint32_t), - std::ios::beg); - - 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; - reader.close(); - - data = aligned_malloc(nptsuint64_t * ndimsuint64_t, ALIGNMENT); - - for (int64_t i = 0; i < (int64_t)nptsuint64_t; i++) { - if (std::find(pts_to_labels[start_id + i].begin(), - pts_to_labels[start_id + i].end(), - filter_label) != pts_to_labels[start_id + i].end() || - std::find(pts_to_labels[start_id + i].begin(), - pts_to_labels[start_id + i].end(), - universal_label) != pts_to_labels[start_id + i].end()) { - rev_map.push_back(start_id + i); - for (int64_t j = 0; j < (int64_t)ndimsuint64_t; j++) { - float cur_val_float = (float)data_T[i * ndimsuint64_t + j]; - std::memcpy((char *)(data + npoints_filt * ndimsuint64_t + j), - (char *)&cur_val_float, sizeof(float)); - } - npoints_filt++; +inline std::vector load_filtered_bin_as_float(const char *filename, float *&data, size_t &npts, size_t &ndims, + int part_num, const char *label_file, + const std::string &filter_label, + const std::string &universal_label, size_t &npoints_filt, + std::vector> &pts_to_labels) +{ + std::ifstream reader(filename, std::ios::binary); + if (reader.fail()) + { + throw diskann::ANNException(std::string("Failed to open file ") + filename, -1); } - } - delete[] data_T; - std::cout << "Finished converting part data to float.. identified " - << npoints_filt << " points matching the filter." << std::endl; - return rev_map; + + std::cout << "Reading bin file " << filename << " ...\n"; + int npts_i32, ndims_i32; + std::vector rev_map; + reader.read((char *)&npts_i32, sizeof(int)); + reader.read((char *)&ndims_i32, sizeof(int)); + uint64_t start_id = part_num * PARTSIZE; + uint64_t end_id = (std::min)(start_id + PARTSIZE, (uint64_t)npts_i32); + npts = end_id - start_id; + ndims = (uint32_t)ndims_i32; + uint64_t nptsuint64_t = (uint64_t)npts; + uint64_t ndimsuint64_t = (uint64_t)ndims; + npoints_filt = 0; + std::cout << "#pts in part = " << npts << ", #dims = " << ndims + << ", size = " << nptsuint64_t * ndimsuint64_t * sizeof(T) << "B" << std::endl; + std::cout << "start and end ids: " << start_id << ", " << end_id << std::endl; + reader.seekg(start_id * ndims * sizeof(T) + 2 * sizeof(uint32_t), std::ios::beg); + + 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; + reader.close(); + + data = aligned_malloc(nptsuint64_t * ndimsuint64_t, ALIGNMENT); + + for (int64_t i = 0; i < (int64_t)nptsuint64_t; i++) + { + if (std::find(pts_to_labels[start_id + i].begin(), pts_to_labels[start_id + i].end(), filter_label) != + pts_to_labels[start_id + i].end() || + std::find(pts_to_labels[start_id + i].begin(), pts_to_labels[start_id + i].end(), universal_label) != + pts_to_labels[start_id + i].end()) + { + rev_map.push_back(start_id + i); + for (int64_t j = 0; j < (int64_t)ndimsuint64_t; j++) + { + float cur_val_float = (float)data_T[i * ndimsuint64_t + j]; + std::memcpy((char *)(data + npoints_filt * ndimsuint64_t + j), (char *)&cur_val_float, sizeof(float)); + } + npoints_filt++; + } + } + delete[] data_T; + std::cout << "Finished converting part data to float.. identified " << npoints_filt + << " points matching the filter." << std::endl; + return rev_map; } -template -inline void save_bin(const std::string filename, T *data, size_t npts, - size_t ndims) { - std::ofstream writer; - writer.exceptions(std::ios::failbit | std::ios::badbit); - writer.open(filename, std::ios::binary | std::ios::out); - std::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; - - writer.write((char *)data, npts * ndims * sizeof(T)); - writer.close(); - std::cout << "Finished writing bin" << std::endl; +template inline void save_bin(const std::string filename, T *data, size_t npts, size_t ndims) +{ + std::ofstream writer; + writer.exceptions(std::ios::failbit | std::ios::badbit); + writer.open(filename, std::ios::binary | std::ios::out); + std::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; + + writer.write((char *)data, npts * ndims * sizeof(T)); + writer.close(); + std::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) { - 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(uint32_t) + 2 * sizeof(int) << "B" - << std::endl; - - writer.write((char *)data, npts * ndims * sizeof(uint32_t)); - writer.write((char *)distances, npts * ndims * sizeof(float)); - writer.close(); - std::cout << "Finished writing truthset" << std::endl; +inline void save_groundtruth_as_one_file(const std::string filename, int32_t *data, float *distances, size_t npts, + size_t ndims) +{ + 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(uint32_t) + 2 * sizeof(int) + << "B" << std::endl; + + writer.write((char *)data, npts * ndims * sizeof(uint32_t)); + writer.write((char *)distances, npts * ndims * sizeof(float)); + writer.close(); + std::cout << "Finished writing truthset" << std::endl; } -inline void parse_label_file_into_vec( - size_t &line_cnt, const std::string &map_file, - std::vector> &pts_to_labels) { - std::ifstream infile(map_file); - std::string line, token; - std::set labels; - infile.clear(); - infile.seekg(0, std::ios::beg); - while (std::getline(infile, line)) { - std::istringstream iss(line); - std::vector lbls(0); - - getline(iss, token, '\t'); - std::istringstream new_iss(token); - while (getline(new_iss, token, ',')) { - token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); - token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); - lbls.push_back(token); - labels.insert(token); - } - if (lbls.size() <= 0) { - std::cout << "No label found"; - exit(-1); +inline void parse_label_file_into_vec(size_t &line_cnt, const std::string &map_file, + std::vector> &pts_to_labels) +{ + std::ifstream infile(map_file); + std::string line, token; + std::set labels; + infile.clear(); + infile.seekg(0, std::ios::beg); + while (std::getline(infile, line)) + { + std::istringstream iss(line); + std::vector lbls(0); + + getline(iss, token, '\t'); + std::istringstream new_iss(token); + while (getline(new_iss, token, ',')) + { + token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); + token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); + lbls.push_back(token); + labels.insert(token); + } + if (lbls.size() <= 0) + { + std::cout << "No label found"; + exit(-1); + } + std::sort(lbls.begin(), lbls.end()); + pts_to_labels.push_back(lbls); } - std::sort(lbls.begin(), lbls.end()); - pts_to_labels.push_back(lbls); - } - std::cout << "Identified " << labels.size() - << " distinct label(s), and populated labels for " - << pts_to_labels.size() << " points" << std::endl; + std::cout << "Identified " << labels.size() << " distinct label(s), and populated labels for " + << pts_to_labels.size() << " points" << std::endl; } template -std::vector>> processUnfilteredParts( - const std::string &base_file, size_t &nqueries, size_t &npoints, - size_t &dim, size_t &k, float *query_data, const diskann::Metric &metric, - std::vector &location_to_tag) { - float *base_data; - int num_parts = get_num_parts(base_file.c_str()); - std::vector>> res(nqueries); - for (int p = 0; p < num_parts; p++) { - size_t start_id = p * PARTSIZE; - load_bin_as_float(base_file.c_str(), base_data, npoints, dim, p); - - int *closest_points_part = new int[nqueries * k]; - float *dist_closest_points_part = new float[nqueries * k]; - - uint32_t part_k; - part_k = k < npoints ? k : npoints; - exact_knn(dim, part_k, closest_points_part, dist_closest_points_part, - npoints, base_data, nqueries, query_data, metric); - - for (size_t i = 0; i < nqueries; i++) { - for (uint64_t j = 0; j < part_k; j++) { - if (!location_to_tag.empty()) - if (location_to_tag[closest_points_part[i * k + j] + start_id] == 0) - continue; - - res[i].push_back(std::make_pair( - (uint32_t)(closest_points_part[i * part_k + j] + start_id), - dist_closest_points_part[i * part_k + j])); - } - } +std::vector>> processUnfilteredParts(const std::string &base_file, + size_t &nqueries, size_t &npoints, + size_t &dim, size_t &k, float *query_data, + const diskann::Metric &metric, + std::vector &location_to_tag) +{ + float *base_data; + int num_parts = get_num_parts(base_file.c_str()); + std::vector>> res(nqueries); + for (int p = 0; p < num_parts; p++) + { + size_t start_id = p * PARTSIZE; + load_bin_as_float(base_file.c_str(), base_data, npoints, dim, p); + + int *closest_points_part = new int[nqueries * k]; + float *dist_closest_points_part = new float[nqueries * k]; + + uint32_t part_k; + part_k = k < npoints ? k : npoints; + exact_knn(dim, part_k, closest_points_part, dist_closest_points_part, npoints, base_data, nqueries, query_data, + metric); + + for (size_t i = 0; i < nqueries; i++) + { + for (uint64_t j = 0; j < part_k; j++) + { + if (!location_to_tag.empty()) + if (location_to_tag[closest_points_part[i * k + j] + start_id] == 0) + continue; + + res[i].push_back(std::make_pair((uint32_t)(closest_points_part[i * part_k + j] + start_id), + dist_closest_points_part[i * part_k + j])); + } + } - delete[] closest_points_part; - delete[] dist_closest_points_part; + delete[] closest_points_part; + delete[] dist_closest_points_part; - diskann::aligned_free(base_data); - } - return res; + diskann::aligned_free(base_data); + } + return res; }; template std::vector>> processFilteredParts( - const std::string &base_file, const std::string &label_file, - const std::string &filter_label, const std::string &universal_label, - size_t &nqueries, size_t &npoints, size_t &dim, size_t &k, - float *query_data, const diskann::Metric &metric, - std::vector &location_to_tag) { - size_t npoints_filt; - float *base_data; - std::vector>> res(nqueries); - int num_parts = get_num_parts(base_file.c_str()); - - std::vector> pts_to_labels; - if (filter_label != "") - parse_label_file_into_vec(npoints, label_file, pts_to_labels); - - for (int p = 0; p < num_parts; p++) { - size_t start_id = p * PARTSIZE; - std::vector rev_map; - if (filter_label != "") - rev_map = load_filtered_bin_as_float( - base_file.c_str(), base_data, npoints, dim, p, label_file.c_str(), - filter_label, universal_label, npoints_filt, pts_to_labels); - int *closest_points_part = new int[nqueries * k]; - float *dist_closest_points_part = new float[nqueries * k]; - - uint32_t part_k; - part_k = k < npoints_filt ? k : npoints_filt; - if (npoints_filt > 0) { - exact_knn(dim, part_k, closest_points_part, dist_closest_points_part, - npoints_filt, base_data, nqueries, query_data, metric); - } + const std::string &base_file, const std::string &label_file, const std::string &filter_label, + const std::string &universal_label, size_t &nqueries, size_t &npoints, size_t &dim, size_t &k, float *query_data, + const diskann::Metric &metric, std::vector &location_to_tag) +{ + size_t npoints_filt; + float *base_data; + std::vector>> res(nqueries); + int num_parts = get_num_parts(base_file.c_str()); - for (size_t i = 0; i < nqueries; i++) { - for (uint64_t j = 0; j < part_k; j++) { - if (!location_to_tag.empty()) - if (location_to_tag[closest_points_part[i * k + j] + start_id] == 0) - continue; + std::vector> pts_to_labels; + if (filter_label != "") + parse_label_file_into_vec(npoints, label_file, pts_to_labels); + + for (int p = 0; p < num_parts; p++) + { + size_t start_id = p * PARTSIZE; + std::vector rev_map; + if (filter_label != "") + rev_map = load_filtered_bin_as_float(base_file.c_str(), base_data, npoints, dim, p, label_file.c_str(), + filter_label, universal_label, npoints_filt, pts_to_labels); + int *closest_points_part = new int[nqueries * k]; + float *dist_closest_points_part = new float[nqueries * k]; + + uint32_t part_k; + part_k = k < npoints_filt ? k : npoints_filt; + if (npoints_filt > 0) + { + exact_knn(dim, part_k, closest_points_part, dist_closest_points_part, npoints_filt, base_data, nqueries, + query_data, metric); + } - res[i].push_back(std::make_pair( - (uint32_t)(rev_map[closest_points_part[i * part_k + j]]), - dist_closest_points_part[i * part_k + j])); - } - } + for (size_t i = 0; i < nqueries; i++) + { + for (uint64_t j = 0; j < part_k; j++) + { + if (!location_to_tag.empty()) + if (location_to_tag[closest_points_part[i * k + j] + start_id] == 0) + continue; + + res[i].push_back(std::make_pair((uint32_t)(rev_map[closest_points_part[i * part_k + j]]), + dist_closest_points_part[i * part_k + j])); + } + } - delete[] closest_points_part; - delete[] dist_closest_points_part; + delete[] closest_points_part; + delete[] dist_closest_points_part; - diskann::aligned_free(base_data); - } - return res; + diskann::aligned_free(base_data); + } + return res; }; template -int aux_main(const std::string &base_file, const std::string &label_file, - const std::string &query_file, const std::string >_file, - size_t k, const std::string &universal_label, - const diskann::Metric &metric, const std::string &filter_label, - const std::string &tags_file = std::string("")) { - size_t npoints, nqueries, dim, npoints_filt; - - float *base_data; - float *query_data; - - load_bin_as_float(query_file.c_str(), query_data, nqueries, dim, 0); - if (nqueries > PARTSIZE) - std::cerr << "WARNING: #Queries provided (" << nqueries - << ") is greater than " << PARTSIZE - << ". Computing GT only for the first " << PARTSIZE << " queries." - << std::endl; - - // load tags - const bool tags_enabled = tags_file.empty() ? false : true; - std::vector location_to_tag = - diskann::loadTags(tags_file, base_file); - - int *closest_points = new int[nqueries * k]; - float *dist_closest_points = new float[nqueries * k]; - - std::vector>> results; - if (filter_label == "") { - results = processUnfilteredParts(base_file, nqueries, npoints, dim, k, - query_data, metric, location_to_tag); - } else { - results = processFilteredParts(base_file, label_file, filter_label, - universal_label, nqueries, npoints, dim, - k, query_data, metric, location_to_tag); - } - - for (size_t i = 0; i < nqueries; i++) { - std::vector> &cur_res = results[i]; - std::sort(cur_res.begin(), cur_res.end(), custom_dist); - size_t j = 0; - for (auto iter : cur_res) { - if (j == k) break; - if (tags_enabled) { - std::uint32_t index_with_tag = location_to_tag[iter.first]; - closest_points[i * k + j] = (int32_t)index_with_tag; - } else { - closest_points[i * k + j] = (int32_t)iter.first; - } - - if (metric == diskann::Metric::INNER_PRODUCT) - dist_closest_points[i * k + j] = -iter.second; - else - dist_closest_points[i * k + j] = iter.second; - - ++j; +int aux_main(const std::string &base_file, const std::string &label_file, const std::string &query_file, + const std::string >_file, size_t k, const std::string &universal_label, const diskann::Metric &metric, + const std::string &filter_label, const std::string &tags_file = std::string("")) +{ + size_t npoints, nqueries, dim, npoints_filt; + + float *base_data; + float *query_data; + + load_bin_as_float(query_file.c_str(), query_data, nqueries, dim, 0); + if (nqueries > PARTSIZE) + std::cerr << "WARNING: #Queries provided (" << nqueries << ") is greater than " << PARTSIZE + << ". Computing GT only for the first " << PARTSIZE << " queries." << std::endl; + + // load tags + const bool tags_enabled = tags_file.empty() ? false : true; + std::vector location_to_tag = diskann::loadTags(tags_file, base_file); + + int *closest_points = new int[nqueries * k]; + float *dist_closest_points = new float[nqueries * k]; + + std::vector>> results; + if (filter_label == "") + { + results = processUnfilteredParts(base_file, nqueries, npoints, dim, k, query_data, metric, location_to_tag); + } + else + { + results = processFilteredParts(base_file, label_file, filter_label, universal_label, nqueries, npoints, dim, + k, query_data, metric, location_to_tag); + } + + for (size_t i = 0; i < nqueries; i++) + { + std::vector> &cur_res = results[i]; + std::sort(cur_res.begin(), cur_res.end(), custom_dist); + size_t j = 0; + for (auto iter : cur_res) + { + if (j == k) + break; + if (tags_enabled) + { + std::uint32_t index_with_tag = location_to_tag[iter.first]; + closest_points[i * k + j] = (int32_t)index_with_tag; + } + else + { + closest_points[i * k + j] = (int32_t)iter.first; + } + + if (metric == diskann::Metric::INNER_PRODUCT) + dist_closest_points[i * k + j] = -iter.second; + else + dist_closest_points[i * k + j] = iter.second; + + ++j; + } + if (j < k) + std::cout << "WARNING: found less than k GT entries for query " << i << std::endl; } - if (j < k) - std::cout << "WARNING: found less than k GT entries for query " << i - << std::endl; - } - - save_groundtruth_as_one_file(gt_file, closest_points, dist_closest_points, - nqueries, k); - delete[] closest_points; - delete[] dist_closest_points; - diskann::aligned_free(query_data); - - return 0; + + save_groundtruth_as_one_file(gt_file, closest_points, dist_closest_points, nqueries, k); + delete[] closest_points; + delete[] dist_closest_points; + diskann::aligned_free(query_data); + + return 0; } -void load_truthset(const std::string &bin_file, uint32_t *&ids, float *&dists, - size_t &npts, size_t &dim) { - size_t read_blk_size = 64 * 1024 * 1024; - cached_ifstream reader(bin_file, read_blk_size); - diskann::cout << "Reading truthset file " << bin_file.c_str() << " ..." - << std::endl; - size_t actual_file_size = reader.get_file_size(); - - int npts_i32, dim_i32; - reader.read((char *)&npts_i32, sizeof(int)); - reader.read((char *)&dim_i32, sizeof(int)); - npts = (uint32_t)npts_i32; - dim = (uint32_t)dim_i32; - - diskann::cout << "Metadata: #pts = " << npts << ", #dims = " << dim << "... " - << std::endl; - - int truthset_type = -1; // 1 means truthset has ids and distances, 2 means - // only ids, -1 is error - size_t expected_file_size_with_dists = - 2 * npts * dim * sizeof(uint32_t) + 2 * sizeof(uint32_t); - - if (actual_file_size == expected_file_size_with_dists) truthset_type = 1; - - size_t expected_file_size_just_ids = - npts * dim * sizeof(uint32_t) + 2 * sizeof(uint32_t); - - if (actual_file_size == expected_file_size_just_ids) truthset_type = 2; - - if (truthset_type == -1) { - std::stringstream stream; - stream << "Error. File size mismatch. File should have bin format, with " - "npts followed by ngt followed by npts*ngt ids and optionally " - "followed by npts*ngt distance values; actual size: " - << actual_file_size - << ", expected: " << expected_file_size_with_dists << " or " - << expected_file_size_just_ids; - diskann::cout << stream.str(); - throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, - __LINE__); - } - - ids = new uint32_t[npts * dim]; - reader.read((char *)ids, npts * dim * sizeof(uint32_t)); - - if (truthset_type == 1) { - dists = new float[npts * dim]; - reader.read((char *)dists, npts * dim * sizeof(float)); - } +void load_truthset(const std::string &bin_file, uint32_t *&ids, float *&dists, size_t &npts, size_t &dim) +{ + size_t read_blk_size = 64 * 1024 * 1024; + cached_ifstream reader(bin_file, read_blk_size); + diskann::cout << "Reading truthset file " << bin_file.c_str() << " ..." << std::endl; + size_t actual_file_size = reader.get_file_size(); + + int npts_i32, dim_i32; + reader.read((char *)&npts_i32, sizeof(int)); + reader.read((char *)&dim_i32, sizeof(int)); + npts = (uint32_t)npts_i32; + dim = (uint32_t)dim_i32; + + diskann::cout << "Metadata: #pts = " << npts << ", #dims = " << dim << "... " << std::endl; + + int truthset_type = -1; // 1 means truthset has ids and distances, 2 means + // only ids, -1 is error + size_t expected_file_size_with_dists = 2 * npts * dim * sizeof(uint32_t) + 2 * sizeof(uint32_t); + + if (actual_file_size == expected_file_size_with_dists) + truthset_type = 1; + + size_t expected_file_size_just_ids = npts * dim * sizeof(uint32_t) + 2 * sizeof(uint32_t); + + if (actual_file_size == expected_file_size_just_ids) + truthset_type = 2; + + if (truthset_type == -1) + { + std::stringstream stream; + stream << "Error. File size mismatch. File should have bin format, with " + "npts followed by ngt followed by npts*ngt ids and optionally " + "followed by npts*ngt distance values; actual size: " + << actual_file_size << ", expected: " << expected_file_size_with_dists << " or " + << expected_file_size_just_ids; + diskann::cout << stream.str(); + throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); + } + + ids = new uint32_t[npts * dim]; + reader.read((char *)ids, npts * dim * sizeof(uint32_t)); + + if (truthset_type == 1) + { + dists = new float[npts * dim]; + reader.read((char *)dists, npts * dim * sizeof(float)); + } } -int main(int argc, char **argv) { - std::string data_type, dist_fn, base_file, query_file, gt_file, tags_file, - label_file, filter_label, universal_label, filter_label_file; - uint64_t K; - - try { - po::options_description desc{"Arguments"}; - - desc.add_options()("help,h", "Print information on arguments"); - - desc.add_options()("data_type", - po::value(&data_type)->required(), - "data type "); - desc.add_options()("dist_fn", po::value(&dist_fn)->required(), - "distance function "); - desc.add_options()("base_file", - po::value(&base_file)->required(), - "File containing the base vectors in binary format"); - desc.add_options()("query_file", - po::value(&query_file)->required(), - "File containing the query vectors in binary format"); - desc.add_options()("label_file", - po::value(&label_file)->default_value(""), - "Input labels file in txt format if present"); - desc.add_options()("filter_label", - po::value(&filter_label)->default_value(""), - "Input filter label if doing filtered groundtruth"); - desc.add_options()( - "universal_label", - po::value(&universal_label)->default_value(""), - "Universal label, if using it, only in conjunction with label_file"); - desc.add_options()("gt_file", po::value(>_file)->required(), - "File name for the writing ground truth in binary " - "format, please don' append .bin at end if " - "no filter_label or filter_label_file is provided it " - "will save the file with '.bin' at end." - "else it will save the file as filename_label.bin"); - desc.add_options()("K", po::value(&K)->required(), - "Number of ground truth nearest neighbors to compute"); - desc.add_options()( - "tags_file", - po::value(&tags_file)->default_value(std::string()), - "File containing the tags in binary format"); - desc.add_options()("filter_label_file", - po::value(&filter_label_file) - ->default_value(std::string("")), - "Filter file for Queries for Filtered Search "); - - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) { - std::cout << desc; - return 0; +int main(int argc, char **argv) +{ + std::string data_type, dist_fn, base_file, query_file, gt_file, tags_file, label_file, filter_label, + universal_label, filter_label_file; + uint64_t K; + + try + { + po::options_description desc{"Arguments"}; + + desc.add_options()("help,h", "Print information on arguments"); + + desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); + desc.add_options()("dist_fn", po::value(&dist_fn)->required(), "distance function "); + desc.add_options()("base_file", po::value(&base_file)->required(), + "File containing the base vectors in binary format"); + desc.add_options()("query_file", po::value(&query_file)->required(), + "File containing the query vectors in binary format"); + desc.add_options()("label_file", po::value(&label_file)->default_value(""), + "Input labels file in txt format if present"); + desc.add_options()("filter_label", po::value(&filter_label)->default_value(""), + "Input filter label if doing filtered groundtruth"); + desc.add_options()("universal_label", po::value(&universal_label)->default_value(""), + "Universal label, if using it, only in conjunction with label_file"); + desc.add_options()("gt_file", po::value(>_file)->required(), + "File name for the writing ground truth in binary " + "format, please don' append .bin at end if " + "no filter_label or filter_label_file is provided it " + "will save the file with '.bin' at end." + "else it will save the file as filename_label.bin"); + desc.add_options()("K", po::value(&K)->required(), + "Number of ground truth nearest neighbors to compute"); + desc.add_options()("tags_file", po::value(&tags_file)->default_value(std::string()), + "File containing the tags in binary format"); + desc.add_options()("filter_label_file", + po::value(&filter_label_file)->default_value(std::string("")), + "Filter file for Queries for Filtered Search "); + + po::variables_map vm; + po::store(po::parse_command_line(argc, argv, desc), vm); + if (vm.count("help")) + { + std::cout << desc; + return 0; + } + po::notify(vm); } - po::notify(vm); - } catch (const std::exception &ex) { - std::cerr << ex.what() << '\n'; - return -1; - } - - if (data_type != std::string("float") && data_type != std::string("int8") && - data_type != std::string("uint8")) { - std::cout << "Unsupported type. float, int8 and uint8 types are supported." - << std::endl; - return -1; - } - - if (filter_label != "" && filter_label_file != "") { - std::cerr - << "Only one of filter_label and query_filters_file should be provided" - << std::endl; - return -1; - } - - diskann::Metric metric; - if (dist_fn == std::string("l2")) { - metric = diskann::Metric::L2; - } else if (dist_fn == std::string("mips")) { - metric = diskann::Metric::INNER_PRODUCT; - } else if (dist_fn == std::string("cosine")) { - metric = diskann::Metric::COSINE; - } else { - std::cerr << "Unsupported distance function. Use l2/mips/cosine." - << std::endl; - return -1; - } - - std::vector filter_labels; - if (filter_label != "") { - filter_labels.push_back(filter_label); - } else if (filter_label_file != "") { - filter_labels = read_file_to_vector_of_strings(filter_label_file, false); - } - - // only when there is no filter label or 1 filter label for all queries - if (filter_labels.size() == 1) { - try { - if (data_type == std::string("float")) - aux_main(base_file, label_file, query_file, gt_file, K, - universal_label, metric, filter_labels[0], tags_file); - if (data_type == std::string("int8")) - aux_main(base_file, label_file, query_file, gt_file, K, - universal_label, metric, filter_labels[0], tags_file); - if (data_type == std::string("uint8")) - aux_main(base_file, label_file, query_file, gt_file, K, - universal_label, metric, filter_labels[0], tags_file); - } catch (const std::exception &e) { - std::cout << std::string(e.what()) << std::endl; - diskann::cerr << "Compute GT failed." << std::endl; - return -1; + catch (const std::exception &ex) + { + std::cerr << ex.what() << '\n'; + return -1; } - } else { // Each query has its own filter label - // Split up data and query bins into label specific ones - tsl::robin_map labels_to_number_of_points; - tsl::robin_map labels_to_number_of_queries; - - label_set all_labels; - for (size_t i = 0; i < filter_labels.size(); i++) { - std::string label = filter_labels[i]; - all_labels.insert(label); - - if (labels_to_number_of_queries.find(label) == - labels_to_number_of_queries.end()) { - labels_to_number_of_queries[label] = 0; - } - labels_to_number_of_queries[label] += 1; + + if (data_type != std::string("float") && data_type != std::string("int8") && data_type != std::string("uint8")) + { + std::cout << "Unsupported type. float, int8 and uint8 types are supported." << std::endl; + return -1; } - size_t npoints; - std::vector> point_to_labels; - parse_label_file_into_vec(npoints, label_file, point_to_labels); - std::vector point_ids_to_labels(point_to_labels.size()); - std::vector query_ids_to_labels(filter_labels.size()); - - for (size_t i = 0; i < point_to_labels.size(); i++) { - for (size_t j = 0; j < point_to_labels[i].size(); j++) { - std::string label = point_to_labels[i][j]; - if (all_labels.find(label) != all_labels.end()) { - point_ids_to_labels[i].insert(point_to_labels[i][j]); - if (labels_to_number_of_points.find(label) == - labels_to_number_of_points.end()) { - labels_to_number_of_points[label] = 0; - } - labels_to_number_of_points[label] += 1; - } - } + if (filter_label != "" && filter_label_file != "") + { + std::cerr << "Only one of filter_label and query_filters_file should be provided" << std::endl; + return -1; } - for (size_t i = 0; i < filter_labels.size(); i++) { - query_ids_to_labels[i].insert(filter_labels[i]); + diskann::Metric metric; + if (dist_fn == std::string("l2")) + { + metric = diskann::Metric::L2; + } + else if (dist_fn == std::string("mips")) + { + metric = diskann::Metric::INNER_PRODUCT; + } + else if (dist_fn == std::string("cosine")) + { + metric = diskann::Metric::COSINE; + } + else + { + std::cerr << "Unsupported distance function. Use l2/mips/cosine." << std::endl; + return -1; + } + + std::vector filter_labels; + if (filter_label != "") + { + filter_labels.push_back(filter_label); + } + else if (filter_label_file != "") + { + filter_labels = read_file_to_vector_of_strings(filter_label_file, false); } - tsl::robin_map> label_id_to_orig_id; - tsl::robin_map> - label_query_id_to_orig_id; - - if (data_type == std::string("float")) { - label_id_to_orig_id = - diskann::generate_label_specific_vector_files_compat( - base_file, labels_to_number_of_points, point_ids_to_labels, - all_labels); - - label_query_id_to_orig_id = - diskann::generate_label_specific_vector_files_compat( - query_file, labels_to_number_of_queries, query_ids_to_labels, - all_labels); // query_filters acts like query_ids_to_labels - } else if (data_type == std::string("int8")) { - label_id_to_orig_id = - diskann::generate_label_specific_vector_files_compat( - base_file, labels_to_number_of_points, point_ids_to_labels, - all_labels); - - label_query_id_to_orig_id = - diskann::generate_label_specific_vector_files_compat( - query_file, labels_to_number_of_queries, query_ids_to_labels, - all_labels); // query_filters acts like query_ids_to_labels - } else if (data_type == std::string("uint8")) { - label_id_to_orig_id = - diskann::generate_label_specific_vector_files_compat( - base_file, labels_to_number_of_points, point_ids_to_labels, - all_labels); - - label_query_id_to_orig_id = - diskann::generate_label_specific_vector_files_compat( - query_file, labels_to_number_of_queries, query_ids_to_labels, - all_labels); // query_filters acts like query_ids_to_labels - } else { - diskann::cerr << "Invalid data type" << std::endl; - return -1; + // only when there is no filter label or 1 filter label for all queries + if (filter_labels.size() == 1) + { + try + { + if (data_type == std::string("float")) + aux_main(base_file, label_file, query_file, gt_file, K, universal_label, metric, + filter_labels[0], tags_file); + if (data_type == std::string("int8")) + aux_main(base_file, label_file, query_file, gt_file, K, universal_label, metric, + filter_labels[0], tags_file); + if (data_type == std::string("uint8")) + aux_main(base_file, label_file, query_file, gt_file, K, universal_label, metric, + filter_labels[0], tags_file); + } + catch (const std::exception &e) + { + std::cout << std::string(e.what()) << std::endl; + diskann::cerr << "Compute GT failed." << std::endl; + return -1; + } } + else + { // Each query has its own filter label + // Split up data and query bins into label specific ones + tsl::robin_map labels_to_number_of_points; + tsl::robin_map labels_to_number_of_queries; + + label_set all_labels; + for (size_t i = 0; i < filter_labels.size(); i++) + { + std::string label = filter_labels[i]; + all_labels.insert(label); + + if (labels_to_number_of_queries.find(label) == labels_to_number_of_queries.end()) + { + labels_to_number_of_queries[label] = 0; + } + labels_to_number_of_queries[label] += 1; + } + + size_t npoints; + std::vector> point_to_labels; + parse_label_file_into_vec(npoints, label_file, point_to_labels); + std::vector point_ids_to_labels(point_to_labels.size()); + std::vector query_ids_to_labels(filter_labels.size()); + + for (size_t i = 0; i < point_to_labels.size(); i++) + { + for (size_t j = 0; j < point_to_labels[i].size(); j++) + { + std::string label = point_to_labels[i][j]; + if (all_labels.find(label) != all_labels.end()) + { + point_ids_to_labels[i].insert(point_to_labels[i][j]); + if (labels_to_number_of_points.find(label) == labels_to_number_of_points.end()) + { + labels_to_number_of_points[label] = 0; + } + labels_to_number_of_points[label] += 1; + } + } + } - // Generate label specific ground truths + for (size_t i = 0; i < filter_labels.size(); i++) + { + query_ids_to_labels[i].insert(filter_labels[i]); + } + + tsl::robin_map> label_id_to_orig_id; + tsl::robin_map> label_query_id_to_orig_id; - try { - for (const auto &label : all_labels) { - std::string filtered_base_file = base_file + "_" + label; - std::string filtered_query_file = query_file + "_" + label; - std::string filtered_gt_file = gt_file + "_" + label; if (data_type == std::string("float")) - aux_main(filtered_base_file, "", filtered_query_file, - filtered_gt_file, K, "", metric, ""); - if (data_type == std::string("int8")) - aux_main(filtered_base_file, "", filtered_query_file, - filtered_gt_file, K, "", metric, ""); - if (data_type == std::string("uint8")) - aux_main(filtered_base_file, "", filtered_query_file, - filtered_gt_file, K, "", metric, ""); - } - } catch (const std::exception &e) { - std::cout << std::string(e.what()) << std::endl; - diskann::cerr << "Compute GT failed." << std::endl; - return -1; - } + { + label_id_to_orig_id = diskann::generate_label_specific_vector_files_compat( + base_file, labels_to_number_of_points, point_ids_to_labels, all_labels); + + label_query_id_to_orig_id = diskann::generate_label_specific_vector_files_compat( + query_file, labels_to_number_of_queries, query_ids_to_labels, + all_labels); // query_filters acts like query_ids_to_labels + } + else if (data_type == std::string("int8")) + { + label_id_to_orig_id = diskann::generate_label_specific_vector_files_compat( + base_file, labels_to_number_of_points, point_ids_to_labels, all_labels); + + label_query_id_to_orig_id = diskann::generate_label_specific_vector_files_compat( + query_file, labels_to_number_of_queries, query_ids_to_labels, + all_labels); // query_filters acts like query_ids_to_labels + } + else if (data_type == std::string("uint8")) + { + label_id_to_orig_id = diskann::generate_label_specific_vector_files_compat( + base_file, labels_to_number_of_points, point_ids_to_labels, all_labels); + + label_query_id_to_orig_id = diskann::generate_label_specific_vector_files_compat( + query_file, labels_to_number_of_queries, query_ids_to_labels, + all_labels); // query_filters acts like query_ids_to_labels + } + else + { + diskann::cerr << "Invalid data type" << std::endl; + return -1; + } - // Combine the label specific ground truths to produce a single GT file + // Generate label specific ground truths + + try + { + for (const auto &label : all_labels) + { + std::string filtered_base_file = base_file + "_" + label; + std::string filtered_query_file = query_file + "_" + label; + std::string filtered_gt_file = gt_file + "_" + label; + if (data_type == std::string("float")) + aux_main(filtered_base_file, "", filtered_query_file, filtered_gt_file, K, "", metric, ""); + if (data_type == std::string("int8")) + aux_main(filtered_base_file, "", filtered_query_file, filtered_gt_file, K, "", metric, ""); + if (data_type == std::string("uint8")) + aux_main(filtered_base_file, "", filtered_query_file, filtered_gt_file, K, "", metric, ""); + } + } + catch (const std::exception &e) + { + std::cout << std::string(e.what()) << std::endl; + diskann::cerr << "Compute GT failed." << std::endl; + return -1; + } - uint32_t *gt_ids = nullptr; - float *gt_dists = nullptr; - size_t gt_num, gt_dim; + // Combine the label specific ground truths to produce a single GT file - std::vector> final_gt_ids; - std::vector> final_gt_dists; + uint32_t *gt_ids = nullptr; + float *gt_dists = nullptr; + size_t gt_num, gt_dim; - int query_num = 0; - for (const auto &lbl : all_labels) { - query_num += labels_to_number_of_queries[lbl]; - } + std::vector> final_gt_ids; + std::vector> final_gt_dists; - for (int i = 0; i < query_num; i++) { - final_gt_ids.push_back(std::vector(K)); - final_gt_dists.push_back(std::vector(K)); - } + int query_num = 0; + for (const auto &lbl : all_labels) + { + query_num += labels_to_number_of_queries[lbl]; + } - for (const auto &lbl : all_labels) { - std::string filtered_gt_file = gt_file + "_" + lbl; - load_truthset(filtered_gt_file, gt_ids, gt_dists, gt_num, gt_dim); + for (int i = 0; i < query_num; i++) + { + final_gt_ids.push_back(std::vector(K)); + final_gt_dists.push_back(std::vector(K)); + } - for (int i = 0; i < labels_to_number_of_queries[lbl]; i++) { - int orig_query_id = label_query_id_to_orig_id[lbl][i]; - for (uint64_t j = 0; j < K; j++) { - final_gt_ids[orig_query_id][j] = - label_id_to_orig_id[lbl][gt_ids[i * K + j]]; - final_gt_dists[orig_query_id][j] = gt_dists[i * K + j]; + for (const auto &lbl : all_labels) + { + std::string filtered_gt_file = gt_file + "_" + lbl; + load_truthset(filtered_gt_file, gt_ids, gt_dists, gt_num, gt_dim); + + for (int i = 0; i < labels_to_number_of_queries[lbl]; i++) + { + int orig_query_id = label_query_id_to_orig_id[lbl][i]; + for (uint64_t j = 0; j < K; j++) + { + final_gt_ids[orig_query_id][j] = label_id_to_orig_id[lbl][gt_ids[i * K + j]]; + final_gt_dists[orig_query_id][j] = gt_dists[i * K + j]; + } + } } - } - } - int32_t *closest_points = new int32_t[query_num * K]; - float *dist_closest_points = new float[query_num * K]; + int32_t *closest_points = new int32_t[query_num * K]; + float *dist_closest_points = new float[query_num * K]; - for (int i = 0; i < query_num; i++) { - for (int j = 0; j < K; j++) { - closest_points[i * K + j] = final_gt_ids[i][j]; - dist_closest_points[i * K + j] = final_gt_dists[i][j]; - } - } + for (int i = 0; i < query_num; i++) + { + for (int j = 0; j < K; j++) + { + closest_points[i * K + j] = final_gt_ids[i][j]; + dist_closest_points[i * K + j] = final_gt_dists[i][j]; + } + } - save_groundtruth_as_one_file(gt_file, closest_points, dist_closest_points, - query_num, K); + save_groundtruth_as_one_file(gt_file, closest_points, dist_closest_points, query_num, K); - // cleanup artifacts - std::cout << "Cleaning up artifacts..." << std::endl; - tsl::robin_set paths_to_clean{gt_file, base_file, query_file}; - clean_up_artifacts(paths_to_clean, all_labels); - } + // cleanup artifacts + std::cout << "Cleaning up artifacts..." << std::endl; + tsl::robin_set paths_to_clean{gt_file, base_file, query_file}; + clean_up_artifacts(paths_to_clean, all_labels); + } } diff --git a/tests/utils/count_bfs_levels.cpp b/tests/utils/count_bfs_levels.cpp index 86156de3b..ddc4eaf0b 100644 --- a/tests/utils/count_bfs_levels.cpp +++ b/tests/utils/count_bfs_levels.cpp @@ -23,56 +23,59 @@ namespace po = boost::program_options; -template -void bfs_count(const std::string &index_path, uint32_t data_dims) { - using TagT = uint32_t; - using LabelT = uint32_t; - diskann::Index index(diskann::Metric::L2, data_dims, 0, - false, false); - std::cout << "Index class instantiated" << std::endl; - index.load(index_path.c_str(), 1, 100); - std::cout << "Index loaded" << std::endl; - index.count_nodes_at_bfs_levels(); +template void bfs_count(const std::string &index_path, uint32_t data_dims) +{ + using TagT = uint32_t; + using LabelT = uint32_t; + diskann::Index index(diskann::Metric::L2, data_dims, 0, false, false); + std::cout << "Index class instantiated" << std::endl; + index.load(index_path.c_str(), 1, 100); + std::cout << "Index loaded" << std::endl; + index.count_nodes_at_bfs_levels(); } -int main(int argc, char **argv) { - std::string data_type, index_path_prefix; - uint32_t data_dims; +int main(int argc, char **argv) +{ + std::string data_type, index_path_prefix; + uint32_t data_dims; - po::options_description desc{"Arguments"}; - try { - desc.add_options()("help,h", "Print information on arguments"); - desc.add_options()("data_type", - po::value(&data_type)->required(), - "data type "); - desc.add_options()("index_path_prefix", - po::value(&index_path_prefix)->required(), - "Path prefix to the index"); - desc.add_options()("data_dims", po::value(&data_dims)->required(), - "Dimensionality of the data"); + po::options_description desc{"Arguments"}; + try + { + desc.add_options()("help,h", "Print information on arguments"); + desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); + desc.add_options()("index_path_prefix", po::value(&index_path_prefix)->required(), + "Path prefix to the index"); + desc.add_options()("data_dims", po::value(&data_dims)->required(), "Dimensionality of the data"); - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) { - std::cout << desc; - return 0; + po::variables_map vm; + po::store(po::parse_command_line(argc, argv, desc), vm); + if (vm.count("help")) + { + std::cout << desc; + return 0; + } + po::notify(vm); + } + catch (const std::exception &ex) + { + std::cerr << ex.what() << '\n'; + return -1; } - po::notify(vm); - } catch (const std::exception &ex) { - std::cerr << ex.what() << '\n'; - return -1; - } - try { - if (data_type == std::string("int8")) - bfs_count(index_path_prefix, data_dims); - else if (data_type == std::string("uint8")) - bfs_count(index_path_prefix, data_dims); - if (data_type == std::string("float")) - bfs_count(index_path_prefix, data_dims); - } catch (std::exception &e) { - std::cout << std::string(e.what()) << std::endl; - diskann::cerr << "Index BFS failed." << std::endl; - return -1; - } + try + { + if (data_type == std::string("int8")) + bfs_count(index_path_prefix, data_dims); + else if (data_type == std::string("uint8")) + bfs_count(index_path_prefix, data_dims); + if (data_type == std::string("float")) + bfs_count(index_path_prefix, data_dims); + } + catch (std::exception &e) + { + std::cout << std::string(e.what()) << std::endl; + diskann::cerr << "Index BFS failed." << std::endl; + return -1; + } } diff --git a/tests/utils/create_disk_layout.cpp b/tests/utils/create_disk_layout.cpp index a6f981db5..f494c1227 100644 --- a/tests/utils/create_disk_layout.cpp +++ b/tests/utils/create_disk_layout.cpp @@ -12,34 +12,37 @@ #include "disk_utils.h" #include "cached_io.h" -template -int create_disk_layout(char **argv) { - 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); - return 0; +template int create_disk_layout(char **argv) +{ + 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); + return 0; } -int main(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); - } +int main(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); + } - int ret_val = -1; - if (std::string(argv[1]) == std::string("float")) - ret_val = create_disk_layout(argv); - else if (std::string(argv[1]) == std::string("int8")) - ret_val = create_disk_layout(argv); - else if (std::string(argv[1]) == std::string("uint8")) - ret_val = create_disk_layout(argv); - else { - std::cout << "unsupported type. use int8/uint8/float " << std::endl; - ret_val = -2; - } - return ret_val; + int ret_val = -1; + if (std::string(argv[1]) == std::string("float")) + ret_val = create_disk_layout(argv); + else if (std::string(argv[1]) == std::string("int8")) + ret_val = create_disk_layout(argv); + else if (std::string(argv[1]) == std::string("uint8")) + ret_val = create_disk_layout(argv); + else + { + std::cout << "unsupported type. use int8/uint8/float " << std::endl; + ret_val = -2; + } + return ret_val; } diff --git a/tests/utils/float_bin_to_int8.cpp b/tests/utils/float_bin_to_int8.cpp index 15d524850..a7632a6cf 100644 --- a/tests/utils/float_bin_to_int8.cpp +++ b/tests/utils/float_bin_to_int8.cpp @@ -4,59 +4,60 @@ #include #include "utils.h" -void block_convert(std::ofstream &writer, int8_t *write_buf, - std::ifstream &reader, float *read_buf, size_t npts, - size_t ndims, float bias, float scale) { - reader.read((char *)read_buf, npts * ndims * sizeof(float)); - - for (size_t i = 0; i < npts; i++) { - for (size_t d = 0; d < ndims; d++) { - write_buf[d + i * ndims] = - (int8_t)((read_buf[d + i * ndims] - bias) * (254.0 / scale)); +void block_convert(std::ofstream &writer, int8_t *write_buf, std::ifstream &reader, float *read_buf, size_t npts, + size_t ndims, float bias, float scale) +{ + reader.read((char *)read_buf, npts * ndims * sizeof(float)); + + for (size_t i = 0; i < npts; i++) + { + for (size_t d = 0; d < ndims; d++) + { + write_buf[d + i * ndims] = (int8_t)((read_buf[d + i * ndims] - bias) * (254.0 / scale)); + } } - } - writer.write((char *)write_buf, npts * ndims); + writer.write((char *)write_buf, npts * ndims); } -int main(int argc, char **argv) { - if (argc != 5) { - std::cout << "Usage: " << argv[0] << " input_bin output_tsv bias scale" - << std::endl; - exit(-1); - } - - std::ifstream reader(argv[1], std::ios::binary); - uint32_t npts_u32; - uint32_t ndims_u32; - reader.read((char *)&npts_u32, sizeof(uint32_t)); - reader.read((char *)&ndims_u32, sizeof(uint32_t)); - size_t npts = npts_u32; - size_t ndims = ndims_u32; - std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims - << std::endl; - - size_t blk_size = 131072; - size_t nblks = ROUND_UP(npts, blk_size) / blk_size; - - std::ofstream writer(argv[2], std::ios::binary); - auto read_buf = new float[blk_size * ndims]; - auto write_buf = new int8_t[blk_size * ndims]; - float bias = atof(argv[3]); - float scale = atof(argv[4]); - - writer.write((char *)(&npts_u32), sizeof(uint32_t)); - writer.write((char *)(&ndims_u32), sizeof(uint32_t)); - - for (size_t i = 0; i < nblks; i++) { - size_t cblk_size = std::min(npts - i * blk_size, blk_size); - block_convert(writer, write_buf, reader, read_buf, cblk_size, ndims, bias, - scale); - std::cout << "Block #" << i << " written" << std::endl; - } - - delete[] read_buf; - delete[] write_buf; - - writer.close(); - reader.close(); +int main(int argc, char **argv) +{ + if (argc != 5) + { + std::cout << "Usage: " << argv[0] << " input_bin output_tsv bias scale" << std::endl; + exit(-1); + } + + std::ifstream reader(argv[1], std::ios::binary); + uint32_t npts_u32; + uint32_t ndims_u32; + reader.read((char *)&npts_u32, sizeof(uint32_t)); + reader.read((char *)&ndims_u32, sizeof(uint32_t)); + size_t npts = npts_u32; + size_t ndims = ndims_u32; + std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims << std::endl; + + size_t blk_size = 131072; + size_t nblks = ROUND_UP(npts, blk_size) / blk_size; + + std::ofstream writer(argv[2], std::ios::binary); + auto read_buf = new float[blk_size * ndims]; + auto write_buf = new int8_t[blk_size * ndims]; + float bias = atof(argv[3]); + float scale = atof(argv[4]); + + writer.write((char *)(&npts_u32), sizeof(uint32_t)); + writer.write((char *)(&ndims_u32), sizeof(uint32_t)); + + for (size_t i = 0; i < nblks; i++) + { + size_t cblk_size = std::min(npts - i * blk_size, blk_size); + block_convert(writer, write_buf, reader, read_buf, cblk_size, ndims, bias, scale); + std::cout << "Block #" << i << " written" << std::endl; + } + + delete[] read_buf; + delete[] write_buf; + + writer.close(); + reader.close(); } diff --git a/tests/utils/fvecs_to_bin.cpp b/tests/utils/fvecs_to_bin.cpp index f645b67c9..873ad3b0c 100644 --- a/tests/utils/fvecs_to_bin.cpp +++ b/tests/utils/fvecs_to_bin.cpp @@ -5,88 +5,91 @@ #include "utils.h" // Convert float types -void block_convert_float(std::ifstream &reader, std::ofstream &writer, - float *read_buf, float *write_buf, size_t npts, - size_t ndims) { - reader.read((char *)read_buf, - npts * (ndims * sizeof(float) + sizeof(uint32_t))); - for (size_t 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)); +void block_convert_float(std::ifstream &reader, std::ofstream &writer, float *read_buf, float *write_buf, size_t npts, + size_t ndims) +{ + reader.read((char *)read_buf, npts * (ndims * sizeof(float) + sizeof(uint32_t))); + for (size_t 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)); } // Convert byte types -void block_convert_byte(std::ifstream &reader, std::ofstream &writer, - uint8_t *read_buf, uint8_t *write_buf, size_t npts, - size_t ndims) { - reader.read((char *)read_buf, - npts * (ndims * sizeof(uint8_t) + sizeof(uint32_t))); - for (size_t i = 0; i < npts; i++) { - memcpy(write_buf + i * ndims, - (read_buf + i * (ndims + sizeof(uint32_t))) + sizeof(uint32_t), - ndims * sizeof(uint8_t)); - } - writer.write((char *)write_buf, npts * ndims * sizeof(uint8_t)); +void block_convert_byte(std::ifstream &reader, std::ofstream &writer, uint8_t *read_buf, uint8_t *write_buf, + size_t npts, size_t ndims) +{ + reader.read((char *)read_buf, npts * (ndims * sizeof(uint8_t) + sizeof(uint32_t))); + for (size_t i = 0; i < npts; i++) + { + memcpy(write_buf + i * ndims, (read_buf + i * (ndims + sizeof(uint32_t))) + sizeof(uint32_t), + ndims * sizeof(uint8_t)); + } + writer.write((char *)write_buf, npts * ndims * sizeof(uint8_t)); } -int main(int argc, char **argv) { - if (argc != 4) { - std::cout << argv[0] << " input_vecs output_bin" - << std::endl; - exit(-1); - } +int main(int argc, char **argv) +{ + if (argc != 4) + { + std::cout << argv[0] << " input_vecs output_bin" << std::endl; + exit(-1); + } - int datasize = sizeof(float); + int datasize = sizeof(float); - if (strcmp(argv[1], "uint8") == 0 || strcmp(argv[1], "int8") == 0) { - datasize = sizeof(uint8_t); - } else if (strcmp(argv[1], "float") != 0) { - std::cout << "Error: type not supported. Use float/int8/uint8" << std::endl; - exit(-1); - } + if (strcmp(argv[1], "uint8") == 0 || strcmp(argv[1], "int8") == 0) + { + datasize = sizeof(uint8_t); + } + else if (strcmp(argv[1], "float") != 0) + { + std::cout << "Error: type not supported. Use float/int8/uint8" << std::endl; + exit(-1); + } - std::ifstream reader(argv[2], std::ios::binary | std::ios::ate); - size_t fsize = reader.tellg(); - reader.seekg(0, std::ios::beg); + std::ifstream reader(argv[2], std::ios::binary | std::ios::ate); + size_t fsize = reader.tellg(); + reader.seekg(0, std::ios::beg); - uint32_t ndims_u32; - reader.read((char *)&ndims_u32, sizeof(uint32_t)); - reader.seekg(0, std::ios::beg); - size_t ndims = (size_t)ndims_u32; - size_t npts = fsize / ((ndims * datasize) + sizeof(uint32_t)); - std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims - << std::endl; + uint32_t ndims_u32; + reader.read((char *)&ndims_u32, sizeof(uint32_t)); + reader.seekg(0, std::ios::beg); + size_t ndims = (size_t)ndims_u32; + size_t npts = fsize / ((ndims * datasize) + sizeof(uint32_t)); + std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims << std::endl; - size_t blk_size = 131072; - size_t nblks = ROUND_UP(npts, blk_size) / blk_size; - std::cout << "# blks: " << nblks << std::endl; - std::ofstream writer(argv[3], std::ios::binary); - int32_t npts_s32 = (int32_t)npts; - int32_t ndims_s32 = (int32_t)ndims; - writer.write((char *)&npts_s32, sizeof(int32_t)); - writer.write((char *)&ndims_s32, sizeof(int32_t)); + size_t blk_size = 131072; + size_t nblks = ROUND_UP(npts, blk_size) / blk_size; + std::cout << "# blks: " << nblks << std::endl; + std::ofstream writer(argv[3], std::ios::binary); + int32_t npts_s32 = (int32_t)npts; + int32_t ndims_s32 = (int32_t)ndims; + writer.write((char *)&npts_s32, sizeof(int32_t)); + writer.write((char *)&ndims_s32, sizeof(int32_t)); - size_t chunknpts = std::min(npts, blk_size); - uint8_t *read_buf = - new uint8_t[chunknpts * ((ndims * datasize) + sizeof(uint32_t))]; - uint8_t *write_buf = new uint8_t[chunknpts * ndims * datasize]; + size_t chunknpts = std::min(npts, blk_size); + uint8_t *read_buf = new uint8_t[chunknpts * ((ndims * datasize) + sizeof(uint32_t))]; + uint8_t *write_buf = new uint8_t[chunknpts * ndims * datasize]; - for (size_t i = 0; i < nblks; i++) { - size_t cblk_size = std::min(npts - i * blk_size, blk_size); - if (datasize == sizeof(float)) { - block_convert_float(reader, writer, (float *)read_buf, (float *)write_buf, - cblk_size, ndims); - } else { - block_convert_byte(reader, writer, read_buf, write_buf, cblk_size, ndims); + for (size_t i = 0; i < nblks; i++) + { + size_t cblk_size = std::min(npts - i * blk_size, blk_size); + if (datasize == sizeof(float)) + { + block_convert_float(reader, writer, (float *)read_buf, (float *)write_buf, cblk_size, ndims); + } + else + { + block_convert_byte(reader, writer, read_buf, write_buf, cblk_size, ndims); + } + std::cout << "Block #" << i << " written" << std::endl; } - std::cout << "Block #" << i << " written" << std::endl; - } - delete[] read_buf; - delete[] write_buf; + delete[] read_buf; + delete[] write_buf; - reader.close(); - writer.close(); + reader.close(); + writer.close(); } diff --git a/tests/utils/fvecs_to_bvecs.cpp b/tests/utils/fvecs_to_bvecs.cpp index 1707bfbfe..f9c2aa71b 100644 --- a/tests/utils/fvecs_to_bvecs.cpp +++ b/tests/utils/fvecs_to_bvecs.cpp @@ -4,53 +4,53 @@ #include #include "utils.h" -void block_convert(std::ifstream &reader, std::ofstream &writer, - float *read_buf, uint8_t *write_buf, size_t npts, - size_t ndims) { - reader.read((char *)read_buf, - npts * (ndims * sizeof(float) + sizeof(uint32_t))); - for (size_t i = 0; i < npts; i++) { - memcpy(write_buf + i * (ndims + 4), read_buf + i * (ndims + 1), - sizeof(uint32_t)); - for (size_t d = 0; d < ndims; d++) - write_buf[i * (ndims + 4) + 4 + d] = - (uint8_t)read_buf[i * (ndims + 1) + 1 + d]; - } - writer.write((char *)write_buf, npts * (ndims * 1 + 4)); +void block_convert(std::ifstream &reader, std::ofstream &writer, float *read_buf, uint8_t *write_buf, size_t npts, + size_t ndims) +{ + reader.read((char *)read_buf, npts * (ndims * sizeof(float) + sizeof(uint32_t))); + for (size_t i = 0; i < npts; i++) + { + memcpy(write_buf + i * (ndims + 4), read_buf + i * (ndims + 1), sizeof(uint32_t)); + for (size_t d = 0; d < ndims; d++) + write_buf[i * (ndims + 4) + 4 + d] = (uint8_t)read_buf[i * (ndims + 1) + 1 + d]; + } + writer.write((char *)write_buf, npts * (ndims * 1 + 4)); } -int main(int argc, char **argv) { - if (argc != 3) { - std::cout << argv[0] << " input_fvecs output_bvecs(uint8)" << std::endl; - exit(-1); - } - std::ifstream reader(argv[1], std::ios::binary | std::ios::ate); - size_t fsize = reader.tellg(); - reader.seekg(0, std::ios::beg); +int main(int argc, char **argv) +{ + if (argc != 3) + { + std::cout << argv[0] << " input_fvecs output_bvecs(uint8)" << std::endl; + exit(-1); + } + std::ifstream reader(argv[1], std::ios::binary | std::ios::ate); + size_t fsize = reader.tellg(); + reader.seekg(0, std::ios::beg); - uint32_t ndims_u32; - reader.read((char *)&ndims_u32, sizeof(uint32_t)); - reader.seekg(0, std::ios::beg); - size_t ndims = (size_t)ndims_u32; - size_t npts = fsize / ((ndims + 1) * sizeof(float)); - std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims - << std::endl; + uint32_t ndims_u32; + reader.read((char *)&ndims_u32, sizeof(uint32_t)); + reader.seekg(0, std::ios::beg); + size_t ndims = (size_t)ndims_u32; + size_t npts = fsize / ((ndims + 1) * sizeof(float)); + std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims << std::endl; - size_t blk_size = 131072; - size_t nblks = ROUND_UP(npts, blk_size) / blk_size; - std::cout << "# blks: " << nblks << std::endl; - std::ofstream writer(argv[2], std::ios::binary); - auto read_buf = new float[npts * (ndims + 1)]; - auto write_buf = new uint8_t[npts * (ndims + 4)]; - for (size_t i = 0; i < nblks; i++) { - size_t 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; - } + size_t blk_size = 131072; + size_t nblks = ROUND_UP(npts, blk_size) / blk_size; + std::cout << "# blks: " << nblks << std::endl; + std::ofstream writer(argv[2], std::ios::binary); + auto read_buf = new float[npts * (ndims + 1)]; + auto write_buf = new uint8_t[npts * (ndims + 4)]; + for (size_t i = 0; i < nblks; i++) + { + size_t 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; + } - delete[] read_buf; - delete[] write_buf; + delete[] read_buf; + delete[] write_buf; - reader.close(); - writer.close(); + reader.close(); + writer.close(); } diff --git a/tests/utils/gen_random_slice.cpp b/tests/utils/gen_random_slice.cpp index 3914473ab..a4cd96e0a 100644 --- a/tests/utils/gen_random_slice.cpp +++ b/tests/utils/gen_random_slice.cpp @@ -20,31 +20,39 @@ #include #include -template -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])); - gen_random_slice(base_file, output_prefix, sampling_rate); - return 0; +template 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])); + gen_random_slice(base_file, output_prefix, sampling_rate); + return 0; } -int main(int argc, char **argv) { - if (argc != 5) { - std::cout << argv[0] - << " data_type [float/int8/uint8] base_bin_file " - "sample_output_prefix sampling_probability" - << std::endl; - exit(-1); - } +int main(int argc, char **argv) +{ + if (argc != 5) + { + std::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(argv); - } else if (std::string(argv[1]) == std::string("int8")) { - aux_main(argv); - } else if (std::string(argv[1]) == std::string("uint8")) { - aux_main(argv); - } else - std::cout << "Unsupported type. Use float/int8/uint8." << std::endl; - return 0; + if (std::string(argv[1]) == std::string("float")) + { + aux_main(argv); + } + else if (std::string(argv[1]) == std::string("int8")) + { + aux_main(argv); + } + else if (std::string(argv[1]) == std::string("uint8")) + { + aux_main(argv); + } + else + std::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 ff4126c68..761983129 100644 --- a/tests/utils/generate_pq.cpp +++ b/tests/utils/generate_pq.cpp @@ -8,66 +8,63 @@ #define KMEANS_ITERS_FOR_PQ 15 template -bool generate_pq(const std::string &data_path, - const std::string &index_prefix_path, - const size_t num_pq_centers, const size_t num_pq_chunks, - const float sampling_rate, const bool opq) { - std::string pq_pivots_path = index_prefix_path + "_pq_pivots.bin"; - std::string pq_compressed_vectors_path = - index_prefix_path + "_pq_compressed.bin"; +bool generate_pq(const std::string &data_path, const std::string &index_prefix_path, const size_t num_pq_centers, + const size_t num_pq_chunks, const float sampling_rate, const bool opq) +{ + std::string pq_pivots_path = index_prefix_path + "_pq_pivots.bin"; + std::string pq_compressed_vectors_path = 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; + // 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; - if (opq) { - diskann::generate_opq_pivots(train_data, train_size, train_dim, - num_pq_centers, num_pq_chunks, pq_pivots_path, - true); - } else { - diskann::generate_pq_pivots(train_data, train_size, train_dim, - num_pq_centers, num_pq_chunks, - KMEANS_ITERS_FOR_PQ, pq_pivots_path); - } - diskann::generate_pq_data_from_pivots(data_path, num_pq_centers, - num_pq_chunks, pq_pivots_path, - pq_compressed_vectors_path, true); + if (opq) + { + diskann::generate_opq_pivots(train_data, train_size, train_dim, num_pq_centers, num_pq_chunks, pq_pivots_path, + true); + } + else + { + diskann::generate_pq_pivots(train_data, train_size, train_dim, num_pq_centers, num_pq_chunks, + KMEANS_ITERS_FOR_PQ, pq_pivots_path); + } + diskann::generate_pq_data_from_pivots(data_path, num_pq_centers, num_pq_chunks, pq_pivots_path, + pq_compressed_vectors_path, true); - delete[] train_data; + delete[] train_data; - return 0; + return 0; } -int main(int argc, char **argv) { - if (argc != 7) { - std::cout << "Usage: \n" - << argv[0] - << " " - " " - " " - << std::endl; - } else { - const std::string data_path(argv[2]); - 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 bool opq = atoi(argv[6]) == 0 ? false : true; - - if (std::string(argv[1]) == std::string("float")) - generate_pq(data_path, index_prefix_path, num_pq_centers, - num_pq_chunks, sampling_rate, opq); - else if (std::string(argv[1]) == std::string("int8")) - generate_pq(data_path, index_prefix_path, num_pq_centers, - num_pq_chunks, sampling_rate, opq); - else if (std::string(argv[1]) == std::string("uint8")) - generate_pq(data_path, index_prefix_path, num_pq_centers, - num_pq_chunks, sampling_rate, opq); +int main(int argc, char **argv) +{ + if (argc != 7) + { + std::cout << "Usage: \n" + << argv[0] + << " " + " " + " " + << std::endl; + } else - std::cout << "Error. wrong file type" << std::endl; - } + { + const std::string data_path(argv[2]); + 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 bool opq = atoi(argv[6]) == 0 ? false : true; + + if (std::string(argv[1]) == std::string("float")) + generate_pq(data_path, index_prefix_path, num_pq_centers, num_pq_chunks, sampling_rate, opq); + else if (std::string(argv[1]) == std::string("int8")) + generate_pq(data_path, index_prefix_path, num_pq_centers, num_pq_chunks, sampling_rate, opq); + else if (std::string(argv[1]) == std::string("uint8")) + generate_pq(data_path, index_prefix_path, num_pq_centers, num_pq_chunks, sampling_rate, opq); + else + std::cout << "Error. wrong file type" << std::endl; + } } diff --git a/tests/utils/generate_synthetic_labels.cpp b/tests/utils/generate_synthetic_labels.cpp index bbbdd55f5..3de2130fb 100644 --- a/tests/utils/generate_synthetic_labels.cpp +++ b/tests/utils/generate_synthetic_labels.cpp @@ -9,169 +9,200 @@ #include "utils.h" namespace po = boost::program_options; -class ZipfDistribution { - public: - ZipfDistribution(int num_points, int num_labels) - : num_labels(num_labels), - num_points(num_points), - uniform_zero_to_one(std::uniform_real_distribution<>(0.0, 1.0)) {} - - std::unordered_map createDistributionMap() { - std::unordered_map map; - int primary_label_freq = ceil(num_points * distribution_factor); - for (int i{1}; i < num_labels + 1; i++) { - map[i] = ceil(primary_label_freq / i); +class ZipfDistribution +{ + public: + ZipfDistribution(int num_points, int num_labels) + : num_labels(num_labels), num_points(num_points), + uniform_zero_to_one(std::uniform_real_distribution<>(0.0, 1.0)) + { } - return map; - } - - int writeDistribution(std::ofstream &outfile) { - auto distribution_map = createDistributionMap(); - for (int i{0}; i < num_points; i++) { - bool label_written = false; - for (auto it = distribution_map.cbegin(), next_it = it; - it != distribution_map.cend(); it = next_it) { - next_it++; - auto label_selection_probability = std::bernoulli_distribution( - distribution_factor / (double)it->first); - if (label_selection_probability(rand_engine)) { - if (label_written) { - outfile << ','; - } - outfile << it->first; - label_written = true; - // remove label from map if we have used all labels - distribution_map[it->first] -= 1; - if (distribution_map[it->first] == 0) { - distribution_map.erase(it); - } + + std::unordered_map createDistributionMap() + { + std::unordered_map map; + int primary_label_freq = ceil(num_points * distribution_factor); + for (int i{1}; i < num_labels + 1; i++) + { + map[i] = ceil(primary_label_freq / i); } - } - if (!label_written) { - outfile << 0; - } - if (i < num_points - 1) { - outfile << '\n'; - } + return map; + } + + int writeDistribution(std::ofstream &outfile) + { + auto distribution_map = createDistributionMap(); + for (int i{0}; i < num_points; i++) + { + bool label_written = false; + for (auto it = distribution_map.cbegin(), next_it = it; it != distribution_map.cend(); it = next_it) + { + next_it++; + auto label_selection_probability = std::bernoulli_distribution(distribution_factor / (double)it->first); + if (label_selection_probability(rand_engine)) + { + if (label_written) + { + outfile << ','; + } + outfile << it->first; + label_written = true; + // remove label from map if we have used all labels + distribution_map[it->first] -= 1; + if (distribution_map[it->first] == 0) + { + distribution_map.erase(it); + } + } + } + if (!label_written) + { + outfile << 0; + } + if (i < num_points - 1) + { + outfile << '\n'; + } + } + return 0; } - return 0; - } - int writeDistribution(std::string filename) { - std::ofstream outfile(filename); - if (!outfile.is_open()) { - std::cerr << "Error: could not open output file " << filename << '\n'; - return -1; + int writeDistribution(std::string filename) + { + std::ofstream outfile(filename); + if (!outfile.is_open()) + { + std::cerr << "Error: could not open output file " << filename << '\n'; + return -1; + } + writeDistribution(outfile); + outfile.close(); } - writeDistribution(outfile); - outfile.close(); - } - - private: - int num_labels; - const int num_points; - const double distribution_factor = 0.7; - std::knuth_b rand_engine; - const std::uniform_real_distribution uniform_zero_to_one; + + private: + int num_labels; + const int num_points; + const double distribution_factor = 0.7; + std::knuth_b rand_engine; + const std::uniform_real_distribution uniform_zero_to_one; }; -int main(int argc, char **argv) { - std::string output_file, distribution_type; - size_t num_labels, num_points; - - try { - po::options_description desc{"Arguments"}; - - desc.add_options()("help,h", "Print information on arguments"); - desc.add_options()("output_file,O", - po::value(&output_file)->required(), - "Filename for saving the label file"); - desc.add_options()("num_points,N", - po::value(&num_points)->required(), - "Number of points in dataset"); - desc.add_options()("num_labels,L", - po::value(&num_labels)->required(), - "Number of unique labels, up to 5000"); - desc.add_options()( - "distribution_type,DT", - po::value(&distribution_type)->default_value("random"), - "Distribution function for labels defaults " - "to random"); - - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) { - std::cout << desc; - return 0; +int main(int argc, char **argv) +{ + std::string output_file, distribution_type; + size_t num_labels, num_points; + + try + { + po::options_description desc{"Arguments"}; + + desc.add_options()("help,h", "Print information on arguments"); + desc.add_options()("output_file,O", po::value(&output_file)->required(), + "Filename for saving the label file"); + desc.add_options()("num_points,N", po::value(&num_points)->required(), "Number of points in dataset"); + desc.add_options()("num_labels,L", po::value(&num_labels)->required(), + "Number of unique labels, up to 5000"); + desc.add_options()("distribution_type,DT", po::value(&distribution_type)->default_value("random"), + "Distribution function for labels defaults " + "to random"); + + po::variables_map vm; + po::store(po::parse_command_line(argc, argv, desc), vm); + if (vm.count("help")) + { + std::cout << desc; + return 0; + } + po::notify(vm); } - po::notify(vm); - } catch (const std::exception &ex) { - std::cerr << ex.what() << '\n'; - return -1; - } - - if (num_labels > 5000) { - std::cerr << "Error: num_labels must be 5000 or less" << '\n'; - return -1; - } - - if (num_points <= 0) { - std::cerr << "Error: num_points must be greater than 0" << '\n'; - return -1; - } - - std::cout << "Generating synthetic labels for " << num_points - << " points with " << num_labels << " unique labels" << '\n'; - - try { - std::ofstream outfile(output_file); - if (!outfile.is_open()) { - std::cerr << "Error: could not open output file " << output_file << '\n'; - return -1; + catch (const std::exception &ex) + { + std::cerr << ex.what() << '\n'; + return -1; } - if (distribution_type == "zipf") { - ZipfDistribution zipf(num_points, num_labels); - zipf.writeDistribution(outfile); - } else if (distribution_type == "random") { - for (size_t i = 0; i < num_points; i++) { - bool label_written = false; - for (size_t j = 1; j <= num_labels; j++) { - // 50% chance to assign each label - if (rand() > (RAND_MAX / 2)) { - if (label_written) { - outfile << ','; + if (num_labels > 5000) + { + std::cerr << "Error: num_labels must be 5000 or less" << '\n'; + return -1; + } + + if (num_points <= 0) + { + std::cerr << "Error: num_points must be greater than 0" << '\n'; + return -1; + } + + std::cout << "Generating synthetic labels for " << num_points << " points with " << num_labels << " unique labels" + << '\n'; + + try + { + std::ofstream outfile(output_file); + if (!outfile.is_open()) + { + std::cerr << "Error: could not open output file " << output_file << '\n'; + return -1; + } + + if (distribution_type == "zipf") + { + ZipfDistribution zipf(num_points, num_labels); + zipf.writeDistribution(outfile); + } + else if (distribution_type == "random") + { + for (size_t i = 0; i < num_points; i++) + { + bool label_written = false; + for (size_t j = 1; j <= num_labels; j++) + { + // 50% chance to assign each label + if (rand() > (RAND_MAX / 2)) + { + if (label_written) + { + outfile << ','; + } + outfile << j; + label_written = true; + } + } + if (!label_written) + { + outfile << 0; + } + if (i < num_points - 1) + { + outfile << '\n'; + } } - outfile << j; - label_written = true; - } } - if (!label_written) { - outfile << 0; + else if (distribution_type == "one_per_point") + { + std::random_device rd; // obtain a random number from hardware + std::mt19937 gen(rd()); // seed the generator + std::uniform_int_distribution<> distr(0, num_labels); // define the range + + for (size_t i = 0; i < num_points; i++) + { + outfile << distr(gen); + if (i != num_points - 1) + outfile << '\n'; + } } - if (i < num_points - 1) { - outfile << '\n'; + if (outfile.is_open()) + { + outfile.close(); } - } - } else if (distribution_type == "one_per_point") { - std::random_device rd; // obtain a random number from hardware - std::mt19937 gen(rd()); // seed the generator - std::uniform_int_distribution<> distr(0, num_labels); // define the range - - for (size_t i = 0; i < num_points; i++) { - outfile << distr(gen); - if (i != num_points - 1) outfile << '\n'; - } + + std::cout << "Labels written to " << output_file << '\n'; } - if (outfile.is_open()) { - outfile.close(); + catch (const std::exception &ex) + { + std::cerr << "Label generation failed: " << ex.what() << '\n'; + return -1; } - std::cout << "Labels written to " << output_file << '\n'; - } catch (const std::exception &ex) { - std::cerr << "Label generation failed: " << ex.what() << '\n'; - return -1; - } - - return 0; + return 0; } \ No newline at end of file diff --git a/tests/utils/int8_to_float.cpp b/tests/utils/int8_to_float.cpp index 05662b2a0..dcdfddc0d 100644 --- a/tests/utils/int8_to_float.cpp +++ b/tests/utils/int8_to_float.cpp @@ -4,18 +4,20 @@ #include #include "utils.h" -int main(int argc, char **argv) { - if (argc != 3) { - std::cout << argv[0] << " input_int8_bin output_float_bin" << std::endl; - exit(-1); - } +int main(int argc, char **argv) +{ + if (argc != 3) + { + std::cout << argv[0] << " input_int8_bin output_float_bin" << std::endl; + exit(-1); + } - int8_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; + int8_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/int8_to_float_scale.cpp b/tests/utils/int8_to_float_scale.cpp index e1217c571..2de1a3a56 100644 --- a/tests/utils/int8_to_float_scale.cpp +++ b/tests/utils/int8_to_float_scale.cpp @@ -4,59 +4,60 @@ #include #include "utils.h" -void block_convert(std::ofstream &writer, float *write_buf, - std::ifstream &reader, int8_t *read_buf, size_t npts, - size_t ndims, float bias, float scale) { - reader.read((char *)read_buf, npts * ndims * sizeof(int8_t)); - - for (size_t i = 0; i < npts; i++) { - for (size_t d = 0; d < ndims; d++) { - write_buf[d + i * ndims] = - (((float)read_buf[d + i * ndims] - bias) * scale); +void block_convert(std::ofstream &writer, float *write_buf, std::ifstream &reader, int8_t *read_buf, size_t npts, + size_t ndims, float bias, float scale) +{ + reader.read((char *)read_buf, npts * ndims * sizeof(int8_t)); + + for (size_t i = 0; i < npts; i++) + { + for (size_t d = 0; d < ndims; d++) + { + write_buf[d + i * ndims] = (((float)read_buf[d + i * ndims] - bias) * scale); + } } - } - writer.write((char *)write_buf, npts * ndims * sizeof(float)); + writer.write((char *)write_buf, npts * ndims * sizeof(float)); } -int main(int argc, char **argv) { - if (argc != 5) { - std::cout << "Usage: " << argv[0] - << " input-int8.bin output-float.bin bias scale" << std::endl; - exit(-1); - } - - std::ifstream reader(argv[1], std::ios::binary); - uint32_t npts_u32; - uint32_t ndims_u32; - reader.read((char *)&npts_u32, sizeof(uint32_t)); - reader.read((char *)&ndims_u32, sizeof(uint32_t)); - size_t npts = npts_u32; - size_t ndims = ndims_u32; - std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims - << std::endl; - - size_t blk_size = 131072; - size_t nblks = ROUND_UP(npts, blk_size) / blk_size; - - std::ofstream writer(argv[2], std::ios::binary); - auto read_buf = new int8_t[blk_size * ndims]; - auto write_buf = new float[blk_size * ndims]; - float bias = atof(argv[3]); - float scale = atof(argv[4]); - - writer.write((char *)(&npts_u32), sizeof(uint32_t)); - writer.write((char *)(&ndims_u32), sizeof(uint32_t)); - - for (size_t i = 0; i < nblks; i++) { - size_t cblk_size = std::min(npts - i * blk_size, blk_size); - block_convert(writer, write_buf, reader, read_buf, cblk_size, ndims, bias, - scale); - std::cout << "Block #" << i << " written" << std::endl; - } - - delete[] read_buf; - delete[] write_buf; - - writer.close(); - reader.close(); +int main(int argc, char **argv) +{ + if (argc != 5) + { + std::cout << "Usage: " << argv[0] << " input-int8.bin output-float.bin bias scale" << std::endl; + exit(-1); + } + + std::ifstream reader(argv[1], std::ios::binary); + uint32_t npts_u32; + uint32_t ndims_u32; + reader.read((char *)&npts_u32, sizeof(uint32_t)); + reader.read((char *)&ndims_u32, sizeof(uint32_t)); + size_t npts = npts_u32; + size_t ndims = ndims_u32; + std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims << std::endl; + + size_t blk_size = 131072; + size_t nblks = ROUND_UP(npts, blk_size) / blk_size; + + std::ofstream writer(argv[2], std::ios::binary); + auto read_buf = new int8_t[blk_size * ndims]; + auto write_buf = new float[blk_size * ndims]; + float bias = atof(argv[3]); + float scale = atof(argv[4]); + + writer.write((char *)(&npts_u32), sizeof(uint32_t)); + writer.write((char *)(&ndims_u32), sizeof(uint32_t)); + + for (size_t i = 0; i < nblks; i++) + { + size_t cblk_size = std::min(npts - i * blk_size, blk_size); + block_convert(writer, write_buf, reader, read_buf, cblk_size, ndims, bias, scale); + std::cout << "Block #" << i << " written" << std::endl; + } + + delete[] read_buf; + delete[] write_buf; + + writer.close(); + reader.close(); } diff --git a/tests/utils/ivecs_to_bin.cpp b/tests/utils/ivecs_to_bin.cpp index 0146330f8..ea8a4a3d2 100644 --- a/tests/utils/ivecs_to_bin.cpp +++ b/tests/utils/ivecs_to_bin.cpp @@ -4,54 +4,55 @@ #include #include "utils.h" -void block_convert(std::ifstream &reader, std::ofstream &writer, - uint32_t *read_buf, uint32_t *write_buf, size_t npts, - size_t ndims) { - reader.read((char *)read_buf, - npts * (ndims * sizeof(uint32_t) + sizeof(uint32_t))); - for (size_t i = 0; i < npts; i++) { - memcpy(write_buf + i * ndims, (read_buf + i * (ndims + 1)) + 1, - ndims * sizeof(uint32_t)); - } - writer.write((char *)write_buf, npts * ndims * sizeof(uint32_t)); +void block_convert(std::ifstream &reader, std::ofstream &writer, uint32_t *read_buf, uint32_t *write_buf, size_t npts, + size_t ndims) +{ + reader.read((char *)read_buf, npts * (ndims * sizeof(uint32_t) + sizeof(uint32_t))); + for (size_t i = 0; i < npts; i++) + { + memcpy(write_buf + i * ndims, (read_buf + i * (ndims + 1)) + 1, ndims * sizeof(uint32_t)); + } + writer.write((char *)write_buf, npts * ndims * sizeof(uint32_t)); } -int main(int argc, char **argv) { - if (argc != 3) { - std::cout << argv[0] << " input_ivecs output_bin" << std::endl; - exit(-1); - } - std::ifstream reader(argv[1], std::ios::binary | std::ios::ate); - size_t fsize = reader.tellg(); - reader.seekg(0, std::ios::beg); +int main(int argc, char **argv) +{ + if (argc != 3) + { + std::cout << argv[0] << " input_ivecs output_bin" << std::endl; + exit(-1); + } + std::ifstream reader(argv[1], std::ios::binary | std::ios::ate); + size_t fsize = reader.tellg(); + reader.seekg(0, std::ios::beg); - uint32_t ndims_u32; - reader.read((char *)&ndims_u32, sizeof(uint32_t)); - reader.seekg(0, std::ios::beg); - size_t ndims = (size_t)ndims_u32; - size_t npts = fsize / ((ndims + 1) * sizeof(uint32_t)); - std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims - << std::endl; + uint32_t ndims_u32; + reader.read((char *)&ndims_u32, sizeof(uint32_t)); + reader.seekg(0, std::ios::beg); + size_t ndims = (size_t)ndims_u32; + size_t npts = fsize / ((ndims + 1) * sizeof(uint32_t)); + std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims << std::endl; - size_t blk_size = 131072; - size_t nblks = ROUND_UP(npts, blk_size) / blk_size; - std::cout << "# blks: " << nblks << std::endl; - std::ofstream writer(argv[2], std::ios::binary); - int npts_s32 = (int)npts; - int ndims_s32 = (int)ndims; - writer.write((char *)&npts_s32, sizeof(int)); - writer.write((char *)&ndims_s32, sizeof(int)); - uint32_t *read_buf = new uint32_t[npts * (ndims + 1)]; - uint32_t *write_buf = new uint32_t[npts * ndims]; - for (size_t i = 0; i < nblks; i++) { - size_t 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; - } + size_t blk_size = 131072; + size_t nblks = ROUND_UP(npts, blk_size) / blk_size; + std::cout << "# blks: " << nblks << std::endl; + std::ofstream writer(argv[2], std::ios::binary); + int npts_s32 = (int)npts; + int ndims_s32 = (int)ndims; + writer.write((char *)&npts_s32, sizeof(int)); + writer.write((char *)&ndims_s32, sizeof(int)); + uint32_t *read_buf = new uint32_t[npts * (ndims + 1)]; + uint32_t *write_buf = new uint32_t[npts * ndims]; + for (size_t i = 0; i < nblks; i++) + { + size_t 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; + } - delete[] read_buf; - delete[] write_buf; + delete[] read_buf; + delete[] write_buf; - reader.close(); - writer.close(); + reader.close(); + writer.close(); } diff --git a/tests/utils/merge_shards.cpp b/tests/utils/merge_shards.cpp index 4ae4fa83d..106c15eef 100644 --- a/tests/utils/merge_shards.cpp +++ b/tests/utils/merge_shards.cpp @@ -14,28 +14,29 @@ #include "cached_io.h" #include "utils.h" -int main(int argc, char **argv) { - if (argc != 9) { - std::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] " - "output_medoids_path[8]" - << std::endl; - exit(-1); - } +int main(int argc, char **argv) +{ + if (argc != 9) + { + std::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] " + "output_medoids_path[8]" + << std::endl; + exit(-1); + } - std::string vamana_prefix(argv[1]); - std::string vamana_suffix(argv[2]); - std::string idmaps_prefix(argv[3]); - std::string idmaps_suffix(argv[4]); - uint64_t nshards = (uint64_t)std::atoi(argv[5]); - uint32_t max_degree = (uint64_t)std::atoi(argv[6]); - std::string output_index(argv[7]); - std::string output_medoids(argv[8]); + std::string vamana_prefix(argv[1]); + std::string vamana_suffix(argv[2]); + std::string idmaps_prefix(argv[3]); + std::string idmaps_suffix(argv[4]); + uint64_t nshards = (uint64_t)std::atoi(argv[5]); + uint32_t max_degree = (uint64_t)std::atoi(argv[6]); + std::string output_index(argv[7]); + std::string output_medoids(argv[8]); - return diskann::merge_shards(vamana_prefix, vamana_suffix, idmaps_prefix, - idmaps_suffix, nshards, max_degree, output_index, - output_medoids); + return diskann::merge_shards(vamana_prefix, vamana_suffix, idmaps_prefix, idmaps_suffix, nshards, max_degree, + output_index, output_medoids); } diff --git a/tests/utils/partition_data.cpp b/tests/utils/partition_data.cpp index 0156e7acd..2c505315c 100644 --- a/tests/utils/partition_data.cpp +++ b/tests/utils/partition_data.cpp @@ -8,33 +8,32 @@ // DEPRECATED: NEED TO REPROGRAM -int main(int argc, char **argv) { - if (argc != 7) { - std::cout << "Usage:\n" - << argv[0] - << " datatype " - " " - " " - << std::endl; - exit(-1); - } +int main(int argc, char **argv) +{ + if (argc != 7) + { + std::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 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]); + const std::string data_path(argv[2]); + const std::string prefix_path(argv[3]); + const float sampling_rate = 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]); - if (std::string(argv[1]) == std::string("float")) - partition(data_path, sampling_rate, num_partitions, max_reps, - prefix_path, k_index); - else if (std::string(argv[1]) == std::string("int8")) - partition(data_path, sampling_rate, num_partitions, max_reps, - prefix_path, k_index); - else if (std::string(argv[1]) == std::string("uint8")) - 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; + if (std::string(argv[1]) == std::string("float")) + partition(data_path, sampling_rate, num_partitions, max_reps, prefix_path, k_index); + else if (std::string(argv[1]) == std::string("int8")) + partition(data_path, sampling_rate, num_partitions, max_reps, prefix_path, k_index); + else if (std::string(argv[1]) == std::string("uint8")) + 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; } diff --git a/tests/utils/partition_with_ram_budget.cpp b/tests/utils/partition_with_ram_budget.cpp index 6ef2c9958..3c546801a 100644 --- a/tests/utils/partition_with_ram_budget.cpp +++ b/tests/utils/partition_with_ram_budget.cpp @@ -8,33 +8,32 @@ // DEPRECATED: NEED TO REPROGRAM -int main(int argc, char **argv) { - if (argc != 8) { - std::cout << "Usage:\n" - << argv[0] - << " datatype " - " " - " " - << std::endl; - exit(-1); - } +int main(int argc, char **argv) +{ + if (argc != 8) + { + std::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 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]); + const std::string data_path(argv[2]); + const std::string prefix_path(argv[3]); + const float sampling_rate = 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]); - if (std::string(argv[1]) == std::string("float")) - partition_with_ram_budget(data_path, sampling_rate, ram_budget, - graph_degree, prefix_path, k_index); - else if (std::string(argv[1]) == std::string("int8")) - partition_with_ram_budget(data_path, sampling_rate, ram_budget, - graph_degree, prefix_path, k_index); - else if (std::string(argv[1]) == std::string("uint8")) - 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; + if (std::string(argv[1]) == std::string("float")) + partition_with_ram_budget(data_path, sampling_rate, ram_budget, graph_degree, prefix_path, k_index); + else if (std::string(argv[1]) == std::string("int8")) + partition_with_ram_budget(data_path, sampling_rate, ram_budget, graph_degree, prefix_path, k_index); + else if (std::string(argv[1]) == std::string("uint8")) + 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; } diff --git a/tests/utils/rand_data_gen.cpp b/tests/utils/rand_data_gen.cpp index d4e81f82e..ea2e67478 100644 --- a/tests/utils/rand_data_gen.cpp +++ b/tests/utils/rand_data_gen.cpp @@ -11,174 +11,199 @@ namespace po = boost::program_options; -int block_write_float(std::ofstream &writer, size_t ndims, size_t npts, - float norm) { - auto vec = new float[ndims]; - - std::random_device rd{}; - std::mt19937 gen{rd()}; - std::normal_distribution<> normal_rand{0, 1}; - - for (size_t i = 0; i < npts; i++) { - float sum = 0; - for (size_t d = 0; d < ndims; ++d) vec[d] = normal_rand(gen); - for (size_t d = 0; d < ndims; ++d) sum += vec[d] * vec[d]; - for (size_t d = 0; d < ndims; ++d) vec[d] = vec[d] * norm / std::sqrt(sum); - - writer.write((char *)vec, ndims * sizeof(float)); - } +int block_write_float(std::ofstream &writer, size_t ndims, size_t npts, float norm) +{ + auto vec = new float[ndims]; + + std::random_device rd{}; + std::mt19937 gen{rd()}; + std::normal_distribution<> normal_rand{0, 1}; + + for (size_t i = 0; i < npts; i++) + { + float sum = 0; + for (size_t d = 0; d < ndims; ++d) + vec[d] = normal_rand(gen); + for (size_t d = 0; d < ndims; ++d) + sum += vec[d] * vec[d]; + for (size_t d = 0; d < ndims; ++d) + vec[d] = vec[d] * norm / std::sqrt(sum); + + writer.write((char *)vec, ndims * sizeof(float)); + } - delete[] vec; - return 0; + delete[] vec; + return 0; } -int block_write_int8(std::ofstream &writer, size_t ndims, size_t npts, - float norm) { - auto vec = new float[ndims]; - auto vec_T = new int8_t[ndims]; - - std::random_device rd{}; - std::mt19937 gen{rd()}; - std::normal_distribution<> normal_rand{0, 1}; - - for (size_t i = 0; i < npts; i++) { - float sum = 0; - for (size_t d = 0; d < ndims; ++d) vec[d] = normal_rand(gen); - for (size_t d = 0; d < ndims; ++d) sum += vec[d] * vec[d]; - for (size_t d = 0; d < ndims; ++d) vec[d] = vec[d] * norm / std::sqrt(sum); - - for (size_t d = 0; d < ndims; ++d) { - vec_T[d] = std::round(vec[d]); +int block_write_int8(std::ofstream &writer, size_t ndims, size_t npts, float norm) +{ + auto vec = new float[ndims]; + auto vec_T = new int8_t[ndims]; + + std::random_device rd{}; + std::mt19937 gen{rd()}; + std::normal_distribution<> normal_rand{0, 1}; + + for (size_t i = 0; i < npts; i++) + { + float sum = 0; + for (size_t d = 0; d < ndims; ++d) + vec[d] = normal_rand(gen); + for (size_t d = 0; d < ndims; ++d) + sum += vec[d] * vec[d]; + for (size_t d = 0; d < ndims; ++d) + vec[d] = vec[d] * norm / std::sqrt(sum); + + for (size_t d = 0; d < ndims; ++d) + { + vec_T[d] = std::round(vec[d]); + } + + writer.write((char *)vec_T, ndims * sizeof(int8_t)); } - writer.write((char *)vec_T, ndims * sizeof(int8_t)); - } - - delete[] vec; - delete[] vec_T; - return 0; + delete[] vec; + delete[] vec_T; + return 0; } -int block_write_uint8(std::ofstream &writer, size_t ndims, size_t npts, - float norm) { - auto vec = new float[ndims]; - auto vec_T = new int8_t[ndims]; - - std::random_device rd{}; - std::mt19937 gen{rd()}; - std::normal_distribution<> normal_rand{0, 1}; +int block_write_uint8(std::ofstream &writer, size_t ndims, size_t npts, float norm) +{ + auto vec = new float[ndims]; + auto vec_T = new int8_t[ndims]; + + std::random_device rd{}; + std::mt19937 gen{rd()}; + std::normal_distribution<> normal_rand{0, 1}; + + for (size_t i = 0; i < npts; i++) + { + float sum = 0; + for (size_t d = 0; d < ndims; ++d) + vec[d] = normal_rand(gen); + for (size_t d = 0; d < ndims; ++d) + sum += vec[d] * vec[d]; + for (size_t d = 0; d < ndims; ++d) + vec[d] = vec[d] * norm / std::sqrt(sum); + + for (size_t d = 0; d < ndims; ++d) + { + vec_T[d] = 128 + std::round(vec[d]); + } + + writer.write((char *)vec_T, ndims * sizeof(uint8_t)); + } - for (size_t i = 0; i < npts; i++) { - float sum = 0; - for (size_t d = 0; d < ndims; ++d) vec[d] = normal_rand(gen); - for (size_t d = 0; d < ndims; ++d) sum += vec[d] * vec[d]; - for (size_t d = 0; d < ndims; ++d) vec[d] = vec[d] * norm / std::sqrt(sum); + delete[] vec; + delete[] vec_T; + return 0; +} - for (size_t d = 0; d < ndims; ++d) { - vec_T[d] = 128 + std::round(vec[d]); +int main(int argc, char **argv) +{ + std::string data_type, output_file; + size_t ndims, npts; + float norm; + + try + { + po::options_description desc{"Arguments"}; + + desc.add_options()("help,h", "Print information on arguments"); + + desc.add_options()("data_type", po::value(&data_type)->required(), "data type "); + desc.add_options()("output_file", po::value(&output_file)->required(), + "File name for saving the random vectors"); + desc.add_options()("ndims,D", po::value(&ndims)->required(), "Dimensoinality of the vector"); + desc.add_options()("npts,N", po::value(&npts)->required(), "Number of vectors"); + desc.add_options()("norm", po::value(&norm)->required(), "Norm of the vectors"); + po::variables_map vm; + po::store(po::parse_command_line(argc, argv, desc), vm); + if (vm.count("help")) + { + std::cout << desc; + return 0; + } + po::notify(vm); + } + catch (const std::exception &ex) + { + std::cerr << ex.what() << '\n'; + return -1; } - writer.write((char *)vec_T, ndims * sizeof(uint8_t)); - } - - delete[] vec; - delete[] vec_T; - return 0; -} + if (data_type != std::string("float") && data_type != std::string("int8") && data_type != std::string("uint8")) + { + std::cout << "Unsupported type. float, int8 and uint8 types are supported." << std::endl; + return -1; + } -int main(int argc, char **argv) { - std::string data_type, output_file; - size_t ndims, npts; - float norm; - - try { - po::options_description desc{"Arguments"}; - - desc.add_options()("help,h", "Print information on arguments"); - - desc.add_options()("data_type", - po::value(&data_type)->required(), - "data type "); - desc.add_options()("output_file", - po::value(&output_file)->required(), - "File name for saving the random vectors"); - desc.add_options()("ndims,D", po::value(&ndims)->required(), - "Dimensoinality of the vector"); - desc.add_options()("npts,N", po::value(&npts)->required(), - "Number of vectors"); - desc.add_options()("norm", po::value(&norm)->required(), - "Norm of the vectors"); - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) { - std::cout << desc; - return 0; + if (norm <= 0.0) + { + std::cerr << "Error: Norm must be a positive number" << std::endl; + return -1; } - po::notify(vm); - } catch (const std::exception &ex) { - std::cerr << ex.what() << '\n'; - return -1; - } - - if (data_type != std::string("float") && data_type != std::string("int8") && - data_type != std::string("uint8")) { - std::cout << "Unsupported type. float, int8 and uint8 types are supported." - << std::endl; - return -1; - } - - if (norm <= 0.0) { - std::cerr << "Error: Norm must be a positive number" << std::endl; - return -1; - } - - if (data_type == std::string("int8") || data_type == std::string("uint8")) { - if (norm > 127) { - std::cerr << "Error: for int8/uint8 datatypes, L2 norm can not be " - "greater " - "than 127" - << std::endl; - return -1; + + if (data_type == std::string("int8") || data_type == std::string("uint8")) + { + if (norm > 127) + { + std::cerr << "Error: for int8/uint8 datatypes, L2 norm can not be " + "greater " + "than 127" + << std::endl; + return -1; + } } - } - - try { - std::ofstream writer; - writer.exceptions(std::ofstream::failbit | std::ofstream::badbit); - writer.open(output_file, std::ios::binary); - auto npts_u32 = (uint32_t)npts; - auto ndims_u32 = (uint32_t)ndims; - writer.write((char *)&npts_u32, sizeof(uint32_t)); - writer.write((char *)&ndims_u32, sizeof(uint32_t)); - - size_t blk_size = 131072; - size_t nblks = ROUND_UP(npts, blk_size) / blk_size; - std::cout << "# blks: " << nblks << std::endl; - - int ret = 0; - for (size_t i = 0; i < nblks; i++) { - size_t cblk_size = std::min(npts - i * blk_size, blk_size); - if (data_type == std::string("float")) { - ret = block_write_float(writer, ndims, cblk_size, norm); - } else if (data_type == std::string("int8")) { - ret = block_write_int8(writer, ndims, cblk_size, norm); - } else if (data_type == std::string("uint8")) { - ret = block_write_uint8(writer, ndims, cblk_size, norm); - } - if (ret == 0) - std::cout << "Block #" << i << " written" << std::endl; - else { + + try + { + std::ofstream writer; + writer.exceptions(std::ofstream::failbit | std::ofstream::badbit); + writer.open(output_file, std::ios::binary); + auto npts_u32 = (uint32_t)npts; + auto ndims_u32 = (uint32_t)ndims; + writer.write((char *)&npts_u32, sizeof(uint32_t)); + writer.write((char *)&ndims_u32, sizeof(uint32_t)); + + size_t blk_size = 131072; + size_t nblks = ROUND_UP(npts, blk_size) / blk_size; + std::cout << "# blks: " << nblks << std::endl; + + int ret = 0; + for (size_t i = 0; i < nblks; i++) + { + size_t cblk_size = std::min(npts - i * blk_size, blk_size); + if (data_type == std::string("float")) + { + ret = block_write_float(writer, ndims, cblk_size, norm); + } + else if (data_type == std::string("int8")) + { + ret = block_write_int8(writer, ndims, cblk_size, norm); + } + else if (data_type == std::string("uint8")) + { + ret = block_write_uint8(writer, ndims, cblk_size, norm); + } + if (ret == 0) + std::cout << "Block #" << i << " written" << std::endl; + else + { + writer.close(); + std::cout << "failed to write" << std::endl; + return -1; + } + } writer.close(); - std::cout << "failed to write" << std::endl; + } + catch (const std::exception &e) + { + std::cout << std::string(e.what()) << std::endl; + diskann::cerr << "Index build failed." << std::endl; return -1; - } } - writer.close(); - } catch (const std::exception &e) { - std::cout << std::string(e.what()) << std::endl; - diskann::cerr << "Index build failed." << std::endl; - return -1; - } - - return 0; + + return 0; } diff --git a/tests/utils/simulate_aggregate_recall.cpp b/tests/utils/simulate_aggregate_recall.cpp index 7965a4854..2e7645ae3 100644 --- a/tests/utils/simulate_aggregate_recall.cpp +++ b/tests/utils/simulate_aggregate_recall.cpp @@ -6,73 +6,80 @@ #include #include -inline float aggregate_recall(const uint32_t k_aggr, const uint32_t k, - const uint32_t npart, uint32_t *count, - const std::vector &recalls) { - float found = 0; - for (uint32_t i = 0; i < npart; ++i) { - size_t max_found = std::min(count[i], k); - found += recalls[max_found - 1] * max_found; - } - return found / (float)k_aggr; +inline float aggregate_recall(const uint32_t k_aggr, const uint32_t k, const uint32_t npart, uint32_t *count, + const std::vector &recalls) +{ + float found = 0; + for (uint32_t i = 0; i < npart; ++i) + { + size_t max_found = std::min(count[i], k); + found += recalls[max_found - 1] * max_found; + } + return found / (float)k_aggr; } -void simulate(const uint32_t k_aggr, const uint32_t k, const uint32_t npart, - const uint32_t nsim, const std::vector &recalls) { - std::random_device r; - std::default_random_engine randeng(r()); - std::uniform_int_distribution uniform_dist(0, npart - 1); +void simulate(const uint32_t k_aggr, const uint32_t k, const uint32_t npart, const uint32_t nsim, + const std::vector &recalls) +{ + std::random_device r; + std::default_random_engine randeng(r()); + std::uniform_int_distribution uniform_dist(0, npart - 1); - uint32_t *count = new uint32_t[npart]; - double aggr_recall = 0; + uint32_t *count = new uint32_t[npart]; + double aggr_recall = 0; - for (uint32_t i = 0; i < nsim; ++i) { - for (uint32_t p = 0; p < npart; ++p) { - count[p] = 0; - } - for (uint32_t t = 0; t < k_aggr; ++t) { - count[uniform_dist(randeng)]++; + for (uint32_t i = 0; i < nsim; ++i) + { + for (uint32_t p = 0; p < npart; ++p) + { + count[p] = 0; + } + for (uint32_t t = 0; t < k_aggr; ++t) + { + count[uniform_dist(randeng)]++; + } + aggr_recall += aggregate_recall(k_aggr, k, npart, count, recalls); } - aggr_recall += aggregate_recall(k_aggr, k, npart, count, recalls); - } - std::cout << "Aggregate recall is " << aggr_recall / (double)nsim - << std::endl; - delete[] count; + std::cout << "Aggregate recall is " << aggr_recall / (double)nsim << std::endl; + delete[] count; } -int main(int argc, char **argv) { - if (argc < 6) { - std::cout << argv[0] - << " k_aggregate k_out npart nsim recall@1 recall@2 ... recall@k" - << std::endl; - exit(-1); - } +int main(int argc, char **argv) +{ + if (argc < 6) + { + std::cout << argv[0] << " k_aggregate k_out npart nsim recall@1 recall@2 ... recall@k" << std::endl; + exit(-1); + } - const uint32_t k_aggr = atoi(argv[1]); - const uint32_t k = atoi(argv[2]); - const uint32_t npart = atoi(argv[3]); - const uint32_t nsim = atoi(argv[4]); + const uint32_t k_aggr = atoi(argv[1]); + const uint32_t k = atoi(argv[2]); + const uint32_t npart = atoi(argv[3]); + const uint32_t nsim = atoi(argv[4]); - std::vector recalls; - for (int ctr = 5; ctr < argc; ctr++) { - recalls.push_back(atof(argv[ctr])); - } + std::vector recalls; + for (int ctr = 5; ctr < argc; ctr++) + { + recalls.push_back(atof(argv[ctr])); + } - if (recalls.size() != k) { - std::cerr << "Please input k numbers for recall@1, recall@2 .. recall@k" - << std::endl; - } - if (k_aggr > npart * k) { - std::cerr << "k_aggr must be <= k * npart" << std::endl; - exit(-1); - } - if (nsim <= npart * k_aggr) { - std::cerr << "Choose nsim > npart*k_aggr" << std::endl; - exit(-1); - } + if (recalls.size() != k) + { + std::cerr << "Please input k numbers for recall@1, recall@2 .. recall@k" << std::endl; + } + if (k_aggr > npart * k) + { + std::cerr << "k_aggr must be <= k * npart" << std::endl; + exit(-1); + } + if (nsim <= npart * k_aggr) + { + std::cerr << "Choose nsim > npart*k_aggr" << std::endl; + exit(-1); + } - simulate(k_aggr, k, npart, nsim, recalls); + simulate(k_aggr, k, npart, nsim, recalls); - return 0; + return 0; } diff --git a/tests/utils/stats_label_data.cpp b/tests/utils/stats_label_data.cpp index 1a11238a1..129d5bcb2 100644 --- a/tests/utils/stats_label_data.cpp +++ b/tests/utils/stats_label_data.cpp @@ -28,123 +28,120 @@ #endif namespace po = boost::program_options; -void stats_analysis(const std::string labels_file, std::string univeral_label, - uint32_t density = 10) { - std::string token, line; - std::ifstream labels_stream(labels_file); - std::unordered_map label_counts; - std::string label_with_max_points; - uint32_t max_points = 0; - long long sum = 0; - long long point_cnt = 0; - float avg_labels_per_pt, mean_label_size; +void stats_analysis(const std::string labels_file, std::string univeral_label, uint32_t density = 10) +{ + std::string token, line; + std::ifstream labels_stream(labels_file); + std::unordered_map label_counts; + std::string label_with_max_points; + uint32_t max_points = 0; + long long sum = 0; + long long point_cnt = 0; + float avg_labels_per_pt, mean_label_size; - std::vector labels_per_point; - uint32_t dense_pts = 0; - if (labels_stream.is_open()) { - while (getline(labels_stream, line)) { - point_cnt++; - std::stringstream iss(line); - uint32_t lbl_cnt = 0; - while (getline(iss, token, ',')) { - lbl_cnt++; - token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); - token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); - if (label_counts.find(token) == label_counts.end()) - label_counts[token] = 0; - label_counts[token]++; - } - if (lbl_cnt >= density) { - dense_pts++; - } - labels_per_point.emplace_back(lbl_cnt); + std::vector labels_per_point; + uint32_t dense_pts = 0; + if (labels_stream.is_open()) + { + while (getline(labels_stream, line)) + { + point_cnt++; + std::stringstream iss(line); + uint32_t lbl_cnt = 0; + while (getline(iss, token, ',')) + { + lbl_cnt++; + token.erase(std::remove(token.begin(), token.end(), '\n'), token.end()); + token.erase(std::remove(token.begin(), token.end(), '\r'), token.end()); + if (label_counts.find(token) == label_counts.end()) + label_counts[token] = 0; + label_counts[token]++; + } + if (lbl_cnt >= density) + { + dense_pts++; + } + labels_per_point.emplace_back(lbl_cnt); + } } - } - std::cout << "fraction of dense points with >= " << density - << " labels = " << (float)dense_pts / (float)labels_per_point.size() - << std::endl; - std::sort(labels_per_point.begin(), labels_per_point.end()); + std::cout << "fraction of dense points with >= " << density + << " labels = " << (float)dense_pts / (float)labels_per_point.size() << std::endl; + std::sort(labels_per_point.begin(), labels_per_point.end()); - std::vector> label_count_vec; + std::vector> label_count_vec; - for (auto it = label_counts.begin(); it != label_counts.end(); it++) { - auto &lbl = *it; - label_count_vec.emplace_back(std::make_pair(lbl.first, lbl.second)); - if (lbl.second > max_points) { - max_points = lbl.second; - label_with_max_points = lbl.first; + for (auto it = label_counts.begin(); it != label_counts.end(); it++) + { + auto &lbl = *it; + label_count_vec.emplace_back(std::make_pair(lbl.first, lbl.second)); + if (lbl.second > max_points) + { + max_points = lbl.second; + label_with_max_points = lbl.first; + } + sum += lbl.second; } - sum += lbl.second; - } - sort(label_count_vec.begin(), label_count_vec.end(), - [](const std::pair &lhs, - const std::pair &rhs) { - return lhs.second < rhs.second; - }); + sort(label_count_vec.begin(), label_count_vec.end(), + [](const std::pair &lhs, const std::pair &rhs) { + return lhs.second < rhs.second; + }); - for (float p = 0; p < 1; p += 0.05) { - std::cout << "Percentile " << (100 * p) << "\t" - << label_count_vec[(size_t)(p * label_count_vec.size())].first - << " with count=" - << label_count_vec[(size_t)(p * label_count_vec.size())].second - << std::endl; - } + for (float p = 0; p < 1; p += 0.05) + { + std::cout << "Percentile " << (100 * p) << "\t" << label_count_vec[(size_t)(p * label_count_vec.size())].first + << " with count=" << label_count_vec[(size_t)(p * label_count_vec.size())].second << std::endl; + } - std::cout << "Most common label " - << "\t" << label_count_vec[label_count_vec.size() - 1].first - << " with count=" - << label_count_vec[label_count_vec.size() - 1].second << std::endl; - if (label_count_vec.size() > 1) - std::cout << "Second common label " - << "\t" << label_count_vec[label_count_vec.size() - 2].first - << " with count=" - << label_count_vec[label_count_vec.size() - 2].second + std::cout << "Most common label " + << "\t" << label_count_vec[label_count_vec.size() - 1].first + << " with count=" << label_count_vec[label_count_vec.size() - 1].second << std::endl; + if (label_count_vec.size() > 1) + std::cout << "Second common label " + << "\t" << label_count_vec[label_count_vec.size() - 2].first + << " with count=" << label_count_vec[label_count_vec.size() - 2].second << std::endl; + if (label_count_vec.size() > 2) + std::cout << "Third common label " + << "\t" << label_count_vec[label_count_vec.size() - 3].first + << " with count=" << label_count_vec[label_count_vec.size() - 3].second << std::endl; + avg_labels_per_pt = (sum) / (float)point_cnt; + mean_label_size = (sum) / label_counts.size(); + std::cout << "Total number of points = " << point_cnt << ", number of labels = " << label_counts.size() << std::endl; - if (label_count_vec.size() > 2) - std::cout << "Third common label " - << "\t" << label_count_vec[label_count_vec.size() - 3].first - << " with count=" - << label_count_vec[label_count_vec.size() - 3].second - << std::endl; - avg_labels_per_pt = (sum) / (float)point_cnt; - mean_label_size = (sum) / label_counts.size(); - std::cout << "Total number of points = " << point_cnt - << ", number of labels = " << label_counts.size() << std::endl; - std::cout << "Average number of labels per point = " << avg_labels_per_pt - << std::endl; - std::cout << "Mean label size excluding 0 = " << mean_label_size << std::endl; - std::cout << "Most popular label is " << label_with_max_points << " with " - << max_points << " pts" << std::endl; + std::cout << "Average number of labels per point = " << avg_labels_per_pt << std::endl; + std::cout << "Mean label size excluding 0 = " << mean_label_size << std::endl; + std::cout << "Most popular label is " << label_with_max_points << " with " << max_points << " pts" << std::endl; } -int main(int argc, char **argv) { - std::string labels_file, universal_label; - uint32_t density; +int main(int argc, char **argv) +{ + std::string labels_file, universal_label; + uint32_t density; - po::options_description desc{"Arguments"}; - try { - desc.add_options()("help,h", "Print information on arguments"); - desc.add_options()("labels_file", - po::value(&labels_file)->required(), - "path to labels data file."); - desc.add_options()("universal_label", - po::value(&universal_label)->required(), - "Universal label used in labels file."); - desc.add_options()( - "density", po::value(&density)->default_value(1), - "Number of labels each point in labels file, defaults to 1"); - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, desc), vm); - if (vm.count("help")) { - std::cout << desc; - return 0; + po::options_description desc{"Arguments"}; + try + { + desc.add_options()("help,h", "Print information on arguments"); + desc.add_options()("labels_file", po::value(&labels_file)->required(), + "path to labels data file."); + desc.add_options()("universal_label", po::value(&universal_label)->required(), + "Universal label used in labels file."); + desc.add_options()("density", po::value(&density)->default_value(1), + "Number of labels each point in labels file, defaults to 1"); + po::variables_map vm; + po::store(po::parse_command_line(argc, argv, desc), vm); + if (vm.count("help")) + { + std::cout << desc; + return 0; + } + po::notify(vm); + } + catch (const std::exception &e) + { + std::cerr << e.what() << '\n'; + return -1; } - po::notify(vm); - } catch (const std::exception &e) { - std::cerr << e.what() << '\n'; - return -1; - } - stats_analysis(labels_file, universal_label, density); + stats_analysis(labels_file, universal_label, density); } diff --git a/tests/utils/tsv_to_bin.cpp b/tests/utils/tsv_to_bin.cpp index 43d304f9e..c590a8f73 100644 --- a/tests/utils/tsv_to_bin.cpp +++ b/tests/utils/tsv_to_bin.cpp @@ -4,105 +4,118 @@ #include #include "utils.h" -void block_convert_float(std::ifstream &reader, std::ofstream &writer, - size_t npts, size_t ndims) { - auto read_buf = new float[npts * (ndims + 1)]; - - auto cursor = read_buf; - float val; - - for (size_t i = 0; i < npts; i++) { - for (size_t d = 0; d < ndims; ++d) { - reader >> val; - *cursor = val; - cursor++; +void block_convert_float(std::ifstream &reader, std::ofstream &writer, size_t npts, size_t ndims) +{ + auto read_buf = new float[npts * (ndims + 1)]; + + auto cursor = read_buf; + float val; + + for (size_t i = 0; i < npts; i++) + { + for (size_t d = 0; d < ndims; ++d) + { + reader >> val; + *cursor = val; + cursor++; + } } - } - writer.write((char *)read_buf, npts * ndims * sizeof(float)); - delete[] read_buf; + writer.write((char *)read_buf, npts * ndims * sizeof(float)); + delete[] read_buf; } -void block_convert_int8(std::ifstream &reader, std::ofstream &writer, - size_t npts, size_t ndims) { - auto read_buf = new int8_t[npts * (ndims + 1)]; - - auto cursor = read_buf; - int val; - - for (size_t i = 0; i < npts; i++) { - for (size_t d = 0; d < ndims; ++d) { - reader >> val; - *cursor = (int8_t)val; - cursor++; +void block_convert_int8(std::ifstream &reader, std::ofstream &writer, size_t npts, size_t ndims) +{ + auto read_buf = new int8_t[npts * (ndims + 1)]; + + auto cursor = read_buf; + int val; + + for (size_t i = 0; i < npts; i++) + { + for (size_t d = 0; d < ndims; ++d) + { + reader >> val; + *cursor = (int8_t)val; + cursor++; + } } - } - writer.write((char *)read_buf, npts * ndims * sizeof(uint8_t)); - delete[] read_buf; + writer.write((char *)read_buf, npts * ndims * sizeof(uint8_t)); + delete[] read_buf; } -void block_convert_uint8(std::ifstream &reader, std::ofstream &writer, - size_t npts, size_t ndims) { - auto read_buf = new uint8_t[npts * (ndims + 1)]; +void block_convert_uint8(std::ifstream &reader, std::ofstream &writer, size_t npts, size_t ndims) +{ + auto read_buf = new uint8_t[npts * (ndims + 1)]; + + auto cursor = read_buf; + int val; + + for (size_t i = 0; i < npts; i++) + { + for (size_t d = 0; d < ndims; ++d) + { + reader >> val; + *cursor = (uint8_t)val; + cursor++; + } + } + writer.write((char *)read_buf, npts * ndims * sizeof(uint8_t)); + delete[] read_buf; +} - auto cursor = read_buf; - int val; +int main(int argc, char **argv) +{ + if (argc != 6) + { + std::cout << argv[0] + << " input_filename.tsv output_filename.bin " + "dim num_pts>" + << std::endl; + exit(-1); + } - for (size_t i = 0; i < npts; i++) { - for (size_t d = 0; d < ndims; ++d) { - reader >> val; - *cursor = (uint8_t)val; - cursor++; + 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; } - } - writer.write((char *)read_buf, npts * ndims * sizeof(uint8_t)); - delete[] read_buf; -} -int main(int argc, char **argv) { - if (argc != 6) { - std::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; - } - - size_t ndims = atoi(argv[4]); - size_t npts = atoi(argv[5]); - - std::ifstream reader(argv[2], std::ios::binary | std::ios::ate); - // size_t fsize = reader.tellg(); - reader.seekg(0, std::ios::beg); - reader.seekg(0, std::ios::beg); - - size_t blk_size = 131072; - size_t nblks = ROUND_UP(npts, blk_size) / blk_size; - std::cout << "# blks: " << nblks << std::endl; - std::ofstream writer(argv[3], std::ios::binary); - auto npts_u32 = (uint32_t)npts; - auto ndims_u32 = (uint32_t)ndims; - writer.write((char *)&npts_u32, sizeof(uint32_t)); - writer.write((char *)&ndims_u32, sizeof(uint32_t)); - - for (size_t i = 0; i < nblks; i++) { - size_t cblk_size = std::min(npts - i * blk_size, blk_size); - if (std::string(argv[1]) == std::string("float")) { - block_convert_float(reader, writer, cblk_size, ndims); - } else if (std::string(argv[1]) == std::string("int8")) { - block_convert_int8(reader, writer, cblk_size, ndims); - } else if (std::string(argv[1]) == std::string("uint8")) { - block_convert_uint8(reader, writer, cblk_size, ndims); + size_t ndims = atoi(argv[4]); + size_t npts = atoi(argv[5]); + + std::ifstream reader(argv[2], std::ios::binary | std::ios::ate); + // size_t fsize = reader.tellg(); + reader.seekg(0, std::ios::beg); + reader.seekg(0, std::ios::beg); + + size_t blk_size = 131072; + size_t nblks = ROUND_UP(npts, blk_size) / blk_size; + std::cout << "# blks: " << nblks << std::endl; + std::ofstream writer(argv[3], std::ios::binary); + auto npts_u32 = (uint32_t)npts; + auto ndims_u32 = (uint32_t)ndims; + writer.write((char *)&npts_u32, sizeof(uint32_t)); + writer.write((char *)&ndims_u32, sizeof(uint32_t)); + + for (size_t i = 0; i < nblks; i++) + { + size_t cblk_size = std::min(npts - i * blk_size, blk_size); + if (std::string(argv[1]) == std::string("float")) + { + block_convert_float(reader, writer, cblk_size, ndims); + } + else if (std::string(argv[1]) == std::string("int8")) + { + block_convert_int8(reader, writer, cblk_size, ndims); + } + else if (std::string(argv[1]) == std::string("uint8")) + { + block_convert_uint8(reader, writer, cblk_size, ndims); + } + std::cout << "Block #" << i << " written" << std::endl; } - std::cout << "Block #" << i << " written" << std::endl; - } - reader.close(); - writer.close(); + reader.close(); + writer.close(); } diff --git a/tests/utils/uint32_to_uint8.cpp b/tests/utils/uint32_to_uint8.cpp index 3b42e77fd..87b6fb8ed 100644 --- a/tests/utils/uint32_to_uint8.cpp +++ b/tests/utils/uint32_to_uint8.cpp @@ -4,18 +4,20 @@ #include #include "utils.h" -int main(int argc, char **argv) { - if (argc != 3) { - std::cout << argv[0] << " input_uint32_bin output_int8_bin" << std::endl; - exit(-1); - } +int main(int argc, char **argv) +{ + if (argc != 3) + { + std::cout << argv[0] << " input_uint32_bin output_int8_bin" << std::endl; + exit(-1); + } - uint32_t *input; - size_t npts, nd; - diskann::load_bin(argv[1], input, npts, nd); - uint8_t *output = new uint8_t[npts * nd]; - diskann::convert_types(input, output, npts, nd); - diskann::save_bin(argv[2], output, npts, nd); - delete[] output; - delete[] input; + uint32_t *input; + size_t npts, nd; + diskann::load_bin(argv[1], input, npts, nd); + uint8_t *output = new uint8_t[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/uint8_to_float.cpp b/tests/utils/uint8_to_float.cpp index 52b4ca083..6415b7c92 100644 --- a/tests/utils/uint8_to_float.cpp +++ b/tests/utils/uint8_to_float.cpp @@ -4,18 +4,20 @@ #include #include "utils.h" -int main(int argc, char **argv) { - if (argc != 3) { - std::cout << argv[0] << " input_uint8_bin output_float_bin" << std::endl; - exit(-1); - } +int main(int argc, char **argv) +{ + if (argc != 3) + { + std::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; + 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/vector_analysis.cpp b/tests/utils/vector_analysis.cpp index 9427a4689..5e4cb9bf4 100644 --- a/tests/utils/vector_analysis.cpp +++ b/tests/utils/vector_analysis.cpp @@ -20,129 +20,144 @@ #include "partition.h" #include "utils.h" -template -int analyze_norm(std::string base_file) { - std::cout << "Analyzing data norms" << std::endl; - T *data; - size_t npts, ndims; - diskann::load_bin(base_file, data, npts, ndims); - std::vector norms(npts, 0); +template int analyze_norm(std::string base_file) +{ + std::cout << "Analyzing data norms" << std::endl; + T *data; + size_t npts, ndims; + diskann::load_bin(base_file, data, npts, ndims); + std::vector norms(npts, 0); #pragma omp parallel for schedule(dynamic) - for (int64_t i = 0; i < (int64_t)npts; i++) { - for (size_t d = 0; d < ndims; d++) - norms[i] += data[i * ndims + d] * data[i * ndims + d]; - norms[i] = std::sqrt(norms[i]); - } - std::sort(norms.begin(), norms.end()); - for (int p = 0; p < 100; p += 5) - std::cout << "percentile " << p << ": " - << norms[std::floor((p / 100.0) * npts)] << std::endl; - std::cout << "percentile 100" - << ": " << norms[npts - 1] << std::endl; - delete[] data; - return 0; + for (int64_t i = 0; i < (int64_t)npts; i++) + { + for (size_t d = 0; d < ndims; d++) + norms[i] += data[i * ndims + d] * data[i * ndims + d]; + norms[i] = std::sqrt(norms[i]); + } + std::sort(norms.begin(), norms.end()); + for (int p = 0; p < 100; p += 5) + std::cout << "percentile " << p << ": " << norms[std::floor((p / 100.0) * npts)] << std::endl; + std::cout << "percentile 100" + << ": " << norms[npts - 1] << std::endl; + delete[] data; + return 0; } -template -int normalize_base(std::string base_file, std::string out_file) { - std::cout << "Normalizing base" << std::endl; - T *data; - size_t npts, ndims; - diskann::load_bin(base_file, data, npts, ndims); - // std::vector norms(npts, 0); +template int normalize_base(std::string base_file, std::string out_file) +{ + std::cout << "Normalizing base" << std::endl; + T *data; + size_t npts, ndims; + diskann::load_bin(base_file, data, npts, ndims); + // std::vector norms(npts, 0); #pragma omp parallel for schedule(dynamic) - for (int64_t i = 0; i < (int64_t)npts; i++) { - float pt_norm = 0; - for (size_t d = 0; d < ndims; d++) - pt_norm += data[i * ndims + d] * data[i * ndims + d]; - pt_norm = std::sqrt(pt_norm); - for (size_t d = 0; d < ndims; d++) - data[i * ndims + d] = data[i * ndims + d] / pt_norm; - } - diskann::save_bin(out_file, data, npts, ndims); - delete[] data; - return 0; + for (int64_t i = 0; i < (int64_t)npts; i++) + { + float pt_norm = 0; + for (size_t d = 0; d < ndims; d++) + pt_norm += data[i * ndims + d] * data[i * ndims + d]; + pt_norm = std::sqrt(pt_norm); + for (size_t d = 0; d < ndims; d++) + data[i * ndims + d] = data[i * ndims + d] / pt_norm; + } + diskann::save_bin(out_file, data, npts, ndims); + delete[] data; + return 0; } -template -int augment_base(std::string base_file, std::string out_file, - bool prep_base = true) { - std::cout << "Analyzing data norms" << std::endl; - T *data; - size_t npts, ndims; - diskann::load_bin(base_file, data, npts, ndims); - std::vector norms(npts, 0); - float max_norm = 0; +template int augment_base(std::string base_file, std::string out_file, bool prep_base = true) +{ + std::cout << "Analyzing data norms" << std::endl; + T *data; + size_t npts, ndims; + diskann::load_bin(base_file, data, npts, ndims); + std::vector norms(npts, 0); + float max_norm = 0; #pragma omp parallel for schedule(dynamic) - for (int64_t i = 0; i < (int64_t)npts; i++) { - for (size_t d = 0; d < ndims; d++) - norms[i] += data[i * ndims + d] * data[i * ndims + d]; - max_norm = norms[i] > max_norm ? norms[i] : max_norm; - } - // std::sort(norms.begin(), norms.end()); - max_norm = std::sqrt(max_norm); - std::cout << "Max norm: " << max_norm << std::endl; - T *new_data; - size_t newdims = ndims + 1; - new_data = new T[npts * newdims]; - for (size_t i = 0; i < npts; i++) { - if (prep_base) { - for (size_t j = 0; j < ndims; j++) { - new_data[i * newdims + j] = data[i * ndims + j] / max_norm; - } - float diff = 1 - (norms[i] / (max_norm * max_norm)); - diff = diff <= 0 ? 0 : std::sqrt(diff); - new_data[i * newdims + ndims] = diff; - if (diff <= 0) { - std::cout << i << " has large max norm, investigate if needed. diff = " - << diff << std::endl; - } - } else { - for (size_t j = 0; j < ndims; j++) { - new_data[i * newdims + j] = data[i * ndims + j] / std::sqrt(norms[i]); - } - new_data[i * newdims + ndims] = 0; + for (int64_t i = 0; i < (int64_t)npts; i++) + { + for (size_t d = 0; d < ndims; d++) + norms[i] += data[i * ndims + d] * data[i * ndims + d]; + max_norm = norms[i] > max_norm ? norms[i] : max_norm; } - } - diskann::save_bin(out_file, new_data, npts, newdims); - delete[] new_data; - delete[] data; - return 0; + // std::sort(norms.begin(), norms.end()); + max_norm = std::sqrt(max_norm); + std::cout << "Max norm: " << max_norm << std::endl; + T *new_data; + size_t newdims = ndims + 1; + new_data = new T[npts * newdims]; + for (size_t i = 0; i < npts; i++) + { + if (prep_base) + { + for (size_t j = 0; j < ndims; j++) + { + new_data[i * newdims + j] = data[i * ndims + j] / max_norm; + } + float diff = 1 - (norms[i] / (max_norm * max_norm)); + diff = diff <= 0 ? 0 : std::sqrt(diff); + new_data[i * newdims + ndims] = diff; + if (diff <= 0) + { + std::cout << i << " has large max norm, investigate if needed. diff = " << diff << std::endl; + } + } + else + { + for (size_t j = 0; j < ndims; j++) + { + new_data[i * newdims + j] = data[i * ndims + j] / std::sqrt(norms[i]); + } + new_data[i * newdims + ndims] = 0; + } + } + diskann::save_bin(out_file, new_data, npts, newdims); + delete[] new_data; + delete[] data; + return 0; } -template -int aux_main(char **argv) { - std::string base_file(argv[2]); - uint32_t option = atoi(argv[3]); - if (option == 1) - analyze_norm(base_file); - else if (option == 2) - augment_base(base_file, std::string(argv[4]), true); - else if (option == 3) - augment_base(base_file, std::string(argv[4]), false); - else if (option == 4) - normalize_base(base_file, std::string(argv[4])); - return 0; +template int aux_main(char **argv) +{ + std::string base_file(argv[2]); + uint32_t option = atoi(argv[3]); + if (option == 1) + analyze_norm(base_file); + else if (option == 2) + augment_base(base_file, std::string(argv[4]), true); + else if (option == 3) + augment_base(base_file, std::string(argv[4]), false); + else if (option == 4) + normalize_base(base_file, std::string(argv[4])); + return 0; } -int main(int argc, char **argv) { - if (argc < 4) { - std::cout << argv[0] - << " data_type [float/int8/uint8] base_bin_file " - "[option: 1-norm analysis, 2-prep_base_for_mip, " - "3-prep_query_for_mip, 4-normalize-vecs] [out_file for " - "options 2/3/4]" - << std::endl; - exit(-1); - } +int main(int argc, char **argv) +{ + if (argc < 4) + { + std::cout << argv[0] + << " data_type [float/int8/uint8] base_bin_file " + "[option: 1-norm analysis, 2-prep_base_for_mip, " + "3-prep_query_for_mip, 4-normalize-vecs] [out_file for " + "options 2/3/4]" + << std::endl; + exit(-1); + } - if (std::string(argv[1]) == std::string("float")) { - aux_main(argv); - } else if (std::string(argv[1]) == std::string("int8")) { - aux_main(argv); - } else if (std::string(argv[1]) == std::string("uint8")) { - aux_main(argv); - } else - std::cout << "Unsupported type. Use float/int8/uint8." << std::endl; - return 0; + if (std::string(argv[1]) == std::string("float")) + { + aux_main(argv); + } + else if (std::string(argv[1]) == std::string("int8")) + { + aux_main(argv); + } + else if (std::string(argv[1]) == std::string("uint8")) + { + aux_main(argv); + } + else + std::cout << "Unsupported type. Use float/int8/uint8." << std::endl; + return 0; } From ca6080e4d8219da6c579600456645258a3cc0062 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Thu, 20 Apr 2023 11:23:34 +0530 Subject: [PATCH 61/69] small change --- src/index.cpp | 4 ---- 1 file changed, 4 deletions(-) diff --git a/src/index.cpp b/src/index.cpp index 58d82c57d..097aae13b 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -1257,10 +1257,6 @@ void Index::prune_neighbors(const uint32_t location, std::vecto pruned_list.clear(); pruned_list.reserve(range); - if (pool.begin()->distance == 0) - { - diskann::cerr << "Warning: a candidate with distance 0 found in prune_neighbors" << std::endl; - } occlude_list(location, pool, alpha, range, max_candidate_size, pruned_list, scratch); assert(pruned_list.size() <= range); From f46f2bdfbf09010d796e1115bbb757ff773cce7c Mon Sep 17 00:00:00 2001 From: Gopal Srinivasa Date: Thu, 20 Apr 2023 11:45:45 +0530 Subject: [PATCH 62/69] Removed slot_manager from this PR --- include/slot_manager.h | 3 --- src/dll/CMakeLists.txt | 2 +- src/slot_manager.cpp | 1 - 3 files changed, 1 insertion(+), 5 deletions(-) delete mode 100644 include/slot_manager.h delete mode 100644 src/slot_manager.cpp diff --git a/include/slot_manager.h b/include/slot_manager.h deleted file mode 100644 index 831382462..000000000 --- a/include/slot_manager.h +++ /dev/null @@ -1,3 +0,0 @@ -#pragma once - -#include \ No newline at end of file diff --git a/src/dll/CMakeLists.txt b/src/dll/CMakeLists.txt index 9e97a35b2..e02996f32 100644 --- a/src/dll/CMakeLists.txt +++ b/src/dll/CMakeLists.txt @@ -4,7 +4,7 @@ add_library(${PROJECT_NAME} SHARED dllmain.cpp ../abstract_data_store.cpp ../partition.cpp ../pq.cpp ../pq_flash_index.cpp ../logger.cpp ../utils.cpp ../windows_aligned_file_reader.cpp ../distance.cpp ../memory_mapper.cpp ../index.cpp ../in_mem_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 ../slot_manager.cpp) + ../ann_exception.cpp ../natural_number_set.cpp ../natural_number_map.cpp ../scratch.cpp) set(TARGET_DIR "$<$:${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}>$<$:${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}>") set(DISKANN_DLL_IMPLIB "${TARGET_DIR}/${PROJECT_NAME}.lib") diff --git a/src/slot_manager.cpp b/src/slot_manager.cpp deleted file mode 100644 index e4dc4a48c..000000000 --- a/src/slot_manager.cpp +++ /dev/null @@ -1 +0,0 @@ -#include "slot_manager.h" \ No newline at end of file From 96c5dc2a9c1b87d1644c9ca6fabf5650ee185f07 Mon Sep 17 00:00:00 2001 From: harsha vardhan simhadri Date: Wed, 19 Apr 2023 15:39:17 -0700 Subject: [PATCH 63/69] newline at EOF in_mem_Graph_store.cpp --- src/in_mem_graph_store.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/in_mem_graph_store.cpp b/src/in_mem_graph_store.cpp index cb83481d3..b14b5fd17 100644 --- a/src/in_mem_graph_store.cpp +++ b/src/in_mem_graph_store.cpp @@ -26,4 +26,4 @@ void InMemGraphStore::set_adj_list(const location_t i, std::vector & { } -} // namespace diskann \ No newline at end of file +} // namespace diskann From d9fc9156335468b074dd229e02084dbec0e31234 Mon Sep 17 00:00:00 2001 From: harsha vardhan simhadri Date: Wed, 19 Apr 2023 15:41:16 -0700 Subject: [PATCH 64/69] remove deleted code in scratch.cpp --- src/scratch.cpp | 5 ----- 1 file changed, 5 deletions(-) diff --git a/src/scratch.cpp b/src/scratch.cpp index 02e594560..79a4b0b00 100644 --- a/src/scratch.cpp +++ b/src/scratch.cpp @@ -24,11 +24,6 @@ InMemQueryScratch::InMemQueryScratch(uint32_t search_l, uint32_t indexing_l, throw diskann::ANNException(ss.str(), -1); } - // REFACTOR - // auto aligned_dim = ROUND_UP(dim, 8); - // alloc_aligned(((void **)&_aligned_query), aligned_dim * sizeof(T), 8 * - // sizeof(T)); memset(_aligned_query, 0, aligned_dim * sizeof(T)); - alloc_aligned(((void **)&_aligned_query), aligned_dim * sizeof(T), alignment_factor * sizeof(T)); memset(_aligned_query, 0, aligned_dim * sizeof(T)); From d2f3bd7ab80d23caf04fcadc7444e71d72500e25 Mon Sep 17 00:00:00 2001 From: harsha vardhan simhadri Date: Sat, 22 Apr 2023 17:53:45 -0700 Subject: [PATCH 65/69] rename distance_metric to distance_fn --- include/in_mem_data_store.h | 2 +- src/in_mem_data_store.cpp | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/include/in_mem_data_store.h b/include/in_mem_data_store.h index f86aa0edc..70d1fa28f 100644 --- a/include/in_mem_data_store.h +++ b/include/in_mem_data_store.h @@ -20,7 +20,7 @@ namespace diskann template class InMemDataStore : public AbstractDataStore { public: - InMemDataStore(const location_t capacity, const size_t dim, std::shared_ptr> distance_metric); + InMemDataStore(const location_t capacity, const size_t dim, std::shared_ptr> distance_fn); virtual ~InMemDataStore(); virtual location_t load(const std::string &filename) override; diff --git a/src/in_mem_data_store.cpp b/src/in_mem_data_store.cpp index 3cc386c4d..fa177bc8f 100644 --- a/src/in_mem_data_store.cpp +++ b/src/in_mem_data_store.cpp @@ -11,8 +11,8 @@ namespace diskann template InMemDataStore::InMemDataStore(const location_t num_points, const size_t dim, - std::shared_ptr> distance_metric) - : AbstractDataStore(num_points, dim), _distance_fn(distance_metric) + std::shared_ptr> distance_fn) + : AbstractDataStore(num_points, dim), _distance_fn(distance_fn) { _aligned_dim = ROUND_UP(dim, _distance_fn->get_required_alignment()); alloc_aligned(((void **)&_data), this->_capacity * _aligned_dim * sizeof(data_t), 8 * sizeof(data_t)); From 4da65dbfbf14ee97deb55328c4b5e3664f9d5722 Mon Sep 17 00:00:00 2001 From: yashpatel007 Date: Mon, 24 Apr 2023 16:07:45 -0400 Subject: [PATCH 66/69] resolving PR comments --- include/utils.h | 2 +- src/in_mem_data_store.cpp | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/include/utils.h b/include/utils.h index 1b9cce924..f53a5e1e8 100644 --- a/include/utils.h +++ b/include/utils.h @@ -1017,7 +1017,7 @@ template inline float get_norm(T *arr, const size_t dim) } // This function is valid only for float data type. -template inline void normalize(T *arr, const size_t dim) +template inline void normalize(T *arr, const size_t dim) { float norm = get_norm(arr, dim); for (uint32_t i = 0; i < dim; i++) diff --git a/src/in_mem_data_store.cpp b/src/in_mem_data_store.cpp index fa177bc8f..6fc901c7c 100644 --- a/src/in_mem_data_store.cpp +++ b/src/in_mem_data_store.cpp @@ -108,9 +108,9 @@ template size_t InMemDataStore::save(const std::string template void InMemDataStore::populate_data(const data_t *vectors, const location_t num_pts) { + memset(_data, 0, _aligned_dim * sizeof(data_t) * num_pts); for (location_t i = 0; i < num_pts; i++) { - memset(_data + i * _aligned_dim, 0, _aligned_dim * sizeof(data_t)); std::memmove(_data + i * _aligned_dim, vectors + i * this->_dim, this->_dim * sizeof(data_t)); } From c9db2de7c9ff7dad12599ed06ea086da7a9d573f Mon Sep 17 00:00:00 2001 From: ravishankar Date: Tue, 25 Apr 2023 11:06:25 +0530 Subject: [PATCH 67/69] minor bug fix for initialization --- src/index.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/index.cpp b/src/index.cpp index 097aae13b..5a3788f7c 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -1540,9 +1540,9 @@ void Index::set_start_points(const T *data, size_t data_count) // memcpy(_data + _aligned_dim * _max_points, data, _aligned_dim * // sizeof(T) * _num_frozen_pts); - for (location_t i = _max_points; i < _max_points + _num_frozen_pts; i++) + for (location_t i = 0; i < _num_frozen_pts; i++) { - _data_store->set_vector(i, data + i * _dim); + _data_store->set_vector(i + _max_points, data + i * _dim); } _has_built = true; diskann::cout << "Index start points set: #" << _num_frozen_pts << std::endl; From 381e95afefb66e51610dd73d6bf26facb23a047a Mon Sep 17 00:00:00 2001 From: yashpatel007 Date: Wed, 26 Apr 2023 11:31:36 -0400 Subject: [PATCH 68/69] clang format fix --- src/in_mem_data_store.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/in_mem_data_store.cpp b/src/in_mem_data_store.cpp index 1dd888aed..b2f708263 100644 --- a/src/in_mem_data_store.cpp +++ b/src/in_mem_data_store.cpp @@ -193,7 +193,8 @@ void InMemDataStore::get_distance(const data_t *query, const location_t template float InMemDataStore::get_distance(const location_t loc1, const location_t loc2) const { - return _distance_fn->compare(_data + loc1 * _aligned_dim, _data + loc2 * _aligned_dim, (uint32_t)this->_aligned_dim); + return _distance_fn->compare(_data + loc1 * _aligned_dim, _data + loc2 * _aligned_dim, + (uint32_t)this->_aligned_dim); } template location_t InMemDataStore::expand(const location_t new_size) From 808c9ab1c8503b98e16cda7a304af25add4ffa0c Mon Sep 17 00:00:00 2001 From: yashpatel007 Date: Wed, 26 Apr 2023 13:33:36 -0400 Subject: [PATCH 69/69] trying to resolve ubuntu build errors --- src/in_mem_graph_store.cpp | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/in_mem_graph_store.cpp b/src/in_mem_graph_store.cpp index b14b5fd17..e9bfd4e9e 100644 --- a/src/in_mem_graph_store.cpp +++ b/src/in_mem_graph_store.cpp @@ -13,9 +13,11 @@ InMemGraphStore::InMemGraphStore(const size_t max_pts) : AbstractGraphStore(max_ int InMemGraphStore::load(const std::string &index_path_prefix) { + return 0; } int InMemGraphStore::store(const std::string &index_path_prefix) { + return 0; } void InMemGraphStore::get_adj_list(const location_t i, std::vector &neighbors)