From a50ce3df11e0235099c4125354e28dbfc729dc2b Mon Sep 17 00:00:00 2001 From: ravishankar Date: Thu, 27 May 2021 20:40:42 +0530 Subject: [PATCH 01/84] working towards inner product in memory indices --- tests/build_memory_index.cpp | 39 +++++++++++++++++++++----------- tests/utils/CMakeLists.txt | 2 +- tests/utils/gen_random_slice.cpp | 14 +++++++----- 3 files changed, 35 insertions(+), 20 deletions(-) diff --git a/tests/build_memory_index.cpp b/tests/build_memory_index.cpp index a5a13595e3..7bbe992c82 100644 --- a/tests/build_memory_index.cpp +++ b/tests/build_memory_index.cpp @@ -16,7 +16,7 @@ #include "memory_mapper.h" template -int build_in_memory_index(const std::string& data_path, const unsigned R, +int build_in_memory_index(const std::string& data_path, _u32 dist_fn, const unsigned R, const unsigned L, const float alpha, const std::string& save_path, const unsigned num_threads) { @@ -29,7 +29,17 @@ int build_in_memory_index(const std::string& data_path, const unsigned R, paras.Set("saturate_graph", 0); paras.Set("num_threads", num_threads); - diskann::Index index(diskann::L2, data_path.c_str()); + diskann::Metric metric; + if (dist_fn == 0) + metric = diskann::L2; + else if (dist_fn == 1) + metric = diskann::INNER_PRODUCT; + else { + std::cout<<"Error. Unsupported distance type. Exitting" << std::endl; + return -1; + } + + diskann::Index index(metric, data_path.c_str()); auto s = std::chrono::high_resolution_clock::now(); index.build(paras); std::chrono::duration diff = @@ -42,9 +52,9 @@ int build_in_memory_index(const std::string& data_path, const unsigned R, } int main(int argc, char** argv) { - if (argc != 8) { + if (argc != 9) { std::cout << "Usage: " << argv[0] - << " [data_type] [data_file.bin] " + << " [data_type] [dist_fn 0 for L2, 1 for inner product] [data_file.bin] " "[output_index_file] " << "[R] [L] [alpha]" << " [num_threads_to_use]. See README for more information on " @@ -53,21 +63,24 @@ int main(int argc, char** argv) { exit(-1); } - const std::string data_path(argv[2]); - const std::string save_path(argv[3]); - const unsigned R = (unsigned) atoi(argv[4]); - const unsigned L = (unsigned) atoi(argv[5]); - const float alpha = (float) atof(argv[6]); - const unsigned num_threads = (unsigned) atoi(argv[7]); + _u32 ctr = 2; + + _u32 dist_fn = (_u32) atoi(argv[ctr++]); + const std::string data_path(argv[ctr++]); + const std::string save_path(argv[ctr++]); + const unsigned R = (unsigned) atoi(argv[ctr++]); + const unsigned L = (unsigned) atoi(argv[ctr++]); + const float alpha = (float) atof(argv[ctr++]); + const unsigned num_threads = (unsigned) atoi(argv[ctr++]); if (std::string(argv[1]) == std::string("int8")) - build_in_memory_index(data_path, R, L, alpha, save_path, + build_in_memory_index(data_path, dist_fn, R, L, alpha, save_path, num_threads); else if (std::string(argv[1]) == std::string("uint8")) - build_in_memory_index(data_path, R, L, alpha, save_path, + build_in_memory_index(data_path, dist_fn, R, L, alpha, save_path, num_threads); else if (std::string(argv[1]) == std::string("float")) - build_in_memory_index(data_path, R, L, alpha, save_path, + build_in_memory_index(data_path, dist_fn, R, L, alpha, save_path, num_threads); else std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; diff --git a/tests/utils/CMakeLists.txt b/tests/utils/CMakeLists.txt index 6f7cb32b8e..164fcffcfd 100644 --- a/tests/utils/CMakeLists.txt +++ b/tests/utils/CMakeLists.txt @@ -122,6 +122,6 @@ endif() # formatter if (LINUX) - add_custom_command(TARGET gen_random_slice PRE_BUILD COMMAND clang-format-4.0 -i ../../../include/*.h ../../../include/dll/*.h ../../../src/*.cpp ../../../tests/*.cpp ../../../src/dll/*.cpp ../../../tests/utils/*.cpp) + add_custom_command(TARGET gen_random_slice PRE_BUILD COMMAND clang-format -i ../../../include/*.h ../../../include/dll/*.h ../../../src/*.cpp ../../../tests/*.cpp ../../../src/dll/*.cpp ../../../tests/utils/*.cpp) endif() diff --git a/tests/utils/gen_random_slice.cpp b/tests/utils/gen_random_slice.cpp index 0417c12c0d..1b102e27c1 100644 --- a/tests/utils/gen_random_slice.cpp +++ b/tests/utils/gen_random_slice.cpp @@ -22,12 +22,6 @@ template int aux_main(int argc, char** argv) { - if (argc != 5) { - std::cout << argv[0] << " data_type [fliat/int8/uint8] base_bin_file " - "sample_output_prefix sampling_probability" - << std::endl; - exit(-1); - } std::string base_file(argv[2]); std::string output_prefix(argv[3]); @@ -37,6 +31,14 @@ int aux_main(int argc, char** argv) { } 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(argc, argv); } else if (std::string(argv[1]) == std::string("int8")) { From 75b4567936f41bd146a36d5412322c2e60e98e1d Mon Sep 17 00:00:00 2001 From: ravishankar Date: Thu, 27 May 2021 21:59:50 +0530 Subject: [PATCH 02/84] done with in-memory code --- include/distance.h | 10 +++- include/index.h | 1 + src/index.cpp | 7 ++- tests/search_memory_index.cpp | 48 +++++++++------- tests/utils/compute_groundtruth.cpp | 89 ++++++++++++++++++++++++----- 5 files changed, 118 insertions(+), 37 deletions(-) diff --git a/include/distance.h b/include/distance.h index 3d403d1e56..c6b8c56d32 100644 --- a/include/distance.h +++ b/include/distance.h @@ -328,7 +328,7 @@ namespace diskann { template class DistanceInnerProduct : public Distance { public: - float compare(const T *a, const T *b, unsigned size) const { + float acompare(const T *a, const T *b, unsigned size) const { float result = 0; #ifdef __GNUC__ #ifdef __AVX__ @@ -426,10 +426,14 @@ namespace diskann { #endif return result; } + float compare(const T *a, const T *b, unsigned size) const { // since we use normally minimization objective for distance comparisons, we are returning 1/x. + float result = acompare(a,b,size); + return 1/result; + } }; template - class DistanceFastL2 : public DistanceInnerProduct { + class DistanceFastL2 : public DistanceInnerProduct { // currently defined only for float. templated for future use. public: float norm(const T *a, unsigned size) const { float result = 0; @@ -522,7 +526,7 @@ namespace diskann { using DistanceInnerProduct::compare; float compare(const T *a, const T *b, float norm, unsigned size) const { // not implement - float result = -2 * DistanceInnerProduct::compare(a, b, size); + float result = -2 * DistanceInnerProduct::acompare(a, b, size); result += norm; return result; } diff --git a/include/index.h b/include/index.h index eedbb1491f..c409b96958 100644 --- a/include/index.h +++ b/include/index.h @@ -167,6 +167,7 @@ namespace diskann { size_t consolidate_deletes(const Parameters ¶meters); private: + Metric _metric = diskann::L2; size_t _dim; size_t _aligned_dim; T * _data; diff --git a/src/index.cpp b/src/index.cpp index 2405878669..7123c661d4 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -61,6 +61,9 @@ namespace { std::cout << "Older CPU. Using slow distance computation" << std::endl; return new diskann::SlowDistanceL2Float(); } + } else if (m == diskann::Metric::INNER_PRODUCT) { + std::cout << "Using Inner Product computation" << std::endl; + return new diskann::DistanceInnerProduct(); } else { std::stringstream stream; stream << "Only L2 metric supported as of now. Email " @@ -129,7 +132,7 @@ namespace diskann { const size_t nd, const size_t num_frozen_pts, const bool enable_tags, const bool store_data, const bool support_eager_delete) - : _num_frozen_pts(num_frozen_pts), _has_built(false), _width(0), + : _metric(m), _num_frozen_pts(num_frozen_pts), _has_built(false), _width(0), _can_delete(false), _eager_done(true), _lazy_done(true), _compacted_order(true), _enable_tags(enable_tags), _consolidated_order(true), _support_eager_delete(support_eager_delete), @@ -1054,6 +1057,8 @@ namespace diskann { for (auto it : best_L_nodes) { indices[pos] = it.id; distances[pos] = it.distance; + if (_metric == diskann::INNER_PRODUCT) + distances[pos] = 1/distances[pos]; pos++; if (pos == K) break; diff --git a/tests/search_memory_index.cpp b/tests/search_memory_index.cpp index 32592710e5..6ad40032a4 100644 --- a/tests/search_memory_index.cpp +++ b/tests/search_memory_index.cpp @@ -27,26 +27,27 @@ int search_memory_index(int argc, char** argv) { size_t query_num, query_dim, query_aligned_dim, gt_num, gt_dim; std::vector<_u64> Lvec; - std::string data_file(argv[2]); - std::string memory_index_file(argv[3]); - _u64 num_threads = std::atoi(argv[4]); - std::string query_bin(argv[5]); - std::string truthset_bin(argv[6]); - _u64 recall_at = std::atoi(argv[7]); - std::string result_output_prefix(argv[8]); - bool use_optimized_search = std::atoi(argv[9]); + _u32 ctr = 2; + _u32 dist_fn = atoi(argv[ctr++]); + std::string data_file(argv[ctr++]); + std::string memory_index_file(argv[ctr++]); + _u64 num_threads = std::atoi(argv[ctr++]); + std::string query_bin(argv[ctr++]); + std::string truthset_bin(argv[ctr++]); + _u64 recall_at = std::atoi(argv[ctr++]); + std::string result_output_prefix(argv[ctr++]); +// bool use_optimized_search = std::atoi(argv[ctr++]); if ((std::string(argv[1]) != std::string("float")) && - (use_optimized_search == true)) { - std::cout << "Error. Optimized search currently only supported for " - "floating point datatypes. Using un-optimized search." + ((dist_fn == 1) || (dist_fn == 2))) { + std::cout << "Error. Inner product and Fast_L2 search currently only supported for " + "floating point datatypes." << std::endl; - use_optimized_search = false; } bool calc_recall_flag = false; - for (int ctr = 10; ctr < argc; ctr++) { + for (; ctr < (_u32) argc; ctr++) { _u64 curL = std::atoi(argv[ctr]); if (curL >= recall_at) Lvec.push_back(curL); @@ -73,14 +74,22 @@ int search_memory_index(int argc, char** argv) { std::cout.setf(std::ios_base::fixed, std::ios_base::floatfield); std::cout.precision(2); - auto metric = diskann::L2; - if (use_optimized_search) + diskann::Metric metric; + if (dist_fn == 0) + metric = diskann::L2; + else if (dist_fn == 1) + metric = diskann::INNER_PRODUCT; + else if(dist_fn == 2) metric = diskann::FAST_L2; + else { + std::cout<<"Error. Unsupported distance function. Exitting"; + return -1; + } diskann::Index index(metric, data_file.c_str()); index.load(memory_index_file.c_str()); // to load NSG std::cout << "Index loaded" << std::endl; - if (use_optimized_search) + if (metric == diskann::FAST_L2) index.optimize_graph(); diskann::Parameters paras; @@ -106,7 +115,7 @@ int search_memory_index(int argc, char** argv) { #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 (use_optimized_search) { + if (metric == diskann::FAST_L2) { index.search_with_opt_graph( query + i * query_aligned_dim, recall_at, L, query_result_ids[test_id].data() + i * recall_at); @@ -160,11 +169,10 @@ int main(int argc, char** argv) { if (argc < 11) { std::cout << "Usage: " << argv[0] - << " [index_type] [data_file.bin] " + << " [index_type] [dist_fn (0 for L2, 1 for Inner Product, 2 for Fast L2 for small datasets)] [data_file.bin] " "[memory_index_path] [num_threads] " "[query_file.bin] [truthset.bin (use \"null\" for none)] " - " [K] [result_output_prefix] [use_optimized_search (for small ~1M " - "data)] " + " [K] [result_output_prefix]" " [L1] [L2] etc. See README for more information on parameters. " << std::endl; exit(-1); diff --git a/tests/utils/compute_groundtruth.cpp b/tests/utils/compute_groundtruth.cpp index 8fef8c929c..cdd8d0c327 100644 --- a/tests/utils/compute_groundtruth.cpp +++ b/tests/utils/compute_groundtruth.cpp @@ -31,10 +31,10 @@ #define ALIGNMENT 512 void command_line_help() { - std::cerr - << " " - << std::endl; + std::cerr << " " + " " + << std::endl; } template @@ -104,6 +104,34 @@ void distsq_to_points( 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, + 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) -1.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 exact_knn(const size_t dim, const size_t k, int *const closest_points, // k * num_queries preallocated, col // major, queries columns @@ -112,14 +140,23 @@ void exact_knn(const size_t dim, const size_t k, // corresponding closes_points size_t npoints, const float *const points, // points in Col major - size_t nqueries, - const float *const queries) // queries in Col major + size_t nqueries, const float *const queries, + bool use_mip = false) // queries in Col major { float *points_l2sq = new float[npoints]; float *queries_l2sq = new float[nqueries]; 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 (use_mip) + std::cout << " inner product "; + 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]; @@ -128,9 +165,15 @@ void exact_knn(const size_t dim, const size_t k, int64_t q_e = ((b + 1) * q_batch_size > nqueries) ? nqueries : (b + 1) * q_batch_size; - 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); + if (!use_mip) { + 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, points_l2sq, q_e - q_b, + queries + (ptrdiff_t) q_b * (ptrdiff_t) dim, queries_l2sq + q_b); + } std::cout << "Computed distances for queries: [" << q_b << "," << q_e << ")" << std::endl; @@ -266,12 +309,12 @@ inline void save_groundtruth_as_one_file(const std::string filename, template int aux_main(int argv, char **argc) { - size_t npoints, nqueries, dim; std::string base_file(argc[2]); std::string query_file(argc[3]); size_t k = atoi(argc[4]); std::string gt_file(argc[5]); + bool use_mip = atoi(argc[6]); float *base_data; float *query_data; @@ -291,7 +334,7 @@ int aux_main(int argv, char **argc) { float *dist_closest_points_part = new float[nqueries * k]; exact_knn(dim, k, closest_points_part, dist_closest_points_part, npoints, - base_data, nqueries, query_data); + base_data, nqueries, query_data, use_mip); for (_u64 i = 0; i < nqueries; i++) { for (_u64 j = 0; j < k; j++) { @@ -303,6 +346,23 @@ int aux_main(int argv, char **argc) { delete[] closest_points_part; delete[] dist_closest_points_part; + + /* + std::cout << "For testing: doing brute force for one point" << + std::endl; std::vector> brute_force; for (_u32 i + = 0; i < npoints; i++) { float cur_pt_dist = 0; for (_u64 k = 0; k < dim; + k++) { cur_pt_dist += base_data[i * dim + k] * query_data[k]; + } + brute_force.push_back(std::make_pair(i, -cur_pt_dist)); + } + + std::sort(brute_force.begin(), brute_force.end(), custom_dist); + for (_u32 i = 0; i < 10; i++) { + std::cout< Date: Fri, 28 May 2021 17:11:13 +0530 Subject: [PATCH 03/84] made the inner product distance function return std::float_max if negative --- include/distance.h | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/include/distance.h b/include/distance.h index c6b8c56d32..b39a1719fb 100644 --- a/include/distance.h +++ b/include/distance.h @@ -428,7 +428,9 @@ namespace diskann { } float compare(const T *a, const T *b, unsigned size) const { // since we use normally minimization objective for distance comparisons, we are returning 1/x. float result = acompare(a,b,size); - return 1/result; + if (result < 0) + return std::numeric_limits::max(); + else return 1/result; } }; From 4fa5f9d9a90b61d1fae679cb5d67c6bda4c7d433 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Sat, 29 May 2021 11:27:44 +0530 Subject: [PATCH 04/84] more changes for disk index support --- include/partition_and_pq.h | 2 +- include/pq_flash_index.h | 3 ++- include/pq_table.h | 23 ++++++++++++++++++----- src/partition_and_pq.cpp | 10 +++++++--- src/pq_flash_index.cpp | 34 ++++++++++++++++++++++++++++------ 5 files changed, 56 insertions(+), 16 deletions(-) diff --git a/include/partition_and_pq.h b/include/partition_and_pq.h index 45b1e26e6d..43a9d84db0 100644 --- a/include/partition_and_pq.h +++ b/include/partition_and_pq.h @@ -54,7 +54,7 @@ DISKANN_DLLEXPORT int generate_pq_pivots(const float *train_data, unsigned num_centers, unsigned num_pq_chunks, unsigned max_k_means_reps, - std::string pq_pivots_path); + std::string pq_pivots_path, bool make_zero_mean = false); template int generate_pq_data_from_pivots(const std::string data_file, diff --git a/include/pq_flash_index.h b/include/pq_flash_index.h index c7601851c6..c20409fb64 100644 --- a/include/pq_flash_index.h +++ b/include/pq_flash_index.h @@ -70,7 +70,7 @@ namespace diskann { // Freeing the reader object is now the client's (DiskANNInterface's) // responsibility. DISKANN_DLLEXPORT PQFlashIndex( - std::shared_ptr &fileReader); + std::shared_ptr &fileReader, diskann::Metric metric = diskann::Metric::L2); DISKANN_DLLEXPORT ~PQFlashIndex(); #ifdef EXEC_ENV_OLS @@ -129,6 +129,7 @@ namespace diskann { // nbrs of node `i`: ((unsigned*)buf) + 1 _u64 max_node_len = 0, nnodes_per_sector = 0, max_degree = 0; + diskann::Metric metric = diskann::Metric::L2; // data info _u64 num_points = 0; _u64 data_dim = 0; diff --git a/include/pq_table.h b/include/pq_table.h index 3cac23c15a..c913271eb1 100644 --- a/include/pq_table.h +++ b/include/pq_table.h @@ -137,11 +137,6 @@ namespace diskann { _u64 permuted_dim_in_query = rearrangement[j]; const float* centers_dim_vec = tables_T + (256 * j); for (_u64 idx = 0; idx < 256; idx++) { - // Gopal. Fixing crash in v14 machines. - // float diff = centers_dim_vec[idx] - - // ((float) query_vec[permuted_dim_in_query] - - // centroid[permuted_dim_in_query]); - // chunk_dists[idx] += (diff * diff); double diff = centers_dim_vec[idx] - (query_vec[permuted_dim_in_query] - centroid[permuted_dim_in_query]); @@ -150,5 +145,23 @@ namespace diskann { } } } + void + populate_chunk_inner_products(const T* query_vec, float* dist_vec) { + memset(dist_vec, 0, 256 * n_chunks * sizeof(float)); + // chunk wise distance computation + for (_u64 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++) { + _u64 permuted_dim_in_query = rearrangement[j]; + const float* centers_dim_vec = tables_T + (256 * j); + for (_u64 idx = 0; idx < 256; idx++) { + double prod = + centers_dim_vec[idx] * query_vec[permuted_dim_in_query]; // 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) + } + } + } + } }; } // namespace diskann diff --git a/src/partition_and_pq.cpp b/src/partition_and_pq.cpp index 9da49b2203..ed6fcf9826 100644 --- a/src/partition_and_pq.cpp +++ b/src/partition_and_pq.cpp @@ -190,13 +190,14 @@ void gen_random_slice(const T *inputdata, size_t npts, size_t ndims, int generate_pq_pivots(const float *passed_train_data, size_t num_train, unsigned dim, unsigned num_centers, unsigned num_pq_chunks, unsigned max_k_means_reps, - std::string pq_pivots_path) { + 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, @@ -221,11 +222,14 @@ int generate_pq_pivots(const float *passed_train_data, size_t num_train, 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) { + for (uint64_t d = 0; d < dim; d++) { for (uint64_t p = 0; p < num_train; p++) { centroid[d] += train_data[p * dim + d]; } @@ -233,12 +237,12 @@ int generate_pq_pivots(const float *passed_train_data, size_t num_train, } // std::memset(centroid, 0 , dim*sizeof(float)); - 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 rearrangement; std::vector chunk_offsets; diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index 21da3521d6..a11379b7d1 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -86,8 +86,8 @@ namespace { namespace diskann { template<> PQFlashIndex<_u8>::PQFlashIndex( - std::shared_ptr &fileReader) - : reader(fileReader) { + std::shared_ptr &fileReader, diskann::Metric metric) + : reader(fileReader), metric(metric) { diskann::cout << "dist_cmp function for _u8 uses slow implementation." " Please contact gopalsr@microsoft.com if you need an AVX/AVX2" @@ -106,12 +106,16 @@ namespace diskann { << std::endl; this->dist_cmp_float = new SlowDistanceL2Float(); } + if (metric != diskann::Metric::L2) { + std::cout<<"Only L2 supported for byte vectors for now. Other distance functions are future work. Falling back to L2 distance." << std::endl; + this->metric = diskann::Metric::L2; + } } template<> PQFlashIndex<_s8>::PQFlashIndex( - std::shared_ptr &fileReader) - : reader(fileReader) { + std::shared_ptr &fileReader, diskann::Metric metric) + : reader(fileReader), metric(metric) { if (Avx2SupportedCPU) { diskann::cout << "Using AVX2 function for dist_cmp and dist_cmp_float" << std::endl; @@ -130,12 +134,18 @@ namespace diskann { this->dist_cmp = new SlowDistanceL2Int(); this->dist_cmp_float = new SlowDistanceL2Float(); } + if (metric != diskann::Metric::L2) { + std::cout<<"Only L2 supported for byte vectors for now. Other distance functions are future work. Falling back to L2 distance." << std::endl; + this->metric = diskann::Metric::L2; + } + } template<> PQFlashIndex::PQFlashIndex( - std::shared_ptr &fileReader) - : reader(fileReader) { + std::shared_ptr &fileReader, diskann::Metric metric) + : reader(fileReader), metric(metric) { + if (metric == diskann::Metric::L2) { if (Avx2SupportedCPU) { diskann::cout << "Using AVX2 functions for dist_cmp and dist_cmp_float" << std::endl; @@ -154,6 +164,15 @@ namespace diskann { this->dist_cmp = new AVXDistanceL2Float(); this->dist_cmp_float = new AVXDistanceL2Float(); } + } else if (metric == diskann::Metric::INNER_PRODUCT) { + this->dist_cmp = new DistanceInnerProduct(); + this->dist_cmp_float = new DistanceInnerProduct(); + } else { + std::cout<<"Unsupported metric type. Reverting to float." << std::endl; + this->dist_cmp = new AVXDistanceL2Float(); + this->dist_cmp_float = new AVXDistanceL2Float(); + this->metric = diskann::Metric::L2; + } } template @@ -801,6 +820,9 @@ namespace diskann { // query <-> PQ chunk centers distances float *pq_dists = query_scratch->aligned_pqtable_dist_scratch; + if (metric==diskann::Metric::INNER_PRODUCT) + pq_table.populate_chunk_inner_products(query, pq_dists); + else if (metric==diskann::Metric::L2) pq_table.populate_chunk_distances(query, pq_dists); // query <-> neighbor list From 661c38be4eda23beb8d5f318fe2e155da69a3b35 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Sat, 29 May 2021 13:41:07 +0530 Subject: [PATCH 05/84] on the way to disk index support for MIPS --- src/aux_utils.cpp | 32 +++++++++++++++++++------------- tests/build_disk_index.cpp | 23 +++++++++++++---------- tests/search_disk_index.cpp | 2 +- 3 files changed, 33 insertions(+), 24 deletions(-) diff --git a/src/aux_utils.cpp b/src/aux_utils.cpp index 6a2990ba66..13a1389ba1 100644 --- a/src/aux_utils.cpp +++ b/src/aux_utils.cpp @@ -343,7 +343,7 @@ namespace diskann { template int build_merged_vamana_index(std::string base_file, - diskann::Metric _compareMetric, unsigned L, + diskann::Metric compareMetric, unsigned L, unsigned R, double sampling_rate, double ram_budget, std::string mem_index_path, std::string medoids_file, @@ -367,7 +367,7 @@ namespace diskann { std::unique_ptr> _pvamanaIndex = std::unique_ptr>( - new diskann::Index(_compareMetric, base_file.c_str())); + new diskann::Index(compareMetric, base_file.c_str())); _pvamanaIndex->build(paras); _pvamanaIndex->save(mem_index_path.c_str()); std::remove(medoids_file.c_str()); @@ -399,7 +399,7 @@ namespace diskann { std::unique_ptr> _pvamanaIndex = std::unique_ptr>( - new diskann::Index(_compareMetric, shard_base_file.c_str())); + new diskann::Index(compareMetric, shard_base_file.c_str())); _pvamanaIndex->build(paras); _pvamanaIndex->save(shard_index_file.c_str()); } @@ -612,7 +612,7 @@ namespace diskann { template bool build_disk_index(const char *dataFilePath, const char *indexFilePath, const char * indexBuildParameters, - diskann::Metric _compareMetric) { + diskann::Metric compareMetric) { std::stringstream parser; parser << std::string(indexBuildParameters); std::string cur_param; @@ -631,6 +631,9 @@ namespace diskann { return false; } + if (compareMetric == diskann::Metric::INNER_PRODUCT) { + std::cout<<"Using Inner Product for PQ and Graph Generation" << std::endl; + } std::string index_prefix_path(indexFilePath); std::string pq_pivots_path = index_prefix_path + "_pq_pivots.bin"; std::string pq_compressed_vectors_path = @@ -697,9 +700,12 @@ namespace diskann { diskann::cout << "Training data loaded of size " << train_size << std::endl; - + + bool make_zero_mean = true; + if (compareMetric == diskann::Metric::INNER_PRODUCT) + make_zero_mean = false; generate_pq_pivots(train_data, train_size, (uint32_t) dim, 256, - (uint32_t) num_pq_chunks, 15, pq_pivots_path); + (uint32_t) num_pq_chunks, 15, pq_pivots_path, make_zero_mean); generate_pq_data_from_pivots(dataFilePath, 256, (uint32_t) num_pq_chunks, pq_pivots_path, @@ -710,7 +716,7 @@ namespace diskann { train_data = nullptr; diskann::build_merged_vamana_index( - dataFilePath, _compareMetric, L, R, p_val, indexing_ram_budget, + dataFilePath, compareMetric, L, R, p_val, indexing_ram_budget, mem_index_path, medoids_path, centroids_path); diskann::create_disk_layout(dataFilePath, mem_index_path, @@ -781,26 +787,26 @@ namespace diskann { template DISKANN_DLLEXPORT bool build_disk_index( const char *dataFilePath, const char *indexFilePath, - const char *indexBuildParameters, diskann::Metric _compareMetric); + const char *indexBuildParameters, diskann::Metric compareMetric); template DISKANN_DLLEXPORT bool build_disk_index( const char *dataFilePath, const char *indexFilePath, - const char *indexBuildParameters, diskann::Metric _compareMetric); + const char *indexBuildParameters, diskann::Metric compareMetric); template DISKANN_DLLEXPORT bool build_disk_index( const char *dataFilePath, const char *indexFilePath, - const char *indexBuildParameters, diskann::Metric _compareMetric); + const char *indexBuildParameters, diskann::Metric compareMetric); template DISKANN_DLLEXPORT int build_merged_vamana_index( - std::string base_file, diskann::Metric _compareMetric, unsigned L, + 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); template DISKANN_DLLEXPORT int build_merged_vamana_index( - std::string base_file, diskann::Metric _compareMetric, unsigned L, + 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); template DISKANN_DLLEXPORT int build_merged_vamana_index( - std::string base_file, diskann::Metric _compareMetric, unsigned L, + 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); diff --git a/tests/build_disk_index.cpp b/tests/build_disk_index.cpp index b29c3f0d1e..8c94d47525 100644 --- a/tests/build_disk_index.cpp +++ b/tests/build_disk_index.cpp @@ -11,29 +11,32 @@ template bool build_index(const char* dataFilePath, const char* indexFilePath, - const char* indexBuildParameters) { + const char* indexBuildParameters, diskann::Metric metric) { return diskann::build_disk_index( - dataFilePath, indexFilePath, indexBuildParameters, diskann::Metric::L2); + dataFilePath, indexFilePath, indexBuildParameters, metric); } int main(int argc, char** argv) { - if (argc != 9) { + if (argc != 10) { std::cout << "Usage: " << argv[0] - << " [data_type] [data_file.bin] " + << " [data_type] [dist_fn: 0 for L2, 1 for MIPS] [data_file.bin] " "[index_prefix_path] " "[R] [L] [B] [M] [T]. See README for more information on " "parameters." << std::endl; } else { - std::string params = std::string(argv[4]) + " " + std::string(argv[5]) + - " " + std::string(argv[6]) + " " + - std::string(argv[7]) + " " + std::string(argv[8]); + diskann::Metric metric = diskann::Metric::L2; + if (atoi(argv[2]) == 1) + metric = diskann::Metric::INNER_PRODUCT; + std::string params = std::string(argv[5]) + " " + std::string(argv[6]) + + " " + std::string(argv[7]) + " " + + std::string(argv[8]) + " " + std::string(argv[9]); if (std::string(argv[1]) == std::string("float")) - build_index(argv[2], argv[3], params.c_str()); + build_index(argv[3], argv[4], params.c_str(), metric); else if (std::string(argv[1]) == std::string("int8")) - build_index(argv[2], argv[3], params.c_str()); + build_index(argv[3], argv[4], params.c_str(), metric); else if (std::string(argv[1]) == std::string("uint8")) - build_index(argv[2], argv[3], params.c_str()); + build_index(argv[3], argv[4], params.c_str(), metric); else std::cout << "Error. wrong file type" << std::endl; } diff --git a/tests/search_disk_index.cpp b/tests/search_disk_index.cpp index 4d5c79ca89..7fa0edeec9 100644 --- a/tests/search_disk_index.cpp +++ b/tests/search_disk_index.cpp @@ -114,7 +114,7 @@ int search_disk_index(int argc, char** argv) { #endif std::unique_ptr> _pFlashIndex( - new diskann::PQFlashIndex(reader)); + new diskann::PQFlashIndex(reader, diskann::Metric::INNER_PRODUCT)); int res = _pFlashIndex->load(num_threads, pq_prefix.c_str(), disk_index_file.c_str()); From a3962d8dd704a8ae8d7ecc01282b90f7251cd06e Mon Sep 17 00:00:00 2001 From: ravishankar Date: Sat, 29 May 2021 16:38:21 +0530 Subject: [PATCH 06/84] works now, need to change the PQ generation for MIPS --- src/aux_utils.cpp | 4 ++-- src/partition_and_pq.cpp | 33 ++++++++++++++++++++++++++++++++- 2 files changed, 34 insertions(+), 3 deletions(-) diff --git a/src/aux_utils.cpp b/src/aux_utils.cpp index 13a1389ba1..b3cc400692 100644 --- a/src/aux_utils.cpp +++ b/src/aux_utils.cpp @@ -56,7 +56,7 @@ namespace diskann { } gt.insert(gt_vec, gt_vec + tie_breaker); - res.insert(res_vec, res_vec + recall_at); + res.insert(res_vec, 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) { if (res.find(v) != res.end()) { @@ -726,7 +726,7 @@ namespace diskann { gen_random_slice(dataFilePath, sample_base_prefix, sample_sampling_rate); - std::remove(mem_index_path.c_str()); + std::remove(mem_index_path.c_str()); auto e = std::chrono::high_resolution_clock::now(); diff --git a/src/partition_and_pq.cpp b/src/partition_and_pq.cpp index ed6fcf9826..be323ad8ff 100644 --- a/src/partition_and_pq.cpp +++ b/src/partition_and_pq.cpp @@ -36,6 +36,7 @@ #endif #define BLOCK_SIZE 5000000 +#define SAVE_INFLATED_PQ true template void gen_random_slice(const std::string base_file, @@ -380,6 +381,8 @@ int generate_pq_data_from_pivots(const std::string data_file, std::unique_ptr rearrangement; std::unique_ptr chunk_offsets; + std::string inflated_pq_file = pq_compressed_vectors_path + "_inflated.bin"; + if (!file_exists(pq_pivots_path)) { diskann::cout << "ERROR: PQ k-means pivot file not found" << std::endl; throw diskann::ANNException("PQ k-means pivot file not found", -1); @@ -446,6 +449,20 @@ int generate_pq_data_from_pivots(const std::string data_file, 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::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<_u32[]> block_compressed_base = std::make_unique<_u32[]>(block_size * (_u64) num_pq_chunks); std::memset(block_compressed_base.get(), 0, @@ -518,6 +535,12 @@ int generate_pq_data_from_pivots(const std::string data_file, #pragma omp parallel for schedule(static, 8192) for (int64_t j = 0; j < (_s64) cur_blk_size; j++) { block_compressed_base[j * num_pq_chunks + i] = closest_center[j]; +#ifdef SAVE_INFLATED_PQ + for (uint64_t k = 0; k < cur_chunk_size; k++) + block_inflated_base[j * dim + chunk_offsets[i] + k] = + cur_pivot_data[closest_center[j] * cur_chunk_size + k] + + centroid[chunk_offsets[i] + k]; +#endif } } @@ -532,8 +555,13 @@ int generate_pq_data_from_pivots(const std::string data_file, 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)); + 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)); +#endif diskann::cout << ".done." << std::endl; } // Gopal. Splittng diskann_dll into separate DLLs for search and build. @@ -542,6 +570,9 @@ int generate_pq_data_from_pivots(const std::string data_file, MallocExtension::instance()->ReleaseFreeMemory(); #endif compressed_file_writer.close(); +#ifdef SAVE_INFLATED_PQ + inflated_file_writer.close(); +#endif return 0; } From 863edb83466553f1570c59ee56014c028c6e9ce0 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Mon, 31 May 2021 11:30:33 +0530 Subject: [PATCH 07/84] now incorporated disk+memory search for inner product --- src/pq_flash_index.cpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index a11379b7d1..c5a1c20840 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -820,9 +820,9 @@ namespace diskann { // query <-> PQ chunk centers distances float *pq_dists = query_scratch->aligned_pqtable_dist_scratch; - if (metric==diskann::Metric::INNER_PRODUCT) - pq_table.populate_chunk_inner_products(query, pq_dists); - else if (metric==diskann::Metric::L2) +// if (metric==diskann::Metric::INNER_PRODUCT) +// pq_table.populate_chunk_inner_products(query, pq_dists); +// else if (metric==diskann::Metric::L2) pq_table.populate_chunk_distances(query, pq_dists); // query <-> neighbor list From 79b1fce67d4118b0872a864a4af6ff0894291cff Mon Sep 17 00:00:00 2001 From: ravishankar Date: Mon, 31 May 2021 15:30:05 +0530 Subject: [PATCH 08/84] support for mips and l2 --- src/pq_flash_index.cpp | 3 +++ tests/search_disk_index.cpp | 35 +++++++++++++++++++++++------------ 2 files changed, 26 insertions(+), 12 deletions(-) diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index c5a1c20840..cb5a51f661 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -165,6 +165,7 @@ namespace diskann { this->dist_cmp_float = new AVXDistanceL2Float(); } } else if (metric == diskann::Metric::INNER_PRODUCT) { + std::cout<<"Using inner product distance function" << std::endl; this->dist_cmp = new DistanceInnerProduct(); this->dist_cmp_float = new DistanceInnerProduct(); } else { @@ -1110,6 +1111,8 @@ namespace diskann { indices[i] = full_retset[i].id; if (distances != nullptr) { distances[i] = full_retset[i].distance; + if (metric == diskann::Metric::INNER_PRODUCT) // flip the sign from convert min to max + distances[i] = 1/distances[i]; } } diff --git a/tests/search_disk_index.cpp b/tests/search_disk_index.cpp index 7fa0edeec9..939c365388 100644 --- a/tests/search_disk_index.cpp +++ b/tests/search_disk_index.cpp @@ -56,26 +56,37 @@ int search_disk_index(int argc, char** argv) { size_t query_num, query_dim, query_aligned_dim, gt_num, gt_dim; std::vector<_u64> Lvec; - std::string index_prefix_path(argv[2]); + _u32 ctr = 2; + _u32 dist_fn = atoi(argv[ctr++]); + std::string index_prefix_path(argv[ctr++]); std::string pq_prefix = index_prefix_path + "_pq"; std::string disk_index_file = index_prefix_path + "_disk.index"; std::string warmup_query_file = index_prefix_path + "_sample_data.bin"; - _u64 num_nodes_to_cache = std::atoi(argv[3]); - _u32 num_threads = std::atoi(argv[4]); - _u32 beamwidth = std::atoi(argv[5]); - std::string query_bin(argv[6]); - std::string truthset_bin(argv[7]); - _u64 recall_at = std::atoi(argv[8]); - std::string result_output_prefix(argv[9]); + _u64 num_nodes_to_cache = std::atoi(argv[ctr++]); + _u32 num_threads = std::atoi(argv[ctr++]); + _u32 beamwidth = std::atoi(argv[ctr++]); + std::string query_bin(argv[ctr++]); + std::string truthset_bin(argv[ctr++]); + _u64 recall_at = std::atoi(argv[ctr++]); + std::string result_output_prefix(argv[ctr++]); bool calc_recall_flag = false; - for (int ctr = 10; ctr < argc; ctr++) { + for (; ctr < (_u32) argc; ctr++) { _u64 curL = std::atoi(argv[ctr]); if (curL >= recall_at) Lvec.push_back(curL); } + diskann::Metric metric; + if (dist_fn == 0) + metric = diskann::Metric::L2; + else if (dist_fn == 1) + metric = diskann::Metric::INNER_PRODUCT; + else { + std::cout<<"Unsupported distance function. Currently only L2/ Inner Product support." << std::endl; + return -1; + } if (Lvec.size() == 0) { diskann::cout << "No valid Lsearch found. Lsearch must be at least recall_at" @@ -114,7 +125,7 @@ int search_disk_index(int argc, char** argv) { #endif std::unique_ptr> _pFlashIndex( - new diskann::PQFlashIndex(reader, diskann::Metric::INNER_PRODUCT)); + new diskann::PQFlashIndex(reader, metric)); int res = _pFlashIndex->load(num_threads, pq_prefix.c_str(), disk_index_file.c_str()); @@ -286,10 +297,10 @@ int search_disk_index(int argc, char** argv) { } int main(int argc, char** argv) { - if (argc < 11) { + if (argc < 12) { diskann::cout << "Usage: " << argv[0] - << " [index_type] [index_prefix_path] " + << " [index_type] [dist_fn 0 for l2/ 1 for inner product] [index_prefix_path] " " [num_nodes_to_cache] [num_threads] [beamwidth (use 0 to " "optimize internally)] " " [query_file.bin] [truthset.bin (use \"null\" for none)] " From a86e2d4a0112306ad239d394f7761e70d485243a Mon Sep 17 00:00:00 2001 From: ravishankar Date: Mon, 31 May 2021 17:13:40 +0530 Subject: [PATCH 09/84] changed inner product to -IP rather than 1/IP --- include/distance.h | 7 ++++--- src/index.cpp | 14 ++++++++++++-- src/pq_flash_index.cpp | 2 +- 3 files changed, 17 insertions(+), 6 deletions(-) diff --git a/include/distance.h b/include/distance.h index b39a1719fb..a65f9e9487 100644 --- a/include/distance.h +++ b/include/distance.h @@ -428,9 +428,10 @@ namespace diskann { } float compare(const T *a, const T *b, unsigned size) const { // since we use normally minimization objective for distance comparisons, we are returning 1/x. float result = acompare(a,b,size); - if (result < 0) - return std::numeric_limits::max(); - else return 1/result; +// if (result < 0) +// return std::numeric_limits::max(); +// else +return -result; } }; diff --git a/src/index.cpp b/src/index.cpp index 7123c661d4..ff00f6dc61 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -538,7 +538,7 @@ namespace diskann { float cur_alpha = 1; while (cur_alpha <= alpha && result.size() < degree) { unsigned start = 0; - + float eps = cur_alpha + 0.01; while (result.size() < degree && (start) < pool.size() && start < maxc) { auto &p = pool[start]; if (occlude_factor[start] > cur_alpha) { @@ -553,8 +553,18 @@ namespace diskann { float djk = _distance->compare( _data + _aligned_dim * (size_t) pool[t].id, _data + _aligned_dim * (size_t) p.id, (unsigned) _aligned_dim); + if (_metric == diskann::Metric::L2) { occlude_factor[t] = (std::max)(occlude_factor[t], pool[t].distance / djk); + } + else if (_metric == diskann::Metric::INNER_PRODUCT) { // stylized rules for inner product since we want max instead of min distance + float x = -pool[t].distance; + float y = -djk; + if (y > cur_alpha * x) { + occlude_factor[t] = + (std::max)(occlude_factor[t], eps); + } + } } start++; } @@ -1058,7 +1068,7 @@ namespace diskann { indices[pos] = it.id; distances[pos] = it.distance; if (_metric == diskann::INNER_PRODUCT) - distances[pos] = 1/distances[pos]; + distances[pos] = -distances[pos]; pos++; if (pos == K) break; diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index cb5a51f661..ef236aa834 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -1112,7 +1112,7 @@ namespace diskann { if (distances != nullptr) { distances[i] = full_retset[i].distance; if (metric == diskann::Metric::INNER_PRODUCT) // flip the sign from convert min to max - distances[i] = 1/distances[i]; + distances[i] = -distances[i]; } } From de9ef76fd892f58a0836de4dbfb8fea9f3e562ec Mon Sep 17 00:00:00 2001 From: ravishankar Date: Tue, 1 Jun 2021 20:59:37 +0530 Subject: [PATCH 10/84] towards adding support for storing PQ vectors in disk index for very large data --- include/pq_flash_index.h | 7 +++++++ include/pq_table.h | 24 ++++++++++++++++++++++++ src/pq_flash_index.cpp | 1 + tests/search_disk_index.cpp | 6 ++++-- 4 files changed, 36 insertions(+), 2 deletions(-) diff --git a/include/pq_flash_index.h b/include/pq_flash_index.h index c20409fb64..f3e8bcde5f 100644 --- a/include/pq_flash_index.h +++ b/include/pq_flash_index.h @@ -152,6 +152,13 @@ namespace diskann { Distance * dist_cmp = nullptr; Distance *dist_cmp_float = nullptr; + // for very large datasets: we use PQ even for the disk resident index + bool use_disk_index_pq = false; + _u64 disk_index_chunk_size; + _u64 disk_index_n_chunks; + FixedChunkPQTable disk_index_pq_table; + + // medoid/start info uint32_t *medoids = nullptr; // by default it is just one entry point of graph, we diff --git a/include/pq_table.h b/include/pq_table.h index c913271eb1..bb71fccba4 100644 --- a/include/pq_table.h +++ b/include/pq_table.h @@ -145,6 +145,30 @@ namespace diskann { } } } + + float compare(const T* query_vec, _u8* base_vec) { + float res = 0; + for (_u64 chunk = 0; chunk < n_chunks; chunk++) { + for (_u64 j = chunk_offsets[chunk]; j < chunk_offsets[chunk + 1]; j++) { + _u64 permuted_dim_in_query = rearrangement[j]; + const float* centers_dim_vec = tables_T + (256 * j); + float diff = centers_dim_vec[base_vec[chunk]] - (query_vec[permuted_dim_in_query] - centroid[permuted_dim_in_query]); + res += diff*diff; + } + } + return res; + } + + void inflate_vector(_u8* base_vec, float* out_vec) { + for (_u64 chunk = 0; chunk < n_chunks; chunk++) { + for (_u64 j = chunk_offsets[chunk]; j < chunk_offsets[chunk + 1]; j++) { + _u64 original_dim = rearrangement[j]; + const float* centers_dim_vec = tables_T + (256 * j); + out_vec[original_dim] = centers_dim_vec[base_vec[chunk]] + centroid[original_dim]; + } + } + } + void populate_chunk_inner_products(const T* query_vec, float* dist_vec) { memset(dist_vec, 0, 256 * n_chunks * sizeof(float)); diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index ef236aa834..ada71f0b7e 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -860,6 +860,7 @@ namespace diskann { } compute_dists(&best_medoid, 1, dist_scratch); + retset[0].id = best_medoid; retset[0].distance = dist_scratch[0]; retset[0].flag = true; diff --git a/tests/search_disk_index.cpp b/tests/search_disk_index.cpp index 939c365388..3a6abddeba 100644 --- a/tests/search_disk_index.cpp +++ b/tests/search_disk_index.cpp @@ -137,13 +137,14 @@ int search_disk_index(int argc, char** argv) { 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->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; @@ -207,7 +208,7 @@ int search_disk_index(int argc, char** argv) { uint32_t optimized_beamwidth = 2; - // query_num = 1; + for (uint32_t test_id = 0; test_id < Lvec.size(); test_id++) { _u64 L = Lvec[test_id]; @@ -225,6 +226,7 @@ int search_disk_index(int argc, char** argv) { diskann::QueryStats* 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) From 4dd8de947960e267a5c6d620fc7fed6efd6eaaab Mon Sep 17 00:00:00 2001 From: ravishankar Date: Tue, 1 Jun 2021 21:00:00 +0530 Subject: [PATCH 11/84] towards adding support for storing PQ vectors in disk index for very large data --- src/pq_flash_index.cpp | 13 ------------- 1 file changed, 13 deletions(-) diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index ada71f0b7e..0bab700ef5 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -897,19 +897,6 @@ namespace diskann { _u32 marker = k; _u32 num_seen = 0; - /* - bool marker_set = false; - diskann::cout << "hop " << hops << ": "; - for (_u32 i = 0; i < cur_list_size; i++) { - diskann::cout << retset[i].id << "( " << retset[i].distance; - if (retset[i].flag && !marker_set) { - diskann::cout << ",*) "; - marker_set = true; - } else - diskann::cout << ") "; - } - diskann::cout << std::endl; - */ while (marker < cur_list_size && frontier.size() < beam_width && num_seen < beam_width + 2) { if (retset[marker].flag) { From d4a658f423cf9f89d4b402b84839197842ca552b Mon Sep 17 00:00:00 2001 From: ravishankar Date: Tue, 1 Jun 2021 22:09:05 +0530 Subject: [PATCH 12/84] halfway through PQ-based disk search option --- include/aux_utils.h | 1 + include/pq_flash_index.h | 8 ++++---- include/pq_table.h | 11 +++++++---- src/aux_utils.cpp | 37 ++++++++++++++++++++++++++++++++++--- src/partition_and_pq.cpp | 2 +- src/pq_flash_index.cpp | 30 ++++++++++++++++++++++++------ tests/build_disk_index.cpp | 6 +++--- 7 files changed, 74 insertions(+), 21 deletions(-) diff --git a/include/aux_utils.h b/include/aux_utils.h index 031a64ba8b..7e946becd4 100644 --- a/include/aux_utils.h +++ b/include/aux_utils.h @@ -36,6 +36,7 @@ namespace diskann { const double THRESHOLD_FOR_CACHING_IN_GB = 1.0; const uint32_t NUM_NODES_TO_CACHE = 250000; const uint32_t WARMUP_L = 20; + const uint32_t NUM_KMEANS_REPS = 12; template class PQFlashIndex; diff --git a/include/pq_flash_index.h b/include/pq_flash_index.h index f3e8bcde5f..489afbed31 100644 --- a/include/pq_flash_index.h +++ b/include/pq_flash_index.h @@ -133,7 +133,9 @@ namespace diskann { // data info _u64 num_points = 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; std::string disk_index_file; std::vector> node_visit_counter; @@ -144,7 +146,6 @@ namespace diskann { // chunk_size = chunk size of each dimension chunk // pq_tables = float* [[2^8 * [chunk_size]] * n_chunks] _u8 * data = nullptr; - _u64 chunk_size; _u64 n_chunks; FixedChunkPQTable pq_table; @@ -154,9 +155,8 @@ namespace diskann { // for very large datasets: we use PQ even for the disk resident index bool use_disk_index_pq = false; - _u64 disk_index_chunk_size; - _u64 disk_index_n_chunks; - FixedChunkPQTable disk_index_pq_table; + _u64 disk_pq_n_chunks; + FixedChunkPQTable disk_pq_table; // medoid/start info diff --git a/include/pq_table.h b/include/pq_table.h index bb71fccba4..825b29c7ec 100644 --- a/include/pq_table.h +++ b/include/pq_table.h @@ -13,8 +13,8 @@ namespace diskann { nullptr; // pq_tables = float* [[2^8 * [chunk_size]] * n_chunks] // _u64 n_chunks; // n_chunks = # of chunks ndims is split into // _u64 chunk_size; // chunk_size = chunk size of each dimension chunk - _u64 ndims; // ndims = chunk_size * n_chunks - _u64 n_chunks; + _u64 ndims = 0; // ndims = chunk_size * n_chunks + _u64 n_chunks = 0; _u32* chunk_offsets = nullptr; _u32* rearrangement = nullptr; float* centroid = nullptr; @@ -79,14 +79,14 @@ namespace diskann { #else diskann::load_bin<_u32>(chunk_offset_file, chunk_offsets, numr, numc); #endif - if (numc != 1 || numr != num_chunks + 1) { + if (numc != 1 || (numr != num_chunks + 1 && num_chunks != 0)) { diskann::cerr << "Error loading chunk offsets file. numc: " << numc << " (should be 1). numr: " << numr << " (should be " << num_chunks + 1 << ")" << std::endl; throw diskann::ANNException("Error loading chunk offsets file", -1, __FUNCSIG__, __FILE__, __LINE__); } - + std::cout<<"PQ data has " << numr - 1 <<" bytes per point." << std::endl; this->n_chunks = numr - 1; #ifdef EXEC_ENV_OLS @@ -126,6 +126,9 @@ namespace diskann { } } +_u32 get_num_chunks() { + return n_chunks; +} void populate_chunk_distances(const T* query_vec, float* dist_vec) { memset(dist_vec, 0, 256 * n_chunks * sizeof(float)); diff --git a/src/aux_utils.cpp b/src/aux_utils.cpp index b3cc400692..98e0e88e57 100644 --- a/src/aux_utils.cpp +++ b/src/aux_utils.cpp @@ -620,17 +620,27 @@ namespace diskann { while (parser >> cur_param) param_list.push_back(cur_param); - if (param_list.size() != 5) { + if (param_list.size() != 5 && param_list.size() != 6) { diskann::cout << "Correct usage of parameters is R (max degree) " "L (indexing list size, better if >= R) B (RAM limit of final " "index in " "GB) M (memory limit while indexing) T (number of threads for " - "indexing)" + "indexing) B' (PQ bytes for disk index: optional parameter for very large dimensional data)" << std::endl; return false; } + _u32 disk_pq_dims = 0; + bool use_disk_pq = false; + + if (param_list.size() == 6) { + disk_pq_dims = atoi(param_list[5].c_str()); + use_disk_pq = true; + if (disk_pq_dims == 0) + use_disk_pq = false; + } + if (compareMetric == diskann::Metric::INNER_PRODUCT) { std::cout<<"Using Inner Product for PQ and Graph Generation" << std::endl; } @@ -643,6 +653,11 @@ namespace diskann { std::string medoids_path = disk_index_path + "_medoids.bin"; std::string centroids_path = disk_index_path + "_centroids.bin"; std::string sample_base_prefix = index_prefix_path + "_sample"; + std::string disk_pq_pivots_path = index_prefix_path + "_disk.index_pq_pivots.bin"; // optional if disk index is also storing pq data + std::string disk_pq_compressed_vectors_path = // optional if disk index is also storing pq data + index_prefix_path + "_disk.index_pq_compressed.bin"; + + unsigned R = (unsigned) atoi(param_list[0].c_str()); unsigned L = (unsigned) atoi(param_list[1].c_str()); @@ -662,6 +677,7 @@ namespace diskann { } _u32 num_threads = (_u32) atoi(param_list[4].c_str()); + if (num_threads != 0) { omp_set_num_threads(num_threads); mkl_set_num_threads(num_threads); @@ -698,6 +714,17 @@ namespace diskann { gen_random_slice(dataFilePath, p_val, train_data, train_size, train_dim); + if (use_disk_pq) { + if (disk_pq_dims > dim) + disk_pq_dims = dim; + + std::cout<<"Compressing base for disk-PQ into " << disk_pq_dims << " chunks " << std::endl; + generate_pq_pivots(train_data, train_size, (uint32_t) dim, 256, + (uint32_t) disk_pq_dims, NUM_KMEANS_REPS, disk_pq_pivots_path, false); + generate_pq_data_from_pivots(dataFilePath, 256, (uint32_t) disk_pq_dims, + disk_pq_pivots_path, + disk_pq_compressed_vectors_path); + } diskann::cout << "Training data loaded of size " << train_size << std::endl; @@ -705,7 +732,7 @@ namespace diskann { if (compareMetric == diskann::Metric::INNER_PRODUCT) make_zero_mean = false; generate_pq_pivots(train_data, train_size, (uint32_t) dim, 256, - (uint32_t) num_pq_chunks, 15, pq_pivots_path, make_zero_mean); + (uint32_t) num_pq_chunks, NUM_KMEANS_REPS, pq_pivots_path, make_zero_mean); generate_pq_data_from_pivots(dataFilePath, 256, (uint32_t) num_pq_chunks, pq_pivots_path, @@ -719,8 +746,12 @@ namespace diskann { dataFilePath, compareMetric, L, R, p_val, indexing_ram_budget, mem_index_path, medoids_path, centroids_path); + if (!use_disk_pq) diskann::create_disk_layout(dataFilePath, mem_index_path, disk_index_path); + else + diskann::create_disk_layout<_u8>(disk_pq_compressed_vectors_path, mem_index_path, + disk_index_path); double sample_sampling_rate = (150000.0 / points_num); gen_random_slice(dataFilePath, sample_base_prefix, diff --git a/src/partition_and_pq.cpp b/src/partition_and_pq.cpp index be323ad8ff..c41f1924b1 100644 --- a/src/partition_and_pq.cpp +++ b/src/partition_and_pq.cpp @@ -36,7 +36,7 @@ #endif #define BLOCK_SIZE 5000000 -#define SAVE_INFLATED_PQ true +//#define SAVE_INFLATED_PQ true template void gen_random_slice(const std::string base_file, diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index 0bab700ef5..484e772a7e 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -49,7 +49,7 @@ // returns region of `node_buf` containing [NNBRS][NBR_ID(_u32)] #define OFFSET_TO_NODE_NHOOD(node_buf) \ - (unsigned *) ((char *) node_buf + data_dim * sizeof(T)) + (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) @@ -320,7 +320,7 @@ namespace diskann { 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, data_dim * sizeof(T)); + 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 @@ -562,7 +562,7 @@ namespace diskann { // 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, data_dim * sizeof(T)); + memcpy(medoid_coords, medoid_disk_coords, disk_bytes_per_point); for (uint32_t i = 0; i < data_dim; i++) centroid_data[cur_m * aligned_dim + i] = medoid_coords[i]; @@ -592,6 +592,7 @@ namespace diskann { std::string centroids_file = std::string(disk_index_file) + "_centroids.bin"; + 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); @@ -608,8 +609,11 @@ namespace diskann { } this->data_dim = pq_file_dim; + this->disk_data_dim = this->data_dim; // will reset later if we use PQ on disk + this->disk_bytes_per_point = this->data_dim * sizeof(T); // will change later if we use PQ on disk this->aligned_dim = ROUND_UP(pq_file_dim, 8); + size_t npts_u64, nchunks_u64; #ifdef EXEC_ENV_OLS diskann::load_bin<_u8>(files, pq_compressed_vectors, this->data, npts_u64, @@ -634,6 +638,20 @@ namespace diskann { << " #aligned_dim: " << aligned_dim << " #chunks: " << n_chunks << std::endl; + +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 + disk_pq_table.load_pq_centroid_bin(files, disk_pq_pivots_path.c_str(), 0); // giving 0 chunks as the pq_table will infer from the chunk_offsets file the correct value +#else + disk_pq_table.load_pq_centroid_bin(disk_pq_pivots_path.c_str(), 0); // giving 0 chunks as the pq_table will infer from the chunk_offsets file the correct value +#endif + disk_pq_n_chunks = disk_pq_table.get_num_chunks(); + disk_bytes_per_point = disk_pq_n_chunks * sizeof(_u8); + std::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 @@ -678,7 +696,7 @@ namespace diskann { READ_U64(index_metadata, medoid_id_on_file); READ_U64(index_metadata, max_node_len); READ_U64(index_metadata, nnodes_per_sector); - max_degree = ((max_node_len - data_dim * sizeof(T)) / sizeof(unsigned)) - 1; + max_degree = ((max_node_len - disk_bytes_per_point) / sizeof(unsigned)) - 1; diskann::cout << "Disk-Index File Meta-data: "; diskann::cout << "# nodes per sector: " << nnodes_per_sector; @@ -796,7 +814,7 @@ namespace diskann { for (uint32_t i = 0; i < this->data_dim; i++) { data.scratch.aligned_query_float[i] = query1[i]; } - memcpy(data.scratch.aligned_query_T, query1, this->data_dim * sizeof(T)); + memcpy(data.scratch.aligned_query_T, query1, disk_bytes_per_point); const T * query = data.scratch.aligned_query_T; const float *query_float = data.scratch.aligned_query_float; @@ -1025,7 +1043,7 @@ namespace diskann { T *node_fp_coords_copy = data_buf + (data_buf_idx * aligned_dim); data_buf_idx++; - memcpy(node_fp_coords_copy, node_fp_coords, data_dim * sizeof(T)); + memcpy(node_fp_coords_copy, node_fp_coords, disk_bytes_per_point); float cur_expanded_dist = dist_cmp->compare(query, node_fp_coords_copy, (unsigned) aligned_dim); diff --git a/tests/build_disk_index.cpp b/tests/build_disk_index.cpp index 8c94d47525..77bdeace37 100644 --- a/tests/build_disk_index.cpp +++ b/tests/build_disk_index.cpp @@ -17,11 +17,11 @@ bool build_index(const char* dataFilePath, const char* indexFilePath, } int main(int argc, char** argv) { - if (argc != 10) { + if (argc != 11) { std::cout << "Usage: " << argv[0] << " [data_type] [dist_fn: 0 for L2, 1 for MIPS] [data_file.bin] " "[index_prefix_path] " - "[R] [L] [B] [M] [T]. See README for more information on " + "[R] [L] [B] [M] [T] [PQ_disk_bytes (for very large dimensionality, use 0 for full vectors)]. See README for more information on " "parameters." << std::endl; } else { @@ -30,7 +30,7 @@ int main(int argc, char** argv) { metric = diskann::Metric::INNER_PRODUCT; std::string params = std::string(argv[5]) + " " + std::string(argv[6]) + " " + std::string(argv[7]) + " " + - std::string(argv[8]) + " " + std::string(argv[9]); + std::string(argv[8]) + " " + std::string(argv[9]) + " " + std::string(argv[10]); if (std::string(argv[1]) == std::string("float")) build_index(argv[3], argv[4], params.c_str(), metric); else if (std::string(argv[1]) == std::string("int8")) From 833d18925407e9246bd2279884ad02c8080624a6 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Wed, 2 Jun 2021 08:17:09 +0530 Subject: [PATCH 13/84] code compiles for disk index pq --- src/pq_flash_index.cpp | 23 ++++++++++++++++++----- 1 file changed, 18 insertions(+), 5 deletions(-) diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index 484e772a7e..24c5a52132 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -564,10 +564,14 @@ namespace diskann { 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((_u8*) medoid_coords, (centroid_data + cur_m*aligned_dim)); + } aligned_free(medoid_buf); + delete[] medoid_coords; } // return ctx @@ -643,9 +647,9 @@ 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 - disk_pq_table.load_pq_centroid_bin(files, disk_pq_pivots_path.c_str(), 0); // giving 0 chunks as the pq_table will 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 #else - disk_pq_table.load_pq_centroid_bin(disk_pq_pivots_path.c_str(), 0); // giving 0 chunks as the pq_table will 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 #endif disk_pq_n_chunks = disk_pq_table.get_num_chunks(); disk_bytes_per_point = disk_pq_n_chunks * sizeof(_u8); @@ -974,8 +978,13 @@ if (file_exists(disk_pq_pivots_path)) { 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 = dist_cmp->compare(query, node_fp_coords_copy, + float cur_expanded_dist; + if (!use_disk_index_pq) + cur_expanded_dist = dist_cmp->compare(query, node_fp_coords_copy, (unsigned) aligned_dim); + else { + cur_expanded_dist = disk_pq_table.compare(query, (_u8*) node_fp_coords_copy); + } full_retset.push_back( Neighbor((unsigned) cached_nhood.first, cur_expanded_dist, true)); @@ -1045,8 +1054,12 @@ if (file_exists(disk_pq_pivots_path)) { data_buf_idx++; memcpy(node_fp_coords_copy, node_fp_coords, disk_bytes_per_point); - float cur_expanded_dist = dist_cmp->compare(query, node_fp_coords_copy, + float cur_expanded_dist; + if (!use_disk_index_pq) + cur_expanded_dist = dist_cmp->compare(query, node_fp_coords_copy, (unsigned) aligned_dim); + else + cur_expanded_dist = disk_pq_table.compare(query, (_u8*) node_fp_coords_copy); full_retset.push_back( Neighbor(frontier_nhood.first, cur_expanded_dist, true)); From b07d0327122644dd0efa0c974010204ed010d83d Mon Sep 17 00:00:00 2001 From: ravishankar Date: Wed, 2 Jun 2021 11:06:33 +0530 Subject: [PATCH 14/84] fixed some bug --- src/aux_utils.cpp | 2 +- src/pq_flash_index.cpp | 28 +++++++++++++++++++++++++++- tests/search_disk_index.cpp | 10 +++++----- tests/utils/create_disk_layout.cpp | 13 +++++++------ 4 files changed, 40 insertions(+), 13 deletions(-) diff --git a/src/aux_utils.cpp b/src/aux_utils.cpp index 98e0e88e57..a9713a6099 100644 --- a/src/aux_utils.cpp +++ b/src/aux_utils.cpp @@ -757,7 +757,7 @@ namespace diskann { gen_random_slice(dataFilePath, sample_base_prefix, sample_sampling_rate); - std::remove(mem_index_path.c_str()); +// std::remove(mem_index_path.c_str()); auto e = std::chrono::high_resolution_clock::now(); diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index 24c5a52132..4a30053346 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -817,8 +817,9 @@ if (file_exists(disk_pq_pivots_path)) { for (uint32_t i = 0; i < this->data_dim; i++) { data.scratch.aligned_query_float[i] = query1[i]; + data.scratch.aligned_query_T[i] = query1[i]; } - memcpy(data.scratch.aligned_query_T, query1, disk_bytes_per_point); +// memcpy(data.scratch.aligned_query_T, query1, disk_bytes_per_point); const T * query = data.scratch.aligned_query_T; const float *query_float = data.scratch.aligned_query_float; @@ -888,6 +889,15 @@ if (file_exists(disk_pq_pivots_path)) { retset[0].flag = true; visited.insert(best_medoid); +/* + std::cout<<"Chose " << retset[0].id<< " as best medoid with distance " << retset[0].distance << std::endl; + + std::cout<<"query from 0 to " << aligned_dim << std::endl; + for (_u32 i = 0; i < aligned_dim-1; i++) { + std::cout< percentiles, std::vector results) { @@ -137,9 +137,9 @@ int search_disk_index(int argc, char** argv) { std::vector node_list; diskann::cout << "Caching " << num_nodes_to_cache << " BFS nodes around medoid(s)" << std::endl; -// _pFlashIndex->cache_bfs_levels(num_nodes_to_cache, node_list); - _pFlashIndex->generate_cache_list_from_sample_queries( - warmup_query_file, 15, 6, num_nodes_to_cache, num_threads, node_list); + _pFlashIndex->cache_bfs_levels(num_nodes_to_cache, node_list); +// _pFlashIndex->generate_cache_list_from_sample_queries( +// warmup_query_file, 15, 6, num_nodes_to_cache, num_threads, node_list); _pFlashIndex->load_cache_list(node_list); node_list.clear(); node_list.shrink_to_fit(); @@ -208,7 +208,7 @@ int search_disk_index(int argc, char** argv) { uint32_t optimized_beamwidth = 2; - +//query_num = 1; for (uint32_t test_id = 0; test_id < Lvec.size(); test_id++) { _u64 L = Lvec[test_id]; diff --git a/tests/utils/create_disk_layout.cpp b/tests/utils/create_disk_layout.cpp index ac378272d8..458e2d6b20 100644 --- a/tests/utils/create_disk_layout.cpp +++ b/tests/utils/create_disk_layout.cpp @@ -14,12 +14,6 @@ template int create_disk_layout(int argc, char **argv) { - if (argc != 5) { - std::cout << argv[0] << " data_type data_bin " - "vamana_index_file output_diskann_index_file" - << std::endl; - exit(-1); - } std::string base_file(argv[2]); std::string vamana_file(argv[3]); std::string output_file(argv[4]); @@ -28,6 +22,13 @@ int create_disk_layout(int argc, char **argv) { } 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(argc, argv); From d6c6b8da57b240525cb393d995bfed015cf74d78 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Wed, 2 Jun 2021 07:02:05 +0000 Subject: [PATCH 15/84] shards are written as and when necessary --- include/partition_and_pq.h | 9 ++ src/aux_utils.cpp | 7 ++ src/partition_and_pq.cpp | 171 ++++++++++++++++++++++++++++++++++++- 3 files changed, 185 insertions(+), 2 deletions(-) diff --git a/include/partition_and_pq.h b/include/partition_and_pq.h index 43a9d84db0..bbae1ed148 100644 --- a/include/partition_and_pq.h +++ b/include/partition_and_pq.h @@ -38,6 +38,15 @@ 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); + +template +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, diff --git a/src/aux_utils.cpp b/src/aux_utils.cpp index a9713a6099..d77d7fa998 100644 --- a/src/aux_utils.cpp +++ b/src/aux_utils.cpp @@ -385,6 +385,12 @@ namespace diskann { 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"; + + 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"; @@ -402,6 +408,7 @@ namespace diskann { new diskann::Index(compareMetric, shard_base_file.c_str())); _pvamanaIndex->build(paras); _pvamanaIndex->save(shard_index_file.c_str()); + std::remove(shard_base_file.c_str()); } diskann::merge_shards(merged_index_prefix + "_subshard-", "_mem.index", diff --git a/src/partition_and_pq.cpp b/src/partition_and_pq.cpp index c41f1924b1..a219dc25a0 100644 --- a/src/partition_and_pq.cpp +++ b/src/partition_and_pq.cpp @@ -35,7 +35,7 @@ #include #endif -#define BLOCK_SIZE 5000000 +#define BLOCK_SIZE 50000 //#define SAVE_INFLATED_PQ true template @@ -741,6 +741,169 @@ int shard_data_into_clusters(const std::string data_file, float *pivots, return 0; } + + + +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; + // create cached reader + writer + cached_ifstream base_reader(data_file, read_blk_size); + _u32 npts32; + _u32 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); + _u32 dummy_size = 0; + _u32 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<_u32[]> block_closest_centers = + std::make_unique<_u32[]>(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++) { + _u32 cur_shard_count = (_u32) 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) { + _u64 read_blk_size = 64 * 1024 * 1024; + // _u64 write_blk_size = 64 * 1024 * 1024; + // create cached reader + writer + cached_ifstream base_reader(data_file, read_blk_size); + _u32 npts32; + _u32 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; + + 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); + + _u32 cur_pos = 0; + _u32 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]) { + 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; + + 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 // 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 @@ -867,7 +1030,7 @@ int partition_with_ram_budget(const std::string data_file, diskann::save_bin(output_file.c_str(), pivot_data, (size_t) num_parts, train_dim); - shard_data_into_clusters(data_file, pivot_data, 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; @@ -926,6 +1089,10 @@ 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); + template DISKANN_DLLEXPORT int generate_pq_data_from_pivots( const std::string data_file, unsigned num_centers, unsigned num_pq_chunks, std::string pq_pivots_path, std::string pq_compressed_vectors_path); From dc5fd39dbb4efab045483f2c7627b7628aac0a4e Mon Sep 17 00:00:00 2001 From: ravishankar Date: Wed, 2 Jun 2021 07:07:26 +0000 Subject: [PATCH 16/84] sharding is now on demand --- src/partition_and_pq.cpp | 1 + 1 file changed, 1 insertion(+) diff --git a/src/partition_and_pq.cpp b/src/partition_and_pq.cpp index a219dc25a0..bf1cb8a81c 100644 --- a/src/partition_and_pq.cpp +++ b/src/partition_and_pq.cpp @@ -881,6 +881,7 @@ int retrieve_shard_data_from_ids(const std::string data_file, std::string idmap_ 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++; From 0a71e5928c32a842c047a2b64598fdf4066b4e3c Mon Sep 17 00:00:00 2001 From: ravishankar Date: Wed, 2 Jun 2021 13:35:06 +0000 Subject: [PATCH 17/84] minor changes --- src/aux_utils.cpp | 8 ++++---- src/partition_and_pq.cpp | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/aux_utils.cpp b/src/aux_utils.cpp index d77d7fa998..e95eff0e30 100644 --- a/src/aux_utils.cpp +++ b/src/aux_utils.cpp @@ -741,9 +741,7 @@ namespace diskann { generate_pq_pivots(train_data, train_size, (uint32_t) dim, 256, (uint32_t) num_pq_chunks, NUM_KMEANS_REPS, pq_pivots_path, make_zero_mean); generate_pq_data_from_pivots(dataFilePath, 256, (uint32_t) - num_pq_chunks, - pq_pivots_path, - pq_compressed_vectors_path); + num_pq_chunks, pq_pivots_path, pq_compressed_vectors_path); delete[] train_data; @@ -764,7 +762,9 @@ namespace diskann { gen_random_slice(dataFilePath, sample_base_prefix, sample_sampling_rate); -// std::remove(mem_index_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(); diff --git a/src/partition_and_pq.cpp b/src/partition_and_pq.cpp index bf1cb8a81c..5e34f35ad1 100644 --- a/src/partition_and_pq.cpp +++ b/src/partition_and_pq.cpp @@ -35,7 +35,7 @@ #include #endif -#define BLOCK_SIZE 50000 +#define BLOCK_SIZE 5000000 //#define SAVE_INFLATED_PQ true template From e4a25e65161764ae18a7d23c2df146fe7c1adf93 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Wed, 2 Jun 2021 15:22:24 +0000 Subject: [PATCH 18/84] fixed one malloc bug in parameters --- include/parameters.h | 3 +++ 1 file changed, 3 insertions(+) diff --git a/include/parameters.h b/include/parameters.h index 42e7d3a463..1cff662850 100644 --- a/include/parameters.h +++ b/include/parameters.h @@ -19,6 +19,9 @@ namespace diskann { template inline void Set(const std::string &name, const ParamType &value) { // ParamType *ptr = (ParamType *) malloc(sizeof(ParamType)); + if (params.find(name) != params.end()) { + free(params[name]); + } ParamType *ptr = new ParamType; *ptr = value; params[name] = (void *) ptr; From 974697b0ec6d15369e47b6d80ef29682535f00ae Mon Sep 17 00:00:00 2001 From: ravishankar Date: Thu, 10 Jun 2021 16:24:59 +0000 Subject: [PATCH 19/84] added a vector analyzer util --- include/aux_utils.h | 2 +- include/utils.h | 5 +++++ src/aux_utils.cpp | 9 +++++++++ tests/utils/CMakeLists.txt | 10 ++++++++++ 4 files changed, 25 insertions(+), 1 deletion(-) diff --git a/include/aux_utils.h b/include/aux_utils.h index 7e946becd4..ac7bad6cc8 100644 --- a/include/aux_utils.h +++ b/include/aux_utils.h @@ -31,7 +31,7 @@ typedef int FileHandle; #include "windows_customizations.h" namespace diskann { - const size_t TRAINING_SET_SIZE = 1500000; + const size_t TRAINING_SET_SIZE = 150000; const double SPACE_FOR_CACHED_NODES_IN_GB = 0.25; const double THRESHOLD_FOR_CACHING_IN_GB = 1.0; const uint32_t NUM_NODES_TO_CACHE = 250000; diff --git a/include/utils.h b/include/utils.h index 6b9db5bf62..3e9bafc088 100644 --- a/include/utils.h +++ b/include/utils.h @@ -218,6 +218,11 @@ namespace diskann { } #endif + inline void wait_for_keystroke() { + int a; + std::cin>> a; + } + template inline void load_bin(const std::string& bin_file, T*& data, size_t& npts, size_t& dim) { diff --git a/src/aux_utils.cpp b/src/aux_utils.cpp index e95eff0e30..a08bb15e1f 100644 --- a/src/aux_utils.cpp +++ b/src/aux_utils.cpp @@ -379,6 +379,9 @@ namespace diskann { partition_with_ram_budget(base_file, sampling_rate, ram_budget, 2 * R / 3, merged_index_prefix, 2); + wait_for_keystroke(); + + std::string cur_centroid_filepath = merged_index_prefix + "_centroids.bin"; std::rename(cur_centroid_filepath.c_str(), centroids_file.c_str()); @@ -409,6 +412,8 @@ namespace diskann { _pvamanaIndex->build(paras); _pvamanaIndex->save(shard_index_file.c_str()); std::remove(shard_base_file.c_str()); + wait_for_keystroke(); + } diskann::merge_shards(merged_index_prefix + "_subshard-", "_mem.index", @@ -735,6 +740,8 @@ namespace diskann { diskann::cout << "Training data loaded of size " << train_size << std::endl; + wait_for_keystroke(); + bool make_zero_mean = true; if (compareMetric == diskann::Metric::INNER_PRODUCT) make_zero_mean = false; @@ -751,6 +758,8 @@ namespace diskann { dataFilePath, compareMetric, L, R, p_val, indexing_ram_budget, mem_index_path, medoids_path, centroids_path); + + if (!use_disk_pq) diskann::create_disk_layout(dataFilePath, mem_index_path, disk_index_path); diff --git a/tests/utils/CMakeLists.txt b/tests/utils/CMakeLists.txt index 164fcffcfd..d4cb63b01f 100644 --- a/tests/utils/CMakeLists.txt +++ b/tests/utils/CMakeLists.txt @@ -42,6 +42,16 @@ else() target_link_libraries(uint32_to_uint8 ${PROJECT_NAME}) endif() + +add_executable(vector_analysis vector_analysis.cpp) +if(MSVC) + target_link_options(vector_analysis PRIVATE /MACHINE:x64) + target_link_libraries(vector_analysis debug ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}/diskann_dll.lib) + target_link_libraries(vector_analysis optimized ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}/diskann_dll.lib) +else() + target_link_libraries(vector_analysis ${PROJECT_NAME} -ltcmalloc) +endif() + add_executable(gen_random_slice gen_random_slice.cpp) if(MSVC) target_link_options(gen_random_slice PRIVATE /MACHINE:x64) From 787819a74d9e3f245cfab154fd1ddabd1e1d3e3d Mon Sep 17 00:00:00 2001 From: ravishankar Date: Thu, 10 Jun 2021 16:29:39 +0000 Subject: [PATCH 20/84] added missing file --- tests/utils/vector_analysis.cpp | 72 +++++++++++++++++++++++++++++++++ 1 file changed, 72 insertions(+) create mode 100644 tests/utils/vector_analysis.cpp diff --git a/tests/utils/vector_analysis.cpp b/tests/utils/vector_analysis.cpp new file mode 100644 index 0000000000..1cd2eb9ff5 --- /dev/null +++ b/tests/utils/vector_analysis.cpp @@ -0,0 +1,72 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "partition_and_pq.h" +#include "utils.h" + +#include +#include +#include +#include + +template +int analyze_norm(std::string base_file) { + std::cout<<"Analyzing data norms" << std::endl; + T* data; + _u64 npts, ndims; + diskann::load_bin(base_file, data, npts, ndims); + std::vector norms(npts, 0); + #pragma omp parallel for schedule(dynamic) + for (_u32 i = 0; i +int aux_main(int argc, char** argv) { + + std::string base_file(argv[2]); + _u32 option = atoi(argv[3]); + if (option == 1) + analyze_norm(base_file); + 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]" + << std::endl; + exit(-1); + } + + if (std::string(argv[1]) == std::string("float")) { + aux_main(argc, argv); + } else if (std::string(argv[1]) == std::string("int8")) { + aux_main(argc, argv); + } else if (std::string(argv[1]) == std::string("uint8")) { + aux_main(argc, argv); + } else + std::cout << "Unsupported type. Use float/int8/uint8." << std::endl; + return 0; +} From d5b3a29c7deb2f1c4d142f8b1314246d87d3a704 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Sat, 12 Jun 2021 03:34:34 +0000 Subject: [PATCH 21/84] fixed a bug which used L2 instead of inner product in cached beam search --- include/pq_table.h | 14 ++++++++++++++ src/pq_flash_index.cpp | 17 ++++++++++++----- 2 files changed, 26 insertions(+), 5 deletions(-) diff --git a/include/pq_table.h b/include/pq_table.h index 825b29c7ec..7c8438e01e 100644 --- a/include/pq_table.h +++ b/include/pq_table.h @@ -162,6 +162,20 @@ _u32 get_num_chunks() { return res; } + float inner_product(const T* query_vec, _u8* base_vec) { + float res = 0; + for (_u64 chunk = 0; chunk < n_chunks; chunk++) { + for (_u64 j = chunk_offsets[chunk]; j < chunk_offsets[chunk + 1]; j++) { + _u64 permuted_dim_in_query = rearrangement[j]; + const float* centers_dim_vec = tables_T + (256 * j); + float diff = centers_dim_vec[base_vec[chunk]]*query_vec[permuted_dim_in_query]; // assumes centroid is 0 to prevent translation errors + res += diff; + } + } + return -res; // returns negative value to simulate distances (max -> min conversion) + } + + void inflate_vector(_u8* base_vec, float* out_vec) { for (_u64 chunk = 0; chunk < n_chunks; chunk++) { for (_u64 j = chunk_offsets[chunk]; j < chunk_offsets[chunk + 1]; j++) { diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index 4a30053346..1943e51241 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -844,9 +844,9 @@ if (file_exists(disk_pq_pivots_path)) { // query <-> PQ chunk centers distances float *pq_dists = query_scratch->aligned_pqtable_dist_scratch; -// if (metric==diskann::Metric::INNER_PRODUCT) -// pq_table.populate_chunk_inner_products(query, pq_dists); -// else if (metric==diskann::Metric::L2) + if (metric==diskann::Metric::INNER_PRODUCT) + pq_table.populate_chunk_inner_products(query, pq_dists); + else if (metric==diskann::Metric::L2) pq_table.populate_chunk_distances(query, pq_dists); // query <-> neighbor list @@ -1000,6 +1000,9 @@ if (file_exists(disk_pq_pivots_path)) { cur_expanded_dist = dist_cmp->compare(query, node_fp_coords_copy, (unsigned) aligned_dim); else { + if (metric == diskann::Metric::INNER_PRODUCT) + cur_expanded_dist = disk_pq_table.inner_product(query, (_u8*) node_fp_coords_copy); + else cur_expanded_dist = disk_pq_table.compare(query, (_u8*) node_fp_coords_copy); } full_retset.push_back( @@ -1075,8 +1078,12 @@ if (file_exists(disk_pq_pivots_path)) { if (!use_disk_index_pq) cur_expanded_dist = dist_cmp->compare(query, node_fp_coords_copy, (unsigned) aligned_dim); - else - cur_expanded_dist = disk_pq_table.compare(query, (_u8*) node_fp_coords_copy); + else { + if (metric == diskann::Metric::INNER_PRODUCT) + cur_expanded_dist = disk_pq_table.inner_product(query, (_u8*) node_fp_coords_copy); + else + cur_expanded_dist = disk_pq_table.compare(query, (_u8*) node_fp_coords_copy); + } full_retset.push_back( Neighbor(frontier_nhood.first, cur_expanded_dist, true)); From b51fea767c7ab66cf366fd7fb4324325073e01cd Mon Sep 17 00:00:00 2001 From: ravishankar Date: Mon, 14 Jun 2021 11:38:46 +0000 Subject: [PATCH 22/84] now setting up the normalizing approach --- tests/utils/vector_analysis.cpp | 50 +++++++++++++++++++++++++++++++-- 1 file changed, 48 insertions(+), 2 deletions(-) diff --git a/tests/utils/vector_analysis.cpp b/tests/utils/vector_analysis.cpp index 1cd2eb9ff5..9f5e5daa9f 100644 --- a/tests/utils/vector_analysis.cpp +++ b/tests/utils/vector_analysis.cpp @@ -36,10 +36,52 @@ int analyze_norm(std::string base_file) { for (_u32 p = 0; p < 100; p+=5) std::cout<<"percentile "< +int augment_base(std::string base_file,std::string out_file, bool prep_base = true) { + std::cout<<"Analyzing data norms" << std::endl; + T* data; + _u64 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 (_u32 i = 0; 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; + _u64 newdims = ndims + 1; + new_data = new T[npts*newdims]; + for (_u64 i = 0;i < npts; i++) { + for (_u64 j = 0; j < ndims; j++) { + new_data[i*newdims + j] = data[i*ndims +j]/ max_norm; + } + if (prep_base) { + 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<(out_file, new_data, npts, newdims); + delete[] new_data; + delete[] data; + return 0; +} + template int aux_main(int argc, char** argv) { @@ -48,14 +90,18 @@ int aux_main(int argc, char** argv) { _u32 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); return 0; } int main(int argc, char** argv) { - if (argc != 4) { + if (argc < 4) { std::cout << argv[0] << " data_type [float/int8/uint8] base_bin_file " - "[option: 1-norm analysis]" + "[option: 1-norm analysis, 2-prep_base_for_mip, 3-prep_query_for_mip] [out_file for options 2/3]" << std::endl; exit(-1); } From a2e4b92858761aa0a70d1d7602b9a2825a70e6fb Mon Sep 17 00:00:00 2001 From: ravishankar Date: Mon, 14 Jun 2021 12:15:48 +0000 Subject: [PATCH 23/84] towards pre-processing data --- include/utils.h | 24 ++++++++++++++++++++++++ src/aux_utils.cpp | 27 ++++++++++++++++----------- src/index.cpp | 4 ++-- src/partition_and_pq.cpp | 2 ++ tests/utils/vector_analysis.cpp | 2 +- 5 files changed, 45 insertions(+), 14 deletions(-) diff --git a/include/utils.h b/include/utils.h index 3e9bafc088..5d76443c99 100644 --- a/include/utils.h +++ b/include/utils.h @@ -52,6 +52,8 @@ typedef int FileHandle; #define IS_512_ALIGNED(X) IS_ALIGNED(X, 512) #define IS_4096_ALIGNED(X) IS_ALIGNED(X, 4096) + + typedef uint64_t _u64; typedef int64_t _s64; typedef uint32_t _u32; @@ -415,6 +417,28 @@ namespace diskann { } } +template +void 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); + _u64 npts, in_dims, out_dims; + float max_norm = 0; + + _u32 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; + + size_t BLOCK_SIZE = 5000000; + size_t block_size = npts <= BLOCK_SIZE ? npts : BLOCK_SIZE; + std::unique_ptr block_data_T = std::make_unique(block_size * out_dims); + +} + // 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) { diff --git a/src/aux_utils.cpp b/src/aux_utils.cpp index a08bb15e1f..97a04d7b09 100644 --- a/src/aux_utils.cpp +++ b/src/aux_utils.cpp @@ -642,7 +642,7 @@ namespace diskann { << std::endl; return false; } - + _u32 disk_pq_dims = 0; bool use_disk_pq = false; @@ -653,9 +653,8 @@ namespace diskann { use_disk_pq = false; } - if (compareMetric == diskann::Metric::INNER_PRODUCT) { - std::cout<<"Using Inner Product for PQ and Graph Generation" << std::endl; - } + std::string base_file(dataFilePath); + std::string data_file_to_use = base_file; std::string index_prefix_path(indexFilePath); std::string pq_pivots_path = index_prefix_path + "_pq_pivots.bin"; std::string pq_compressed_vectors_path = @@ -670,6 +669,12 @@ namespace diskann { index_prefix_path + "_disk.index_pq_compressed.bin"; + if (compareMetric == diskann::Metric::INNER_PRODUCT) { + std::cout<<"Using Inner Product search, so need to pre-process base data into temp file. Please ensure there is sufficient space!!" << std::endl; + std::string prepped_base = index_prefix_path + "_prepped_base.bin"; + data_file_to_use = prepped_base; + diskann::prepare_base_for_inner_products(base_file, prepped_base); + } unsigned R = (unsigned) atoi(param_list[0].c_str()); unsigned L = (unsigned) atoi(param_list[1].c_str()); @@ -704,7 +709,7 @@ namespace diskann { size_t points_num, dim; - diskann::get_bin_metadata(dataFilePath, points_num, dim); + diskann::get_bin_metadata(data_file_to_use.c_str(), points_num, dim); size_t num_pq_chunks = (size_t)(std::floor)(_u64(final_index_ram_limit / points_num)); @@ -723,7 +728,7 @@ namespace diskann { double p_val = ((double) TRAINING_SET_SIZE / (double) points_num); // generates random sample and sets it to train_data and updates // train_size - gen_random_slice(dataFilePath, p_val, train_data, train_size, + gen_random_slice(data_file_to_use.c_str(), p_val, train_data, train_size, train_dim); if (use_disk_pq) { @@ -733,7 +738,7 @@ namespace diskann { std::cout<<"Compressing base for disk-PQ into " << disk_pq_dims << " chunks " << std::endl; generate_pq_pivots(train_data, train_size, (uint32_t) dim, 256, (uint32_t) disk_pq_dims, NUM_KMEANS_REPS, disk_pq_pivots_path, false); - generate_pq_data_from_pivots(dataFilePath, 256, (uint32_t) disk_pq_dims, + generate_pq_data_from_pivots(data_file_to_use.c_str(), 256, (uint32_t) disk_pq_dims, disk_pq_pivots_path, disk_pq_compressed_vectors_path); } @@ -747,7 +752,7 @@ namespace diskann { make_zero_mean = false; generate_pq_pivots(train_data, train_size, (uint32_t) dim, 256, (uint32_t) num_pq_chunks, NUM_KMEANS_REPS, pq_pivots_path, make_zero_mean); - generate_pq_data_from_pivots(dataFilePath, 256, (uint32_t) + generate_pq_data_from_pivots(data_file_to_use.c_str(), 256, (uint32_t) num_pq_chunks, pq_pivots_path, pq_compressed_vectors_path); delete[] train_data; @@ -755,20 +760,20 @@ namespace diskann { train_data = nullptr; diskann::build_merged_vamana_index( - dataFilePath, compareMetric, L, R, p_val, indexing_ram_budget, + data_file_to_use.c_str(), compareMetric, L, R, p_val, indexing_ram_budget, mem_index_path, medoids_path, centroids_path); if (!use_disk_pq) - diskann::create_disk_layout(dataFilePath, mem_index_path, + diskann::create_disk_layout(data_file_to_use.c_str(), mem_index_path, disk_index_path); else diskann::create_disk_layout<_u8>(disk_pq_compressed_vectors_path, mem_index_path, disk_index_path); double sample_sampling_rate = (150000.0 / points_num); - gen_random_slice(dataFilePath, sample_base_prefix, + gen_random_slice(data_file_to_use.c_str(), sample_base_prefix, sample_sampling_rate); std::remove(mem_index_path.c_str()); diff --git a/src/index.cpp b/src/index.cpp index ff00f6dc61..b8d5c26642 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -504,9 +504,9 @@ namespace diskann { std::vector init_ids, std::vector & expanded_nodes_info, tsl::robin_set &expanded_nodes_ids) { - const T * node_coords = _data + _aligned_dim * node_id; + T * node_coords = _data + _aligned_dim * node_id; std::vector best_L_nodes; - + if (init_ids.size() == 0) init_ids.emplace_back(_ep); diff --git a/src/partition_and_pq.cpp b/src/partition_and_pq.cpp index 5e34f35ad1..6c8b0a431c 100644 --- a/src/partition_and_pq.cpp +++ b/src/partition_and_pq.cpp @@ -35,7 +35,9 @@ #include #endif +// block size for reading/ processing large files and matrices in blocks #define BLOCK_SIZE 5000000 + //#define SAVE_INFLATED_PQ true template diff --git a/tests/utils/vector_analysis.cpp b/tests/utils/vector_analysis.cpp index 9f5e5daa9f..41ee323c7f 100644 --- a/tests/utils/vector_analysis.cpp +++ b/tests/utils/vector_analysis.cpp @@ -70,7 +70,7 @@ std::cout<<"Max norm: " << max_norm << std::endl; diff = diff <= 0 ? 0 : std::sqrt(diff); new_data[i*newdims + ndims] = diff; if (diff <= 0) { - std::cout< Date: Mon, 14 Jun 2021 16:41:45 +0000 Subject: [PATCH 24/84] working towards newer inner product --- include/pq_flash_index.h | 4 +-- include/pq_table.h | 10 +++--- include/utils.h | 47 +++++++++++++++++++++++-- src/aux_utils.cpp | 37 +++++++++++++++----- src/pq_flash_index.cpp | 68 ++++++++++++++++++++++--------------- tests/search_disk_index.cpp | 6 ++-- 6 files changed, 125 insertions(+), 47 deletions(-) diff --git a/include/pq_flash_index.h b/include/pq_flash_index.h index 489afbed31..fbc6fc62b3 100644 --- a/include/pq_flash_index.h +++ b/include/pq_flash_index.h @@ -147,7 +147,7 @@ namespace diskann { // pq_tables = float* [[2^8 * [chunk_size]] * n_chunks] _u8 * data = nullptr; _u64 n_chunks; - FixedChunkPQTable pq_table; + FixedChunkPQTable pq_table; // distance comparator Distance * dist_cmp = nullptr; @@ -156,7 +156,7 @@ namespace diskann { // for very large datasets: we use PQ even for the disk resident index bool use_disk_index_pq = false; _u64 disk_pq_n_chunks; - FixedChunkPQTable disk_pq_table; + FixedChunkPQTable disk_pq_table; // medoid/start info diff --git a/include/pq_table.h b/include/pq_table.h index 7c8438e01e..927e5bc7db 100644 --- a/include/pq_table.h +++ b/include/pq_table.h @@ -6,7 +6,7 @@ #include "utils.h" namespace diskann { - template +// template class FixedChunkPQTable { // data_dim = n_chunks * chunk_size; float* tables = @@ -130,7 +130,7 @@ _u32 get_num_chunks() { return n_chunks; } void - populate_chunk_distances(const T* query_vec, float* dist_vec) { + populate_chunk_distances(const float* query_vec, float* dist_vec) { memset(dist_vec, 0, 256 * n_chunks * sizeof(float)); // chunk wise distance computation for (_u64 chunk = 0; chunk < n_chunks; chunk++) { @@ -149,7 +149,7 @@ _u32 get_num_chunks() { } } - float compare(const T* query_vec, _u8* base_vec) { + float compare(const float* query_vec, _u8* base_vec) { float res = 0; for (_u64 chunk = 0; chunk < n_chunks; chunk++) { for (_u64 j = chunk_offsets[chunk]; j < chunk_offsets[chunk + 1]; j++) { @@ -162,7 +162,7 @@ _u32 get_num_chunks() { return res; } - float inner_product(const T* query_vec, _u8* base_vec) { + float inner_product(const float* query_vec, _u8* base_vec) { float res = 0; for (_u64 chunk = 0; chunk < n_chunks; chunk++) { for (_u64 j = chunk_offsets[chunk]; j < chunk_offsets[chunk + 1]; j++) { @@ -187,7 +187,7 @@ _u32 get_num_chunks() { } void - populate_chunk_inner_products(const T* query_vec, float* dist_vec) { + 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 (_u64 chunk = 0; chunk < n_chunks; chunk++) { diff --git a/include/utils.h b/include/utils.h index 5d76443c99..ae90cacdd8 100644 --- a/include/utils.h +++ b/include/utils.h @@ -432,11 +432,54 @@ void prepare_base_for_inner_products(const std::string in_file, const std::strin npts = npts32; in_dims = dims32; out_dims = in_dims+1; + _u32 outdims32 = (_u32) out_dims; - size_t BLOCK_SIZE = 5000000; + 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 block_data_T = std::make_unique(block_size * out_dims); + 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); + _u64 num_blocks = DIV_ROUND_UP(npts, block_size); + + std::vector norms(npts, 0); + + for (_u64 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; + in_reader.read((char *) in_block_data.get(), block_pts * in_dims * sizeof(T)); + for (_u64 p = 0; p < block_pts; p++) { + for (_u64 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(_u32), std::ios::beg); + for (_u64 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; + in_reader.read((char *) in_block_data.get(), block_pts * in_dims * sizeof(T)); + for (_u64 p = 0; p < block_pts; p++) { + for (_u64 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.close(); } // plain saves data as npts X ndims array into filename diff --git a/src/aux_utils.cpp b/src/aux_utils.cpp index 97a04d7b09..d57fa0cff6 100644 --- a/src/aux_utils.cpp +++ b/src/aux_utils.cpp @@ -360,7 +360,7 @@ namespace diskann { paras.Set("L", (unsigned) L); paras.Set("R", (unsigned) R); paras.Set("C", 750); - paras.Set("alpha", 2.0f); + paras.Set("alpha", 1.2f); paras.Set("num_rnds", 2); paras.Set("saturate_graph", 1); paras.Set("save_path", mem_index_path); @@ -738,6 +738,11 @@ namespace diskann { std::cout<<"Compressing base for disk-PQ into " << disk_pq_dims << " chunks " << std::endl; generate_pq_pivots(train_data, train_size, (uint32_t) dim, 256, (uint32_t) disk_pq_dims, NUM_KMEANS_REPS, disk_pq_pivots_path, false); + if (compareMetric == diskann::Metric::INNER_PRODUCT) + generate_pq_data_from_pivots(data_file_to_use.c_str(), 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.c_str(), 256, (uint32_t) disk_pq_dims, disk_pq_pivots_path, disk_pq_compressed_vectors_path); @@ -745,13 +750,17 @@ namespace diskann { diskann::cout << "Training data loaded of size " << train_size << std::endl; - wait_for_keystroke(); +// wait_for_keystroke(); bool make_zero_mean = true; if (compareMetric == diskann::Metric::INNER_PRODUCT) make_zero_mean = false; generate_pq_pivots(train_data, train_size, (uint32_t) dim, 256, (uint32_t) num_pq_chunks, NUM_KMEANS_REPS, pq_pivots_path, make_zero_mean); + if (compareMetric == diskann::Metric::INNER_PRODUCT) + generate_pq_data_from_pivots(data_file_to_use.c_str(), 256, (uint32_t) + num_pq_chunks, pq_pivots_path, pq_compressed_vectors_path); + else generate_pq_data_from_pivots(data_file_to_use.c_str(), 256, (uint32_t) num_pq_chunks, pq_pivots_path, pq_compressed_vectors_path); @@ -759,24 +768,36 @@ namespace diskann { train_data = nullptr; + if (compareMetric == diskann::Metric::INNER_PRODUCT) + 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); + else diskann::build_merged_vamana_index( - data_file_to_use.c_str(), compareMetric, L, R, p_val, indexing_ram_budget, + data_file_to_use.c_str(), diskann::Metric::L2, L, R, p_val, indexing_ram_budget, mem_index_path, medoids_path, centroids_path); - - - if (!use_disk_pq) - diskann::create_disk_layout(data_file_to_use.c_str(), mem_index_path, + if (!use_disk_pq) { + if (compareMetric == diskann::Metric::INNER_PRODUCT) + diskann::create_disk_layout(data_file_to_use.c_str(), mem_index_path, disk_index_path); + else + diskann::create_disk_layout(data_file_to_use.c_str(), mem_index_path, + disk_index_path); + } else diskann::create_disk_layout<_u8>(disk_pq_compressed_vectors_path, mem_index_path, disk_index_path); double sample_sampling_rate = (150000.0 / points_num); + if (compareMetric == diskann::Metric::INNER_PRODUCT) + gen_random_slice(data_file_to_use.c_str(), sample_base_prefix, + sample_sampling_rate); + else gen_random_slice(data_file_to_use.c_str(), sample_base_prefix, sample_sampling_rate); - std::remove(mem_index_path.c_str()); +// std::remove(mem_index_path.c_str()); if (use_disk_pq) std::remove(disk_pq_compressed_vectors_path.c_str()); diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index 1943e51241..d140b0f77d 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -614,7 +614,10 @@ namespace diskann { this->data_dim = pq_file_dim; this->disk_data_dim = this->data_dim; // will reset later if we use PQ on disk - this->disk_bytes_per_point = this->data_dim * sizeof(T); // will change later if we use PQ on disk + this->disk_bytes_per_point = this->data_dim * sizeof(T); // will change later if we use PQ on disk or if we are using inner product without PQ + if (metric == diskann::Metric::INNER_PRODUCT) { + this->disk_bytes_per_point = this->data_dim * sizeof(float); // because we normalize the data and store it as float if no PQ + } this->aligned_dim = ROUND_UP(pq_file_dim, 8); @@ -815,10 +818,24 @@ if (file_exists(disk_pq_pivots_path)) { data = this->thread_data.pop(); } + float query_norm = 0; + if (metric == diskann::Metric::L2) { for (uint32_t i = 0; i < this->data_dim; i++) { data.scratch.aligned_query_float[i] = query1[i]; data.scratch.aligned_query_T[i] = query1[i]; } + } else if (metric == diskann::Metric::INNER_PRODUCT) { + for (uint32_t i = 0; i < this->data_dim - 1; i++) { + data.scratch.aligned_query_float[i] = query1[i]; + query_norm += query1[i]*query1[i]; + } + query_norm = std::sqrt(query_norm); + data.scratch.aligned_query_float[this->data_dim -1] = 0; + for (uint32_t i = 0; i < this->data_dim - 1; i++) { + data.scratch.aligned_query_float[i] /= query_norm; + } + } + // memcpy(data.scratch.aligned_query_T, query1, disk_bytes_per_point); const T * query = data.scratch.aligned_query_T; const float *query_float = data.scratch.aligned_query_float; @@ -844,10 +861,10 @@ if (file_exists(disk_pq_pivots_path)) { // query <-> PQ chunk centers distances float *pq_dists = query_scratch->aligned_pqtable_dist_scratch; - if (metric==diskann::Metric::INNER_PRODUCT) - pq_table.populate_chunk_inner_products(query, pq_dists); - else if (metric==diskann::Metric::L2) - pq_table.populate_chunk_distances(query, pq_dists); +// if (metric==diskann::Metric::INNER_PRODUCT) +// pq_table.populate_chunk_inner_products(query, pq_dists); +// else if (metric==diskann::Metric::L2) + pq_table.populate_chunk_distances(query_float, pq_dists); // query <-> neighbor list float *dist_scratch = query_scratch->aligned_dist_scratch; @@ -889,15 +906,6 @@ if (file_exists(disk_pq_pivots_path)) { retset[0].flag = true; visited.insert(best_medoid); -/* - std::cout<<"Chose " << retset[0].id<< " as best medoid with distance " << retset[0].distance << std::endl; - - std::cout<<"query from 0 to " << aligned_dim << std::endl; - for (_u32 i = 0; i < aligned_dim-1; i++) { - std::cout<second; float cur_expanded_dist; - if (!use_disk_index_pq) + if (!use_disk_index_pq) { + if (metric == diskann::Metric::INNER_PRODUCT) + cur_expanded_dist = dist_cmp_float->compare(query_float, (float*) node_fp_coords_copy, + (unsigned) aligned_dim); + else cur_expanded_dist = dist_cmp->compare(query, node_fp_coords_copy, (unsigned) aligned_dim); + } else { if (metric == diskann::Metric::INNER_PRODUCT) - cur_expanded_dist = disk_pq_table.inner_product(query, (_u8*) node_fp_coords_copy); + cur_expanded_dist = disk_pq_table.inner_product(query_float, (_u8*) node_fp_coords_copy); else - cur_expanded_dist = disk_pq_table.compare(query, (_u8*) node_fp_coords_copy); + cur_expanded_dist = disk_pq_table.compare(query_float, (_u8*) node_fp_coords_copy); } full_retset.push_back( Neighbor((unsigned) cached_nhood.first, cur_expanded_dist, true)); @@ -1075,14 +1082,21 @@ if (file_exists(disk_pq_pivots_path)) { 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(query, node_fp_coords_copy, + if (!use_disk_index_pq) { + if (metric == diskann::Metric::INNER_PRODUCT) + cur_expanded_dist = dist_cmp_float->compare(query_float, (float*) node_fp_coords_copy, (unsigned) aligned_dim); + else + cur_expanded_dist = dist_cmp->compare(query, node_fp_coords_copy, + (unsigned) aligned_dim); + + + } else { if (metric == diskann::Metric::INNER_PRODUCT) - cur_expanded_dist = disk_pq_table.inner_product(query, (_u8*) node_fp_coords_copy); + cur_expanded_dist = disk_pq_table.inner_product(query_float, (_u8*) node_fp_coords_copy); else - cur_expanded_dist = disk_pq_table.compare(query, (_u8*) node_fp_coords_copy); + cur_expanded_dist = disk_pq_table.compare(query_float, (_u8*) node_fp_coords_copy); } full_retset.push_back( Neighbor(frontier_nhood.first, cur_expanded_dist, true)); diff --git a/tests/search_disk_index.cpp b/tests/search_disk_index.cpp index fe2fe5b9e5..e89c562a63 100644 --- a/tests/search_disk_index.cpp +++ b/tests/search_disk_index.cpp @@ -137,9 +137,9 @@ int search_disk_index(int argc, char** argv) { std::vector node_list; diskann::cout << "Caching " << num_nodes_to_cache << " BFS nodes around medoid(s)" << std::endl; - _pFlashIndex->cache_bfs_levels(num_nodes_to_cache, node_list); -// _pFlashIndex->generate_cache_list_from_sample_queries( -// warmup_query_file, 15, 6, num_nodes_to_cache, num_threads, node_list); +// _pFlashIndex->cache_bfs_levels(num_nodes_to_cache, node_list); + _pFlashIndex->generate_cache_list_from_sample_queries( + warmup_query_file, 15, 6, num_nodes_to_cache, num_threads, node_list); _pFlashIndex->load_cache_list(node_list); node_list.clear(); node_list.shrink_to_fit(); From f5e55d17da0af357374f5c0e27e867a827dbd9e9 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Tue, 15 Jun 2021 07:58:47 +0000 Subject: [PATCH 25/84] more changes to do MIPS by reducing to L2 with extra coordinate --- src/aux_utils.cpp | 8 ++++---- tests/search_disk_index.cpp | 6 +++--- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/src/aux_utils.cpp b/src/aux_utils.cpp index d57fa0cff6..34cde93cc2 100644 --- a/src/aux_utils.cpp +++ b/src/aux_utils.cpp @@ -362,7 +362,7 @@ namespace diskann { paras.Set("C", 750); paras.Set("alpha", 1.2f); paras.Set("num_rnds", 2); - paras.Set("saturate_graph", 1); + paras.Set("saturate_graph", 0); paras.Set("save_path", mem_index_path); std::unique_ptr> _pvamanaIndex = @@ -379,7 +379,7 @@ namespace diskann { partition_with_ram_budget(base_file, sampling_rate, ram_budget, 2 * R / 3, merged_index_prefix, 2); - wait_for_keystroke(); +// wait_for_keystroke(); std::string cur_centroid_filepath = merged_index_prefix + "_centroids.bin"; @@ -401,9 +401,9 @@ namespace diskann { paras.Set("L", L); paras.Set("R", (2 * (R / 3))); paras.Set("C", 750); - paras.Set("alpha", 2.0f); + paras.Set("alpha", 1.2f); paras.Set("num_rnds", 2); - paras.Set("saturate_graph", 1); + paras.Set("saturate_graph", 0); paras.Set("save_path", shard_index_file); std::unique_ptr> _pvamanaIndex = diff --git a/tests/search_disk_index.cpp b/tests/search_disk_index.cpp index e89c562a63..fe2fe5b9e5 100644 --- a/tests/search_disk_index.cpp +++ b/tests/search_disk_index.cpp @@ -137,9 +137,9 @@ int search_disk_index(int argc, char** argv) { std::vector node_list; diskann::cout << "Caching " << num_nodes_to_cache << " BFS nodes around medoid(s)" << std::endl; -// _pFlashIndex->cache_bfs_levels(num_nodes_to_cache, node_list); - _pFlashIndex->generate_cache_list_from_sample_queries( - warmup_query_file, 15, 6, num_nodes_to_cache, num_threads, node_list); + _pFlashIndex->cache_bfs_levels(num_nodes_to_cache, node_list); +// _pFlashIndex->generate_cache_list_from_sample_queries( +// warmup_query_file, 15, 6, num_nodes_to_cache, num_threads, node_list); _pFlashIndex->load_cache_list(node_list); node_list.clear(); node_list.shrink_to_fit(); From 6fe07e01174161fd83863950c7174ea20ef1d68a Mon Sep 17 00:00:00 2001 From: ravishankar Date: Tue, 15 Jun 2021 16:18:55 +0000 Subject: [PATCH 26/84] cleaned up code a bit, need to test everything again --- src/aux_utils.cpp | 37 +++++++++++++++---------------------- src/partition_and_pq.cpp | 2 +- src/pq_flash_index.cpp | 28 ++++++---------------------- 3 files changed, 22 insertions(+), 45 deletions(-) diff --git a/src/aux_utils.cpp b/src/aux_utils.cpp index 34cde93cc2..ef362351e8 100644 --- a/src/aux_utils.cpp +++ b/src/aux_utils.cpp @@ -379,9 +379,7 @@ namespace diskann { partition_with_ram_budget(base_file, sampling_rate, ram_budget, 2 * R / 3, merged_index_prefix, 2); -// wait_for_keystroke(); - - + wait_for_keystroke(); std::string cur_centroid_filepath = merged_index_prefix + "_centroids.bin"; std::rename(cur_centroid_filepath.c_str(), centroids_file.c_str()); @@ -643,9 +641,18 @@ namespace diskann { return false; } + + 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. Please contact us if you need other scenarios." << std::endl; + throw diskann::ANNException(stream.str(), -1); + + } + _u32 disk_pq_dims = 0; bool use_disk_pq = false; +// 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() == 6) { disk_pq_dims = atoi(param_list[5].c_str()); use_disk_pq = true; @@ -668,7 +675,7 @@ namespace diskann { std::string disk_pq_compressed_vectors_path = // optional if disk index is also storing pq data 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) { std::cout<<"Using Inner Product search, so need to pre-process base data into temp file. Please ensure there is sufficient space!!" << std::endl; std::string prepped_base = index_prefix_path + "_prepped_base.bin"; @@ -750,17 +757,16 @@ namespace diskann { diskann::cout << "Training data loaded of size " << train_size << std::endl; -// wait_for_keystroke(); + wait_for_keystroke(); +// don't translate data to make zero mean for PQ compression. We must not translate for inner product search. bool make_zero_mean = true; if (compareMetric == diskann::Metric::INNER_PRODUCT) make_zero_mean = false; + generate_pq_pivots(train_data, train_size, (uint32_t) dim, 256, (uint32_t) num_pq_chunks, NUM_KMEANS_REPS, pq_pivots_path, make_zero_mean); - if (compareMetric == diskann::Metric::INNER_PRODUCT) - generate_pq_data_from_pivots(data_file_to_use.c_str(), 256, (uint32_t) - num_pq_chunks, pq_pivots_path, pq_compressed_vectors_path); - else + generate_pq_data_from_pivots(data_file_to_use.c_str(), 256, (uint32_t) num_pq_chunks, pq_pivots_path, pq_compressed_vectors_path); @@ -768,20 +774,11 @@ namespace diskann { train_data = nullptr; - if (compareMetric == diskann::Metric::INNER_PRODUCT) - 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); - else 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); if (!use_disk_pq) { - if (compareMetric == diskann::Metric::INNER_PRODUCT) - diskann::create_disk_layout(data_file_to_use.c_str(), mem_index_path, - disk_index_path); - else diskann::create_disk_layout(data_file_to_use.c_str(), mem_index_path, disk_index_path); } @@ -790,10 +787,6 @@ namespace diskann { disk_index_path); double sample_sampling_rate = (150000.0 / points_num); - if (compareMetric == diskann::Metric::INNER_PRODUCT) - gen_random_slice(data_file_to_use.c_str(), sample_base_prefix, - sample_sampling_rate); - else gen_random_slice(data_file_to_use.c_str(), sample_base_prefix, sample_sampling_rate); diff --git a/src/partition_and_pq.cpp b/src/partition_and_pq.cpp index 6c8b0a431c..6b047f7316 100644 --- a/src/partition_and_pq.cpp +++ b/src/partition_and_pq.cpp @@ -745,7 +745,7 @@ int shard_data_into_clusters(const std::string data_file, float *pivots, - +// 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, diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index d140b0f77d..864f009b82 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -615,9 +615,6 @@ namespace diskann { this->data_dim = pq_file_dim; this->disk_data_dim = this->data_dim; // will reset later if we use PQ on disk this->disk_bytes_per_point = this->data_dim * sizeof(T); // will change later if we use PQ on disk or if we are using inner product without PQ - if (metric == diskann::Metric::INNER_PRODUCT) { - this->disk_bytes_per_point = this->data_dim * sizeof(float); // because we normalize the data and store it as float if no PQ - } this->aligned_dim = ROUND_UP(pq_file_dim, 8); @@ -818,17 +815,18 @@ if (file_exists(disk_pq_pivots_path)) { data = this->thread_data.pop(); } +// copy query to thread specific aligned and allocated memory (for distance calculations we need aligned data) + float query_norm = 0; - if (metric == diskann::Metric::L2) { + for (uint32_t i = 0; i < this->data_dim; i++) { data.scratch.aligned_query_float[i] = query1[i]; data.scratch.aligned_query_T[i] = query1[i]; - } - } else if (metric == diskann::Metric::INNER_PRODUCT) { - for (uint32_t i = 0; i < this->data_dim - 1; i++) { - data.scratch.aligned_query_float[i] = query1[i]; query_norm += query1[i]*query1[i]; } + +// 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) { query_norm = std::sqrt(query_norm); data.scratch.aligned_query_float[this->data_dim -1] = 0; for (uint32_t i = 0; i < this->data_dim - 1; i++) { @@ -836,7 +834,6 @@ if (file_exists(disk_pq_pivots_path)) { } } -// memcpy(data.scratch.aligned_query_T, query1, disk_bytes_per_point); const T * query = data.scratch.aligned_query_T; const float *query_float = data.scratch.aligned_query_float; @@ -861,9 +858,6 @@ if (file_exists(disk_pq_pivots_path)) { // query <-> PQ chunk centers distances float *pq_dists = query_scratch->aligned_pqtable_dist_scratch; -// if (metric==diskann::Metric::INNER_PRODUCT) -// pq_table.populate_chunk_inner_products(query, pq_dists); -// else if (metric==diskann::Metric::L2) pq_table.populate_chunk_distances(query_float, pq_dists); // query <-> neighbor list @@ -999,10 +993,6 @@ if (file_exists(disk_pq_pivots_path)) { T * node_fp_coords_copy = global_cache_iter->second; float cur_expanded_dist; if (!use_disk_index_pq) { - if (metric == diskann::Metric::INNER_PRODUCT) - cur_expanded_dist = dist_cmp_float->compare(query_float, (float*) node_fp_coords_copy, - (unsigned) aligned_dim); - else cur_expanded_dist = dist_cmp->compare(query, node_fp_coords_copy, (unsigned) aligned_dim); } @@ -1083,14 +1073,8 @@ if (file_exists(disk_pq_pivots_path)) { float cur_expanded_dist; if (!use_disk_index_pq) { - if (metric == diskann::Metric::INNER_PRODUCT) - cur_expanded_dist = dist_cmp_float->compare(query_float, (float*) node_fp_coords_copy, - (unsigned) aligned_dim); - else cur_expanded_dist = dist_cmp->compare(query, node_fp_coords_copy, (unsigned) aligned_dim); - - } else { if (metric == diskann::Metric::INNER_PRODUCT) From aa707f221755d51a5a3cb270ea79996072467eff Mon Sep 17 00:00:00 2001 From: ravishankar Date: Wed, 16 Jun 2021 09:16:28 +0000 Subject: [PATCH 27/84] testing underway --- include/aux_utils.h | 2 ++ include/utils.h | 1 + src/aux_utils.cpp | 6 +++--- src/index.cpp | 3 ++- 4 files changed, 8 insertions(+), 4 deletions(-) diff --git a/include/aux_utils.h b/include/aux_utils.h index ac7bad6cc8..ca0dbce57c 100644 --- a/include/aux_utils.h +++ b/include/aux_utils.h @@ -29,6 +29,8 @@ typedef int FileHandle; #include "common_includes.h" #include "utils.h" #include "windows_customizations.h" +#include "gperftools/malloc_extension.h" + namespace diskann { const size_t TRAINING_SET_SIZE = 150000; diff --git a/include/utils.h b/include/utils.h index ae90cacdd8..ae487b05ce 100644 --- a/include/utils.h +++ b/include/utils.h @@ -222,6 +222,7 @@ namespace diskann { inline void wait_for_keystroke() { int a; + std::cout<<"Press any number to continue.." << std::endl; std::cin>> a; } diff --git a/src/aux_utils.cpp b/src/aux_utils.cpp index ef362351e8..c059343e5d 100644 --- a/src/aux_utils.cpp +++ b/src/aux_utils.cpp @@ -354,7 +354,7 @@ namespace diskann { double full_index_ram = ESTIMATE_RAM_USAGE(base_num, base_dim, sizeof(T), R); if (full_index_ram < ram_budget * 1024 * 1024 * 1024) { - diskann::cout << "Full index fits in RAM, building in one shot" + diskann::cout << "Full index fits in RAM budget, should consume at most " << full_index_ram/(1024*1024*1024) <<"GBs, so building in one shot" << std::endl; diskann::Parameters paras; paras.Set("L", (unsigned) L); @@ -757,8 +757,6 @@ namespace diskann { diskann::cout << "Training data loaded of size " << train_size << std::endl; - wait_for_keystroke(); - // don't translate data to make zero mean for PQ compression. We must not translate for inner product search. bool make_zero_mean = true; if (compareMetric == diskann::Metric::INNER_PRODUCT) @@ -773,6 +771,8 @@ namespace diskann { delete[] train_data; train_data = nullptr; + MallocExtension::instance()->ReleaseFreeMemory(); + diskann::build_merged_vamana_index( data_file_to_use.c_str(), diskann::Metric::L2, L, R, p_val, indexing_ram_budget, diff --git a/src/index.cpp b/src/index.cpp index b8d5c26642..9bdc2785c0 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -137,6 +137,7 @@ namespace diskann { _compacted_order(true), _enable_tags(enable_tags), _consolidated_order(true), _support_eager_delete(support_eager_delete), _store_data(store_data) { + // data is stored to _nd * aligned_dim matrix with necessary // zero-padding diskann::cout << "Number of frozen points = " << _num_frozen_pts @@ -182,7 +183,6 @@ namespace diskann { this->_distance = ::get_distance_function(m); _locks = std::vector(_max_points + _num_frozen_pts); - _width = 0; } @@ -789,6 +789,7 @@ namespace diskann { _final_graph[p].reserve((size_t)(std::ceil(range * SLACK_FACTOR * 1.05))); } + std::random_device rd; std::mt19937 gen(rd()); std::uniform_real_distribution<> dis(0, 1); From 54630f17247fbae5cb055ac68684ef07869300e7 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Wed, 16 Jun 2021 10:05:33 +0000 Subject: [PATCH 28/84] added back saturate graph to create denser indices --- src/aux_utils.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/aux_utils.cpp b/src/aux_utils.cpp index c059343e5d..25e139ffd2 100644 --- a/src/aux_utils.cpp +++ b/src/aux_utils.cpp @@ -362,7 +362,7 @@ namespace diskann { paras.Set("C", 750); paras.Set("alpha", 1.2f); paras.Set("num_rnds", 2); - paras.Set("saturate_graph", 0); + paras.Set("saturate_graph", 1); paras.Set("save_path", mem_index_path); std::unique_ptr> _pvamanaIndex = @@ -401,7 +401,7 @@ namespace diskann { paras.Set("C", 750); paras.Set("alpha", 1.2f); paras.Set("num_rnds", 2); - paras.Set("saturate_graph", 0); + paras.Set("saturate_graph", 1); paras.Set("save_path", shard_index_file); std::unique_ptr> _pvamanaIndex = From dd210bf6917da73b78d35c719cbc7a1255ae7d3f Mon Sep 17 00:00:00 2001 From: ravishankar Date: Wed, 16 Jun 2021 11:36:10 +0000 Subject: [PATCH 29/84] now we dont sample a new test dataset every iteration for estimating sharding --- include/partition_and_pq.h | 3 +-- src/aux_utils.cpp | 3 +-- src/partition_and_pq.cpp | 48 +++++++++++++++++--------------------- 3 files changed, 23 insertions(+), 31 deletions(-) diff --git a/include/partition_and_pq.h b/include/partition_and_pq.h index bbae1ed148..0afc85410b 100644 --- a/include/partition_and_pq.h +++ b/include/partition_and_pq.h @@ -27,8 +27,7 @@ template void gen_random_slice(const T *inputdata, size_t npts, size_t ndims, double p_val, float *&sampled_data, size_t &slice_size); -template -int estimate_cluster_sizes(const std::string data_file, float *pivots, +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); diff --git a/src/aux_utils.cpp b/src/aux_utils.cpp index 25e139ffd2..5685bebee6 100644 --- a/src/aux_utils.cpp +++ b/src/aux_utils.cpp @@ -379,7 +379,6 @@ namespace diskann { partition_with_ram_budget(base_file, sampling_rate, ram_budget, 2 * R / 3, merged_index_prefix, 2); - wait_for_keystroke(); std::string cur_centroid_filepath = merged_index_prefix + "_centroids.bin"; std::rename(cur_centroid_filepath.c_str(), centroids_file.c_str()); @@ -410,7 +409,7 @@ namespace diskann { _pvamanaIndex->build(paras); _pvamanaIndex->save(shard_index_file.c_str()); std::remove(shard_base_file.c_str()); - wait_for_keystroke(); +// wait_for_keystroke(); } diff --git a/src/partition_and_pq.cpp b/src/partition_and_pq.cpp index 6b047f7316..9210c1371a 100644 --- a/src/partition_and_pq.cpp +++ b/src/partition_and_pq.cpp @@ -578,35 +578,20 @@ int generate_pq_data_from_pivots(const std::string data_file, return 0; } -template -int estimate_cluster_sizes(const std::string data_file, float *pivots, - const size_t num_centers, const size_t dim, +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 num_test, test_dim; - float *test_data_float; - double sampling_rate = 0.01; - - gen_random_slice(data_file, sampling_rate, test_data_float, num_test, - test_dim); - - if (test_dim != dim) { - diskann::cout << "Error. dimensions dont match for pivot set and base set" - << std::endl; - return -1; - } size_t *shard_counts = new size_t[num_centers]; for (size_t i = 0; i < num_centers; i++) { shard_counts[i] = 0; } - - size_t num_points = 0, num_dim = 0; - diskann::get_bin_metadata(data_file, num_points, num_dim); - size_t block_size = num_points <= BLOCK_SIZE ? num_points : BLOCK_SIZE; + + size_t block_size = num_test <= BLOCK_SIZE ? num_test : BLOCK_SIZE; _u32 * block_closest_centers = new _u32[block_size * k_base]; float *block_data_float; @@ -619,7 +604,7 @@ int estimate_cluster_sizes(const std::string data_file, float *pivots, block_data_float = test_data_float + start_id * test_dim; - math_utils::compute_closest_centers(block_data_float, cur_blk_size, dim, + math_utils::compute_closest_centers(block_data_float, cur_blk_size, test_dim, pivots, num_centers, k_base, block_closest_centers); @@ -635,8 +620,8 @@ int estimate_cluster_sizes(const std::string data_file, float *pivots, for (size_t i = 0; i < num_centers; i++) { _u32 cur_shard_count = (_u32) shard_counts[i]; cluster_sizes.push_back( - size_t(((double) cur_shard_count) * (1.0 / sampling_rate))); - diskann::cout << cur_shard_count * (1.0 / sampling_rate) << " "; + (size_t)cur_shard_count); + diskann::cout << cur_shard_count << " "; } diskann::cout << std::endl; delete[] shard_counts; @@ -952,9 +937,9 @@ int partition(const std::string data_file, const float sampling_rate, // now pivots are ready. need to stream base points and assign them to // closest clusters. - std::vector cluster_sizes; - estimate_cluster_sizes(data_file, pivot_data, num_parts, train_dim, k_base, - cluster_sizes); + //std::vector cluster_sizes; + //estimate_cluster_sizes(data_file, pivot_data, num_parts, train_dim, k_base, + // cluster_sizes); shard_data_into_clusters(data_file, pivot_data, num_parts, train_dim, k_base, prefix_path); @@ -971,7 +956,7 @@ int partition_with_ram_budget(const std::string data_file, size_t train_dim; size_t num_train; float *train_data_float; - size_t max_k_means_reps = 20; + size_t max_k_means_reps = 10; int num_parts = 3; bool fit_in_ram = false; @@ -979,6 +964,13 @@ int partition_with_ram_budget(const std::string data_file, 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); @@ -1010,10 +1002,11 @@ int partition_with_ram_budget(const std::string data_file, // closest clusters. std::vector cluster_sizes; - estimate_cluster_sizes(data_file, pivot_data, num_parts, train_dim, + estimate_cluster_sizes(test_data_float, num_test, pivot_data, num_parts, train_dim, k_base, cluster_sizes); for (auto &p : cluster_sizes) { + p = (_u64) (p/ sampling_rate); // to account for the fact that p is the size of the shard over the testing sample. double cur_shard_ram_estimate = ESTIMATE_RAM_USAGE(p, train_dim, sizeof(T), graph_degree); @@ -1037,6 +1030,7 @@ int partition_with_ram_budget(const std::string data_file, k_base, prefix_path); delete[] pivot_data; delete[] train_data_float; + delete[] test_data_float; return num_parts; } From dac9d4bda792976648e088fcba7b3aa37a4eb55c Mon Sep 17 00:00:00 2001 From: ravishankar Date: Wed, 16 Jun 2021 11:37:06 +0000 Subject: [PATCH 30/84] now num_parts increases by 2 --- src/partition_and_pq.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/partition_and_pq.cpp b/src/partition_and_pq.cpp index 9210c1371a..dd9c00a56c 100644 --- a/src/partition_and_pq.cpp +++ b/src/partition_and_pq.cpp @@ -1018,7 +1018,7 @@ int partition_with_ram_budget(const std::string data_file, << "GB, budget given is " << ram_budget << std::endl; if (max_ram_usage > 1024 * 1024 * 1024 * ram_budget) { fit_in_ram = false; - num_parts++; + num_parts+=2; } } From 2ad09538eba7e143d22c4b11af7b40f33e1d8efc Mon Sep 17 00:00:00 2001 From: ravishankar Date: Wed, 16 Jun 2021 11:50:56 +0000 Subject: [PATCH 31/84] cleaned up warnings in Debug mode compiler --- include/distance.h | 22 ++++++++++++++++------ include/index.h | 2 +- include/pq_flash_index.h | 3 +-- src/index.cpp | 2 +- src/pq_flash_index.cpp | 3 +-- tests/utils/compute_groundtruth.cpp | 14 ++++++-------- tests/utils/create_disk_layout.cpp | 8 ++++---- tests/utils/gen_random_slice.cpp | 8 ++++---- tests/utils/vector_analysis.cpp | 8 ++++---- 9 files changed, 38 insertions(+), 32 deletions(-) diff --git a/include/distance.h b/include/distance.h index a65f9e9487..1de05fefb0 100644 --- a/include/distance.h +++ b/include/distance.h @@ -255,9 +255,14 @@ namespace diskann { virtual float compare(const int8_t *a, const int8_t *b, unsigned int length) const { #ifndef _WINDOWS - std::cout << "AVX only supported in Windows build."; - return 0; - } +int32_t result = 0; +#pragma omp simd reduction(+ : result) aligned(a, b : 8) + for (_s32 i = 0; i < (_s32) 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; + } #else __m128 r = _mm_setzero_ps(); __m128i r1; @@ -302,9 +307,14 @@ namespace diskann { virtual float compare(const float *a, const float *b, unsigned int length) const { #ifndef _WINDOWS - std::cout << "AVX only supported in Windows build."; - return 0; - } +float result = 0; +#pragma omp simd reduction(+ : result) aligned(a, b : 8) + for (_s32 i = 0; i < (_s32) length; i++) { + result += (a[i] - b[i]) * + (a[i] - b[i]); + } + return result; + } #else __m128 diff, v1, v2; __m128 sum = _mm_set1_ps(0); diff --git a/include/index.h b/include/index.h index c409b96958..b800781d38 100644 --- a/include/index.h +++ b/include/index.h @@ -63,7 +63,7 @@ namespace diskann { DISKANN_DLLEXPORT std::pair search_with_tags( const T *query, const size_t K, const unsigned L, TagT *tags, - unsigned frozen_pts, unsigned *indices_buffer = NULL); + unsigned *indices_buffer = NULL); // repositions frozen points to the end of _data - if they have been moved // during deletion diff --git a/include/pq_flash_index.h b/include/pq_flash_index.h index fbc6fc62b3..f17448e7d4 100644 --- a/include/pq_flash_index.h +++ b/include/pq_flash_index.h @@ -112,8 +112,7 @@ namespace diskann { // implemented DISKANN_DLLEXPORT void cached_beam_search( const T *query, const _u64 k_search, const _u64 l_search, _u64 *res_ids, - float *res_dists, const _u64 beam_width, QueryStats *stats = nullptr, - Distance *output_dist_func = nullptr); + float *res_dists, const _u64 beam_width, QueryStats *stats = nullptr); std::shared_ptr &reader; protected: diff --git a/src/index.cpp b/src/index.cpp index 9bdc2785c0..0d01c92322 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -1080,7 +1080,7 @@ namespace diskann { template std::pair Index::search_with_tags( const T *query, const size_t K, const unsigned L, TagT *tags, - unsigned frozen_pts, unsigned *indices_buffer) { + unsigned *indices_buffer) { const bool alloc = indices_buffer == NULL; auto indices = alloc ? new unsigned[K] : indices_buffer; auto ret = search(query, K, L, indices); diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index 864f009b82..6d0176b896 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -807,8 +807,7 @@ if (file_exists(disk_pq_pivots_path)) { const _u64 l_search, _u64 *indices, float * distances, const _u64 beam_width, - QueryStats * stats, - Distance *output_dist_func) { + QueryStats * stats) { ThreadData data = this->thread_data.pop(); while (data.scratch.sector_scratch == nullptr) { this->thread_data.wait_for_push_notify(); diff --git a/tests/utils/compute_groundtruth.cpp b/tests/utils/compute_groundtruth.cpp index cdd8d0c327..8f9f77399a 100644 --- a/tests/utils/compute_groundtruth.cpp +++ b/tests/utils/compute_groundtruth.cpp @@ -108,9 +108,7 @@ 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, - 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; @@ -171,8 +169,8 @@ void exact_knn(const size_t dim, const size_t k, queries_l2sq + q_b); } else { inner_prod_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); + 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; @@ -308,7 +306,7 @@ inline void save_groundtruth_as_one_file(const std::string filename, } template -int aux_main(int argv, char **argc) { +int aux_main(char **argc) { size_t npoints, nqueries, dim; std::string base_file(argc[2]); std::string query_file(argc[3]); @@ -393,9 +391,9 @@ int main(int argc, char **argv) { } if (std::string(argv[1]) == std::string("float")) - aux_main(argc, argv); + aux_main(argv); if (std::string(argv[1]) == std::string("int8")) - aux_main(argc, argv); + aux_main(argv); if (std::string(argv[1]) == std::string("uint8")) - aux_main(argc, argv); + aux_main(argv); } diff --git a/tests/utils/create_disk_layout.cpp b/tests/utils/create_disk_layout.cpp index 458e2d6b20..21c6cacede 100644 --- a/tests/utils/create_disk_layout.cpp +++ b/tests/utils/create_disk_layout.cpp @@ -13,7 +13,7 @@ #include "utils.h" template -int create_disk_layout(int argc, char **argv) { +int create_disk_layout(char **argv) { std::string base_file(argv[2]); std::string vamana_file(argv[3]); std::string output_file(argv[4]); @@ -31,11 +31,11 @@ int main(int argc, char **argv) { int ret_val = -1; if (std::string(argv[1]) == std::string("float")) - ret_val = create_disk_layout(argc, argv); + ret_val = create_disk_layout(argv); else if (std::string(argv[1]) == std::string("int8")) - ret_val = create_disk_layout(argc, argv); + ret_val = create_disk_layout(argv); else if (std::string(argv[1]) == std::string("uint8")) - ret_val = create_disk_layout(argc, argv); + ret_val = create_disk_layout(argv); else { std::cout << "unsupported type. use int8/uint8/float " << std::endl; ret_val = -2; diff --git a/tests/utils/gen_random_slice.cpp b/tests/utils/gen_random_slice.cpp index 1b102e27c1..dccb50a13c 100644 --- a/tests/utils/gen_random_slice.cpp +++ b/tests/utils/gen_random_slice.cpp @@ -21,7 +21,7 @@ #include template -int aux_main(int argc, char** argv) { +int aux_main(char** argv) { std::string base_file(argv[2]); std::string output_prefix(argv[3]); @@ -40,11 +40,11 @@ int main(int argc, char** argv) { } if (std::string(argv[1]) == std::string("float")) { - aux_main(argc, argv); + aux_main(argv); } else if (std::string(argv[1]) == std::string("int8")) { - aux_main(argc, argv); + aux_main(argv); } else if (std::string(argv[1]) == std::string("uint8")) { - aux_main(argc, argv); + aux_main(argv); } else std::cout << "Unsupported type. Use float/int8/uint8." << std::endl; return 0; diff --git a/tests/utils/vector_analysis.cpp b/tests/utils/vector_analysis.cpp index 41ee323c7f..43d333b532 100644 --- a/tests/utils/vector_analysis.cpp +++ b/tests/utils/vector_analysis.cpp @@ -84,7 +84,7 @@ std::cout<<"Max norm: " << max_norm << std::endl; template -int aux_main(int argc, char** argv) { +int aux_main(char** argv) { std::string base_file(argv[2]); _u32 option = atoi(argv[3]); @@ -107,11 +107,11 @@ int main(int argc, char** argv) { } if (std::string(argv[1]) == std::string("float")) { - aux_main(argc, argv); + aux_main(argv); } else if (std::string(argv[1]) == std::string("int8")) { - aux_main(argc, argv); + aux_main(argv); } else if (std::string(argv[1]) == std::string("uint8")) { - aux_main(argc, argv); + aux_main(argv); } else std::cout << "Unsupported type. Use float/int8/uint8." << std::endl; return 0; From 247f83e4682a1dabad8ee6e519bf4c6dd27eb4be Mon Sep 17 00:00:00 2001 From: ravishankar Date: Thu, 27 May 2021 20:40:42 +0530 Subject: [PATCH 32/84] working towards inner product in memory indices --- tests/build_memory_index.cpp | 39 +++++++++++++++++++++----------- tests/utils/CMakeLists.txt | 2 +- tests/utils/gen_random_slice.cpp | 14 +++++++----- 3 files changed, 35 insertions(+), 20 deletions(-) diff --git a/tests/build_memory_index.cpp b/tests/build_memory_index.cpp index a5a13595e3..7bbe992c82 100644 --- a/tests/build_memory_index.cpp +++ b/tests/build_memory_index.cpp @@ -16,7 +16,7 @@ #include "memory_mapper.h" template -int build_in_memory_index(const std::string& data_path, const unsigned R, +int build_in_memory_index(const std::string& data_path, _u32 dist_fn, const unsigned R, const unsigned L, const float alpha, const std::string& save_path, const unsigned num_threads) { @@ -29,7 +29,17 @@ int build_in_memory_index(const std::string& data_path, const unsigned R, paras.Set("saturate_graph", 0); paras.Set("num_threads", num_threads); - diskann::Index index(diskann::L2, data_path.c_str()); + diskann::Metric metric; + if (dist_fn == 0) + metric = diskann::L2; + else if (dist_fn == 1) + metric = diskann::INNER_PRODUCT; + else { + std::cout<<"Error. Unsupported distance type. Exitting" << std::endl; + return -1; + } + + diskann::Index index(metric, data_path.c_str()); auto s = std::chrono::high_resolution_clock::now(); index.build(paras); std::chrono::duration diff = @@ -42,9 +52,9 @@ int build_in_memory_index(const std::string& data_path, const unsigned R, } int main(int argc, char** argv) { - if (argc != 8) { + if (argc != 9) { std::cout << "Usage: " << argv[0] - << " [data_type] [data_file.bin] " + << " [data_type] [dist_fn 0 for L2, 1 for inner product] [data_file.bin] " "[output_index_file] " << "[R] [L] [alpha]" << " [num_threads_to_use]. See README for more information on " @@ -53,21 +63,24 @@ int main(int argc, char** argv) { exit(-1); } - const std::string data_path(argv[2]); - const std::string save_path(argv[3]); - const unsigned R = (unsigned) atoi(argv[4]); - const unsigned L = (unsigned) atoi(argv[5]); - const float alpha = (float) atof(argv[6]); - const unsigned num_threads = (unsigned) atoi(argv[7]); + _u32 ctr = 2; + + _u32 dist_fn = (_u32) atoi(argv[ctr++]); + const std::string data_path(argv[ctr++]); + const std::string save_path(argv[ctr++]); + const unsigned R = (unsigned) atoi(argv[ctr++]); + const unsigned L = (unsigned) atoi(argv[ctr++]); + const float alpha = (float) atof(argv[ctr++]); + const unsigned num_threads = (unsigned) atoi(argv[ctr++]); if (std::string(argv[1]) == std::string("int8")) - build_in_memory_index(data_path, R, L, alpha, save_path, + build_in_memory_index(data_path, dist_fn, R, L, alpha, save_path, num_threads); else if (std::string(argv[1]) == std::string("uint8")) - build_in_memory_index(data_path, R, L, alpha, save_path, + build_in_memory_index(data_path, dist_fn, R, L, alpha, save_path, num_threads); else if (std::string(argv[1]) == std::string("float")) - build_in_memory_index(data_path, R, L, alpha, save_path, + build_in_memory_index(data_path, dist_fn, R, L, alpha, save_path, num_threads); else std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; diff --git a/tests/utils/CMakeLists.txt b/tests/utils/CMakeLists.txt index e69722dcfc..0955d57b97 100644 --- a/tests/utils/CMakeLists.txt +++ b/tests/utils/CMakeLists.txt @@ -136,6 +136,6 @@ endif() # formatter if (LINUX) - add_custom_command(TARGET gen_random_slice PRE_BUILD COMMAND clang-format-4.0 -i ../../../include/*.h ../../../include/dll/*.h ../../../src/*.cpp ../../../tests/*.cpp ../../../src/dll/*.cpp ../../../tests/utils/*.cpp) + add_custom_command(TARGET gen_random_slice PRE_BUILD COMMAND clang-format -i ../../../include/*.h ../../../include/dll/*.h ../../../src/*.cpp ../../../tests/*.cpp ../../../src/dll/*.cpp ../../../tests/utils/*.cpp) endif() diff --git a/tests/utils/gen_random_slice.cpp b/tests/utils/gen_random_slice.cpp index 0417c12c0d..1b102e27c1 100644 --- a/tests/utils/gen_random_slice.cpp +++ b/tests/utils/gen_random_slice.cpp @@ -22,12 +22,6 @@ template int aux_main(int argc, char** argv) { - if (argc != 5) { - std::cout << argv[0] << " data_type [fliat/int8/uint8] base_bin_file " - "sample_output_prefix sampling_probability" - << std::endl; - exit(-1); - } std::string base_file(argv[2]); std::string output_prefix(argv[3]); @@ -37,6 +31,14 @@ int aux_main(int argc, char** argv) { } 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(argc, argv); } else if (std::string(argv[1]) == std::string("int8")) { From e339efd1331934b399431c85dce0bfb827c47b36 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Thu, 27 May 2021 21:59:50 +0530 Subject: [PATCH 33/84] done with in-memory code --- include/distance.h | 10 +++- include/index.h | 1 + src/index.cpp | 7 ++- tests/search_memory_index.cpp | 48 +++++++++------- tests/utils/compute_groundtruth.cpp | 89 ++++++++++++++++++++++++----- 5 files changed, 118 insertions(+), 37 deletions(-) diff --git a/include/distance.h b/include/distance.h index 3d403d1e56..c6b8c56d32 100644 --- a/include/distance.h +++ b/include/distance.h @@ -328,7 +328,7 @@ namespace diskann { template class DistanceInnerProduct : public Distance { public: - float compare(const T *a, const T *b, unsigned size) const { + float acompare(const T *a, const T *b, unsigned size) const { float result = 0; #ifdef __GNUC__ #ifdef __AVX__ @@ -426,10 +426,14 @@ namespace diskann { #endif return result; } + float compare(const T *a, const T *b, unsigned size) const { // since we use normally minimization objective for distance comparisons, we are returning 1/x. + float result = acompare(a,b,size); + return 1/result; + } }; template - class DistanceFastL2 : public DistanceInnerProduct { + class DistanceFastL2 : public DistanceInnerProduct { // currently defined only for float. templated for future use. public: float norm(const T *a, unsigned size) const { float result = 0; @@ -522,7 +526,7 @@ namespace diskann { using DistanceInnerProduct::compare; float compare(const T *a, const T *b, float norm, unsigned size) const { // not implement - float result = -2 * DistanceInnerProduct::compare(a, b, size); + float result = -2 * DistanceInnerProduct::acompare(a, b, size); result += norm; return result; } diff --git a/include/index.h b/include/index.h index eedbb1491f..c409b96958 100644 --- a/include/index.h +++ b/include/index.h @@ -167,6 +167,7 @@ namespace diskann { size_t consolidate_deletes(const Parameters ¶meters); private: + Metric _metric = diskann::L2; size_t _dim; size_t _aligned_dim; T * _data; diff --git a/src/index.cpp b/src/index.cpp index 2405878669..7123c661d4 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -61,6 +61,9 @@ namespace { std::cout << "Older CPU. Using slow distance computation" << std::endl; return new diskann::SlowDistanceL2Float(); } + } else if (m == diskann::Metric::INNER_PRODUCT) { + std::cout << "Using Inner Product computation" << std::endl; + return new diskann::DistanceInnerProduct(); } else { std::stringstream stream; stream << "Only L2 metric supported as of now. Email " @@ -129,7 +132,7 @@ namespace diskann { const size_t nd, const size_t num_frozen_pts, const bool enable_tags, const bool store_data, const bool support_eager_delete) - : _num_frozen_pts(num_frozen_pts), _has_built(false), _width(0), + : _metric(m), _num_frozen_pts(num_frozen_pts), _has_built(false), _width(0), _can_delete(false), _eager_done(true), _lazy_done(true), _compacted_order(true), _enable_tags(enable_tags), _consolidated_order(true), _support_eager_delete(support_eager_delete), @@ -1054,6 +1057,8 @@ namespace diskann { for (auto it : best_L_nodes) { indices[pos] = it.id; distances[pos] = it.distance; + if (_metric == diskann::INNER_PRODUCT) + distances[pos] = 1/distances[pos]; pos++; if (pos == K) break; diff --git a/tests/search_memory_index.cpp b/tests/search_memory_index.cpp index 32592710e5..6ad40032a4 100644 --- a/tests/search_memory_index.cpp +++ b/tests/search_memory_index.cpp @@ -27,26 +27,27 @@ int search_memory_index(int argc, char** argv) { size_t query_num, query_dim, query_aligned_dim, gt_num, gt_dim; std::vector<_u64> Lvec; - std::string data_file(argv[2]); - std::string memory_index_file(argv[3]); - _u64 num_threads = std::atoi(argv[4]); - std::string query_bin(argv[5]); - std::string truthset_bin(argv[6]); - _u64 recall_at = std::atoi(argv[7]); - std::string result_output_prefix(argv[8]); - bool use_optimized_search = std::atoi(argv[9]); + _u32 ctr = 2; + _u32 dist_fn = atoi(argv[ctr++]); + std::string data_file(argv[ctr++]); + std::string memory_index_file(argv[ctr++]); + _u64 num_threads = std::atoi(argv[ctr++]); + std::string query_bin(argv[ctr++]); + std::string truthset_bin(argv[ctr++]); + _u64 recall_at = std::atoi(argv[ctr++]); + std::string result_output_prefix(argv[ctr++]); +// bool use_optimized_search = std::atoi(argv[ctr++]); if ((std::string(argv[1]) != std::string("float")) && - (use_optimized_search == true)) { - std::cout << "Error. Optimized search currently only supported for " - "floating point datatypes. Using un-optimized search." + ((dist_fn == 1) || (dist_fn == 2))) { + std::cout << "Error. Inner product and Fast_L2 search currently only supported for " + "floating point datatypes." << std::endl; - use_optimized_search = false; } bool calc_recall_flag = false; - for (int ctr = 10; ctr < argc; ctr++) { + for (; ctr < (_u32) argc; ctr++) { _u64 curL = std::atoi(argv[ctr]); if (curL >= recall_at) Lvec.push_back(curL); @@ -73,14 +74,22 @@ int search_memory_index(int argc, char** argv) { std::cout.setf(std::ios_base::fixed, std::ios_base::floatfield); std::cout.precision(2); - auto metric = diskann::L2; - if (use_optimized_search) + diskann::Metric metric; + if (dist_fn == 0) + metric = diskann::L2; + else if (dist_fn == 1) + metric = diskann::INNER_PRODUCT; + else if(dist_fn == 2) metric = diskann::FAST_L2; + else { + std::cout<<"Error. Unsupported distance function. Exitting"; + return -1; + } diskann::Index index(metric, data_file.c_str()); index.load(memory_index_file.c_str()); // to load NSG std::cout << "Index loaded" << std::endl; - if (use_optimized_search) + if (metric == diskann::FAST_L2) index.optimize_graph(); diskann::Parameters paras; @@ -106,7 +115,7 @@ int search_memory_index(int argc, char** argv) { #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 (use_optimized_search) { + if (metric == diskann::FAST_L2) { index.search_with_opt_graph( query + i * query_aligned_dim, recall_at, L, query_result_ids[test_id].data() + i * recall_at); @@ -160,11 +169,10 @@ int main(int argc, char** argv) { if (argc < 11) { std::cout << "Usage: " << argv[0] - << " [index_type] [data_file.bin] " + << " [index_type] [dist_fn (0 for L2, 1 for Inner Product, 2 for Fast L2 for small datasets)] [data_file.bin] " "[memory_index_path] [num_threads] " "[query_file.bin] [truthset.bin (use \"null\" for none)] " - " [K] [result_output_prefix] [use_optimized_search (for small ~1M " - "data)] " + " [K] [result_output_prefix]" " [L1] [L2] etc. See README for more information on parameters. " << std::endl; exit(-1); diff --git a/tests/utils/compute_groundtruth.cpp b/tests/utils/compute_groundtruth.cpp index 8fef8c929c..cdd8d0c327 100644 --- a/tests/utils/compute_groundtruth.cpp +++ b/tests/utils/compute_groundtruth.cpp @@ -31,10 +31,10 @@ #define ALIGNMENT 512 void command_line_help() { - std::cerr - << " " - << std::endl; + std::cerr << " " + " " + << std::endl; } template @@ -104,6 +104,34 @@ void distsq_to_points( 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, + 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) -1.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 exact_knn(const size_t dim, const size_t k, int *const closest_points, // k * num_queries preallocated, col // major, queries columns @@ -112,14 +140,23 @@ void exact_knn(const size_t dim, const size_t k, // corresponding closes_points size_t npoints, const float *const points, // points in Col major - size_t nqueries, - const float *const queries) // queries in Col major + size_t nqueries, const float *const queries, + bool use_mip = false) // queries in Col major { float *points_l2sq = new float[npoints]; float *queries_l2sq = new float[nqueries]; 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 (use_mip) + std::cout << " inner product "; + 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]; @@ -128,9 +165,15 @@ void exact_knn(const size_t dim, const size_t k, int64_t q_e = ((b + 1) * q_batch_size > nqueries) ? nqueries : (b + 1) * q_batch_size; - 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); + if (!use_mip) { + 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, points_l2sq, q_e - q_b, + queries + (ptrdiff_t) q_b * (ptrdiff_t) dim, queries_l2sq + q_b); + } std::cout << "Computed distances for queries: [" << q_b << "," << q_e << ")" << std::endl; @@ -266,12 +309,12 @@ inline void save_groundtruth_as_one_file(const std::string filename, template int aux_main(int argv, char **argc) { - size_t npoints, nqueries, dim; std::string base_file(argc[2]); std::string query_file(argc[3]); size_t k = atoi(argc[4]); std::string gt_file(argc[5]); + bool use_mip = atoi(argc[6]); float *base_data; float *query_data; @@ -291,7 +334,7 @@ int aux_main(int argv, char **argc) { float *dist_closest_points_part = new float[nqueries * k]; exact_knn(dim, k, closest_points_part, dist_closest_points_part, npoints, - base_data, nqueries, query_data); + base_data, nqueries, query_data, use_mip); for (_u64 i = 0; i < nqueries; i++) { for (_u64 j = 0; j < k; j++) { @@ -303,6 +346,23 @@ int aux_main(int argv, char **argc) { delete[] closest_points_part; delete[] dist_closest_points_part; + + /* + std::cout << "For testing: doing brute force for one point" << + std::endl; std::vector> brute_force; for (_u32 i + = 0; i < npoints; i++) { float cur_pt_dist = 0; for (_u64 k = 0; k < dim; + k++) { cur_pt_dist += base_data[i * dim + k] * query_data[k]; + } + brute_force.push_back(std::make_pair(i, -cur_pt_dist)); + } + + std::sort(brute_force.begin(), brute_force.end(), custom_dist); + for (_u32 i = 0; i < 10; i++) { + std::cout< Date: Fri, 28 May 2021 17:11:13 +0530 Subject: [PATCH 34/84] made the inner product distance function return std::float_max if negative --- include/distance.h | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/include/distance.h b/include/distance.h index c6b8c56d32..b39a1719fb 100644 --- a/include/distance.h +++ b/include/distance.h @@ -428,7 +428,9 @@ namespace diskann { } float compare(const T *a, const T *b, unsigned size) const { // since we use normally minimization objective for distance comparisons, we are returning 1/x. float result = acompare(a,b,size); - return 1/result; + if (result < 0) + return std::numeric_limits::max(); + else return 1/result; } }; From 34bfa3f83a6ca7d7f5acbd96c1e412d57035a0ec Mon Sep 17 00:00:00 2001 From: ravishankar Date: Sat, 29 May 2021 11:27:44 +0530 Subject: [PATCH 35/84] more changes for disk index support --- include/partition_and_pq.h | 2 +- include/pq_flash_index.h | 3 ++- include/pq_table.h | 23 ++++++++++++++++++----- src/partition_and_pq.cpp | 10 +++++++--- src/pq_flash_index.cpp | 34 ++++++++++++++++++++++++++++------ 5 files changed, 56 insertions(+), 16 deletions(-) diff --git a/include/partition_and_pq.h b/include/partition_and_pq.h index 45b1e26e6d..43a9d84db0 100644 --- a/include/partition_and_pq.h +++ b/include/partition_and_pq.h @@ -54,7 +54,7 @@ DISKANN_DLLEXPORT int generate_pq_pivots(const float *train_data, unsigned num_centers, unsigned num_pq_chunks, unsigned max_k_means_reps, - std::string pq_pivots_path); + std::string pq_pivots_path, bool make_zero_mean = false); template int generate_pq_data_from_pivots(const std::string data_file, diff --git a/include/pq_flash_index.h b/include/pq_flash_index.h index c7601851c6..c20409fb64 100644 --- a/include/pq_flash_index.h +++ b/include/pq_flash_index.h @@ -70,7 +70,7 @@ namespace diskann { // Freeing the reader object is now the client's (DiskANNInterface's) // responsibility. DISKANN_DLLEXPORT PQFlashIndex( - std::shared_ptr &fileReader); + std::shared_ptr &fileReader, diskann::Metric metric = diskann::Metric::L2); DISKANN_DLLEXPORT ~PQFlashIndex(); #ifdef EXEC_ENV_OLS @@ -129,6 +129,7 @@ namespace diskann { // nbrs of node `i`: ((unsigned*)buf) + 1 _u64 max_node_len = 0, nnodes_per_sector = 0, max_degree = 0; + diskann::Metric metric = diskann::Metric::L2; // data info _u64 num_points = 0; _u64 data_dim = 0; diff --git a/include/pq_table.h b/include/pq_table.h index 3cac23c15a..c913271eb1 100644 --- a/include/pq_table.h +++ b/include/pq_table.h @@ -137,11 +137,6 @@ namespace diskann { _u64 permuted_dim_in_query = rearrangement[j]; const float* centers_dim_vec = tables_T + (256 * j); for (_u64 idx = 0; idx < 256; idx++) { - // Gopal. Fixing crash in v14 machines. - // float diff = centers_dim_vec[idx] - - // ((float) query_vec[permuted_dim_in_query] - - // centroid[permuted_dim_in_query]); - // chunk_dists[idx] += (diff * diff); double diff = centers_dim_vec[idx] - (query_vec[permuted_dim_in_query] - centroid[permuted_dim_in_query]); @@ -150,5 +145,23 @@ namespace diskann { } } } + void + populate_chunk_inner_products(const T* query_vec, float* dist_vec) { + memset(dist_vec, 0, 256 * n_chunks * sizeof(float)); + // chunk wise distance computation + for (_u64 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++) { + _u64 permuted_dim_in_query = rearrangement[j]; + const float* centers_dim_vec = tables_T + (256 * j); + for (_u64 idx = 0; idx < 256; idx++) { + double prod = + centers_dim_vec[idx] * query_vec[permuted_dim_in_query]; // 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) + } + } + } + } }; } // namespace diskann diff --git a/src/partition_and_pq.cpp b/src/partition_and_pq.cpp index 9da49b2203..ed6fcf9826 100644 --- a/src/partition_and_pq.cpp +++ b/src/partition_and_pq.cpp @@ -190,13 +190,14 @@ void gen_random_slice(const T *inputdata, size_t npts, size_t ndims, int generate_pq_pivots(const float *passed_train_data, size_t num_train, unsigned dim, unsigned num_centers, unsigned num_pq_chunks, unsigned max_k_means_reps, - std::string pq_pivots_path) { + 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, @@ -221,11 +222,14 @@ int generate_pq_pivots(const float *passed_train_data, size_t num_train, 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) { + for (uint64_t d = 0; d < dim; d++) { for (uint64_t p = 0; p < num_train; p++) { centroid[d] += train_data[p * dim + d]; } @@ -233,12 +237,12 @@ int generate_pq_pivots(const float *passed_train_data, size_t num_train, } // std::memset(centroid, 0 , dim*sizeof(float)); - 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 rearrangement; std::vector chunk_offsets; diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index 21da3521d6..a11379b7d1 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -86,8 +86,8 @@ namespace { namespace diskann { template<> PQFlashIndex<_u8>::PQFlashIndex( - std::shared_ptr &fileReader) - : reader(fileReader) { + std::shared_ptr &fileReader, diskann::Metric metric) + : reader(fileReader), metric(metric) { diskann::cout << "dist_cmp function for _u8 uses slow implementation." " Please contact gopalsr@microsoft.com if you need an AVX/AVX2" @@ -106,12 +106,16 @@ namespace diskann { << std::endl; this->dist_cmp_float = new SlowDistanceL2Float(); } + if (metric != diskann::Metric::L2) { + std::cout<<"Only L2 supported for byte vectors for now. Other distance functions are future work. Falling back to L2 distance." << std::endl; + this->metric = diskann::Metric::L2; + } } template<> PQFlashIndex<_s8>::PQFlashIndex( - std::shared_ptr &fileReader) - : reader(fileReader) { + std::shared_ptr &fileReader, diskann::Metric metric) + : reader(fileReader), metric(metric) { if (Avx2SupportedCPU) { diskann::cout << "Using AVX2 function for dist_cmp and dist_cmp_float" << std::endl; @@ -130,12 +134,18 @@ namespace diskann { this->dist_cmp = new SlowDistanceL2Int(); this->dist_cmp_float = new SlowDistanceL2Float(); } + if (metric != diskann::Metric::L2) { + std::cout<<"Only L2 supported for byte vectors for now. Other distance functions are future work. Falling back to L2 distance." << std::endl; + this->metric = diskann::Metric::L2; + } + } template<> PQFlashIndex::PQFlashIndex( - std::shared_ptr &fileReader) - : reader(fileReader) { + std::shared_ptr &fileReader, diskann::Metric metric) + : reader(fileReader), metric(metric) { + if (metric == diskann::Metric::L2) { if (Avx2SupportedCPU) { diskann::cout << "Using AVX2 functions for dist_cmp and dist_cmp_float" << std::endl; @@ -154,6 +164,15 @@ namespace diskann { this->dist_cmp = new AVXDistanceL2Float(); this->dist_cmp_float = new AVXDistanceL2Float(); } + } else if (metric == diskann::Metric::INNER_PRODUCT) { + this->dist_cmp = new DistanceInnerProduct(); + this->dist_cmp_float = new DistanceInnerProduct(); + } else { + std::cout<<"Unsupported metric type. Reverting to float." << std::endl; + this->dist_cmp = new AVXDistanceL2Float(); + this->dist_cmp_float = new AVXDistanceL2Float(); + this->metric = diskann::Metric::L2; + } } template @@ -801,6 +820,9 @@ namespace diskann { // query <-> PQ chunk centers distances float *pq_dists = query_scratch->aligned_pqtable_dist_scratch; + if (metric==diskann::Metric::INNER_PRODUCT) + pq_table.populate_chunk_inner_products(query, pq_dists); + else if (metric==diskann::Metric::L2) pq_table.populate_chunk_distances(query, pq_dists); // query <-> neighbor list From b81b90ae0d5213d2c5c1a8fe8734d677e3eb9624 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Sat, 29 May 2021 13:41:07 +0530 Subject: [PATCH 36/84] on the way to disk index support for MIPS --- src/aux_utils.cpp | 32 +++++++++++++++++++------------- tests/build_disk_index.cpp | 23 +++++++++++++---------- tests/search_disk_index.cpp | 2 +- 3 files changed, 33 insertions(+), 24 deletions(-) diff --git a/src/aux_utils.cpp b/src/aux_utils.cpp index 6a2990ba66..13a1389ba1 100644 --- a/src/aux_utils.cpp +++ b/src/aux_utils.cpp @@ -343,7 +343,7 @@ namespace diskann { template int build_merged_vamana_index(std::string base_file, - diskann::Metric _compareMetric, unsigned L, + diskann::Metric compareMetric, unsigned L, unsigned R, double sampling_rate, double ram_budget, std::string mem_index_path, std::string medoids_file, @@ -367,7 +367,7 @@ namespace diskann { std::unique_ptr> _pvamanaIndex = std::unique_ptr>( - new diskann::Index(_compareMetric, base_file.c_str())); + new diskann::Index(compareMetric, base_file.c_str())); _pvamanaIndex->build(paras); _pvamanaIndex->save(mem_index_path.c_str()); std::remove(medoids_file.c_str()); @@ -399,7 +399,7 @@ namespace diskann { std::unique_ptr> _pvamanaIndex = std::unique_ptr>( - new diskann::Index(_compareMetric, shard_base_file.c_str())); + new diskann::Index(compareMetric, shard_base_file.c_str())); _pvamanaIndex->build(paras); _pvamanaIndex->save(shard_index_file.c_str()); } @@ -612,7 +612,7 @@ namespace diskann { template bool build_disk_index(const char *dataFilePath, const char *indexFilePath, const char * indexBuildParameters, - diskann::Metric _compareMetric) { + diskann::Metric compareMetric) { std::stringstream parser; parser << std::string(indexBuildParameters); std::string cur_param; @@ -631,6 +631,9 @@ namespace diskann { return false; } + if (compareMetric == diskann::Metric::INNER_PRODUCT) { + std::cout<<"Using Inner Product for PQ and Graph Generation" << std::endl; + } std::string index_prefix_path(indexFilePath); std::string pq_pivots_path = index_prefix_path + "_pq_pivots.bin"; std::string pq_compressed_vectors_path = @@ -697,9 +700,12 @@ namespace diskann { diskann::cout << "Training data loaded of size " << train_size << std::endl; - + + bool make_zero_mean = true; + if (compareMetric == diskann::Metric::INNER_PRODUCT) + make_zero_mean = false; generate_pq_pivots(train_data, train_size, (uint32_t) dim, 256, - (uint32_t) num_pq_chunks, 15, pq_pivots_path); + (uint32_t) num_pq_chunks, 15, pq_pivots_path, make_zero_mean); generate_pq_data_from_pivots(dataFilePath, 256, (uint32_t) num_pq_chunks, pq_pivots_path, @@ -710,7 +716,7 @@ namespace diskann { train_data = nullptr; diskann::build_merged_vamana_index( - dataFilePath, _compareMetric, L, R, p_val, indexing_ram_budget, + dataFilePath, compareMetric, L, R, p_val, indexing_ram_budget, mem_index_path, medoids_path, centroids_path); diskann::create_disk_layout(dataFilePath, mem_index_path, @@ -781,26 +787,26 @@ namespace diskann { template DISKANN_DLLEXPORT bool build_disk_index( const char *dataFilePath, const char *indexFilePath, - const char *indexBuildParameters, diskann::Metric _compareMetric); + const char *indexBuildParameters, diskann::Metric compareMetric); template DISKANN_DLLEXPORT bool build_disk_index( const char *dataFilePath, const char *indexFilePath, - const char *indexBuildParameters, diskann::Metric _compareMetric); + const char *indexBuildParameters, diskann::Metric compareMetric); template DISKANN_DLLEXPORT bool build_disk_index( const char *dataFilePath, const char *indexFilePath, - const char *indexBuildParameters, diskann::Metric _compareMetric); + const char *indexBuildParameters, diskann::Metric compareMetric); template DISKANN_DLLEXPORT int build_merged_vamana_index( - std::string base_file, diskann::Metric _compareMetric, unsigned L, + 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); template DISKANN_DLLEXPORT int build_merged_vamana_index( - std::string base_file, diskann::Metric _compareMetric, unsigned L, + 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); template DISKANN_DLLEXPORT int build_merged_vamana_index( - std::string base_file, diskann::Metric _compareMetric, unsigned L, + 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); diff --git a/tests/build_disk_index.cpp b/tests/build_disk_index.cpp index b29c3f0d1e..8c94d47525 100644 --- a/tests/build_disk_index.cpp +++ b/tests/build_disk_index.cpp @@ -11,29 +11,32 @@ template bool build_index(const char* dataFilePath, const char* indexFilePath, - const char* indexBuildParameters) { + const char* indexBuildParameters, diskann::Metric metric) { return diskann::build_disk_index( - dataFilePath, indexFilePath, indexBuildParameters, diskann::Metric::L2); + dataFilePath, indexFilePath, indexBuildParameters, metric); } int main(int argc, char** argv) { - if (argc != 9) { + if (argc != 10) { std::cout << "Usage: " << argv[0] - << " [data_type] [data_file.bin] " + << " [data_type] [dist_fn: 0 for L2, 1 for MIPS] [data_file.bin] " "[index_prefix_path] " "[R] [L] [B] [M] [T]. See README for more information on " "parameters." << std::endl; } else { - std::string params = std::string(argv[4]) + " " + std::string(argv[5]) + - " " + std::string(argv[6]) + " " + - std::string(argv[7]) + " " + std::string(argv[8]); + diskann::Metric metric = diskann::Metric::L2; + if (atoi(argv[2]) == 1) + metric = diskann::Metric::INNER_PRODUCT; + std::string params = std::string(argv[5]) + " " + std::string(argv[6]) + + " " + std::string(argv[7]) + " " + + std::string(argv[8]) + " " + std::string(argv[9]); if (std::string(argv[1]) == std::string("float")) - build_index(argv[2], argv[3], params.c_str()); + build_index(argv[3], argv[4], params.c_str(), metric); else if (std::string(argv[1]) == std::string("int8")) - build_index(argv[2], argv[3], params.c_str()); + build_index(argv[3], argv[4], params.c_str(), metric); else if (std::string(argv[1]) == std::string("uint8")) - build_index(argv[2], argv[3], params.c_str()); + build_index(argv[3], argv[4], params.c_str(), metric); else std::cout << "Error. wrong file type" << std::endl; } diff --git a/tests/search_disk_index.cpp b/tests/search_disk_index.cpp index 4d5c79ca89..7fa0edeec9 100644 --- a/tests/search_disk_index.cpp +++ b/tests/search_disk_index.cpp @@ -114,7 +114,7 @@ int search_disk_index(int argc, char** argv) { #endif std::unique_ptr> _pFlashIndex( - new diskann::PQFlashIndex(reader)); + new diskann::PQFlashIndex(reader, diskann::Metric::INNER_PRODUCT)); int res = _pFlashIndex->load(num_threads, pq_prefix.c_str(), disk_index_file.c_str()); From 6f2aa4e926a7f014e66a262b542e51a0759c8592 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Sat, 29 May 2021 16:38:21 +0530 Subject: [PATCH 37/84] works now, need to change the PQ generation for MIPS --- src/aux_utils.cpp | 4 ++-- src/partition_and_pq.cpp | 33 ++++++++++++++++++++++++++++++++- 2 files changed, 34 insertions(+), 3 deletions(-) diff --git a/src/aux_utils.cpp b/src/aux_utils.cpp index 13a1389ba1..b3cc400692 100644 --- a/src/aux_utils.cpp +++ b/src/aux_utils.cpp @@ -56,7 +56,7 @@ namespace diskann { } gt.insert(gt_vec, gt_vec + tie_breaker); - res.insert(res_vec, res_vec + recall_at); + res.insert(res_vec, 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) { if (res.find(v) != res.end()) { @@ -726,7 +726,7 @@ namespace diskann { gen_random_slice(dataFilePath, sample_base_prefix, sample_sampling_rate); - std::remove(mem_index_path.c_str()); + std::remove(mem_index_path.c_str()); auto e = std::chrono::high_resolution_clock::now(); diff --git a/src/partition_and_pq.cpp b/src/partition_and_pq.cpp index ed6fcf9826..be323ad8ff 100644 --- a/src/partition_and_pq.cpp +++ b/src/partition_and_pq.cpp @@ -36,6 +36,7 @@ #endif #define BLOCK_SIZE 5000000 +#define SAVE_INFLATED_PQ true template void gen_random_slice(const std::string base_file, @@ -380,6 +381,8 @@ int generate_pq_data_from_pivots(const std::string data_file, std::unique_ptr rearrangement; std::unique_ptr chunk_offsets; + std::string inflated_pq_file = pq_compressed_vectors_path + "_inflated.bin"; + if (!file_exists(pq_pivots_path)) { diskann::cout << "ERROR: PQ k-means pivot file not found" << std::endl; throw diskann::ANNException("PQ k-means pivot file not found", -1); @@ -446,6 +449,20 @@ int generate_pq_data_from_pivots(const std::string data_file, 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::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<_u32[]> block_compressed_base = std::make_unique<_u32[]>(block_size * (_u64) num_pq_chunks); std::memset(block_compressed_base.get(), 0, @@ -518,6 +535,12 @@ int generate_pq_data_from_pivots(const std::string data_file, #pragma omp parallel for schedule(static, 8192) for (int64_t j = 0; j < (_s64) cur_blk_size; j++) { block_compressed_base[j * num_pq_chunks + i] = closest_center[j]; +#ifdef SAVE_INFLATED_PQ + for (uint64_t k = 0; k < cur_chunk_size; k++) + block_inflated_base[j * dim + chunk_offsets[i] + k] = + cur_pivot_data[closest_center[j] * cur_chunk_size + k] + + centroid[chunk_offsets[i] + k]; +#endif } } @@ -532,8 +555,13 @@ int generate_pq_data_from_pivots(const std::string data_file, 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)); + 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)); +#endif diskann::cout << ".done." << std::endl; } // Gopal. Splittng diskann_dll into separate DLLs for search and build. @@ -542,6 +570,9 @@ int generate_pq_data_from_pivots(const std::string data_file, MallocExtension::instance()->ReleaseFreeMemory(); #endif compressed_file_writer.close(); +#ifdef SAVE_INFLATED_PQ + inflated_file_writer.close(); +#endif return 0; } From 0a21b457a49d2b7a8760eff74c2e880aa56d9ee6 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Mon, 31 May 2021 11:30:33 +0530 Subject: [PATCH 38/84] now incorporated disk+memory search for inner product --- src/pq_flash_index.cpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index a11379b7d1..c5a1c20840 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -820,9 +820,9 @@ namespace diskann { // query <-> PQ chunk centers distances float *pq_dists = query_scratch->aligned_pqtable_dist_scratch; - if (metric==diskann::Metric::INNER_PRODUCT) - pq_table.populate_chunk_inner_products(query, pq_dists); - else if (metric==diskann::Metric::L2) +// if (metric==diskann::Metric::INNER_PRODUCT) +// pq_table.populate_chunk_inner_products(query, pq_dists); +// else if (metric==diskann::Metric::L2) pq_table.populate_chunk_distances(query, pq_dists); // query <-> neighbor list From 1e1e14cf4a9d2191cca04f111beb1384b51d1a33 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Mon, 31 May 2021 15:30:05 +0530 Subject: [PATCH 39/84] support for mips and l2 --- src/pq_flash_index.cpp | 3 +++ tests/search_disk_index.cpp | 35 +++++++++++++++++++++++------------ 2 files changed, 26 insertions(+), 12 deletions(-) diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index c5a1c20840..cb5a51f661 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -165,6 +165,7 @@ namespace diskann { this->dist_cmp_float = new AVXDistanceL2Float(); } } else if (metric == diskann::Metric::INNER_PRODUCT) { + std::cout<<"Using inner product distance function" << std::endl; this->dist_cmp = new DistanceInnerProduct(); this->dist_cmp_float = new DistanceInnerProduct(); } else { @@ -1110,6 +1111,8 @@ namespace diskann { indices[i] = full_retset[i].id; if (distances != nullptr) { distances[i] = full_retset[i].distance; + if (metric == diskann::Metric::INNER_PRODUCT) // flip the sign from convert min to max + distances[i] = 1/distances[i]; } } diff --git a/tests/search_disk_index.cpp b/tests/search_disk_index.cpp index 7fa0edeec9..939c365388 100644 --- a/tests/search_disk_index.cpp +++ b/tests/search_disk_index.cpp @@ -56,26 +56,37 @@ int search_disk_index(int argc, char** argv) { size_t query_num, query_dim, query_aligned_dim, gt_num, gt_dim; std::vector<_u64> Lvec; - std::string index_prefix_path(argv[2]); + _u32 ctr = 2; + _u32 dist_fn = atoi(argv[ctr++]); + std::string index_prefix_path(argv[ctr++]); std::string pq_prefix = index_prefix_path + "_pq"; std::string disk_index_file = index_prefix_path + "_disk.index"; std::string warmup_query_file = index_prefix_path + "_sample_data.bin"; - _u64 num_nodes_to_cache = std::atoi(argv[3]); - _u32 num_threads = std::atoi(argv[4]); - _u32 beamwidth = std::atoi(argv[5]); - std::string query_bin(argv[6]); - std::string truthset_bin(argv[7]); - _u64 recall_at = std::atoi(argv[8]); - std::string result_output_prefix(argv[9]); + _u64 num_nodes_to_cache = std::atoi(argv[ctr++]); + _u32 num_threads = std::atoi(argv[ctr++]); + _u32 beamwidth = std::atoi(argv[ctr++]); + std::string query_bin(argv[ctr++]); + std::string truthset_bin(argv[ctr++]); + _u64 recall_at = std::atoi(argv[ctr++]); + std::string result_output_prefix(argv[ctr++]); bool calc_recall_flag = false; - for (int ctr = 10; ctr < argc; ctr++) { + for (; ctr < (_u32) argc; ctr++) { _u64 curL = std::atoi(argv[ctr]); if (curL >= recall_at) Lvec.push_back(curL); } + diskann::Metric metric; + if (dist_fn == 0) + metric = diskann::Metric::L2; + else if (dist_fn == 1) + metric = diskann::Metric::INNER_PRODUCT; + else { + std::cout<<"Unsupported distance function. Currently only L2/ Inner Product support." << std::endl; + return -1; + } if (Lvec.size() == 0) { diskann::cout << "No valid Lsearch found. Lsearch must be at least recall_at" @@ -114,7 +125,7 @@ int search_disk_index(int argc, char** argv) { #endif std::unique_ptr> _pFlashIndex( - new diskann::PQFlashIndex(reader, diskann::Metric::INNER_PRODUCT)); + new diskann::PQFlashIndex(reader, metric)); int res = _pFlashIndex->load(num_threads, pq_prefix.c_str(), disk_index_file.c_str()); @@ -286,10 +297,10 @@ int search_disk_index(int argc, char** argv) { } int main(int argc, char** argv) { - if (argc < 11) { + if (argc < 12) { diskann::cout << "Usage: " << argv[0] - << " [index_type] [index_prefix_path] " + << " [index_type] [dist_fn 0 for l2/ 1 for inner product] [index_prefix_path] " " [num_nodes_to_cache] [num_threads] [beamwidth (use 0 to " "optimize internally)] " " [query_file.bin] [truthset.bin (use \"null\" for none)] " From 92f583d021614a0551b9ed353637285c7241bc12 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Mon, 31 May 2021 17:13:40 +0530 Subject: [PATCH 40/84] changed inner product to -IP rather than 1/IP --- include/distance.h | 7 ++++--- src/index.cpp | 14 ++++++++++++-- src/pq_flash_index.cpp | 2 +- 3 files changed, 17 insertions(+), 6 deletions(-) diff --git a/include/distance.h b/include/distance.h index b39a1719fb..a65f9e9487 100644 --- a/include/distance.h +++ b/include/distance.h @@ -428,9 +428,10 @@ namespace diskann { } float compare(const T *a, const T *b, unsigned size) const { // since we use normally minimization objective for distance comparisons, we are returning 1/x. float result = acompare(a,b,size); - if (result < 0) - return std::numeric_limits::max(); - else return 1/result; +// if (result < 0) +// return std::numeric_limits::max(); +// else +return -result; } }; diff --git a/src/index.cpp b/src/index.cpp index 7123c661d4..ff00f6dc61 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -538,7 +538,7 @@ namespace diskann { float cur_alpha = 1; while (cur_alpha <= alpha && result.size() < degree) { unsigned start = 0; - + float eps = cur_alpha + 0.01; while (result.size() < degree && (start) < pool.size() && start < maxc) { auto &p = pool[start]; if (occlude_factor[start] > cur_alpha) { @@ -553,8 +553,18 @@ namespace diskann { float djk = _distance->compare( _data + _aligned_dim * (size_t) pool[t].id, _data + _aligned_dim * (size_t) p.id, (unsigned) _aligned_dim); + if (_metric == diskann::Metric::L2) { occlude_factor[t] = (std::max)(occlude_factor[t], pool[t].distance / djk); + } + else if (_metric == diskann::Metric::INNER_PRODUCT) { // stylized rules for inner product since we want max instead of min distance + float x = -pool[t].distance; + float y = -djk; + if (y > cur_alpha * x) { + occlude_factor[t] = + (std::max)(occlude_factor[t], eps); + } + } } start++; } @@ -1058,7 +1068,7 @@ namespace diskann { indices[pos] = it.id; distances[pos] = it.distance; if (_metric == diskann::INNER_PRODUCT) - distances[pos] = 1/distances[pos]; + distances[pos] = -distances[pos]; pos++; if (pos == K) break; diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index cb5a51f661..ef236aa834 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -1112,7 +1112,7 @@ namespace diskann { if (distances != nullptr) { distances[i] = full_retset[i].distance; if (metric == diskann::Metric::INNER_PRODUCT) // flip the sign from convert min to max - distances[i] = 1/distances[i]; + distances[i] = -distances[i]; } } From 035579ecb3c727ccc2f2a838a531af4613603323 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Tue, 1 Jun 2021 20:59:37 +0530 Subject: [PATCH 41/84] towards adding support for storing PQ vectors in disk index for very large data --- include/pq_flash_index.h | 7 +++++++ include/pq_table.h | 24 ++++++++++++++++++++++++ src/pq_flash_index.cpp | 1 + tests/search_disk_index.cpp | 6 ++++-- 4 files changed, 36 insertions(+), 2 deletions(-) diff --git a/include/pq_flash_index.h b/include/pq_flash_index.h index c20409fb64..f3e8bcde5f 100644 --- a/include/pq_flash_index.h +++ b/include/pq_flash_index.h @@ -152,6 +152,13 @@ namespace diskann { Distance * dist_cmp = nullptr; Distance *dist_cmp_float = nullptr; + // for very large datasets: we use PQ even for the disk resident index + bool use_disk_index_pq = false; + _u64 disk_index_chunk_size; + _u64 disk_index_n_chunks; + FixedChunkPQTable disk_index_pq_table; + + // medoid/start info uint32_t *medoids = nullptr; // by default it is just one entry point of graph, we diff --git a/include/pq_table.h b/include/pq_table.h index c913271eb1..bb71fccba4 100644 --- a/include/pq_table.h +++ b/include/pq_table.h @@ -145,6 +145,30 @@ namespace diskann { } } } + + float compare(const T* query_vec, _u8* base_vec) { + float res = 0; + for (_u64 chunk = 0; chunk < n_chunks; chunk++) { + for (_u64 j = chunk_offsets[chunk]; j < chunk_offsets[chunk + 1]; j++) { + _u64 permuted_dim_in_query = rearrangement[j]; + const float* centers_dim_vec = tables_T + (256 * j); + float diff = centers_dim_vec[base_vec[chunk]] - (query_vec[permuted_dim_in_query] - centroid[permuted_dim_in_query]); + res += diff*diff; + } + } + return res; + } + + void inflate_vector(_u8* base_vec, float* out_vec) { + for (_u64 chunk = 0; chunk < n_chunks; chunk++) { + for (_u64 j = chunk_offsets[chunk]; j < chunk_offsets[chunk + 1]; j++) { + _u64 original_dim = rearrangement[j]; + const float* centers_dim_vec = tables_T + (256 * j); + out_vec[original_dim] = centers_dim_vec[base_vec[chunk]] + centroid[original_dim]; + } + } + } + void populate_chunk_inner_products(const T* query_vec, float* dist_vec) { memset(dist_vec, 0, 256 * n_chunks * sizeof(float)); diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index ef236aa834..ada71f0b7e 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -860,6 +860,7 @@ namespace diskann { } compute_dists(&best_medoid, 1, dist_scratch); + retset[0].id = best_medoid; retset[0].distance = dist_scratch[0]; retset[0].flag = true; diff --git a/tests/search_disk_index.cpp b/tests/search_disk_index.cpp index 939c365388..3a6abddeba 100644 --- a/tests/search_disk_index.cpp +++ b/tests/search_disk_index.cpp @@ -137,13 +137,14 @@ int search_disk_index(int argc, char** argv) { 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->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; @@ -207,7 +208,7 @@ int search_disk_index(int argc, char** argv) { uint32_t optimized_beamwidth = 2; - // query_num = 1; + for (uint32_t test_id = 0; test_id < Lvec.size(); test_id++) { _u64 L = Lvec[test_id]; @@ -225,6 +226,7 @@ int search_disk_index(int argc, char** argv) { diskann::QueryStats* 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) From af0dfabd29dd9d64c5079f1f6e8f21fd0dc2116c Mon Sep 17 00:00:00 2001 From: ravishankar Date: Tue, 1 Jun 2021 21:00:00 +0530 Subject: [PATCH 42/84] towards adding support for storing PQ vectors in disk index for very large data --- src/pq_flash_index.cpp | 13 ------------- 1 file changed, 13 deletions(-) diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index ada71f0b7e..0bab700ef5 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -897,19 +897,6 @@ namespace diskann { _u32 marker = k; _u32 num_seen = 0; - /* - bool marker_set = false; - diskann::cout << "hop " << hops << ": "; - for (_u32 i = 0; i < cur_list_size; i++) { - diskann::cout << retset[i].id << "( " << retset[i].distance; - if (retset[i].flag && !marker_set) { - diskann::cout << ",*) "; - marker_set = true; - } else - diskann::cout << ") "; - } - diskann::cout << std::endl; - */ while (marker < cur_list_size && frontier.size() < beam_width && num_seen < beam_width + 2) { if (retset[marker].flag) { From 52d2180f5cbbe04d91aca1ecc22ff6aa7b224fa5 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Tue, 1 Jun 2021 22:09:05 +0530 Subject: [PATCH 43/84] halfway through PQ-based disk search option --- include/aux_utils.h | 1 + include/pq_flash_index.h | 8 ++++---- include/pq_table.h | 11 +++++++---- src/aux_utils.cpp | 37 ++++++++++++++++++++++++++++++++++--- src/partition_and_pq.cpp | 2 +- src/pq_flash_index.cpp | 30 ++++++++++++++++++++++++------ tests/build_disk_index.cpp | 6 +++--- 7 files changed, 74 insertions(+), 21 deletions(-) diff --git a/include/aux_utils.h b/include/aux_utils.h index 031a64ba8b..7e946becd4 100644 --- a/include/aux_utils.h +++ b/include/aux_utils.h @@ -36,6 +36,7 @@ namespace diskann { const double THRESHOLD_FOR_CACHING_IN_GB = 1.0; const uint32_t NUM_NODES_TO_CACHE = 250000; const uint32_t WARMUP_L = 20; + const uint32_t NUM_KMEANS_REPS = 12; template class PQFlashIndex; diff --git a/include/pq_flash_index.h b/include/pq_flash_index.h index f3e8bcde5f..489afbed31 100644 --- a/include/pq_flash_index.h +++ b/include/pq_flash_index.h @@ -133,7 +133,9 @@ namespace diskann { // data info _u64 num_points = 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; std::string disk_index_file; std::vector> node_visit_counter; @@ -144,7 +146,6 @@ namespace diskann { // chunk_size = chunk size of each dimension chunk // pq_tables = float* [[2^8 * [chunk_size]] * n_chunks] _u8 * data = nullptr; - _u64 chunk_size; _u64 n_chunks; FixedChunkPQTable pq_table; @@ -154,9 +155,8 @@ namespace diskann { // for very large datasets: we use PQ even for the disk resident index bool use_disk_index_pq = false; - _u64 disk_index_chunk_size; - _u64 disk_index_n_chunks; - FixedChunkPQTable disk_index_pq_table; + _u64 disk_pq_n_chunks; + FixedChunkPQTable disk_pq_table; // medoid/start info diff --git a/include/pq_table.h b/include/pq_table.h index bb71fccba4..825b29c7ec 100644 --- a/include/pq_table.h +++ b/include/pq_table.h @@ -13,8 +13,8 @@ namespace diskann { nullptr; // pq_tables = float* [[2^8 * [chunk_size]] * n_chunks] // _u64 n_chunks; // n_chunks = # of chunks ndims is split into // _u64 chunk_size; // chunk_size = chunk size of each dimension chunk - _u64 ndims; // ndims = chunk_size * n_chunks - _u64 n_chunks; + _u64 ndims = 0; // ndims = chunk_size * n_chunks + _u64 n_chunks = 0; _u32* chunk_offsets = nullptr; _u32* rearrangement = nullptr; float* centroid = nullptr; @@ -79,14 +79,14 @@ namespace diskann { #else diskann::load_bin<_u32>(chunk_offset_file, chunk_offsets, numr, numc); #endif - if (numc != 1 || numr != num_chunks + 1) { + if (numc != 1 || (numr != num_chunks + 1 && num_chunks != 0)) { diskann::cerr << "Error loading chunk offsets file. numc: " << numc << " (should be 1). numr: " << numr << " (should be " << num_chunks + 1 << ")" << std::endl; throw diskann::ANNException("Error loading chunk offsets file", -1, __FUNCSIG__, __FILE__, __LINE__); } - + std::cout<<"PQ data has " << numr - 1 <<" bytes per point." << std::endl; this->n_chunks = numr - 1; #ifdef EXEC_ENV_OLS @@ -126,6 +126,9 @@ namespace diskann { } } +_u32 get_num_chunks() { + return n_chunks; +} void populate_chunk_distances(const T* query_vec, float* dist_vec) { memset(dist_vec, 0, 256 * n_chunks * sizeof(float)); diff --git a/src/aux_utils.cpp b/src/aux_utils.cpp index b3cc400692..98e0e88e57 100644 --- a/src/aux_utils.cpp +++ b/src/aux_utils.cpp @@ -620,17 +620,27 @@ namespace diskann { while (parser >> cur_param) param_list.push_back(cur_param); - if (param_list.size() != 5) { + if (param_list.size() != 5 && param_list.size() != 6) { diskann::cout << "Correct usage of parameters is R (max degree) " "L (indexing list size, better if >= R) B (RAM limit of final " "index in " "GB) M (memory limit while indexing) T (number of threads for " - "indexing)" + "indexing) B' (PQ bytes for disk index: optional parameter for very large dimensional data)" << std::endl; return false; } + _u32 disk_pq_dims = 0; + bool use_disk_pq = false; + + if (param_list.size() == 6) { + disk_pq_dims = atoi(param_list[5].c_str()); + use_disk_pq = true; + if (disk_pq_dims == 0) + use_disk_pq = false; + } + if (compareMetric == diskann::Metric::INNER_PRODUCT) { std::cout<<"Using Inner Product for PQ and Graph Generation" << std::endl; } @@ -643,6 +653,11 @@ namespace diskann { std::string medoids_path = disk_index_path + "_medoids.bin"; std::string centroids_path = disk_index_path + "_centroids.bin"; std::string sample_base_prefix = index_prefix_path + "_sample"; + std::string disk_pq_pivots_path = index_prefix_path + "_disk.index_pq_pivots.bin"; // optional if disk index is also storing pq data + std::string disk_pq_compressed_vectors_path = // optional if disk index is also storing pq data + index_prefix_path + "_disk.index_pq_compressed.bin"; + + unsigned R = (unsigned) atoi(param_list[0].c_str()); unsigned L = (unsigned) atoi(param_list[1].c_str()); @@ -662,6 +677,7 @@ namespace diskann { } _u32 num_threads = (_u32) atoi(param_list[4].c_str()); + if (num_threads != 0) { omp_set_num_threads(num_threads); mkl_set_num_threads(num_threads); @@ -698,6 +714,17 @@ namespace diskann { gen_random_slice(dataFilePath, p_val, train_data, train_size, train_dim); + if (use_disk_pq) { + if (disk_pq_dims > dim) + disk_pq_dims = dim; + + std::cout<<"Compressing base for disk-PQ into " << disk_pq_dims << " chunks " << std::endl; + generate_pq_pivots(train_data, train_size, (uint32_t) dim, 256, + (uint32_t) disk_pq_dims, NUM_KMEANS_REPS, disk_pq_pivots_path, false); + generate_pq_data_from_pivots(dataFilePath, 256, (uint32_t) disk_pq_dims, + disk_pq_pivots_path, + disk_pq_compressed_vectors_path); + } diskann::cout << "Training data loaded of size " << train_size << std::endl; @@ -705,7 +732,7 @@ namespace diskann { if (compareMetric == diskann::Metric::INNER_PRODUCT) make_zero_mean = false; generate_pq_pivots(train_data, train_size, (uint32_t) dim, 256, - (uint32_t) num_pq_chunks, 15, pq_pivots_path, make_zero_mean); + (uint32_t) num_pq_chunks, NUM_KMEANS_REPS, pq_pivots_path, make_zero_mean); generate_pq_data_from_pivots(dataFilePath, 256, (uint32_t) num_pq_chunks, pq_pivots_path, @@ -719,8 +746,12 @@ namespace diskann { dataFilePath, compareMetric, L, R, p_val, indexing_ram_budget, mem_index_path, medoids_path, centroids_path); + if (!use_disk_pq) diskann::create_disk_layout(dataFilePath, mem_index_path, disk_index_path); + else + diskann::create_disk_layout<_u8>(disk_pq_compressed_vectors_path, mem_index_path, + disk_index_path); double sample_sampling_rate = (150000.0 / points_num); gen_random_slice(dataFilePath, sample_base_prefix, diff --git a/src/partition_and_pq.cpp b/src/partition_and_pq.cpp index be323ad8ff..c41f1924b1 100644 --- a/src/partition_and_pq.cpp +++ b/src/partition_and_pq.cpp @@ -36,7 +36,7 @@ #endif #define BLOCK_SIZE 5000000 -#define SAVE_INFLATED_PQ true +//#define SAVE_INFLATED_PQ true template void gen_random_slice(const std::string base_file, diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index 0bab700ef5..484e772a7e 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -49,7 +49,7 @@ // returns region of `node_buf` containing [NNBRS][NBR_ID(_u32)] #define OFFSET_TO_NODE_NHOOD(node_buf) \ - (unsigned *) ((char *) node_buf + data_dim * sizeof(T)) + (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) @@ -320,7 +320,7 @@ namespace diskann { 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, data_dim * sizeof(T)); + 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 @@ -562,7 +562,7 @@ namespace diskann { // 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, data_dim * sizeof(T)); + memcpy(medoid_coords, medoid_disk_coords, disk_bytes_per_point); for (uint32_t i = 0; i < data_dim; i++) centroid_data[cur_m * aligned_dim + i] = medoid_coords[i]; @@ -592,6 +592,7 @@ namespace diskann { std::string centroids_file = std::string(disk_index_file) + "_centroids.bin"; + 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); @@ -608,8 +609,11 @@ namespace diskann { } this->data_dim = pq_file_dim; + this->disk_data_dim = this->data_dim; // will reset later if we use PQ on disk + this->disk_bytes_per_point = this->data_dim * sizeof(T); // will change later if we use PQ on disk this->aligned_dim = ROUND_UP(pq_file_dim, 8); + size_t npts_u64, nchunks_u64; #ifdef EXEC_ENV_OLS diskann::load_bin<_u8>(files, pq_compressed_vectors, this->data, npts_u64, @@ -634,6 +638,20 @@ namespace diskann { << " #aligned_dim: " << aligned_dim << " #chunks: " << n_chunks << std::endl; + +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 + disk_pq_table.load_pq_centroid_bin(files, disk_pq_pivots_path.c_str(), 0); // giving 0 chunks as the pq_table will infer from the chunk_offsets file the correct value +#else + disk_pq_table.load_pq_centroid_bin(disk_pq_pivots_path.c_str(), 0); // giving 0 chunks as the pq_table will infer from the chunk_offsets file the correct value +#endif + disk_pq_n_chunks = disk_pq_table.get_num_chunks(); + disk_bytes_per_point = disk_pq_n_chunks * sizeof(_u8); + std::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 @@ -678,7 +696,7 @@ namespace diskann { READ_U64(index_metadata, medoid_id_on_file); READ_U64(index_metadata, max_node_len); READ_U64(index_metadata, nnodes_per_sector); - max_degree = ((max_node_len - data_dim * sizeof(T)) / sizeof(unsigned)) - 1; + max_degree = ((max_node_len - disk_bytes_per_point) / sizeof(unsigned)) - 1; diskann::cout << "Disk-Index File Meta-data: "; diskann::cout << "# nodes per sector: " << nnodes_per_sector; @@ -796,7 +814,7 @@ namespace diskann { for (uint32_t i = 0; i < this->data_dim; i++) { data.scratch.aligned_query_float[i] = query1[i]; } - memcpy(data.scratch.aligned_query_T, query1, this->data_dim * sizeof(T)); + memcpy(data.scratch.aligned_query_T, query1, disk_bytes_per_point); const T * query = data.scratch.aligned_query_T; const float *query_float = data.scratch.aligned_query_float; @@ -1025,7 +1043,7 @@ namespace diskann { T *node_fp_coords_copy = data_buf + (data_buf_idx * aligned_dim); data_buf_idx++; - memcpy(node_fp_coords_copy, node_fp_coords, data_dim * sizeof(T)); + memcpy(node_fp_coords_copy, node_fp_coords, disk_bytes_per_point); float cur_expanded_dist = dist_cmp->compare(query, node_fp_coords_copy, (unsigned) aligned_dim); diff --git a/tests/build_disk_index.cpp b/tests/build_disk_index.cpp index 8c94d47525..77bdeace37 100644 --- a/tests/build_disk_index.cpp +++ b/tests/build_disk_index.cpp @@ -17,11 +17,11 @@ bool build_index(const char* dataFilePath, const char* indexFilePath, } int main(int argc, char** argv) { - if (argc != 10) { + if (argc != 11) { std::cout << "Usage: " << argv[0] << " [data_type] [dist_fn: 0 for L2, 1 for MIPS] [data_file.bin] " "[index_prefix_path] " - "[R] [L] [B] [M] [T]. See README for more information on " + "[R] [L] [B] [M] [T] [PQ_disk_bytes (for very large dimensionality, use 0 for full vectors)]. See README for more information on " "parameters." << std::endl; } else { @@ -30,7 +30,7 @@ int main(int argc, char** argv) { metric = diskann::Metric::INNER_PRODUCT; std::string params = std::string(argv[5]) + " " + std::string(argv[6]) + " " + std::string(argv[7]) + " " + - std::string(argv[8]) + " " + std::string(argv[9]); + std::string(argv[8]) + " " + std::string(argv[9]) + " " + std::string(argv[10]); if (std::string(argv[1]) == std::string("float")) build_index(argv[3], argv[4], params.c_str(), metric); else if (std::string(argv[1]) == std::string("int8")) From f4fbee21076754807539c1f5fb7a46f83989adcc Mon Sep 17 00:00:00 2001 From: ravishankar Date: Wed, 2 Jun 2021 08:17:09 +0530 Subject: [PATCH 44/84] code compiles for disk index pq --- src/pq_flash_index.cpp | 23 ++++++++++++++++++----- 1 file changed, 18 insertions(+), 5 deletions(-) diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index 484e772a7e..24c5a52132 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -564,10 +564,14 @@ namespace diskann { 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((_u8*) medoid_coords, (centroid_data + cur_m*aligned_dim)); + } aligned_free(medoid_buf); + delete[] medoid_coords; } // return ctx @@ -643,9 +647,9 @@ 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 - disk_pq_table.load_pq_centroid_bin(files, disk_pq_pivots_path.c_str(), 0); // giving 0 chunks as the pq_table will 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 #else - disk_pq_table.load_pq_centroid_bin(disk_pq_pivots_path.c_str(), 0); // giving 0 chunks as the pq_table will 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 #endif disk_pq_n_chunks = disk_pq_table.get_num_chunks(); disk_bytes_per_point = disk_pq_n_chunks * sizeof(_u8); @@ -974,8 +978,13 @@ if (file_exists(disk_pq_pivots_path)) { 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 = dist_cmp->compare(query, node_fp_coords_copy, + float cur_expanded_dist; + if (!use_disk_index_pq) + cur_expanded_dist = dist_cmp->compare(query, node_fp_coords_copy, (unsigned) aligned_dim); + else { + cur_expanded_dist = disk_pq_table.compare(query, (_u8*) node_fp_coords_copy); + } full_retset.push_back( Neighbor((unsigned) cached_nhood.first, cur_expanded_dist, true)); @@ -1045,8 +1054,12 @@ if (file_exists(disk_pq_pivots_path)) { data_buf_idx++; memcpy(node_fp_coords_copy, node_fp_coords, disk_bytes_per_point); - float cur_expanded_dist = dist_cmp->compare(query, node_fp_coords_copy, + float cur_expanded_dist; + if (!use_disk_index_pq) + cur_expanded_dist = dist_cmp->compare(query, node_fp_coords_copy, (unsigned) aligned_dim); + else + cur_expanded_dist = disk_pq_table.compare(query, (_u8*) node_fp_coords_copy); full_retset.push_back( Neighbor(frontier_nhood.first, cur_expanded_dist, true)); From cb966dc2439e05a0b3a140512b9db78d78a859ed Mon Sep 17 00:00:00 2001 From: ravishankar Date: Wed, 2 Jun 2021 11:06:33 +0530 Subject: [PATCH 45/84] fixed some bug --- src/aux_utils.cpp | 2 +- src/pq_flash_index.cpp | 28 +++++++++++++++++++++++++++- tests/search_disk_index.cpp | 10 +++++----- tests/utils/create_disk_layout.cpp | 13 +++++++------ 4 files changed, 40 insertions(+), 13 deletions(-) diff --git a/src/aux_utils.cpp b/src/aux_utils.cpp index 98e0e88e57..a9713a6099 100644 --- a/src/aux_utils.cpp +++ b/src/aux_utils.cpp @@ -757,7 +757,7 @@ namespace diskann { gen_random_slice(dataFilePath, sample_base_prefix, sample_sampling_rate); - std::remove(mem_index_path.c_str()); +// std::remove(mem_index_path.c_str()); auto e = std::chrono::high_resolution_clock::now(); diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index 24c5a52132..4a30053346 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -817,8 +817,9 @@ if (file_exists(disk_pq_pivots_path)) { for (uint32_t i = 0; i < this->data_dim; i++) { data.scratch.aligned_query_float[i] = query1[i]; + data.scratch.aligned_query_T[i] = query1[i]; } - memcpy(data.scratch.aligned_query_T, query1, disk_bytes_per_point); +// memcpy(data.scratch.aligned_query_T, query1, disk_bytes_per_point); const T * query = data.scratch.aligned_query_T; const float *query_float = data.scratch.aligned_query_float; @@ -888,6 +889,15 @@ if (file_exists(disk_pq_pivots_path)) { retset[0].flag = true; visited.insert(best_medoid); +/* + std::cout<<"Chose " << retset[0].id<< " as best medoid with distance " << retset[0].distance << std::endl; + + std::cout<<"query from 0 to " << aligned_dim << std::endl; + for (_u32 i = 0; i < aligned_dim-1; i++) { + std::cout< percentiles, std::vector results) { @@ -137,9 +137,9 @@ int search_disk_index(int argc, char** argv) { std::vector node_list; diskann::cout << "Caching " << num_nodes_to_cache << " BFS nodes around medoid(s)" << std::endl; -// _pFlashIndex->cache_bfs_levels(num_nodes_to_cache, node_list); - _pFlashIndex->generate_cache_list_from_sample_queries( - warmup_query_file, 15, 6, num_nodes_to_cache, num_threads, node_list); + _pFlashIndex->cache_bfs_levels(num_nodes_to_cache, node_list); +// _pFlashIndex->generate_cache_list_from_sample_queries( +// warmup_query_file, 15, 6, num_nodes_to_cache, num_threads, node_list); _pFlashIndex->load_cache_list(node_list); node_list.clear(); node_list.shrink_to_fit(); @@ -208,7 +208,7 @@ int search_disk_index(int argc, char** argv) { uint32_t optimized_beamwidth = 2; - +//query_num = 1; for (uint32_t test_id = 0; test_id < Lvec.size(); test_id++) { _u64 L = Lvec[test_id]; diff --git a/tests/utils/create_disk_layout.cpp b/tests/utils/create_disk_layout.cpp index ac378272d8..458e2d6b20 100644 --- a/tests/utils/create_disk_layout.cpp +++ b/tests/utils/create_disk_layout.cpp @@ -14,12 +14,6 @@ template int create_disk_layout(int argc, char **argv) { - if (argc != 5) { - std::cout << argv[0] << " data_type data_bin " - "vamana_index_file output_diskann_index_file" - << std::endl; - exit(-1); - } std::string base_file(argv[2]); std::string vamana_file(argv[3]); std::string output_file(argv[4]); @@ -28,6 +22,13 @@ int create_disk_layout(int argc, char **argv) { } 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(argc, argv); From 3070a29b8acfb3a55e1743d7f4afadf61c77d769 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Wed, 2 Jun 2021 07:02:05 +0000 Subject: [PATCH 46/84] shards are written as and when necessary --- include/partition_and_pq.h | 9 ++ src/aux_utils.cpp | 7 ++ src/partition_and_pq.cpp | 171 ++++++++++++++++++++++++++++++++++++- 3 files changed, 185 insertions(+), 2 deletions(-) diff --git a/include/partition_and_pq.h b/include/partition_and_pq.h index 43a9d84db0..bbae1ed148 100644 --- a/include/partition_and_pq.h +++ b/include/partition_and_pq.h @@ -38,6 +38,15 @@ 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); + +template +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, diff --git a/src/aux_utils.cpp b/src/aux_utils.cpp index a9713a6099..d77d7fa998 100644 --- a/src/aux_utils.cpp +++ b/src/aux_utils.cpp @@ -385,6 +385,12 @@ namespace diskann { 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"; + + 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"; @@ -402,6 +408,7 @@ namespace diskann { new diskann::Index(compareMetric, shard_base_file.c_str())); _pvamanaIndex->build(paras); _pvamanaIndex->save(shard_index_file.c_str()); + std::remove(shard_base_file.c_str()); } diskann::merge_shards(merged_index_prefix + "_subshard-", "_mem.index", diff --git a/src/partition_and_pq.cpp b/src/partition_and_pq.cpp index c41f1924b1..a219dc25a0 100644 --- a/src/partition_and_pq.cpp +++ b/src/partition_and_pq.cpp @@ -35,7 +35,7 @@ #include #endif -#define BLOCK_SIZE 5000000 +#define BLOCK_SIZE 50000 //#define SAVE_INFLATED_PQ true template @@ -741,6 +741,169 @@ int shard_data_into_clusters(const std::string data_file, float *pivots, return 0; } + + + +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; + // create cached reader + writer + cached_ifstream base_reader(data_file, read_blk_size); + _u32 npts32; + _u32 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); + _u32 dummy_size = 0; + _u32 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<_u32[]> block_closest_centers = + std::make_unique<_u32[]>(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++) { + _u32 cur_shard_count = (_u32) 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) { + _u64 read_blk_size = 64 * 1024 * 1024; + // _u64 write_blk_size = 64 * 1024 * 1024; + // create cached reader + writer + cached_ifstream base_reader(data_file, read_blk_size); + _u32 npts32; + _u32 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; + + 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); + + _u32 cur_pos = 0; + _u32 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]) { + 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; + + 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 // 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 @@ -867,7 +1030,7 @@ int partition_with_ram_budget(const std::string data_file, diskann::save_bin(output_file.c_str(), pivot_data, (size_t) num_parts, train_dim); - shard_data_into_clusters(data_file, pivot_data, 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; @@ -926,6 +1089,10 @@ 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); + template DISKANN_DLLEXPORT int generate_pq_data_from_pivots( const std::string data_file, unsigned num_centers, unsigned num_pq_chunks, std::string pq_pivots_path, std::string pq_compressed_vectors_path); From e916751487d0aef87d2696ab16a2c6a542ae7ead Mon Sep 17 00:00:00 2001 From: ravishankar Date: Wed, 2 Jun 2021 07:07:26 +0000 Subject: [PATCH 47/84] sharding is now on demand --- src/partition_and_pq.cpp | 1 + 1 file changed, 1 insertion(+) diff --git a/src/partition_and_pq.cpp b/src/partition_and_pq.cpp index a219dc25a0..bf1cb8a81c 100644 --- a/src/partition_and_pq.cpp +++ b/src/partition_and_pq.cpp @@ -881,6 +881,7 @@ int retrieve_shard_data_from_ids(const std::string data_file, std::string idmap_ 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++; From 063336aaa95ffb8da1e82d2d4a336ad10888acfe Mon Sep 17 00:00:00 2001 From: ravishankar Date: Wed, 2 Jun 2021 13:35:06 +0000 Subject: [PATCH 48/84] minor changes --- src/aux_utils.cpp | 8 ++++---- src/partition_and_pq.cpp | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/aux_utils.cpp b/src/aux_utils.cpp index d77d7fa998..e95eff0e30 100644 --- a/src/aux_utils.cpp +++ b/src/aux_utils.cpp @@ -741,9 +741,7 @@ namespace diskann { generate_pq_pivots(train_data, train_size, (uint32_t) dim, 256, (uint32_t) num_pq_chunks, NUM_KMEANS_REPS, pq_pivots_path, make_zero_mean); generate_pq_data_from_pivots(dataFilePath, 256, (uint32_t) - num_pq_chunks, - pq_pivots_path, - pq_compressed_vectors_path); + num_pq_chunks, pq_pivots_path, pq_compressed_vectors_path); delete[] train_data; @@ -764,7 +762,9 @@ namespace diskann { gen_random_slice(dataFilePath, sample_base_prefix, sample_sampling_rate); -// std::remove(mem_index_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(); diff --git a/src/partition_and_pq.cpp b/src/partition_and_pq.cpp index bf1cb8a81c..5e34f35ad1 100644 --- a/src/partition_and_pq.cpp +++ b/src/partition_and_pq.cpp @@ -35,7 +35,7 @@ #include #endif -#define BLOCK_SIZE 50000 +#define BLOCK_SIZE 5000000 //#define SAVE_INFLATED_PQ true template From 6ded4aa0977315f560c104c4eb552724dc6247b9 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Wed, 2 Jun 2021 15:22:24 +0000 Subject: [PATCH 49/84] fixed one malloc bug in parameters --- include/parameters.h | 3 +++ 1 file changed, 3 insertions(+) diff --git a/include/parameters.h b/include/parameters.h index 42e7d3a463..1cff662850 100644 --- a/include/parameters.h +++ b/include/parameters.h @@ -19,6 +19,9 @@ namespace diskann { template inline void Set(const std::string &name, const ParamType &value) { // ParamType *ptr = (ParamType *) malloc(sizeof(ParamType)); + if (params.find(name) != params.end()) { + free(params[name]); + } ParamType *ptr = new ParamType; *ptr = value; params[name] = (void *) ptr; From 66393e7f7cad4a5a4f7a180cdfc875ca6ea2dabb Mon Sep 17 00:00:00 2001 From: ravishankar Date: Thu, 10 Jun 2021 16:24:59 +0000 Subject: [PATCH 50/84] added a vector analyzer util --- include/aux_utils.h | 2 +- include/utils.h | 5 +++++ src/aux_utils.cpp | 9 +++++++++ tests/utils/CMakeLists.txt | 10 ++++++++++ 4 files changed, 25 insertions(+), 1 deletion(-) diff --git a/include/aux_utils.h b/include/aux_utils.h index 7e946becd4..ac7bad6cc8 100644 --- a/include/aux_utils.h +++ b/include/aux_utils.h @@ -31,7 +31,7 @@ typedef int FileHandle; #include "windows_customizations.h" namespace diskann { - const size_t TRAINING_SET_SIZE = 1500000; + const size_t TRAINING_SET_SIZE = 150000; const double SPACE_FOR_CACHED_NODES_IN_GB = 0.25; const double THRESHOLD_FOR_CACHING_IN_GB = 1.0; const uint32_t NUM_NODES_TO_CACHE = 250000; diff --git a/include/utils.h b/include/utils.h index 6b9db5bf62..3e9bafc088 100644 --- a/include/utils.h +++ b/include/utils.h @@ -218,6 +218,11 @@ namespace diskann { } #endif + inline void wait_for_keystroke() { + int a; + std::cin>> a; + } + template inline void load_bin(const std::string& bin_file, T*& data, size_t& npts, size_t& dim) { diff --git a/src/aux_utils.cpp b/src/aux_utils.cpp index e95eff0e30..a08bb15e1f 100644 --- a/src/aux_utils.cpp +++ b/src/aux_utils.cpp @@ -379,6 +379,9 @@ namespace diskann { partition_with_ram_budget(base_file, sampling_rate, ram_budget, 2 * R / 3, merged_index_prefix, 2); + wait_for_keystroke(); + + std::string cur_centroid_filepath = merged_index_prefix + "_centroids.bin"; std::rename(cur_centroid_filepath.c_str(), centroids_file.c_str()); @@ -409,6 +412,8 @@ namespace diskann { _pvamanaIndex->build(paras); _pvamanaIndex->save(shard_index_file.c_str()); std::remove(shard_base_file.c_str()); + wait_for_keystroke(); + } diskann::merge_shards(merged_index_prefix + "_subshard-", "_mem.index", @@ -735,6 +740,8 @@ namespace diskann { diskann::cout << "Training data loaded of size " << train_size << std::endl; + wait_for_keystroke(); + bool make_zero_mean = true; if (compareMetric == diskann::Metric::INNER_PRODUCT) make_zero_mean = false; @@ -751,6 +758,8 @@ namespace diskann { dataFilePath, compareMetric, L, R, p_val, indexing_ram_budget, mem_index_path, medoids_path, centroids_path); + + if (!use_disk_pq) diskann::create_disk_layout(dataFilePath, mem_index_path, disk_index_path); diff --git a/tests/utils/CMakeLists.txt b/tests/utils/CMakeLists.txt index 0955d57b97..7769229235 100644 --- a/tests/utils/CMakeLists.txt +++ b/tests/utils/CMakeLists.txt @@ -56,6 +56,16 @@ else() target_link_libraries(uint32_to_uint8 ${PROJECT_NAME}) endif() + +add_executable(vector_analysis vector_analysis.cpp) +if(MSVC) + target_link_options(vector_analysis PRIVATE /MACHINE:x64) + target_link_libraries(vector_analysis debug ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}/diskann_dll.lib) + target_link_libraries(vector_analysis optimized ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}/diskann_dll.lib) +else() + target_link_libraries(vector_analysis ${PROJECT_NAME} -ltcmalloc) +endif() + add_executable(gen_random_slice gen_random_slice.cpp) if(MSVC) target_link_options(gen_random_slice PRIVATE /MACHINE:x64) From 56cf277a6f724fefa1b8732d522672cf5ba8ed72 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Thu, 10 Jun 2021 16:29:39 +0000 Subject: [PATCH 51/84] added missing file --- tests/utils/vector_analysis.cpp | 72 +++++++++++++++++++++++++++++++++ 1 file changed, 72 insertions(+) create mode 100644 tests/utils/vector_analysis.cpp diff --git a/tests/utils/vector_analysis.cpp b/tests/utils/vector_analysis.cpp new file mode 100644 index 0000000000..1cd2eb9ff5 --- /dev/null +++ b/tests/utils/vector_analysis.cpp @@ -0,0 +1,72 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "partition_and_pq.h" +#include "utils.h" + +#include +#include +#include +#include + +template +int analyze_norm(std::string base_file) { + std::cout<<"Analyzing data norms" << std::endl; + T* data; + _u64 npts, ndims; + diskann::load_bin(base_file, data, npts, ndims); + std::vector norms(npts, 0); + #pragma omp parallel for schedule(dynamic) + for (_u32 i = 0; i +int aux_main(int argc, char** argv) { + + std::string base_file(argv[2]); + _u32 option = atoi(argv[3]); + if (option == 1) + analyze_norm(base_file); + 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]" + << std::endl; + exit(-1); + } + + if (std::string(argv[1]) == std::string("float")) { + aux_main(argc, argv); + } else if (std::string(argv[1]) == std::string("int8")) { + aux_main(argc, argv); + } else if (std::string(argv[1]) == std::string("uint8")) { + aux_main(argc, argv); + } else + std::cout << "Unsupported type. Use float/int8/uint8." << std::endl; + return 0; +} From 43d94c61026b9802f2290bd328c244312b61649b Mon Sep 17 00:00:00 2001 From: ravishankar Date: Sat, 12 Jun 2021 03:34:34 +0000 Subject: [PATCH 52/84] fixed a bug which used L2 instead of inner product in cached beam search --- include/pq_table.h | 14 ++++++++++++++ src/pq_flash_index.cpp | 17 ++++++++++++----- 2 files changed, 26 insertions(+), 5 deletions(-) diff --git a/include/pq_table.h b/include/pq_table.h index 825b29c7ec..7c8438e01e 100644 --- a/include/pq_table.h +++ b/include/pq_table.h @@ -162,6 +162,20 @@ _u32 get_num_chunks() { return res; } + float inner_product(const T* query_vec, _u8* base_vec) { + float res = 0; + for (_u64 chunk = 0; chunk < n_chunks; chunk++) { + for (_u64 j = chunk_offsets[chunk]; j < chunk_offsets[chunk + 1]; j++) { + _u64 permuted_dim_in_query = rearrangement[j]; + const float* centers_dim_vec = tables_T + (256 * j); + float diff = centers_dim_vec[base_vec[chunk]]*query_vec[permuted_dim_in_query]; // assumes centroid is 0 to prevent translation errors + res += diff; + } + } + return -res; // returns negative value to simulate distances (max -> min conversion) + } + + void inflate_vector(_u8* base_vec, float* out_vec) { for (_u64 chunk = 0; chunk < n_chunks; chunk++) { for (_u64 j = chunk_offsets[chunk]; j < chunk_offsets[chunk + 1]; j++) { diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index 4a30053346..1943e51241 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -844,9 +844,9 @@ if (file_exists(disk_pq_pivots_path)) { // query <-> PQ chunk centers distances float *pq_dists = query_scratch->aligned_pqtable_dist_scratch; -// if (metric==diskann::Metric::INNER_PRODUCT) -// pq_table.populate_chunk_inner_products(query, pq_dists); -// else if (metric==diskann::Metric::L2) + if (metric==diskann::Metric::INNER_PRODUCT) + pq_table.populate_chunk_inner_products(query, pq_dists); + else if (metric==diskann::Metric::L2) pq_table.populate_chunk_distances(query, pq_dists); // query <-> neighbor list @@ -1000,6 +1000,9 @@ if (file_exists(disk_pq_pivots_path)) { cur_expanded_dist = dist_cmp->compare(query, node_fp_coords_copy, (unsigned) aligned_dim); else { + if (metric == diskann::Metric::INNER_PRODUCT) + cur_expanded_dist = disk_pq_table.inner_product(query, (_u8*) node_fp_coords_copy); + else cur_expanded_dist = disk_pq_table.compare(query, (_u8*) node_fp_coords_copy); } full_retset.push_back( @@ -1075,8 +1078,12 @@ if (file_exists(disk_pq_pivots_path)) { if (!use_disk_index_pq) cur_expanded_dist = dist_cmp->compare(query, node_fp_coords_copy, (unsigned) aligned_dim); - else - cur_expanded_dist = disk_pq_table.compare(query, (_u8*) node_fp_coords_copy); + else { + if (metric == diskann::Metric::INNER_PRODUCT) + cur_expanded_dist = disk_pq_table.inner_product(query, (_u8*) node_fp_coords_copy); + else + cur_expanded_dist = disk_pq_table.compare(query, (_u8*) node_fp_coords_copy); + } full_retset.push_back( Neighbor(frontier_nhood.first, cur_expanded_dist, true)); From 31b2ae19122927144c5ac3fa828de18c769d57d5 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Mon, 14 Jun 2021 11:38:46 +0000 Subject: [PATCH 53/84] now setting up the normalizing approach --- tests/utils/vector_analysis.cpp | 50 +++++++++++++++++++++++++++++++-- 1 file changed, 48 insertions(+), 2 deletions(-) diff --git a/tests/utils/vector_analysis.cpp b/tests/utils/vector_analysis.cpp index 1cd2eb9ff5..9f5e5daa9f 100644 --- a/tests/utils/vector_analysis.cpp +++ b/tests/utils/vector_analysis.cpp @@ -36,10 +36,52 @@ int analyze_norm(std::string base_file) { for (_u32 p = 0; p < 100; p+=5) std::cout<<"percentile "< +int augment_base(std::string base_file,std::string out_file, bool prep_base = true) { + std::cout<<"Analyzing data norms" << std::endl; + T* data; + _u64 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 (_u32 i = 0; 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; + _u64 newdims = ndims + 1; + new_data = new T[npts*newdims]; + for (_u64 i = 0;i < npts; i++) { + for (_u64 j = 0; j < ndims; j++) { + new_data[i*newdims + j] = data[i*ndims +j]/ max_norm; + } + if (prep_base) { + 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<(out_file, new_data, npts, newdims); + delete[] new_data; + delete[] data; + return 0; +} + template int aux_main(int argc, char** argv) { @@ -48,14 +90,18 @@ int aux_main(int argc, char** argv) { _u32 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); return 0; } int main(int argc, char** argv) { - if (argc != 4) { + if (argc < 4) { std::cout << argv[0] << " data_type [float/int8/uint8] base_bin_file " - "[option: 1-norm analysis]" + "[option: 1-norm analysis, 2-prep_base_for_mip, 3-prep_query_for_mip] [out_file for options 2/3]" << std::endl; exit(-1); } From c36f77a33dd5dcaca4aae40b8a7aaa97919a1169 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Mon, 14 Jun 2021 12:15:48 +0000 Subject: [PATCH 54/84] towards pre-processing data --- include/utils.h | 24 ++++++++++++++++++++++++ src/aux_utils.cpp | 27 ++++++++++++++++----------- src/index.cpp | 4 ++-- src/partition_and_pq.cpp | 2 ++ tests/utils/vector_analysis.cpp | 2 +- 5 files changed, 45 insertions(+), 14 deletions(-) diff --git a/include/utils.h b/include/utils.h index 3e9bafc088..5d76443c99 100644 --- a/include/utils.h +++ b/include/utils.h @@ -52,6 +52,8 @@ typedef int FileHandle; #define IS_512_ALIGNED(X) IS_ALIGNED(X, 512) #define IS_4096_ALIGNED(X) IS_ALIGNED(X, 4096) + + typedef uint64_t _u64; typedef int64_t _s64; typedef uint32_t _u32; @@ -415,6 +417,28 @@ namespace diskann { } } +template +void 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); + _u64 npts, in_dims, out_dims; + float max_norm = 0; + + _u32 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; + + size_t BLOCK_SIZE = 5000000; + size_t block_size = npts <= BLOCK_SIZE ? npts : BLOCK_SIZE; + std::unique_ptr block_data_T = std::make_unique(block_size * out_dims); + +} + // 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) { diff --git a/src/aux_utils.cpp b/src/aux_utils.cpp index a08bb15e1f..97a04d7b09 100644 --- a/src/aux_utils.cpp +++ b/src/aux_utils.cpp @@ -642,7 +642,7 @@ namespace diskann { << std::endl; return false; } - + _u32 disk_pq_dims = 0; bool use_disk_pq = false; @@ -653,9 +653,8 @@ namespace diskann { use_disk_pq = false; } - if (compareMetric == diskann::Metric::INNER_PRODUCT) { - std::cout<<"Using Inner Product for PQ and Graph Generation" << std::endl; - } + std::string base_file(dataFilePath); + std::string data_file_to_use = base_file; std::string index_prefix_path(indexFilePath); std::string pq_pivots_path = index_prefix_path + "_pq_pivots.bin"; std::string pq_compressed_vectors_path = @@ -670,6 +669,12 @@ namespace diskann { index_prefix_path + "_disk.index_pq_compressed.bin"; + if (compareMetric == diskann::Metric::INNER_PRODUCT) { + std::cout<<"Using Inner Product search, so need to pre-process base data into temp file. Please ensure there is sufficient space!!" << std::endl; + std::string prepped_base = index_prefix_path + "_prepped_base.bin"; + data_file_to_use = prepped_base; + diskann::prepare_base_for_inner_products(base_file, prepped_base); + } unsigned R = (unsigned) atoi(param_list[0].c_str()); unsigned L = (unsigned) atoi(param_list[1].c_str()); @@ -704,7 +709,7 @@ namespace diskann { size_t points_num, dim; - diskann::get_bin_metadata(dataFilePath, points_num, dim); + diskann::get_bin_metadata(data_file_to_use.c_str(), points_num, dim); size_t num_pq_chunks = (size_t)(std::floor)(_u64(final_index_ram_limit / points_num)); @@ -723,7 +728,7 @@ namespace diskann { double p_val = ((double) TRAINING_SET_SIZE / (double) points_num); // generates random sample and sets it to train_data and updates // train_size - gen_random_slice(dataFilePath, p_val, train_data, train_size, + gen_random_slice(data_file_to_use.c_str(), p_val, train_data, train_size, train_dim); if (use_disk_pq) { @@ -733,7 +738,7 @@ namespace diskann { std::cout<<"Compressing base for disk-PQ into " << disk_pq_dims << " chunks " << std::endl; generate_pq_pivots(train_data, train_size, (uint32_t) dim, 256, (uint32_t) disk_pq_dims, NUM_KMEANS_REPS, disk_pq_pivots_path, false); - generate_pq_data_from_pivots(dataFilePath, 256, (uint32_t) disk_pq_dims, + generate_pq_data_from_pivots(data_file_to_use.c_str(), 256, (uint32_t) disk_pq_dims, disk_pq_pivots_path, disk_pq_compressed_vectors_path); } @@ -747,7 +752,7 @@ namespace diskann { make_zero_mean = false; generate_pq_pivots(train_data, train_size, (uint32_t) dim, 256, (uint32_t) num_pq_chunks, NUM_KMEANS_REPS, pq_pivots_path, make_zero_mean); - generate_pq_data_from_pivots(dataFilePath, 256, (uint32_t) + generate_pq_data_from_pivots(data_file_to_use.c_str(), 256, (uint32_t) num_pq_chunks, pq_pivots_path, pq_compressed_vectors_path); delete[] train_data; @@ -755,20 +760,20 @@ namespace diskann { train_data = nullptr; diskann::build_merged_vamana_index( - dataFilePath, compareMetric, L, R, p_val, indexing_ram_budget, + data_file_to_use.c_str(), compareMetric, L, R, p_val, indexing_ram_budget, mem_index_path, medoids_path, centroids_path); if (!use_disk_pq) - diskann::create_disk_layout(dataFilePath, mem_index_path, + diskann::create_disk_layout(data_file_to_use.c_str(), mem_index_path, disk_index_path); else diskann::create_disk_layout<_u8>(disk_pq_compressed_vectors_path, mem_index_path, disk_index_path); double sample_sampling_rate = (150000.0 / points_num); - gen_random_slice(dataFilePath, sample_base_prefix, + gen_random_slice(data_file_to_use.c_str(), sample_base_prefix, sample_sampling_rate); std::remove(mem_index_path.c_str()); diff --git a/src/index.cpp b/src/index.cpp index ff00f6dc61..b8d5c26642 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -504,9 +504,9 @@ namespace diskann { std::vector init_ids, std::vector & expanded_nodes_info, tsl::robin_set &expanded_nodes_ids) { - const T * node_coords = _data + _aligned_dim * node_id; + T * node_coords = _data + _aligned_dim * node_id; std::vector best_L_nodes; - + if (init_ids.size() == 0) init_ids.emplace_back(_ep); diff --git a/src/partition_and_pq.cpp b/src/partition_and_pq.cpp index 5e34f35ad1..6c8b0a431c 100644 --- a/src/partition_and_pq.cpp +++ b/src/partition_and_pq.cpp @@ -35,7 +35,9 @@ #include #endif +// block size for reading/ processing large files and matrices in blocks #define BLOCK_SIZE 5000000 + //#define SAVE_INFLATED_PQ true template diff --git a/tests/utils/vector_analysis.cpp b/tests/utils/vector_analysis.cpp index 9f5e5daa9f..41ee323c7f 100644 --- a/tests/utils/vector_analysis.cpp +++ b/tests/utils/vector_analysis.cpp @@ -70,7 +70,7 @@ std::cout<<"Max norm: " << max_norm << std::endl; diff = diff <= 0 ? 0 : std::sqrt(diff); new_data[i*newdims + ndims] = diff; if (diff <= 0) { - std::cout< Date: Mon, 14 Jun 2021 16:41:45 +0000 Subject: [PATCH 55/84] working towards newer inner product --- include/pq_flash_index.h | 4 +-- include/pq_table.h | 10 +++--- include/utils.h | 47 +++++++++++++++++++++++-- src/aux_utils.cpp | 37 +++++++++++++++----- src/pq_flash_index.cpp | 68 ++++++++++++++++++++++--------------- tests/search_disk_index.cpp | 6 ++-- 6 files changed, 125 insertions(+), 47 deletions(-) diff --git a/include/pq_flash_index.h b/include/pq_flash_index.h index 489afbed31..fbc6fc62b3 100644 --- a/include/pq_flash_index.h +++ b/include/pq_flash_index.h @@ -147,7 +147,7 @@ namespace diskann { // pq_tables = float* [[2^8 * [chunk_size]] * n_chunks] _u8 * data = nullptr; _u64 n_chunks; - FixedChunkPQTable pq_table; + FixedChunkPQTable pq_table; // distance comparator Distance * dist_cmp = nullptr; @@ -156,7 +156,7 @@ namespace diskann { // for very large datasets: we use PQ even for the disk resident index bool use_disk_index_pq = false; _u64 disk_pq_n_chunks; - FixedChunkPQTable disk_pq_table; + FixedChunkPQTable disk_pq_table; // medoid/start info diff --git a/include/pq_table.h b/include/pq_table.h index 7c8438e01e..927e5bc7db 100644 --- a/include/pq_table.h +++ b/include/pq_table.h @@ -6,7 +6,7 @@ #include "utils.h" namespace diskann { - template +// template class FixedChunkPQTable { // data_dim = n_chunks * chunk_size; float* tables = @@ -130,7 +130,7 @@ _u32 get_num_chunks() { return n_chunks; } void - populate_chunk_distances(const T* query_vec, float* dist_vec) { + populate_chunk_distances(const float* query_vec, float* dist_vec) { memset(dist_vec, 0, 256 * n_chunks * sizeof(float)); // chunk wise distance computation for (_u64 chunk = 0; chunk < n_chunks; chunk++) { @@ -149,7 +149,7 @@ _u32 get_num_chunks() { } } - float compare(const T* query_vec, _u8* base_vec) { + float compare(const float* query_vec, _u8* base_vec) { float res = 0; for (_u64 chunk = 0; chunk < n_chunks; chunk++) { for (_u64 j = chunk_offsets[chunk]; j < chunk_offsets[chunk + 1]; j++) { @@ -162,7 +162,7 @@ _u32 get_num_chunks() { return res; } - float inner_product(const T* query_vec, _u8* base_vec) { + float inner_product(const float* query_vec, _u8* base_vec) { float res = 0; for (_u64 chunk = 0; chunk < n_chunks; chunk++) { for (_u64 j = chunk_offsets[chunk]; j < chunk_offsets[chunk + 1]; j++) { @@ -187,7 +187,7 @@ _u32 get_num_chunks() { } void - populate_chunk_inner_products(const T* query_vec, float* dist_vec) { + 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 (_u64 chunk = 0; chunk < n_chunks; chunk++) { diff --git a/include/utils.h b/include/utils.h index 5d76443c99..ae90cacdd8 100644 --- a/include/utils.h +++ b/include/utils.h @@ -432,11 +432,54 @@ void prepare_base_for_inner_products(const std::string in_file, const std::strin npts = npts32; in_dims = dims32; out_dims = in_dims+1; + _u32 outdims32 = (_u32) out_dims; - size_t BLOCK_SIZE = 5000000; + 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 block_data_T = std::make_unique(block_size * out_dims); + 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); + _u64 num_blocks = DIV_ROUND_UP(npts, block_size); + + std::vector norms(npts, 0); + + for (_u64 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; + in_reader.read((char *) in_block_data.get(), block_pts * in_dims * sizeof(T)); + for (_u64 p = 0; p < block_pts; p++) { + for (_u64 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(_u32), std::ios::beg); + for (_u64 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; + in_reader.read((char *) in_block_data.get(), block_pts * in_dims * sizeof(T)); + for (_u64 p = 0; p < block_pts; p++) { + for (_u64 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.close(); } // plain saves data as npts X ndims array into filename diff --git a/src/aux_utils.cpp b/src/aux_utils.cpp index 97a04d7b09..d57fa0cff6 100644 --- a/src/aux_utils.cpp +++ b/src/aux_utils.cpp @@ -360,7 +360,7 @@ namespace diskann { paras.Set("L", (unsigned) L); paras.Set("R", (unsigned) R); paras.Set("C", 750); - paras.Set("alpha", 2.0f); + paras.Set("alpha", 1.2f); paras.Set("num_rnds", 2); paras.Set("saturate_graph", 1); paras.Set("save_path", mem_index_path); @@ -738,6 +738,11 @@ namespace diskann { std::cout<<"Compressing base for disk-PQ into " << disk_pq_dims << " chunks " << std::endl; generate_pq_pivots(train_data, train_size, (uint32_t) dim, 256, (uint32_t) disk_pq_dims, NUM_KMEANS_REPS, disk_pq_pivots_path, false); + if (compareMetric == diskann::Metric::INNER_PRODUCT) + generate_pq_data_from_pivots(data_file_to_use.c_str(), 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.c_str(), 256, (uint32_t) disk_pq_dims, disk_pq_pivots_path, disk_pq_compressed_vectors_path); @@ -745,13 +750,17 @@ namespace diskann { diskann::cout << "Training data loaded of size " << train_size << std::endl; - wait_for_keystroke(); +// wait_for_keystroke(); bool make_zero_mean = true; if (compareMetric == diskann::Metric::INNER_PRODUCT) make_zero_mean = false; generate_pq_pivots(train_data, train_size, (uint32_t) dim, 256, (uint32_t) num_pq_chunks, NUM_KMEANS_REPS, pq_pivots_path, make_zero_mean); + if (compareMetric == diskann::Metric::INNER_PRODUCT) + generate_pq_data_from_pivots(data_file_to_use.c_str(), 256, (uint32_t) + num_pq_chunks, pq_pivots_path, pq_compressed_vectors_path); + else generate_pq_data_from_pivots(data_file_to_use.c_str(), 256, (uint32_t) num_pq_chunks, pq_pivots_path, pq_compressed_vectors_path); @@ -759,24 +768,36 @@ namespace diskann { train_data = nullptr; + if (compareMetric == diskann::Metric::INNER_PRODUCT) + 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); + else diskann::build_merged_vamana_index( - data_file_to_use.c_str(), compareMetric, L, R, p_val, indexing_ram_budget, + data_file_to_use.c_str(), diskann::Metric::L2, L, R, p_val, indexing_ram_budget, mem_index_path, medoids_path, centroids_path); - - - if (!use_disk_pq) - diskann::create_disk_layout(data_file_to_use.c_str(), mem_index_path, + if (!use_disk_pq) { + if (compareMetric == diskann::Metric::INNER_PRODUCT) + diskann::create_disk_layout(data_file_to_use.c_str(), mem_index_path, disk_index_path); + else + diskann::create_disk_layout(data_file_to_use.c_str(), mem_index_path, + disk_index_path); + } else diskann::create_disk_layout<_u8>(disk_pq_compressed_vectors_path, mem_index_path, disk_index_path); double sample_sampling_rate = (150000.0 / points_num); + if (compareMetric == diskann::Metric::INNER_PRODUCT) + gen_random_slice(data_file_to_use.c_str(), sample_base_prefix, + sample_sampling_rate); + else gen_random_slice(data_file_to_use.c_str(), sample_base_prefix, sample_sampling_rate); - std::remove(mem_index_path.c_str()); +// std::remove(mem_index_path.c_str()); if (use_disk_pq) std::remove(disk_pq_compressed_vectors_path.c_str()); diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index 1943e51241..d140b0f77d 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -614,7 +614,10 @@ namespace diskann { this->data_dim = pq_file_dim; this->disk_data_dim = this->data_dim; // will reset later if we use PQ on disk - this->disk_bytes_per_point = this->data_dim * sizeof(T); // will change later if we use PQ on disk + this->disk_bytes_per_point = this->data_dim * sizeof(T); // will change later if we use PQ on disk or if we are using inner product without PQ + if (metric == diskann::Metric::INNER_PRODUCT) { + this->disk_bytes_per_point = this->data_dim * sizeof(float); // because we normalize the data and store it as float if no PQ + } this->aligned_dim = ROUND_UP(pq_file_dim, 8); @@ -815,10 +818,24 @@ if (file_exists(disk_pq_pivots_path)) { data = this->thread_data.pop(); } + float query_norm = 0; + if (metric == diskann::Metric::L2) { for (uint32_t i = 0; i < this->data_dim; i++) { data.scratch.aligned_query_float[i] = query1[i]; data.scratch.aligned_query_T[i] = query1[i]; } + } else if (metric == diskann::Metric::INNER_PRODUCT) { + for (uint32_t i = 0; i < this->data_dim - 1; i++) { + data.scratch.aligned_query_float[i] = query1[i]; + query_norm += query1[i]*query1[i]; + } + query_norm = std::sqrt(query_norm); + data.scratch.aligned_query_float[this->data_dim -1] = 0; + for (uint32_t i = 0; i < this->data_dim - 1; i++) { + data.scratch.aligned_query_float[i] /= query_norm; + } + } + // memcpy(data.scratch.aligned_query_T, query1, disk_bytes_per_point); const T * query = data.scratch.aligned_query_T; const float *query_float = data.scratch.aligned_query_float; @@ -844,10 +861,10 @@ if (file_exists(disk_pq_pivots_path)) { // query <-> PQ chunk centers distances float *pq_dists = query_scratch->aligned_pqtable_dist_scratch; - if (metric==diskann::Metric::INNER_PRODUCT) - pq_table.populate_chunk_inner_products(query, pq_dists); - else if (metric==diskann::Metric::L2) - pq_table.populate_chunk_distances(query, pq_dists); +// if (metric==diskann::Metric::INNER_PRODUCT) +// pq_table.populate_chunk_inner_products(query, pq_dists); +// else if (metric==diskann::Metric::L2) + pq_table.populate_chunk_distances(query_float, pq_dists); // query <-> neighbor list float *dist_scratch = query_scratch->aligned_dist_scratch; @@ -889,15 +906,6 @@ if (file_exists(disk_pq_pivots_path)) { retset[0].flag = true; visited.insert(best_medoid); -/* - std::cout<<"Chose " << retset[0].id<< " as best medoid with distance " << retset[0].distance << std::endl; - - std::cout<<"query from 0 to " << aligned_dim << std::endl; - for (_u32 i = 0; i < aligned_dim-1; i++) { - std::cout<second; float cur_expanded_dist; - if (!use_disk_index_pq) + if (!use_disk_index_pq) { + if (metric == diskann::Metric::INNER_PRODUCT) + cur_expanded_dist = dist_cmp_float->compare(query_float, (float*) node_fp_coords_copy, + (unsigned) aligned_dim); + else cur_expanded_dist = dist_cmp->compare(query, node_fp_coords_copy, (unsigned) aligned_dim); + } else { if (metric == diskann::Metric::INNER_PRODUCT) - cur_expanded_dist = disk_pq_table.inner_product(query, (_u8*) node_fp_coords_copy); + cur_expanded_dist = disk_pq_table.inner_product(query_float, (_u8*) node_fp_coords_copy); else - cur_expanded_dist = disk_pq_table.compare(query, (_u8*) node_fp_coords_copy); + cur_expanded_dist = disk_pq_table.compare(query_float, (_u8*) node_fp_coords_copy); } full_retset.push_back( Neighbor((unsigned) cached_nhood.first, cur_expanded_dist, true)); @@ -1075,14 +1082,21 @@ if (file_exists(disk_pq_pivots_path)) { 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(query, node_fp_coords_copy, + if (!use_disk_index_pq) { + if (metric == diskann::Metric::INNER_PRODUCT) + cur_expanded_dist = dist_cmp_float->compare(query_float, (float*) node_fp_coords_copy, (unsigned) aligned_dim); + else + cur_expanded_dist = dist_cmp->compare(query, node_fp_coords_copy, + (unsigned) aligned_dim); + + + } else { if (metric == diskann::Metric::INNER_PRODUCT) - cur_expanded_dist = disk_pq_table.inner_product(query, (_u8*) node_fp_coords_copy); + cur_expanded_dist = disk_pq_table.inner_product(query_float, (_u8*) node_fp_coords_copy); else - cur_expanded_dist = disk_pq_table.compare(query, (_u8*) node_fp_coords_copy); + cur_expanded_dist = disk_pq_table.compare(query_float, (_u8*) node_fp_coords_copy); } full_retset.push_back( Neighbor(frontier_nhood.first, cur_expanded_dist, true)); diff --git a/tests/search_disk_index.cpp b/tests/search_disk_index.cpp index fe2fe5b9e5..e89c562a63 100644 --- a/tests/search_disk_index.cpp +++ b/tests/search_disk_index.cpp @@ -137,9 +137,9 @@ int search_disk_index(int argc, char** argv) { std::vector node_list; diskann::cout << "Caching " << num_nodes_to_cache << " BFS nodes around medoid(s)" << std::endl; - _pFlashIndex->cache_bfs_levels(num_nodes_to_cache, node_list); -// _pFlashIndex->generate_cache_list_from_sample_queries( -// warmup_query_file, 15, 6, num_nodes_to_cache, num_threads, node_list); +// _pFlashIndex->cache_bfs_levels(num_nodes_to_cache, node_list); + _pFlashIndex->generate_cache_list_from_sample_queries( + warmup_query_file, 15, 6, num_nodes_to_cache, num_threads, node_list); _pFlashIndex->load_cache_list(node_list); node_list.clear(); node_list.shrink_to_fit(); From 8c314e16c37c275db2fe7c01100f17bd28e38295 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Tue, 15 Jun 2021 07:58:47 +0000 Subject: [PATCH 56/84] more changes to do MIPS by reducing to L2 with extra coordinate --- src/aux_utils.cpp | 8 ++++---- tests/search_disk_index.cpp | 6 +++--- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/src/aux_utils.cpp b/src/aux_utils.cpp index d57fa0cff6..34cde93cc2 100644 --- a/src/aux_utils.cpp +++ b/src/aux_utils.cpp @@ -362,7 +362,7 @@ namespace diskann { paras.Set("C", 750); paras.Set("alpha", 1.2f); paras.Set("num_rnds", 2); - paras.Set("saturate_graph", 1); + paras.Set("saturate_graph", 0); paras.Set("save_path", mem_index_path); std::unique_ptr> _pvamanaIndex = @@ -379,7 +379,7 @@ namespace diskann { partition_with_ram_budget(base_file, sampling_rate, ram_budget, 2 * R / 3, merged_index_prefix, 2); - wait_for_keystroke(); +// wait_for_keystroke(); std::string cur_centroid_filepath = merged_index_prefix + "_centroids.bin"; @@ -401,9 +401,9 @@ namespace diskann { paras.Set("L", L); paras.Set("R", (2 * (R / 3))); paras.Set("C", 750); - paras.Set("alpha", 2.0f); + paras.Set("alpha", 1.2f); paras.Set("num_rnds", 2); - paras.Set("saturate_graph", 1); + paras.Set("saturate_graph", 0); paras.Set("save_path", shard_index_file); std::unique_ptr> _pvamanaIndex = diff --git a/tests/search_disk_index.cpp b/tests/search_disk_index.cpp index e89c562a63..fe2fe5b9e5 100644 --- a/tests/search_disk_index.cpp +++ b/tests/search_disk_index.cpp @@ -137,9 +137,9 @@ int search_disk_index(int argc, char** argv) { std::vector node_list; diskann::cout << "Caching " << num_nodes_to_cache << " BFS nodes around medoid(s)" << std::endl; -// _pFlashIndex->cache_bfs_levels(num_nodes_to_cache, node_list); - _pFlashIndex->generate_cache_list_from_sample_queries( - warmup_query_file, 15, 6, num_nodes_to_cache, num_threads, node_list); + _pFlashIndex->cache_bfs_levels(num_nodes_to_cache, node_list); +// _pFlashIndex->generate_cache_list_from_sample_queries( +// warmup_query_file, 15, 6, num_nodes_to_cache, num_threads, node_list); _pFlashIndex->load_cache_list(node_list); node_list.clear(); node_list.shrink_to_fit(); From fc7efffc90f70c31318cac945ac6d4027868550b Mon Sep 17 00:00:00 2001 From: ravishankar Date: Tue, 15 Jun 2021 16:18:55 +0000 Subject: [PATCH 57/84] cleaned up code a bit, need to test everything again --- src/aux_utils.cpp | 37 +++++++++++++++---------------------- src/partition_and_pq.cpp | 2 +- src/pq_flash_index.cpp | 28 ++++++---------------------- 3 files changed, 22 insertions(+), 45 deletions(-) diff --git a/src/aux_utils.cpp b/src/aux_utils.cpp index 34cde93cc2..ef362351e8 100644 --- a/src/aux_utils.cpp +++ b/src/aux_utils.cpp @@ -379,9 +379,7 @@ namespace diskann { partition_with_ram_budget(base_file, sampling_rate, ram_budget, 2 * R / 3, merged_index_prefix, 2); -// wait_for_keystroke(); - - + wait_for_keystroke(); std::string cur_centroid_filepath = merged_index_prefix + "_centroids.bin"; std::rename(cur_centroid_filepath.c_str(), centroids_file.c_str()); @@ -643,9 +641,18 @@ namespace diskann { return false; } + + 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. Please contact us if you need other scenarios." << std::endl; + throw diskann::ANNException(stream.str(), -1); + + } + _u32 disk_pq_dims = 0; bool use_disk_pq = false; +// 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() == 6) { disk_pq_dims = atoi(param_list[5].c_str()); use_disk_pq = true; @@ -668,7 +675,7 @@ namespace diskann { std::string disk_pq_compressed_vectors_path = // optional if disk index is also storing pq data 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) { std::cout<<"Using Inner Product search, so need to pre-process base data into temp file. Please ensure there is sufficient space!!" << std::endl; std::string prepped_base = index_prefix_path + "_prepped_base.bin"; @@ -750,17 +757,16 @@ namespace diskann { diskann::cout << "Training data loaded of size " << train_size << std::endl; -// wait_for_keystroke(); + wait_for_keystroke(); +// don't translate data to make zero mean for PQ compression. We must not translate for inner product search. bool make_zero_mean = true; if (compareMetric == diskann::Metric::INNER_PRODUCT) make_zero_mean = false; + generate_pq_pivots(train_data, train_size, (uint32_t) dim, 256, (uint32_t) num_pq_chunks, NUM_KMEANS_REPS, pq_pivots_path, make_zero_mean); - if (compareMetric == diskann::Metric::INNER_PRODUCT) - generate_pq_data_from_pivots(data_file_to_use.c_str(), 256, (uint32_t) - num_pq_chunks, pq_pivots_path, pq_compressed_vectors_path); - else + generate_pq_data_from_pivots(data_file_to_use.c_str(), 256, (uint32_t) num_pq_chunks, pq_pivots_path, pq_compressed_vectors_path); @@ -768,20 +774,11 @@ namespace diskann { train_data = nullptr; - if (compareMetric == diskann::Metric::INNER_PRODUCT) - 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); - else 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); if (!use_disk_pq) { - if (compareMetric == diskann::Metric::INNER_PRODUCT) - diskann::create_disk_layout(data_file_to_use.c_str(), mem_index_path, - disk_index_path); - else diskann::create_disk_layout(data_file_to_use.c_str(), mem_index_path, disk_index_path); } @@ -790,10 +787,6 @@ namespace diskann { disk_index_path); double sample_sampling_rate = (150000.0 / points_num); - if (compareMetric == diskann::Metric::INNER_PRODUCT) - gen_random_slice(data_file_to_use.c_str(), sample_base_prefix, - sample_sampling_rate); - else gen_random_slice(data_file_to_use.c_str(), sample_base_prefix, sample_sampling_rate); diff --git a/src/partition_and_pq.cpp b/src/partition_and_pq.cpp index 6c8b0a431c..6b047f7316 100644 --- a/src/partition_and_pq.cpp +++ b/src/partition_and_pq.cpp @@ -745,7 +745,7 @@ int shard_data_into_clusters(const std::string data_file, float *pivots, - +// 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, diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index d140b0f77d..864f009b82 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -615,9 +615,6 @@ namespace diskann { this->data_dim = pq_file_dim; this->disk_data_dim = this->data_dim; // will reset later if we use PQ on disk this->disk_bytes_per_point = this->data_dim * sizeof(T); // will change later if we use PQ on disk or if we are using inner product without PQ - if (metric == diskann::Metric::INNER_PRODUCT) { - this->disk_bytes_per_point = this->data_dim * sizeof(float); // because we normalize the data and store it as float if no PQ - } this->aligned_dim = ROUND_UP(pq_file_dim, 8); @@ -818,17 +815,18 @@ if (file_exists(disk_pq_pivots_path)) { data = this->thread_data.pop(); } +// copy query to thread specific aligned and allocated memory (for distance calculations we need aligned data) + float query_norm = 0; - if (metric == diskann::Metric::L2) { + for (uint32_t i = 0; i < this->data_dim; i++) { data.scratch.aligned_query_float[i] = query1[i]; data.scratch.aligned_query_T[i] = query1[i]; - } - } else if (metric == diskann::Metric::INNER_PRODUCT) { - for (uint32_t i = 0; i < this->data_dim - 1; i++) { - data.scratch.aligned_query_float[i] = query1[i]; query_norm += query1[i]*query1[i]; } + +// 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) { query_norm = std::sqrt(query_norm); data.scratch.aligned_query_float[this->data_dim -1] = 0; for (uint32_t i = 0; i < this->data_dim - 1; i++) { @@ -836,7 +834,6 @@ if (file_exists(disk_pq_pivots_path)) { } } -// memcpy(data.scratch.aligned_query_T, query1, disk_bytes_per_point); const T * query = data.scratch.aligned_query_T; const float *query_float = data.scratch.aligned_query_float; @@ -861,9 +858,6 @@ if (file_exists(disk_pq_pivots_path)) { // query <-> PQ chunk centers distances float *pq_dists = query_scratch->aligned_pqtable_dist_scratch; -// if (metric==diskann::Metric::INNER_PRODUCT) -// pq_table.populate_chunk_inner_products(query, pq_dists); -// else if (metric==diskann::Metric::L2) pq_table.populate_chunk_distances(query_float, pq_dists); // query <-> neighbor list @@ -999,10 +993,6 @@ if (file_exists(disk_pq_pivots_path)) { T * node_fp_coords_copy = global_cache_iter->second; float cur_expanded_dist; if (!use_disk_index_pq) { - if (metric == diskann::Metric::INNER_PRODUCT) - cur_expanded_dist = dist_cmp_float->compare(query_float, (float*) node_fp_coords_copy, - (unsigned) aligned_dim); - else cur_expanded_dist = dist_cmp->compare(query, node_fp_coords_copy, (unsigned) aligned_dim); } @@ -1083,14 +1073,8 @@ if (file_exists(disk_pq_pivots_path)) { float cur_expanded_dist; if (!use_disk_index_pq) { - if (metric == diskann::Metric::INNER_PRODUCT) - cur_expanded_dist = dist_cmp_float->compare(query_float, (float*) node_fp_coords_copy, - (unsigned) aligned_dim); - else cur_expanded_dist = dist_cmp->compare(query, node_fp_coords_copy, (unsigned) aligned_dim); - - } else { if (metric == diskann::Metric::INNER_PRODUCT) From d1f9fcc62f54e886048508ec8073fd358752b458 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Wed, 16 Jun 2021 09:16:28 +0000 Subject: [PATCH 58/84] testing underway --- include/aux_utils.h | 2 ++ include/utils.h | 1 + src/aux_utils.cpp | 6 +++--- src/index.cpp | 3 ++- 4 files changed, 8 insertions(+), 4 deletions(-) diff --git a/include/aux_utils.h b/include/aux_utils.h index ac7bad6cc8..ca0dbce57c 100644 --- a/include/aux_utils.h +++ b/include/aux_utils.h @@ -29,6 +29,8 @@ typedef int FileHandle; #include "common_includes.h" #include "utils.h" #include "windows_customizations.h" +#include "gperftools/malloc_extension.h" + namespace diskann { const size_t TRAINING_SET_SIZE = 150000; diff --git a/include/utils.h b/include/utils.h index ae90cacdd8..ae487b05ce 100644 --- a/include/utils.h +++ b/include/utils.h @@ -222,6 +222,7 @@ namespace diskann { inline void wait_for_keystroke() { int a; + std::cout<<"Press any number to continue.." << std::endl; std::cin>> a; } diff --git a/src/aux_utils.cpp b/src/aux_utils.cpp index ef362351e8..c059343e5d 100644 --- a/src/aux_utils.cpp +++ b/src/aux_utils.cpp @@ -354,7 +354,7 @@ namespace diskann { double full_index_ram = ESTIMATE_RAM_USAGE(base_num, base_dim, sizeof(T), R); if (full_index_ram < ram_budget * 1024 * 1024 * 1024) { - diskann::cout << "Full index fits in RAM, building in one shot" + diskann::cout << "Full index fits in RAM budget, should consume at most " << full_index_ram/(1024*1024*1024) <<"GBs, so building in one shot" << std::endl; diskann::Parameters paras; paras.Set("L", (unsigned) L); @@ -757,8 +757,6 @@ namespace diskann { diskann::cout << "Training data loaded of size " << train_size << std::endl; - wait_for_keystroke(); - // don't translate data to make zero mean for PQ compression. We must not translate for inner product search. bool make_zero_mean = true; if (compareMetric == diskann::Metric::INNER_PRODUCT) @@ -773,6 +771,8 @@ namespace diskann { delete[] train_data; train_data = nullptr; + MallocExtension::instance()->ReleaseFreeMemory(); + diskann::build_merged_vamana_index( data_file_to_use.c_str(), diskann::Metric::L2, L, R, p_val, indexing_ram_budget, diff --git a/src/index.cpp b/src/index.cpp index b8d5c26642..9bdc2785c0 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -137,6 +137,7 @@ namespace diskann { _compacted_order(true), _enable_tags(enable_tags), _consolidated_order(true), _support_eager_delete(support_eager_delete), _store_data(store_data) { + // data is stored to _nd * aligned_dim matrix with necessary // zero-padding diskann::cout << "Number of frozen points = " << _num_frozen_pts @@ -182,7 +183,6 @@ namespace diskann { this->_distance = ::get_distance_function(m); _locks = std::vector(_max_points + _num_frozen_pts); - _width = 0; } @@ -789,6 +789,7 @@ namespace diskann { _final_graph[p].reserve((size_t)(std::ceil(range * SLACK_FACTOR * 1.05))); } + std::random_device rd; std::mt19937 gen(rd()); std::uniform_real_distribution<> dis(0, 1); From d7edf6ccc4716f7508edd96aaa0312e305c98dcf Mon Sep 17 00:00:00 2001 From: ravishankar Date: Wed, 16 Jun 2021 10:05:33 +0000 Subject: [PATCH 59/84] added back saturate graph to create denser indices --- src/aux_utils.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/aux_utils.cpp b/src/aux_utils.cpp index c059343e5d..25e139ffd2 100644 --- a/src/aux_utils.cpp +++ b/src/aux_utils.cpp @@ -362,7 +362,7 @@ namespace diskann { paras.Set("C", 750); paras.Set("alpha", 1.2f); paras.Set("num_rnds", 2); - paras.Set("saturate_graph", 0); + paras.Set("saturate_graph", 1); paras.Set("save_path", mem_index_path); std::unique_ptr> _pvamanaIndex = @@ -401,7 +401,7 @@ namespace diskann { paras.Set("C", 750); paras.Set("alpha", 1.2f); paras.Set("num_rnds", 2); - paras.Set("saturate_graph", 0); + paras.Set("saturate_graph", 1); paras.Set("save_path", shard_index_file); std::unique_ptr> _pvamanaIndex = From 510481316fcc773c4585f3a76501bbdaa5b4a05c Mon Sep 17 00:00:00 2001 From: ravishankar Date: Wed, 16 Jun 2021 11:36:10 +0000 Subject: [PATCH 60/84] now we dont sample a new test dataset every iteration for estimating sharding --- include/partition_and_pq.h | 3 +-- src/aux_utils.cpp | 3 +-- src/partition_and_pq.cpp | 48 +++++++++++++++++--------------------- 3 files changed, 23 insertions(+), 31 deletions(-) diff --git a/include/partition_and_pq.h b/include/partition_and_pq.h index bbae1ed148..0afc85410b 100644 --- a/include/partition_and_pq.h +++ b/include/partition_and_pq.h @@ -27,8 +27,7 @@ template void gen_random_slice(const T *inputdata, size_t npts, size_t ndims, double p_val, float *&sampled_data, size_t &slice_size); -template -int estimate_cluster_sizes(const std::string data_file, float *pivots, +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); diff --git a/src/aux_utils.cpp b/src/aux_utils.cpp index 25e139ffd2..5685bebee6 100644 --- a/src/aux_utils.cpp +++ b/src/aux_utils.cpp @@ -379,7 +379,6 @@ namespace diskann { partition_with_ram_budget(base_file, sampling_rate, ram_budget, 2 * R / 3, merged_index_prefix, 2); - wait_for_keystroke(); std::string cur_centroid_filepath = merged_index_prefix + "_centroids.bin"; std::rename(cur_centroid_filepath.c_str(), centroids_file.c_str()); @@ -410,7 +409,7 @@ namespace diskann { _pvamanaIndex->build(paras); _pvamanaIndex->save(shard_index_file.c_str()); std::remove(shard_base_file.c_str()); - wait_for_keystroke(); +// wait_for_keystroke(); } diff --git a/src/partition_and_pq.cpp b/src/partition_and_pq.cpp index 6b047f7316..9210c1371a 100644 --- a/src/partition_and_pq.cpp +++ b/src/partition_and_pq.cpp @@ -578,35 +578,20 @@ int generate_pq_data_from_pivots(const std::string data_file, return 0; } -template -int estimate_cluster_sizes(const std::string data_file, float *pivots, - const size_t num_centers, const size_t dim, +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 num_test, test_dim; - float *test_data_float; - double sampling_rate = 0.01; - - gen_random_slice(data_file, sampling_rate, test_data_float, num_test, - test_dim); - - if (test_dim != dim) { - diskann::cout << "Error. dimensions dont match for pivot set and base set" - << std::endl; - return -1; - } size_t *shard_counts = new size_t[num_centers]; for (size_t i = 0; i < num_centers; i++) { shard_counts[i] = 0; } - - size_t num_points = 0, num_dim = 0; - diskann::get_bin_metadata(data_file, num_points, num_dim); - size_t block_size = num_points <= BLOCK_SIZE ? num_points : BLOCK_SIZE; + + size_t block_size = num_test <= BLOCK_SIZE ? num_test : BLOCK_SIZE; _u32 * block_closest_centers = new _u32[block_size * k_base]; float *block_data_float; @@ -619,7 +604,7 @@ int estimate_cluster_sizes(const std::string data_file, float *pivots, block_data_float = test_data_float + start_id * test_dim; - math_utils::compute_closest_centers(block_data_float, cur_blk_size, dim, + math_utils::compute_closest_centers(block_data_float, cur_blk_size, test_dim, pivots, num_centers, k_base, block_closest_centers); @@ -635,8 +620,8 @@ int estimate_cluster_sizes(const std::string data_file, float *pivots, for (size_t i = 0; i < num_centers; i++) { _u32 cur_shard_count = (_u32) shard_counts[i]; cluster_sizes.push_back( - size_t(((double) cur_shard_count) * (1.0 / sampling_rate))); - diskann::cout << cur_shard_count * (1.0 / sampling_rate) << " "; + (size_t)cur_shard_count); + diskann::cout << cur_shard_count << " "; } diskann::cout << std::endl; delete[] shard_counts; @@ -952,9 +937,9 @@ int partition(const std::string data_file, const float sampling_rate, // now pivots are ready. need to stream base points and assign them to // closest clusters. - std::vector cluster_sizes; - estimate_cluster_sizes(data_file, pivot_data, num_parts, train_dim, k_base, - cluster_sizes); + //std::vector cluster_sizes; + //estimate_cluster_sizes(data_file, pivot_data, num_parts, train_dim, k_base, + // cluster_sizes); shard_data_into_clusters(data_file, pivot_data, num_parts, train_dim, k_base, prefix_path); @@ -971,7 +956,7 @@ int partition_with_ram_budget(const std::string data_file, size_t train_dim; size_t num_train; float *train_data_float; - size_t max_k_means_reps = 20; + size_t max_k_means_reps = 10; int num_parts = 3; bool fit_in_ram = false; @@ -979,6 +964,13 @@ int partition_with_ram_budget(const std::string data_file, 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); @@ -1010,10 +1002,11 @@ int partition_with_ram_budget(const std::string data_file, // closest clusters. std::vector cluster_sizes; - estimate_cluster_sizes(data_file, pivot_data, num_parts, train_dim, + estimate_cluster_sizes(test_data_float, num_test, pivot_data, num_parts, train_dim, k_base, cluster_sizes); for (auto &p : cluster_sizes) { + p = (_u64) (p/ sampling_rate); // to account for the fact that p is the size of the shard over the testing sample. double cur_shard_ram_estimate = ESTIMATE_RAM_USAGE(p, train_dim, sizeof(T), graph_degree); @@ -1037,6 +1030,7 @@ int partition_with_ram_budget(const std::string data_file, k_base, prefix_path); delete[] pivot_data; delete[] train_data_float; + delete[] test_data_float; return num_parts; } From 971a90c4e21fd48ed0654a27a435d6538403e074 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Wed, 16 Jun 2021 11:37:06 +0000 Subject: [PATCH 61/84] now num_parts increases by 2 --- src/partition_and_pq.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/partition_and_pq.cpp b/src/partition_and_pq.cpp index 9210c1371a..dd9c00a56c 100644 --- a/src/partition_and_pq.cpp +++ b/src/partition_and_pq.cpp @@ -1018,7 +1018,7 @@ int partition_with_ram_budget(const std::string data_file, << "GB, budget given is " << ram_budget << std::endl; if (max_ram_usage > 1024 * 1024 * 1024 * ram_budget) { fit_in_ram = false; - num_parts++; + num_parts+=2; } } From 552e0f1d086ab6da7046995922e8f45ef5189b02 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Wed, 16 Jun 2021 11:50:56 +0000 Subject: [PATCH 62/84] cleaned up warnings in Debug mode compiler --- include/distance.h | 22 ++++++++++++++++------ include/index.h | 2 +- include/pq_flash_index.h | 3 +-- src/index.cpp | 2 +- src/pq_flash_index.cpp | 3 +-- tests/utils/compute_groundtruth.cpp | 14 ++++++-------- tests/utils/create_disk_layout.cpp | 8 ++++---- tests/utils/gen_random_slice.cpp | 8 ++++---- tests/utils/vector_analysis.cpp | 8 ++++---- 9 files changed, 38 insertions(+), 32 deletions(-) diff --git a/include/distance.h b/include/distance.h index a65f9e9487..1de05fefb0 100644 --- a/include/distance.h +++ b/include/distance.h @@ -255,9 +255,14 @@ namespace diskann { virtual float compare(const int8_t *a, const int8_t *b, unsigned int length) const { #ifndef _WINDOWS - std::cout << "AVX only supported in Windows build."; - return 0; - } +int32_t result = 0; +#pragma omp simd reduction(+ : result) aligned(a, b : 8) + for (_s32 i = 0; i < (_s32) 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; + } #else __m128 r = _mm_setzero_ps(); __m128i r1; @@ -302,9 +307,14 @@ namespace diskann { virtual float compare(const float *a, const float *b, unsigned int length) const { #ifndef _WINDOWS - std::cout << "AVX only supported in Windows build."; - return 0; - } +float result = 0; +#pragma omp simd reduction(+ : result) aligned(a, b : 8) + for (_s32 i = 0; i < (_s32) length; i++) { + result += (a[i] - b[i]) * + (a[i] - b[i]); + } + return result; + } #else __m128 diff, v1, v2; __m128 sum = _mm_set1_ps(0); diff --git a/include/index.h b/include/index.h index c409b96958..b800781d38 100644 --- a/include/index.h +++ b/include/index.h @@ -63,7 +63,7 @@ namespace diskann { DISKANN_DLLEXPORT std::pair search_with_tags( const T *query, const size_t K, const unsigned L, TagT *tags, - unsigned frozen_pts, unsigned *indices_buffer = NULL); + unsigned *indices_buffer = NULL); // repositions frozen points to the end of _data - if they have been moved // during deletion diff --git a/include/pq_flash_index.h b/include/pq_flash_index.h index fbc6fc62b3..f17448e7d4 100644 --- a/include/pq_flash_index.h +++ b/include/pq_flash_index.h @@ -112,8 +112,7 @@ namespace diskann { // implemented DISKANN_DLLEXPORT void cached_beam_search( const T *query, const _u64 k_search, const _u64 l_search, _u64 *res_ids, - float *res_dists, const _u64 beam_width, QueryStats *stats = nullptr, - Distance *output_dist_func = nullptr); + float *res_dists, const _u64 beam_width, QueryStats *stats = nullptr); std::shared_ptr &reader; protected: diff --git a/src/index.cpp b/src/index.cpp index 9bdc2785c0..0d01c92322 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -1080,7 +1080,7 @@ namespace diskann { template std::pair Index::search_with_tags( const T *query, const size_t K, const unsigned L, TagT *tags, - unsigned frozen_pts, unsigned *indices_buffer) { + unsigned *indices_buffer) { const bool alloc = indices_buffer == NULL; auto indices = alloc ? new unsigned[K] : indices_buffer; auto ret = search(query, K, L, indices); diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index 864f009b82..6d0176b896 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -807,8 +807,7 @@ if (file_exists(disk_pq_pivots_path)) { const _u64 l_search, _u64 *indices, float * distances, const _u64 beam_width, - QueryStats * stats, - Distance *output_dist_func) { + QueryStats * stats) { ThreadData data = this->thread_data.pop(); while (data.scratch.sector_scratch == nullptr) { this->thread_data.wait_for_push_notify(); diff --git a/tests/utils/compute_groundtruth.cpp b/tests/utils/compute_groundtruth.cpp index cdd8d0c327..8f9f77399a 100644 --- a/tests/utils/compute_groundtruth.cpp +++ b/tests/utils/compute_groundtruth.cpp @@ -108,9 +108,7 @@ 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, - 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; @@ -171,8 +169,8 @@ void exact_knn(const size_t dim, const size_t k, queries_l2sq + q_b); } else { inner_prod_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); + 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; @@ -308,7 +306,7 @@ inline void save_groundtruth_as_one_file(const std::string filename, } template -int aux_main(int argv, char **argc) { +int aux_main(char **argc) { size_t npoints, nqueries, dim; std::string base_file(argc[2]); std::string query_file(argc[3]); @@ -393,9 +391,9 @@ int main(int argc, char **argv) { } if (std::string(argv[1]) == std::string("float")) - aux_main(argc, argv); + aux_main(argv); if (std::string(argv[1]) == std::string("int8")) - aux_main(argc, argv); + aux_main(argv); if (std::string(argv[1]) == std::string("uint8")) - aux_main(argc, argv); + aux_main(argv); } diff --git a/tests/utils/create_disk_layout.cpp b/tests/utils/create_disk_layout.cpp index 458e2d6b20..21c6cacede 100644 --- a/tests/utils/create_disk_layout.cpp +++ b/tests/utils/create_disk_layout.cpp @@ -13,7 +13,7 @@ #include "utils.h" template -int create_disk_layout(int argc, char **argv) { +int create_disk_layout(char **argv) { std::string base_file(argv[2]); std::string vamana_file(argv[3]); std::string output_file(argv[4]); @@ -31,11 +31,11 @@ int main(int argc, char **argv) { int ret_val = -1; if (std::string(argv[1]) == std::string("float")) - ret_val = create_disk_layout(argc, argv); + ret_val = create_disk_layout(argv); else if (std::string(argv[1]) == std::string("int8")) - ret_val = create_disk_layout(argc, argv); + ret_val = create_disk_layout(argv); else if (std::string(argv[1]) == std::string("uint8")) - ret_val = create_disk_layout(argc, argv); + ret_val = create_disk_layout(argv); else { std::cout << "unsupported type. use int8/uint8/float " << std::endl; ret_val = -2; diff --git a/tests/utils/gen_random_slice.cpp b/tests/utils/gen_random_slice.cpp index 1b102e27c1..dccb50a13c 100644 --- a/tests/utils/gen_random_slice.cpp +++ b/tests/utils/gen_random_slice.cpp @@ -21,7 +21,7 @@ #include template -int aux_main(int argc, char** argv) { +int aux_main(char** argv) { std::string base_file(argv[2]); std::string output_prefix(argv[3]); @@ -40,11 +40,11 @@ int main(int argc, char** argv) { } if (std::string(argv[1]) == std::string("float")) { - aux_main(argc, argv); + aux_main(argv); } else if (std::string(argv[1]) == std::string("int8")) { - aux_main(argc, argv); + aux_main(argv); } else if (std::string(argv[1]) == std::string("uint8")) { - aux_main(argc, argv); + aux_main(argv); } else std::cout << "Unsupported type. Use float/int8/uint8." << std::endl; return 0; diff --git a/tests/utils/vector_analysis.cpp b/tests/utils/vector_analysis.cpp index 41ee323c7f..43d333b532 100644 --- a/tests/utils/vector_analysis.cpp +++ b/tests/utils/vector_analysis.cpp @@ -84,7 +84,7 @@ std::cout<<"Max norm: " << max_norm << std::endl; template -int aux_main(int argc, char** argv) { +int aux_main(char** argv) { std::string base_file(argv[2]); _u32 option = atoi(argv[3]); @@ -107,11 +107,11 @@ int main(int argc, char** argv) { } if (std::string(argv[1]) == std::string("float")) { - aux_main(argc, argv); + aux_main(argv); } else if (std::string(argv[1]) == std::string("int8")) { - aux_main(argc, argv); + aux_main(argv); } else if (std::string(argv[1]) == std::string("uint8")) { - aux_main(argc, argv); + aux_main(argv); } else std::cout << "Unsupported type. Use float/int8/uint8." << std::endl; return 0; From 39fc6d953fc5921203ce166ea5d90f5ad4885198 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Tue, 6 Jul 2021 17:21:25 +0000 Subject: [PATCH 63/84] added a normalizer to vector analysis --- tests/utils/vector_analysis.cpp | 25 ++++++++++++++++++++++++- 1 file changed, 24 insertions(+), 1 deletion(-) diff --git a/tests/utils/vector_analysis.cpp b/tests/utils/vector_analysis.cpp index 43d333b532..65f42771c5 100644 --- a/tests/utils/vector_analysis.cpp +++ b/tests/utils/vector_analysis.cpp @@ -40,6 +40,27 @@ int analyze_norm(std::string base_file) { return 0; } +template +int normalize_base(std::string base_file, std::string out_file) { + std::cout<<"Normalizing base" << std::endl; + T* data; + _u64 npts, ndims; + diskann::load_bin(base_file, data, npts, ndims); +// std::vector norms(npts, 0); + #pragma omp parallel for schedule(dynamic) + for (_u32 i = 0; i(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) { @@ -94,6 +115,8 @@ int aux_main(char** argv) { 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; } @@ -101,7 +124,7 @@ 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] [out_file for options 2/3]" + "[option: 1-norm analysis, 2-prep_base_for_mip, 3-prep_query_for_mip, 4-normalize-vecs] [out_file for options 2/3]" << std::endl; exit(-1); } From 313cbdaa43ad33530d6e52818a20e56d36f45eb0 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Tue, 13 Jul 2021 11:16:43 +0000 Subject: [PATCH 64/84] fixed one bug for MIPS --- src/pq_flash_index.cpp | 7 +++++-- tests/utils/CMakeLists.txt | 10 ++++++++++ tests/utils/uint8_to_float.cpp | 21 +++++++++++++++++++++ tests/utils/vector_analysis.cpp | 10 +++++++--- 4 files changed, 43 insertions(+), 5 deletions(-) create mode 100644 tests/utils/uint8_to_float.cpp diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index 6d0176b896..bd6b62ea7b 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -814,9 +814,12 @@ if (file_exists(disk_pq_pivots_path)) { data = this->thread_data.pop(); } + // copy query to thread specific aligned and allocated memory (for distance calculations we need aligned data) float query_norm = 0; + const T * query = data.scratch.aligned_query_T; + const float *query_float = data.scratch.aligned_query_float; for (uint32_t i = 0; i < this->data_dim; i++) { data.scratch.aligned_query_float[i] = query1[i]; @@ -827,14 +830,14 @@ if (file_exists(disk_pq_pivots_path)) { // 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) { query_norm = std::sqrt(query_norm); + data.scratch.aligned_query_T[this->data_dim -1] = 0; data.scratch.aligned_query_float[this->data_dim -1] = 0; for (uint32_t i = 0; i < this->data_dim - 1; i++) { + data.scratch.aligned_query_T[i] /= query_norm; data.scratch.aligned_query_float[i] /= query_norm; } } - const T * query = data.scratch.aligned_query_T; - const float *query_float = data.scratch.aligned_query_float; IOContext &ctx = data.ctx; auto query_scratch = &(data.scratch); diff --git a/tests/utils/CMakeLists.txt b/tests/utils/CMakeLists.txt index 7769229235..a8171e739c 100644 --- a/tests/utils/CMakeLists.txt +++ b/tests/utils/CMakeLists.txt @@ -47,6 +47,16 @@ else() target_link_libraries(int8_to_float ${PROJECT_NAME}) endif() +add_executable(uint8_to_float uint8_to_float.cpp) +if(MSVC) + target_link_options(uint8_to_float PRIVATE /MACHINE:x64) + target_link_libraries(uint8_to_float debug ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}/diskann_dll.lib) + target_link_libraries(uint8_to_float optimized ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}/diskann_dll.lib) +else() + target_link_libraries(uint8_to_float ${PROJECT_NAME}) +endif() + + add_executable(uint32_to_uint8 uint32_to_uint8.cpp) if(MSVC) target_link_options(uint32_to_uint8 PRIVATE /MACHINE:x64) diff --git a/tests/utils/uint8_to_float.cpp b/tests/utils/uint8_to_float.cpp new file mode 100644 index 0000000000..e383489fd4 --- /dev/null +++ b/tests/utils/uint8_to_float.cpp @@ -0,0 +1,21 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#include +#include "utils.h" + +int main(int argc, char** argv) { + if (argc != 3) { + 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; +} diff --git a/tests/utils/vector_analysis.cpp b/tests/utils/vector_analysis.cpp index 65f42771c5..e370db2774 100644 --- a/tests/utils/vector_analysis.cpp +++ b/tests/utils/vector_analysis.cpp @@ -83,19 +83,23 @@ std::cout<<"Max norm: " << max_norm << std::endl; _u64 newdims = ndims + 1; new_data = new T[npts*newdims]; for (_u64 i = 0;i < npts; i++) { + if (prep_base) { for (_u64 j = 0; j < ndims; j++) { new_data[i*newdims + j] = data[i*ndims +j]/ max_norm; } - if (prep_base) { 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< void 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; diff --git a/src/aux_utils.cpp b/src/aux_utils.cpp index 5685bebee6..1a8e2387fd 100644 --- a/src/aux_utils.cpp +++ b/src/aux_utils.cpp @@ -354,7 +354,7 @@ namespace diskann { double full_index_ram = ESTIMATE_RAM_USAGE(base_num, base_dim, sizeof(T), R); 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) <<"GBs, so building in one shot" + 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::Parameters paras; paras.Set("L", (unsigned) L); @@ -643,7 +643,7 @@ namespace diskann { 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. Please contact us if you need other scenarios." << std::endl; + stream << "DiskANN currently only supports floating point data for Max Inner Product Search. " << std::endl; throw diskann::ANNException(stream.str(), -1); } @@ -676,7 +676,7 @@ namespace diskann { // 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) { - std::cout<<"Using Inner Product search, so need to pre-process base data into temp file. Please ensure there is sufficient space!!" << std::endl; + 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; diskann::prepare_base_for_inner_products(base_file, prepped_base); diff --git a/src/index.cpp b/src/index.cpp index 0d01c92322..061760da47 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -504,7 +504,7 @@ namespace diskann { std::vector init_ids, std::vector & expanded_nodes_info, tsl::robin_set &expanded_nodes_ids) { - T * node_coords = _data + _aligned_dim * node_id; + const T * node_coords = _data + _aligned_dim * node_id; std::vector best_L_nodes; if (init_ids.size() == 0) @@ -538,7 +538,7 @@ namespace diskann { float cur_alpha = 1; while (cur_alpha <= alpha && result.size() < degree) { unsigned start = 0; - float eps = cur_alpha + 0.01; + float eps = cur_alpha + 0.01; // 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. while (result.size() < degree && (start) < pool.size() && start < maxc) { auto &p = pool[start]; if (occlude_factor[start] > cur_alpha) { diff --git a/src/partition_and_pq.cpp b/src/partition_and_pq.cpp index dd9c00a56c..c2accdfbd0 100644 --- a/src/partition_and_pq.cpp +++ b/src/partition_and_pq.cpp @@ -231,7 +231,7 @@ int generate_pq_pivots(const float *passed_train_data, size_t num_train, for (uint64_t d = 0; d < dim; d++) { centroid[d] = 0; } - if (make_zero_mean) { + 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]; @@ -239,7 +239,7 @@ int generate_pq_pivots(const float *passed_train_data, size_t num_train, centroid[d] /= num_train; } - // std::memset(centroid, 0 , dim*sizeof(float)); + for (uint64_t d = 0; d < dim; d++) { for (uint64_t p = 0; p < num_train; p++) { train_data[p * dim + d] -= centroid[d]; @@ -937,10 +937,7 @@ int partition(const std::string data_file, const float sampling_rate, // now pivots are ready. need to stream base points and assign them to // closest clusters. - //std::vector cluster_sizes; - //estimate_cluster_sizes(data_file, pivot_data, num_parts, train_dim, k_base, - // cluster_sizes); - + shard_data_into_clusters(data_file, pivot_data, num_parts, train_dim, k_base, prefix_path); delete[] pivot_data; diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index bd6b62ea7b..a8eae0986f 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -1002,7 +1002,7 @@ if (file_exists(disk_pq_pivots_path)) { if (metric == diskann::Metric::INNER_PRODUCT) cur_expanded_dist = disk_pq_table.inner_product(query_float, (_u8*) node_fp_coords_copy); else - cur_expanded_dist = disk_pq_table.compare(query_float, (_u8*) node_fp_coords_copy); + cur_expanded_dist = disk_pq_table.l2_distance(query_float, (_u8*) node_fp_coords_copy); } full_retset.push_back( Neighbor((unsigned) cached_nhood.first, cur_expanded_dist, true)); @@ -1082,7 +1082,7 @@ if (file_exists(disk_pq_pivots_path)) { if (metric == diskann::Metric::INNER_PRODUCT) cur_expanded_dist = disk_pq_table.inner_product(query_float, (_u8*) node_fp_coords_copy); else - cur_expanded_dist = disk_pq_table.compare(query_float, (_u8*) node_fp_coords_copy); + cur_expanded_dist = disk_pq_table.l2_distance(query_float, (_u8*) node_fp_coords_copy); } full_retset.push_back( Neighbor(frontier_nhood.first, cur_expanded_dist, true)); diff --git a/tests/build_disk_index.cpp b/tests/build_disk_index.cpp index 77bdeace37..ed0b71b8e8 100644 --- a/tests/build_disk_index.cpp +++ b/tests/build_disk_index.cpp @@ -19,15 +19,17 @@ bool build_index(const char* dataFilePath, const char* indexFilePath, int main(int argc, char** argv) { if (argc != 11) { std::cout << "Usage: " << argv[0] - << " [data_type] [dist_fn: 0 for L2, 1 for MIPS] [data_file.bin] " + << " [data_type] [dist_fn: l2/mips] [data_file.bin] " "[index_prefix_path] " "[R] [L] [B] [M] [T] [PQ_disk_bytes (for very large dimensionality, use 0 for full vectors)]. See README for more information on " "parameters." << std::endl; } else { diskann::Metric metric = diskann::Metric::L2; - if (atoi(argv[2]) == 1) + + if (std::string(argv[2]) == std::string("mips")) metric = diskann::Metric::INNER_PRODUCT; + std::string params = std::string(argv[5]) + " " + std::string(argv[6]) + " " + std::string(argv[7]) + " " + std::string(argv[8]) + " " + std::string(argv[9]) + " " + std::string(argv[10]); diff --git a/tests/build_memory_index.cpp b/tests/build_memory_index.cpp index 7bbe992c82..8faf39dcc8 100644 --- a/tests/build_memory_index.cpp +++ b/tests/build_memory_index.cpp @@ -16,7 +16,7 @@ #include "memory_mapper.h" template -int build_in_memory_index(const std::string& data_path, _u32 dist_fn, const unsigned R, +int build_in_memory_index(const std::string& data_path, const diskann::Metric &metric, const unsigned R, const unsigned L, const float alpha, const std::string& save_path, const unsigned num_threads) { @@ -29,16 +29,6 @@ int build_in_memory_index(const std::string& data_path, _u32 dist_fn, const unsi paras.Set("saturate_graph", 0); paras.Set("num_threads", num_threads); - diskann::Metric metric; - if (dist_fn == 0) - metric = diskann::L2; - else if (dist_fn == 1) - metric = diskann::INNER_PRODUCT; - else { - std::cout<<"Error. Unsupported distance type. Exitting" << std::endl; - return -1; - } - diskann::Index index(metric, data_path.c_str()); auto s = std::chrono::high_resolution_clock::now(); index.build(paras); @@ -54,7 +44,7 @@ int build_in_memory_index(const std::string& data_path, _u32 dist_fn, const unsi int main(int argc, char** argv) { if (argc != 9) { std::cout << "Usage: " << argv[0] - << " [data_type] [dist_fn 0 for L2, 1 for inner product] [data_file.bin] " + << " [data_type] [l2/mips] [data_file.bin] " "[output_index_file] " << "[R] [L] [alpha]" << " [num_threads_to_use]. See README for more information on " @@ -65,7 +55,17 @@ int main(int argc, char** argv) { _u32 ctr = 2; - _u32 dist_fn = (_u32) atoi(argv[ctr++]); +diskann::Metric metric; + if (std::string(argv[ctr]) == std::string("mips")) + metric = diskann::Metric::INNER_PRODUCT; + else if (std::string(argv[ctr]) == std::string("l2")) + metric = diskann::Metric::L2; + else { + std::cout<<"Unsupported distance function. Currently only L2/ Inner Product support." << std::endl; + return -1; + } +ctr++; + const std::string data_path(argv[ctr++]); const std::string save_path(argv[ctr++]); const unsigned R = (unsigned) atoi(argv[ctr++]); @@ -74,13 +74,13 @@ int main(int argc, char** argv) { const unsigned num_threads = (unsigned) atoi(argv[ctr++]); if (std::string(argv[1]) == std::string("int8")) - build_in_memory_index(data_path, dist_fn, R, L, alpha, save_path, + build_in_memory_index(data_path, metric, R, L, alpha, save_path, num_threads); else if (std::string(argv[1]) == std::string("uint8")) - build_in_memory_index(data_path, dist_fn, R, L, alpha, save_path, + build_in_memory_index(data_path, metric, R, L, alpha, save_path, num_threads); else if (std::string(argv[1]) == std::string("float")) - build_in_memory_index(data_path, dist_fn, R, L, alpha, save_path, + build_in_memory_index(data_path, metric, R, L, alpha, save_path, num_threads); else std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; diff --git a/tests/search_disk_index.cpp b/tests/search_disk_index.cpp index fe2fe5b9e5..1eb35d2bd7 100644 --- a/tests/search_disk_index.cpp +++ b/tests/search_disk_index.cpp @@ -57,7 +57,19 @@ int search_disk_index(int argc, char** argv) { std::vector<_u64> Lvec; _u32 ctr = 2; - _u32 dist_fn = atoi(argv[ctr++]); + diskann::Metric metric; + + if (std::string(argv[ctr]) == std::string("mips")) + metric = diskann::Metric::INNER_PRODUCT; + else if (std::string(argv[ctr]) == std::string("l2")) + metric = diskann::Metric::L2; + else { + std::cout<<"Unsupported distance function. Currently only L2/ Inner Product support." << std::endl; + return -1; + } + + ctr++; + std::string index_prefix_path(argv[ctr++]); std::string pq_prefix = index_prefix_path + "_pq"; std::string disk_index_file = index_prefix_path + "_disk.index"; @@ -78,15 +90,6 @@ int search_disk_index(int argc, char** argv) { Lvec.push_back(curL); } - diskann::Metric metric; - if (dist_fn == 0) - metric = diskann::Metric::L2; - else if (dist_fn == 1) - metric = diskann::Metric::INNER_PRODUCT; - else { - std::cout<<"Unsupported distance function. Currently only L2/ Inner Product support." << std::endl; - return -1; - } if (Lvec.size() == 0) { diskann::cout << "No valid Lsearch found. Lsearch must be at least recall_at" @@ -208,7 +211,6 @@ int search_disk_index(int argc, char** argv) { uint32_t optimized_beamwidth = 2; -//query_num = 1; for (uint32_t test_id = 0; test_id < Lvec.size(); test_id++) { _u64 L = Lvec[test_id]; diff --git a/tests/search_memory_index.cpp b/tests/search_memory_index.cpp index 6ad40032a4..334da66f98 100644 --- a/tests/search_memory_index.cpp +++ b/tests/search_memory_index.cpp @@ -28,7 +28,28 @@ int search_memory_index(int argc, char** argv) { std::vector<_u64> Lvec; _u32 ctr = 2; - _u32 dist_fn = atoi(argv[ctr++]); + diskann::Metric metric; + + if (std::string(argv[ctr]) == std::string("mips")) + metric = diskann::Metric::INNER_PRODUCT; + else if (std::string(argv[ctr]) == std::string("l2")) + metric = diskann::Metric::L2; + else if (std::string(argv[ctr]) == std::string("fast_l2")) + metric = diskann::Metric::FAST_L2; + else { + std::cout<<"Unsupported distance function. Currently only L2/ Inner Product/FAST_L2 support." << std::endl; + return -1; + } +ctr++; + + if ((std::string(argv[1]) != std::string("float")) && + ((metric == diskann::Metric::INNER_PRODUCT) || (metric == diskann::Metric::FAST_L2))) { + std::cout << "Error. Inner product and Fast_L2 search currently only supported for " + "floating point datatypes." + << std::endl; + } + + std::string data_file(argv[ctr++]); std::string memory_index_file(argv[ctr++]); _u64 num_threads = std::atoi(argv[ctr++]); @@ -38,12 +59,6 @@ int search_memory_index(int argc, char** argv) { std::string result_output_prefix(argv[ctr++]); // bool use_optimized_search = std::atoi(argv[ctr++]); - if ((std::string(argv[1]) != std::string("float")) && - ((dist_fn == 1) || (dist_fn == 2))) { - std::cout << "Error. Inner product and Fast_L2 search currently only supported for " - "floating point datatypes." - << std::endl; - } bool calc_recall_flag = false; @@ -74,17 +89,7 @@ int search_memory_index(int argc, char** argv) { std::cout.setf(std::ios_base::fixed, std::ios_base::floatfield); std::cout.precision(2); - diskann::Metric metric; - if (dist_fn == 0) - metric = diskann::L2; - else if (dist_fn == 1) - metric = diskann::INNER_PRODUCT; - else if(dist_fn == 2) - metric = diskann::FAST_L2; - else { - std::cout<<"Error. Unsupported distance function. Exitting"; - return -1; - } + diskann::Index index(metric, data_file.c_str()); index.load(memory_index_file.c_str()); // to load NSG std::cout << "Index loaded" << std::endl; diff --git a/tests/utils/compute_groundtruth.cpp b/tests/utils/compute_groundtruth.cpp index 8f9f77399a..de46cab904 100644 --- a/tests/utils/compute_groundtruth.cpp +++ b/tests/utils/compute_groundtruth.cpp @@ -31,7 +31,7 @@ #define ALIGNMENT 512 void command_line_help() { - std::cerr << " " " " << std::endl; @@ -120,12 +120,7 @@ void inner_prod_to_points( cblas_sgemm(CblasColMajor, CblasTrans, CblasNoTrans, npoints, nqueries, dim, (float) -1.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; } @@ -345,21 +340,7 @@ int aux_main(char **argc) { delete[] closest_points_part; delete[] dist_closest_points_part; - /* - std::cout << "For testing: doing brute force for one point" << - std::endl; std::vector> brute_force; for (_u32 i - = 0; i < npoints; i++) { float cur_pt_dist = 0; for (_u64 k = 0; k < dim; - k++) { cur_pt_dist += base_data[i * dim + k] * query_data[k]; - } - brute_force.push_back(std::make_pair(i, -cur_pt_dist)); - } - - std::sort(brute_force.begin(), brute_force.end(), custom_dist); - for (_u32 i = 0; i < 10; i++) { - std::cout< Date: Thu, 15 Jul 2021 10:14:07 +0000 Subject: [PATCH 66/84] fixed minor typos. now running unit tests --- tests/search_disk_index.cpp | 7 ++++++- tests/search_memory_index.cpp | 2 +- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/tests/search_disk_index.cpp b/tests/search_disk_index.cpp index 1eb35d2bd7..241968a7cf 100644 --- a/tests/search_disk_index.cpp +++ b/tests/search_disk_index.cpp @@ -67,6 +67,11 @@ int search_disk_index(int argc, char** argv) { std::cout<<"Unsupported distance function. Currently only L2/ Inner Product support." << std::endl; return -1; } + + if ((std::string(argv[1]) != std::string("float")) && (metric == diskann::Metric::INNER_PRODUCT)) { + std::cout<<"Currently support only floating point data for Inner Product." << std::endl; + return -1; + } ctr++; @@ -304,7 +309,7 @@ int main(int argc, char** argv) { if (argc < 12) { diskann::cout << "Usage: " << argv[0] - << " [index_type] [dist_fn 0 for l2/ 1 for inner product] [index_prefix_path] " + << " [index_type] [dist_fn] [index_prefix_path] " " [num_nodes_to_cache] [num_threads] [beamwidth (use 0 to " "optimize internally)] " " [query_file.bin] [truthset.bin (use \"null\" for none)] " diff --git a/tests/search_memory_index.cpp b/tests/search_memory_index.cpp index 334da66f98..64fbcd6909 100644 --- a/tests/search_memory_index.cpp +++ b/tests/search_memory_index.cpp @@ -174,7 +174,7 @@ int main(int argc, char** argv) { if (argc < 11) { std::cout << "Usage: " << argv[0] - << " [index_type] [dist_fn (0 for L2, 1 for Inner Product, 2 for Fast L2 for small datasets)] [data_file.bin] " + << " [index_type] [dist_fn (l2/mips/fast_l2)] [data_file.bin] " "[memory_index_path] [num_threads] " "[query_file.bin] [truthset.bin (use \"null\" for none)] " " [K] [result_output_prefix]" From d018a20067280a41ad4c59e04bb4fea65c6eb65b Mon Sep 17 00:00:00 2001 From: ravishankar Date: Thu, 15 Jul 2021 10:40:44 +0000 Subject: [PATCH 67/84] ran clang-format as it doesnt run by default due to LINUX flag not set anywhere --- include/aligned_file_reader.h | 4 +- include/aux_utils.h | 1 - include/distance.h | 37 +-- include/exceptions.h | 2 +- include/index.h | 2 +- include/memory_mapper.h | 2 +- include/partition_and_pq.h | 28 +-- include/percentile_stats.h | 2 +- include/pq_flash_index.h | 25 +- include/pq_table.h | 70 +++--- include/timer.h | 2 +- include/utils.h | 154 +++++++------ include/windows_aligned_file_reader.h | 5 +- src/aux_utils.cpp | 263 +++++++++++----------- src/index.cpp | 61 ++--- src/linux_aligned_file_reader.cpp | 2 +- src/partition_and_pq.cpp | 166 +++++++------- src/pq_flash_index.cpp | 191 ++++++++-------- src/windows_aligned_file_reader.cpp | 2 +- tests/build_disk_index.cpp | 18 +- tests/build_memory_index.cpp | 17 +- tests/search_disk_index.cpp | 34 +-- tests/search_memory_index.cpp | 28 +-- tests/test_incremental_index.cpp | 2 +- tests/utils/bin_to_tsv.cpp | 8 +- tests/utils/compute_groundtruth.cpp | 16 +- tests/utils/create_disk_layout.cpp | 7 +- tests/utils/float_bin_to_int8.cpp | 1 - tests/utils/gen_random_slice.cpp | 9 +- tests/utils/partition_data.cpp | 7 +- tests/utils/partition_with_ram_budget.cpp | 7 +- tests/utils/tsv_to_bin.cpp | 7 +- tests/utils/uint8_to_float.cpp | 2 +- tests/utils/vector_analysis.cpp | 124 +++++----- 34 files changed, 683 insertions(+), 623 deletions(-) diff --git a/include/aligned_file_reader.h b/include/aligned_file_reader.h index c33bacd666..b6b56a012e 100644 --- a/include/aligned_file_reader.h +++ b/include/aligned_file_reader.h @@ -18,7 +18,7 @@ typedef io_context_t IOContext; #include #ifndef USE_BING_INFRA -struct IOContext{ +struct IOContext { HANDLE fhandle = NULL; HANDLE iocp = NULL; std::vector reqs; @@ -77,7 +77,7 @@ struct AlignedRead { class AlignedFileReader { protected: tsl::robin_map ctx_map; - std::mutex ctx_mut; + std::mutex ctx_mut; public: // returns the thread-specific context diff --git a/include/aux_utils.h b/include/aux_utils.h index d3bffcd2e4..1698b9ed54 100644 --- a/include/aux_utils.h +++ b/include/aux_utils.h @@ -31,7 +31,6 @@ typedef int FileHandle; #include "windows_customizations.h" #include "gperftools/malloc_extension.h" - namespace diskann { const size_t TRAINING_SET_SIZE = 1000000; const double SPACE_FOR_CACHED_NODES_IN_GB = 0.25; diff --git a/include/distance.h b/include/distance.h index 3afb56e236..ddfa9bb3b4 100644 --- a/include/distance.h +++ b/include/distance.h @@ -255,16 +255,16 @@ namespace diskann { virtual float compare(const int8_t *a, const int8_t *b, unsigned int length) const { #ifndef _WINDOWS -int32_t result = 0; + int32_t result = 0; #pragma omp simd reduction(+ : result) aligned(a, b : 8) for (_s32 i = 0; i < (_s32) 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; - } + } #else - __m128 r = _mm_setzero_ps(); + __m128 r = _mm_setzero_ps(); __m128i r1; while (length >= 16) { r1 = _mm_subs_epi8(_mm_load_si128((__m128i *) a), @@ -278,7 +278,7 @@ int32_t result = 0; float res = r.m128_f32[0]; if (length >= 8) { - __m128 r2 = _mm_setzero_ps(); + __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)); @@ -290,7 +290,7 @@ int32_t result = 0; } if (length >= 4) { - __m128 r2 = _mm_setzero_ps(); + __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)); @@ -307,14 +307,13 @@ int32_t result = 0; virtual float compare(const float *a, const float *b, unsigned int length) const { #ifndef _WINDOWS -float result = 0; + float result = 0; #pragma omp simd reduction(+ : result) aligned(a, b : 8) for (_s32 i = 0; i < (_s32) length; i++) { - result += (a[i] - b[i]) * - (a[i] - b[i]); + result += (a[i] - b[i]) * (a[i] - b[i]); } return result; - } + } #else __m128 diff, v1, v2; __m128 sum = _mm_set1_ps(0); @@ -436,17 +435,21 @@ float result = 0; #endif return result; } - float compare(const T *a, const T *b, unsigned size) const { // since we use normally minimization objective for distance comparisons, we are returning 1/x. - float result = inner_product(a,b,size); -// if (result < 0) -// return std::numeric_limits::max(); -// else -return -result; + float compare(const T *a, const T *b, unsigned size) + const { // since we use normally minimization objective for distance + // comparisons, we are returning 1/x. + 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. + class DistanceFastL2 + : public DistanceInnerProduct { // currently defined only for float. + // templated for future use. public: float norm(const T *a, unsigned size) const { float result = 0; @@ -539,7 +542,7 @@ return -result; using DistanceInnerProduct::compare; float compare(const T *a, const T *b, float norm, unsigned size) const { // not implement - float result = -2 * DistanceInnerProduct::inner_product(a, b, size); + floatresult = -2 * DistanceInnerProduct::inner_product(a, b, size); result += norm; return result; } diff --git a/include/exceptions.h b/include/exceptions.h index 0323ac3dcb..eefb0f69cc 100644 --- a/include/exceptions.h +++ b/include/exceptions.h @@ -12,4 +12,4 @@ namespace diskann { : std::logic_error("Function not yet implemented.") { } }; -} +} // namespace diskann diff --git a/include/index.h b/include/index.h index b800781d38..fdcac5fafd 100644 --- a/include/index.h +++ b/include/index.h @@ -52,7 +52,7 @@ namespace diskann { // Gopal. Added search overload that takes L as parameter, so that we // can customize L on a per-query basis without tampering with "Parameters" - DISKANN_DLLEXPORT std::pair search(const T *query, + DISKANN_DLLEXPORT std::pair search(const T * query, const size_t K, const unsigned L, unsigned *indices); diff --git a/include/memory_mapper.h b/include/memory_mapper.h index 4ebe6ec62e..a0ec974b98 100644 --- a/include/memory_mapper.h +++ b/include/memory_mapper.h @@ -38,4 +38,4 @@ namespace diskann { ~MemoryMapper(); }; -} \ No newline at end of file +} // namespace diskann \ No newline at end of file diff --git a/include/partition_and_pq.h b/include/partition_and_pq.h index 0afc85410b..7273dcf62f 100644 --- a/include/partition_and_pq.h +++ b/include/partition_and_pq.h @@ -27,9 +27,9 @@ template 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, +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 @@ -38,13 +38,15 @@ int shard_data_into_clusters(const std::string data_file, float *pivots, 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, @@ -57,12 +59,10 @@ int partition_with_ram_budget(const std::string data_file, size_t graph_degree, const std::string prefix_path, size_t k_base); -DISKANN_DLLEXPORT int generate_pq_pivots(const float *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 *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); template int generate_pq_data_from_pivots(const std::string data_file, diff --git a/include/percentile_stats.h b/include/percentile_stats.h index 808546c165..6a7b7cec7f 100644 --- a/include/percentile_stats.h +++ b/include/percentile_stats.h @@ -58,4 +58,4 @@ namespace diskann { } return avg / len; } -} +} // namespace diskann diff --git a/include/pq_flash_index.h b/include/pq_flash_index.h index f17448e7d4..4d9cb3a7f7 100644 --- a/include/pq_flash_index.h +++ b/include/pq_flash_index.h @@ -70,7 +70,8 @@ namespace diskann { // Freeing the reader object is now the client's (DiskANNInterface's) // responsibility. DISKANN_DLLEXPORT PQFlashIndex( - std::shared_ptr &fileReader, diskann::Metric metric = diskann::Metric::L2); + std::shared_ptr &fileReader, + diskann::Metric metric = diskann::Metric::L2); DISKANN_DLLEXPORT ~PQFlashIndex(); #ifdef EXEC_ENV_OLS @@ -79,8 +80,8 @@ namespace diskann { const char *disk_index_file); #else // load compressed data, and obtains the handle to the disk-resident index - DISKANN_DLLEXPORT int load(uint32_t num_threads, const char *pq_prefix, - const char *disk_index_file); + DISKANN_DLLEXPORT int load(uint32_t num_threads, const char *pq_prefix, + const char *disk_index_file); #endif DISKANN_DLLEXPORT void load_cache_list(std::vector &node_list); @@ -132,11 +133,12 @@ namespace diskann { // data info _u64 num_points = 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 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; - std::string disk_index_file; + std::string disk_index_file; std::vector> node_visit_counter; // PQ data @@ -144,8 +146,8 @@ namespace diskann { // data: _u8 * 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; + _u8 * data = nullptr; + _u64 n_chunks; FixedChunkPQTable pq_table; // distance comparator @@ -153,11 +155,10 @@ namespace diskann { Distance *dist_cmp_float = nullptr; // for very large datasets: we use PQ even for the disk resident index - bool use_disk_index_pq = false; - _u64 disk_pq_n_chunks; + bool use_disk_index_pq = false; + _u64 disk_pq_n_chunks; FixedChunkPQTable disk_pq_table; - // medoid/start info uint32_t *medoids = nullptr; // by default it is just one entry point of graph, we @@ -169,11 +170,11 @@ namespace diskann { // closest centroid as the starting point of search // nhood_cache - unsigned *nhood_cache_buf = nullptr; + unsigned * nhood_cache_buf = nullptr; tsl::robin_map<_u32, std::pair<_u32, _u32 *>> nhood_cache; // coord_cache - T *coord_cache_buf = nullptr; + T * coord_cache_buf = nullptr; tsl::robin_map<_u32, T *> coord_cache; // thread-specific scratch diff --git a/include/pq_table.h b/include/pq_table.h index 1c86ca9c31..9525c2f5c2 100644 --- a/include/pq_table.h +++ b/include/pq_table.h @@ -6,7 +6,7 @@ #include "utils.h" namespace diskann { -// template + // template class FixedChunkPQTable { // data_dim = n_chunks * chunk_size; float* tables = @@ -86,7 +86,8 @@ namespace diskann { throw diskann::ANNException("Error loading chunk offsets file", -1, __FUNCSIG__, __FILE__, __LINE__); } - std::cout<<"PQ data has " << numr - 1 <<" bytes per point." << std::endl; + std::cout << "PQ data has " << numr - 1 << " bytes per point." + << std::endl; this->n_chunks = numr - 1; #ifdef EXEC_ENV_OLS @@ -126,11 +127,11 @@ namespace diskann { } } -_u32 get_num_chunks() { - return n_chunks; -} - void - populate_chunk_distances(const float* query_vec, float* dist_vec) { + _u32 + get_num_chunks() { + return n_chunks; + } + void populate_chunk_distances(const float* query_vec, float* dist_vec) { memset(dist_vec, 0, 256 * n_chunks * sizeof(float)); // chunk wise distance computation for (_u64 chunk = 0; chunk < n_chunks; chunk++) { @@ -149,45 +150,50 @@ _u32 get_num_chunks() { } } - float l2_distance(const float* query_vec, _u8* base_vec) { - float res = 0; - for (_u64 chunk = 0; chunk < n_chunks; chunk++) { + float l2_distance(const float* query_vec, _u8* base_vec) { + float res = 0; + for (_u64 chunk = 0; chunk < n_chunks; chunk++) { for (_u64 j = chunk_offsets[chunk]; j < chunk_offsets[chunk + 1]; j++) { _u64 permuted_dim_in_query = rearrangement[j]; const float* centers_dim_vec = tables_T + (256 * j); - float diff = centers_dim_vec[base_vec[chunk]] - (query_vec[permuted_dim_in_query] - centroid[permuted_dim_in_query]); - res += diff*diff; + float diff = centers_dim_vec[base_vec[chunk]] - + (query_vec[permuted_dim_in_query] - + centroid[permuted_dim_in_query]); + res += diff * diff; } } return res; - } + } - float inner_product(const float* query_vec, _u8* base_vec) { - float res = 0; - for (_u64 chunk = 0; chunk < n_chunks; chunk++) { + float inner_product(const float* query_vec, _u8* base_vec) { + float res = 0; + for (_u64 chunk = 0; chunk < n_chunks; chunk++) { for (_u64 j = chunk_offsets[chunk]; j < chunk_offsets[chunk + 1]; j++) { _u64 permuted_dim_in_query = rearrangement[j]; const float* centers_dim_vec = tables_T + (256 * j); - float diff = centers_dim_vec[base_vec[chunk]]*query_vec[permuted_dim_in_query]; // assumes centroid is 0 to prevent translation errors + float diff = + centers_dim_vec[base_vec[chunk]] * + query_vec[permuted_dim_in_query]; // 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) + } - void inflate_vector(_u8* base_vec, float* out_vec) { - for (_u64 chunk = 0; chunk < n_chunks; chunk++) { + void inflate_vector(_u8* base_vec, float* out_vec) { + for (_u64 chunk = 0; chunk < n_chunks; chunk++) { for (_u64 j = chunk_offsets[chunk]; j < chunk_offsets[chunk + 1]; j++) { _u64 original_dim = rearrangement[j]; const float* centers_dim_vec = tables_T + (256 * j); - out_vec[original_dim] = centers_dim_vec[base_vec[chunk]] + centroid[original_dim]; + out_vec[original_dim] = + centers_dim_vec[base_vec[chunk]] + centroid[original_dim]; } } - } + } - void - populate_chunk_inner_products(const float* query_vec, float* dist_vec) { + void 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 (_u64 chunk = 0; chunk < n_chunks; chunk++) { @@ -198,11 +204,17 @@ _u32 get_num_chunks() { const float* centers_dim_vec = tables_T + (256 * j); for (_u64 idx = 0; idx < 256; idx++) { double prod = - centers_dim_vec[idx] * query_vec[permuted_dim_in_query]; // 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) + centers_dim_vec[idx] * + query_vec[permuted_dim_in_query]; // 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) } } } } -}; +}; // namespace diskann } // namespace diskann diff --git a/include/timer.h b/include/timer.h index 4671c33beb..bf52ed8835 100644 --- a/include/timer.h +++ b/include/timer.h @@ -22,4 +22,4 @@ namespace diskann { .count(); } }; -} +} // namespace diskann diff --git a/include/utils.h b/include/utils.h index 138f05e332..86c5f1fd3c 100644 --- a/include/utils.h +++ b/include/utils.h @@ -52,8 +52,6 @@ typedef int FileHandle; #define IS_512_ALIGNED(X) IS_ALIGNED(X, 512) #define IS_4096_ALIGNED(X) IS_ALIGNED(X, 4096) - - typedef uint64_t _u64; typedef int64_t _s64; typedef uint32_t _u32; @@ -222,10 +220,10 @@ namespace diskann { inline void wait_for_keystroke() { int a; - std::cout<<"Press any number to continue.." << std::endl; - std::cin>> a; + std::cout << "Press any number to continue.." << std::endl; + std::cin >> a; } - + template inline void load_bin(const std::string& bin_file, T*& data, size_t& npts, size_t& dim) { @@ -418,76 +416,87 @@ namespace diskann { } } -//this function will take in_file of n*d dimensions and save the output as a floating point matrix -// with n*(d+1) dimensions. All vectors are scaled by a large value M so that the norms are <=1 -// and the final coordinate is set so that the resulting norm (in d+1 coordinates) is equal to 1 -// this is a classical transformation from MIPS to L2 search from "On Symmetric and Asymmetric LSHs for Inner Product Search" -// by Neyshabur and Srebro - -template -void 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); - _u64 npts, in_dims, out_dims; - float max_norm = 0; - - _u32 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; - - 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); - _u64 num_blocks = DIV_ROUND_UP(npts, block_size); - - std::vector norms(npts, 0); - - for (_u64 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; - in_reader.read((char *) in_block_data.get(), block_pts * in_dims * sizeof(T)); - for (_u64 p = 0; p < block_pts; p++) { - for (_u64 j = 0; j < in_dims; j++) { - norms[start_id + p] += in_block_data[p*in_dims + j]*in_block_data[p*in_dims + j]; + // this function will take in_file of n*d dimensions and save the output as a + // floating point matrix + // with n*(d+1) dimensions. All vectors are scaled by a large value M so that + // the norms are <=1 and the final coordinate is set so that the resulting + // norm (in d+1 coordinates) is equal to 1 this is a classical transformation + // from MIPS to L2 search from "On Symmetric and Asymmetric LSHs for Inner + // Product Search" by Neyshabur and Srebro + + template + void 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); + _u64 npts, in_dims, out_dims; + float max_norm = 0; + + _u32 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; + + 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); + _u64 num_blocks = DIV_ROUND_UP(npts, block_size); + + std::vector norms(npts, 0); + + for (_u64 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; + in_reader.read((char*) in_block_data.get(), + block_pts * in_dims * sizeof(T)); + for (_u64 p = 0; p < block_pts; p++) { + for (_u64 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 = max_norm > norms[start_id + p] ? max_norm : norms[start_id + p]; } - } - max_norm = std::sqrt(max_norm); - - in_reader.seekg(2*sizeof(_u32), std::ios::beg); - for (_u64 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; - in_reader.read((char *) in_block_data.get(), block_pts * in_dims * sizeof(T)); - for (_u64 p = 0; p < block_pts; p++) { - for (_u64 j = 0; j < in_dims; j++) { - out_block_data[p*out_dims + j] = in_block_data[p*in_dims + j] / max_norm; + max_norm = std::sqrt(max_norm); + + in_reader.seekg(2 * sizeof(_u32), std::ios::beg); + for (_u64 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; + in_reader.read((char*) in_block_data.get(), + block_pts * in_dims * sizeof(T)); + for (_u64 p = 0; p < block_pts; p++) { + for (_u64 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; } - 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(); -} + out_writer.close(); + } // plain saves data as npts X ndims array into filename template @@ -574,8 +583,9 @@ inline bool validate_file_size(const std::string& name) { size_t expected_file_size; in.read((char*) &expected_file_size, sizeof(uint64_t)); if (actual_file_size != expected_file_size) { - diskann::cout << "Error loading" << name << ". Expected " - "size (metadata): " + diskann::cout << "Error loading" << name + << ". Expected " + "size (metadata): " << expected_file_size << ", actual file size : " << actual_file_size << ". Exitting." << std::endl; diff --git a/include/windows_aligned_file_reader.h b/include/windows_aligned_file_reader.h index 8fec3d4f0e..433d3c0bf7 100644 --- a/include/windows_aligned_file_reader.h +++ b/include/windows_aligned_file_reader.h @@ -31,7 +31,7 @@ class WindowsAlignedFileReader : public AlignedFileReader { // Open & close ops // Blocking calls DISKANN_DLLEXPORT virtual void open(const std::string &fname); - DISKANN_DLLEXPORT virtual void close(); + DISKANN_DLLEXPORT virtual void close(); DISKANN_DLLEXPORT virtual void register_thread(); DISKANN_DLLEXPORT virtual void deregister_thread() { @@ -41,8 +41,7 @@ class WindowsAlignedFileReader : public AlignedFileReader { // process batch of aligned requests in parallel // NOTE :: blocking call for the calling thread, but can thread-safe DISKANN_DLLEXPORT virtual void read(std::vector &read_reqs, - IOContext &ctx, - bool async); + IOContext &ctx, bool async); }; #endif // USE_BING_INFRA #endif //_WINDOWS diff --git a/src/aux_utils.cpp b/src/aux_utils.cpp index 1a8e2387fd..e1f8205796 100644 --- a/src/aux_utils.cpp +++ b/src/aux_utils.cpp @@ -56,7 +56,9 @@ namespace diskann { } 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.insert(res_vec, + 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) { if (res.find(v) != res.end()) { @@ -152,9 +154,8 @@ namespace diskann { std::ifstream reader(fname.c_str(), std::ios::binary); reader.read((char *) &npts32, sizeof(uint32_t)); reader.read((char *) &dim, sizeof(uint32_t)); - if (dim != 1 || - actual_file_size != - ((size_t) npts32) * sizeof(uint32_t) + 2 * sizeof(uint32_t)) { + if (dim != 1 || actual_file_size != ((size_t) npts32) * sizeof(uint32_t) + + 2 * sizeof(uint32_t)) { std::stringstream stream; stream << "Error reading idmap file. Check if the file is bin file with " "1 dimensional data. Actual: " @@ -209,11 +210,11 @@ namespace diskann { node_shard.push_back(std::make_pair((_u32) node_id, (_u32) shard)); } } - std::sort(node_shard.begin(), node_shard.end(), [](const auto &left, - const auto &right) { - return left.first < right.first || - (left.first == right.first && left.second < right.second); - }); + std::sort(node_shard.begin(), node_shard.end(), + [](const auto &left, const auto &right) { + return left.first < right.first || (left.first == right.first && + left.second < right.second); + }); diskann::cout << "Finished computing node -> shards map" << std::endl; // create cached vamana readers @@ -354,8 +355,9 @@ namespace diskann { double full_index_ram = ESTIMATE_RAM_USAGE(base_num, base_dim, sizeof(T), R); 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::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::Parameters paras; paras.Set("L", (unsigned) L); paras.Set("R", (unsigned) R); @@ -386,10 +388,11 @@ namespace diskann { 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_ids_file = merged_index_prefix + "_subshard-" + + std::to_string(p) + "_ids_uint32.bin"; - retrieve_shard_data_from_ids(base_file, shard_ids_file, shard_base_file); + 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"; @@ -409,8 +412,7 @@ namespace diskann { _pvamanaIndex->build(paras); _pvamanaIndex->save(shard_index_file.c_str()); std::remove(shard_base_file.c_str()); -// wait_for_keystroke(); - + // wait_for_keystroke(); } diskann::merge_shards(merged_index_prefix + "_subshard-", "_mem.index", @@ -438,60 +440,55 @@ namespace diskann { // optimizes the beamwidth to maximize QPS for a given L_search subject to // 99.9 latency not blowing up - template - uint32_t optimize_beamwidth( - std::unique_ptr> &pFlashIndex, T - *tuning_sample, - _u64 tuning_sample_num, _u64 tuning_sample_aligned_dim, uint32_t L, - uint32_t nthreads, uint32_t start_bw) { - uint32_t cur_bw = start_bw; - double max_qps = 0; - uint32_t best_bw = start_bw; - bool stop_flag = false; - - while (!stop_flag) { - std::vector tuning_sample_result_ids_64(tuning_sample_num, - 0); - std::vector tuning_sample_result_dists(tuning_sample_num, - 0); - diskann::QueryStats * stats = new - diskann::QueryStats[tuning_sample_num]; - - auto s = std::chrono::high_resolution_clock::now(); - #pragma omp parallel for schedule(dynamic, 1) num_threads(nthreads) - for (_s64 i = 0; i < (int64_t) tuning_sample_num; i++) { - pFlashIndex->cached_beam_search( - tuning_sample + (i * tuning_sample_aligned_dim), 1, L, - tuning_sample_result_ids_64.data() + (i * 1), - tuning_sample_result_dists.data() + (i * 1), cur_bw, stats + - i); - } - auto e = std::chrono::high_resolution_clock::now(); - std::chrono::duration diff = e - s; - double qps = (1.0f * tuning_sample_num) / (1.0f * diff.count()); - - double lat_999 = diskann::get_percentile_stats( - stats, tuning_sample_num, 0.999, - [](const diskann::QueryStats &stats) { return stats.total_us; }); - - double mean_latency = diskann::get_mean_stats( - stats, tuning_sample_num, - [](const diskann::QueryStats &stats) { return stats.total_us; }); - - if (qps > max_qps && lat_999 < (15000) + mean_latency * 2) { - max_qps = qps; - best_bw = cur_bw; - cur_bw = (uint32_t)(std::ceil)((float) cur_bw * 1.1); - } else { - stop_flag = true; - } - if (cur_bw > 64) - stop_flag = true; - - delete[] stats; + template + uint32_t optimize_beamwidth( + std::unique_ptr> &pFlashIndex, T *tuning_sample, + _u64 tuning_sample_num, _u64 tuning_sample_aligned_dim, uint32_t L, + uint32_t nthreads, uint32_t start_bw) { + uint32_t cur_bw = start_bw; + double max_qps = 0; + uint32_t best_bw = start_bw; + bool stop_flag = false; + + while (!stop_flag) { + std::vector tuning_sample_result_ids_64(tuning_sample_num, 0); + std::vector tuning_sample_result_dists(tuning_sample_num, 0); + diskann::QueryStats * stats = new diskann::QueryStats[tuning_sample_num]; + + auto s = std::chrono::high_resolution_clock::now(); +#pragma omp parallel for schedule(dynamic, 1) num_threads(nthreads) + for (_s64 i = 0; i < (int64_t) tuning_sample_num; i++) { + pFlashIndex->cached_beam_search( + tuning_sample + (i * tuning_sample_aligned_dim), 1, L, + tuning_sample_result_ids_64.data() + (i * 1), + tuning_sample_result_dists.data() + (i * 1), cur_bw, stats + i); } - return best_bw; + auto e = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = e - s; + double qps = (1.0f * tuning_sample_num) / (1.0f * diff.count()); + + double lat_999 = diskann::get_percentile_stats( + stats, tuning_sample_num, 0.999, + [](const diskann::QueryStats &stats) { return stats.total_us; }); + + double mean_latency = diskann::get_mean_stats( + stats, tuning_sample_num, + [](const diskann::QueryStats &stats) { return stats.total_us; }); + + if (qps > max_qps && lat_999 < (15000) + mean_latency * 2) { + max_qps = qps; + best_bw = cur_bw; + cur_bw = (uint32_t)(std::ceil)((float) cur_bw * 1.1); + } else { + stop_flag = true; + } + if (cur_bw > 64) + stop_flag = true; + + delete[] stats; } + return best_bw; + } template void create_disk_layout(const std::string base_file, @@ -548,18 +545,15 @@ namespace diskann { << std::endl; // SECTOR_LEN buffer for each sector - std::unique_ptr sector_buf = - std::make_unique(SECTOR_LEN); - std::unique_ptr node_buf = - std::make_unique(max_node_len); + std::unique_ptr sector_buf = std::make_unique(SECTOR_LEN); + std::unique_ptr node_buf = std::make_unique(max_node_len); unsigned &nnbrs = *(unsigned *) (node_buf.get() + ndims_64 * sizeof(T)); unsigned *nhood_buf = (unsigned *) (node_buf.get() + (ndims_64 * sizeof(T)) + sizeof(unsigned)); // number of sectors (1 for meta data) - _u64 n_sectors = ROUND_UP(npts_64, nnodes_per_sector) / - nnodes_per_sector; + _u64 n_sectors = ROUND_UP(npts_64, nnodes_per_sector) / nnodes_per_sector; _u64 disk_index_file_size = (n_sectors + 1) * SECTOR_LEN; // write first sector with metadata *(_u64 *) (sector_buf.get() + 0 * sizeof(_u64)) = disk_index_file_size; @@ -593,8 +587,7 @@ namespace diskann { // write coords of node first // T *node_coords = data + ((_u64) ndims_64 * cur_node_id); - base_reader.read((char *) cur_node_coords.get(), sizeof(T) * - ndims_64); + base_reader.read((char *) cur_node_coords.get(), sizeof(T) * ndims_64); memcpy(node_buf.get(), cur_node_coords.get(), ndims_64 * sizeof(T)); // write nnbrs @@ -635,28 +628,32 @@ namespace diskann { "L (indexing list size, better if >= R) B (RAM limit of final " "index in " "GB) M (memory limit while indexing) T (number of threads for " - "indexing) B' (PQ bytes for disk index: optional parameter for very large dimensional data)" + "indexing) B' (PQ bytes for disk index: optional parameter for " + "very large dimensional data)" << std::endl; return false; } - - - 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); + 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); } _u32 disk_pq_dims = 0; bool use_disk_pq = false; -// 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 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() == 6) { - disk_pq_dims = atoi(param_list[5].c_str()); - use_disk_pq = true; - if (disk_pq_dims == 0) - use_disk_pq = false; + disk_pq_dims = atoi(param_list[5].c_str()); + use_disk_pq = true; + if (disk_pq_dims == 0) + use_disk_pq = false; } std::string base_file(dataFilePath); @@ -670,13 +667,23 @@ namespace diskann { std::string medoids_path = disk_index_path + "_medoids.bin"; std::string centroids_path = disk_index_path + "_centroids.bin"; std::string sample_base_prefix = index_prefix_path + "_sample"; - std::string disk_pq_pivots_path = index_prefix_path + "_disk.index_pq_pivots.bin"; // optional if disk index is also storing pq data - std::string disk_pq_compressed_vectors_path = // optional if disk index is also storing pq data + std::string disk_pq_pivots_path = + index_prefix_path + + "_disk.index_pq_pivots.bin"; // optional if disk index is also storing + // pq data + std::string disk_pq_compressed_vectors_path = // optional if disk index is + // also storing pq data 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! + // 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) { - 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::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; diskann::prepare_base_for_inner_products(base_file, prepped_base); @@ -700,7 +707,6 @@ namespace diskann { } _u32 num_threads = (_u32) atoi(param_list[4].c_str()); - if (num_threads != 0) { omp_set_num_threads(num_threads); mkl_set_num_threads(num_threads); @@ -735,66 +741,67 @@ namespace diskann { // generates random sample and sets it to train_data and updates // train_size gen_random_slice(data_file_to_use.c_str(), p_val, train_data, train_size, - train_dim); + train_dim); if (use_disk_pq) { if (disk_pq_dims > dim) - disk_pq_dims = dim; + disk_pq_dims = dim; - std::cout<<"Compressing base for disk-PQ into " << disk_pq_dims << " chunks " << std::endl; + std::cout << "Compressing base for disk-PQ into " << disk_pq_dims + << " chunks " << std::endl; generate_pq_pivots(train_data, train_size, (uint32_t) dim, 256, - (uint32_t) disk_pq_dims, NUM_KMEANS_REPS, disk_pq_pivots_path, false); - if (compareMetric == diskann::Metric::INNER_PRODUCT) - generate_pq_data_from_pivots(data_file_to_use.c_str(), 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.c_str(), 256, (uint32_t) disk_pq_dims, - disk_pq_pivots_path, - disk_pq_compressed_vectors_path); + (uint32_t) disk_pq_dims, NUM_KMEANS_REPS, + disk_pq_pivots_path, false); + if (compareMetric == diskann::Metric::INNER_PRODUCT) + generate_pq_data_from_pivots( + data_file_to_use.c_str(), 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.c_str(), 256, (uint32_t) disk_pq_dims, + disk_pq_pivots_path, disk_pq_compressed_vectors_path); } - diskann::cout << "Training data loaded of size " << train_size << - std::endl; - -// don't translate data to make zero mean for PQ compression. We must not translate for inner product search. + diskann::cout << "Training data loaded of size " << train_size << std::endl; + + // don't translate data to make zero mean for PQ compression. We must not + // translate for inner product search. bool make_zero_mean = true; if (compareMetric == diskann::Metric::INNER_PRODUCT) - make_zero_mean = false; + make_zero_mean = false; generate_pq_pivots(train_data, train_size, (uint32_t) dim, 256, - (uint32_t) num_pq_chunks, NUM_KMEANS_REPS, pq_pivots_path, make_zero_mean); - - generate_pq_data_from_pivots(data_file_to_use.c_str(), 256, (uint32_t) - num_pq_chunks, pq_pivots_path, pq_compressed_vectors_path); + (uint32_t) num_pq_chunks, NUM_KMEANS_REPS, + pq_pivots_path, make_zero_mean); + + generate_pq_data_from_pivots(data_file_to_use.c_str(), 256, + (uint32_t) num_pq_chunks, pq_pivots_path, + pq_compressed_vectors_path); delete[] train_data; train_data = nullptr; MallocExtension::instance()->ReleaseFreeMemory(); - 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); + data_file_to_use.c_str(), diskann::Metric::L2, L, R, p_val, + indexing_ram_budget, mem_index_path, medoids_path, centroids_path); 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<_u8>(disk_pq_compressed_vectors_path, mem_index_path, - disk_index_path); + diskann::create_disk_layout(data_file_to_use.c_str(), mem_index_path, + disk_index_path); + } else + diskann::create_disk_layout<_u8>(disk_pq_compressed_vectors_path, + mem_index_path, disk_index_path); double sample_sampling_rate = (150000.0 / points_num); gen_random_slice(data_file_to_use.c_str(), sample_base_prefix, - sample_sampling_rate); + sample_sampling_rate); -// std::remove(mem_index_path.c_str()); + // std::remove(mem_index_path.c_str()); if (use_disk_pq) - std::remove(disk_pq_compressed_vectors_path.c_str()); + std::remove(disk_pq_compressed_vectors_path.c_str()); - auto e = - std::chrono::high_resolution_clock::now(); + auto e = std::chrono::high_resolution_clock::now(); std::chrono::duration diff = e - s; diskann::cout << "Indexing time: " << diff.count() << std::endl; diff --git a/src/index.cpp b/src/index.cpp index 061760da47..38b58da964 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -61,9 +61,9 @@ namespace { std::cout << "Older CPU. Using slow distance computation" << std::endl; return new diskann::SlowDistanceL2Float(); } - } else if (m == diskann::Metric::INNER_PRODUCT) { - std::cout << "Using Inner Product computation" << std::endl; - return new diskann::DistanceInnerProduct(); + } else if (m == diskann::Metric::INNER_PRODUCT) { + std::cout << "Using Inner Product computation" << std::endl; + return new diskann::DistanceInnerProduct(); } else { std::stringstream stream; stream << "Only L2 metric supported as of now. Email " @@ -132,12 +132,11 @@ namespace diskann { const size_t nd, const size_t num_frozen_pts, const bool enable_tags, const bool store_data, const bool support_eager_delete) - : _metric(m), _num_frozen_pts(num_frozen_pts), _has_built(false), _width(0), - _can_delete(false), _eager_done(true), _lazy_done(true), + : _metric(m), _num_frozen_pts(num_frozen_pts), _has_built(false), + _width(0), _can_delete(false), _eager_done(true), _lazy_done(true), _compacted_order(true), _enable_tags(enable_tags), _consolidated_order(true), _support_eager_delete(support_eager_delete), _store_data(store_data) { - // data is stored to _nd * aligned_dim matrix with necessary // zero-padding diskann::cout << "Number of frozen points = " << _num_frozen_pts @@ -375,7 +374,7 @@ namespace diskann { center[j] /= _nd; // compute all to one distance - float * distances = new float[_nd](); + float *distances = new float[_nd](); #pragma omp parallel for schedule(static, 65536) for (_s64 i = 0; i < (_s64) _nd; i++) { // extract point and distance reference @@ -506,7 +505,7 @@ namespace diskann { tsl::robin_set &expanded_nodes_ids) { const T * node_coords = _data + _aligned_dim * node_id; std::vector best_L_nodes; - + if (init_ids.size() == 0) init_ids.emplace_back(_ep); @@ -538,7 +537,10 @@ namespace diskann { float cur_alpha = 1; while (cur_alpha <= alpha && result.size() < degree) { unsigned start = 0; - float eps = cur_alpha + 0.01; // 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.01; // 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. while (result.size() < degree && (start) < pool.size() && start < maxc) { auto &p = pool[start]; if (occlude_factor[start] > cur_alpha) { @@ -553,18 +555,20 @@ namespace diskann { float djk = _distance->compare( _data + _aligned_dim * (size_t) pool[t].id, _data + _aligned_dim * (size_t) p.id, (unsigned) _aligned_dim); - if (_metric == diskann::Metric::L2) { - occlude_factor[t] = - (std::max)(occlude_factor[t], pool[t].distance / djk); - } - else if (_metric == diskann::Metric::INNER_PRODUCT) { // stylized rules for inner product since we want max instead of min distance - float x = -pool[t].distance; - float y = -djk; - if (y > cur_alpha * x) { - occlude_factor[t] = - (std::max)(occlude_factor[t], eps); - } - } + if (_metric == diskann::Metric::L2) { + occlude_factor[t] = + (std::max)(occlude_factor[t], pool[t].distance / djk); + } else if (_metric == + diskann::Metric::INNER_PRODUCT) { // stylized rules for + // inner product since + // we want max instead + // of min distance + float x = -pool[t].distance; + float y = -djk; + if (y > cur_alpha * x) { + occlude_factor[t] = (std::max)(occlude_factor[t], eps); + } + } } start++; } @@ -573,7 +577,7 @@ namespace diskann { } template - void Index::prune_neighbors(const unsigned location, + void Index::prune_neighbors(const unsigned location, std::vector &pool, const Parameters & parameter, std::vector &pruned_list) { @@ -653,7 +657,7 @@ namespace diskann { * the current node n. */ template - void Index::inter_insert(unsigned n, + void Index::inter_insert(unsigned n, std::vector &pruned_list, const Parameters & parameter, bool update_in_graph) { @@ -789,7 +793,6 @@ namespace diskann { _final_graph[p].reserve((size_t)(std::ceil(range * SLACK_FACTOR * 1.05))); } - std::random_device rd; std::mt19937 gen(rd()); std::uniform_real_distribution<> dis(0, 1); @@ -985,7 +988,7 @@ namespace diskann { } template - void Index::build(Parameters ¶meters, + void Index::build(Parameters & parameters, const std::vector &tags) { if (_enable_tags) { if (tags.size() != _nd) { @@ -1023,7 +1026,7 @@ namespace diskann { } template - std::pair Index::search(const T *query, + std::pair Index::search(const T * query, const size_t K, const unsigned L, unsigned * indices) { @@ -1069,7 +1072,7 @@ namespace diskann { indices[pos] = it.id; distances[pos] = it.distance; if (_metric == diskann::INNER_PRODUCT) - distances[pos] = -distances[pos]; + distances[pos] = -distances[pos]; pos++; if (pos == K) break; @@ -1289,7 +1292,7 @@ namespace diskann { } template - int Index::eager_delete(const TagT tag, + int Index::eager_delete(const TagT tag, const Parameters ¶meters) { if (_lazy_done && (!_consolidated_order)) { diskann::cout << "Lazy delete reuests issued but data not consolidated, " @@ -1765,7 +1768,7 @@ namespace diskann { template int Index::disable_delete(const Parameters ¶meters, - const bool consolidate) { + const bool consolidate) { LockGuard guard(_change_lock); if (!_can_delete) { diskann::cerr << "Delete not currently enabled" << std::endl; diff --git a/src/linux_aligned_file_reader.cpp b/src/linux_aligned_file_reader.cpp index 69eae9b60e..bb95201d41 100644 --- a/src/linux_aligned_file_reader.cpp +++ b/src/linux_aligned_file_reader.cpp @@ -89,7 +89,7 @@ namespace { } std::cout << std::endl;*/ } -} +} // namespace LinuxAlignedFileReader::LinuxAlignedFileReader() { this->file_desc = -1; diff --git a/src/partition_and_pq.cpp b/src/partition_and_pq.cpp index c2accdfbd0..044bdd70f5 100644 --- a/src/partition_and_pq.cpp +++ b/src/partition_and_pq.cpp @@ -200,7 +200,6 @@ int generate_pq_pivots(const float *passed_train_data, size_t num_train, return -1; } - std::unique_ptr train_data = std::make_unique(num_train * dim); std::memcpy(train_data.get(), passed_train_data, @@ -225,27 +224,30 @@ int generate_pq_pivots(const float *passed_train_data, size_t num_train, 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. + 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]; + for (uint64_t p = 0; p < num_train; p++) { + centroid[d] += train_data[p * dim + d]; + } + centroid[d] /= num_train; } - 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 rearrangement; std::vector chunk_offsets; @@ -258,7 +260,7 @@ int generate_pq_pivots(const float *passed_train_data, size_t num_train, std::vector> bin_to_dims(num_pq_chunks); tsl::robin_map dim_to_bin; - std::vector bin_loads(num_pq_chunks, 0); + std::vector bin_loads(num_pq_chunks, 0); // Process dimensions not inserted by previous loop for (uint32_t d = 0; d < dim; d++) { @@ -445,25 +447,22 @@ int generate_pq_data_from_pivots(const std::string data_file, std::ofstream compressed_file_writer(pq_compressed_vectors_path, std::ios::binary); - _u32 num_pq_chunks_u32 = num_pq_chunks; + _u32 num_pq_chunks_u32 = num_pq_chunks; compressed_file_writer.write((char *) &num_points, sizeof(uint32_t)); compressed_file_writer.write((char *) &num_pq_chunks_u32, sizeof(uint32_t)); 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); + 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)); + 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)); -#endif + 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); @@ -538,7 +537,7 @@ int generate_pq_data_from_pivots(const std::string data_file, for (int64_t j = 0; j < (_s64) cur_blk_size; j++) { block_compressed_base[j * num_pq_chunks + i] = closest_center[j]; #ifdef SAVE_INFLATED_PQ - for (uint64_t k = 0; k < cur_chunk_size; k++) + for (uint64_t k = 0; k < cur_chunk_size; k++) block_inflated_base[j * dim + chunk_offsets[i] + k] = cur_pivot_data[closest_center[j] * cur_chunk_size + k] + centroid[chunk_offsets[i] + k]; @@ -557,12 +556,11 @@ int generate_pq_data_from_pivots(const std::string data_file, 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)); + 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; } @@ -578,19 +576,18 @@ int generate_pq_data_from_pivots(const std::string data_file, return 0; } -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, +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]; 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; _u32 * block_closest_centers = new _u32[block_size * k_base]; float *block_data_float; @@ -604,8 +601,8 @@ int estimate_cluster_sizes(float* test_data_float, size_t num_test, float *pivot 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, + 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++) { @@ -619,8 +616,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]; - cluster_sizes.push_back( - (size_t)cur_shard_count); + cluster_sizes.push_back((size_t) cur_shard_count); diskann::cout << cur_shard_count << " "; } diskann::cout << std::endl; @@ -728,13 +724,13 @@ int shard_data_into_clusters(const std::string data_file, float *pivots, return 0; } - - -// useful for partitioning large dataset. we first generate only the IDS for each shard, and retrieve the actual vectors on demand. +// 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) { +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; // create cached reader + writer @@ -819,10 +815,10 @@ int shard_data_into_clusters_only_ids(const std::string data_file, float *pivots return 0; } - - 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) { _u64 read_blk_size = 64 * 1024 * 1024; // _u64 write_blk_size = 64 * 1024 * 1024; // create cached reader + writer @@ -834,21 +830,19 @@ int retrieve_shard_data_from_ids(const std::string data_file, std::string idmap_ size_t num_points = npts32; size_t dim = basedim32; - - _u32 dummy_size = 0; - + _u32 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; + _u32 *shard_ids; + _u64 shard_size, tmp; diskann::load_bin<_u32>(idmap_filename, shard_ids, shard_size, tmp); _u32 cur_pos = 0; _u32 num_written = 0; - std::cout<<"Shard has " << shard_size<< " points" << std::endl; + 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); @@ -864,34 +858,30 @@ int retrieve_shard_data_from_ids(const std::string data_file, std::string idmap_ 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++; - } + 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; + 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; + 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 // 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 @@ -937,7 +927,6 @@ int partition(const std::string data_file, const float sampling_rate, // 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; @@ -967,7 +956,6 @@ int partition_with_ram_budget(const std::string data_file, 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); @@ -999,11 +987,13 @@ int partition_with_ram_budget(const std::string data_file, // closest clusters. std::vector cluster_sizes; - estimate_cluster_sizes(test_data_float, num_test, pivot_data, num_parts, train_dim, - k_base, 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) { - p = (_u64) (p/ sampling_rate); // to account for the fact that p is the size of the shard over the testing sample. + p = (_u64)(p / + sampling_rate); // to account for the fact that p is the size + // of the shard over the testing sample. double cur_shard_ram_estimate = ESTIMATE_RAM_USAGE(p, train_dim, sizeof(T), graph_degree); @@ -1015,7 +1005,7 @@ int partition_with_ram_budget(const std::string data_file, << "GB, budget given is " << ram_budget << std::endl; if (max_ram_usage > 1024 * 1024 * 1024 * ram_budget) { fit_in_ram = false; - num_parts+=2; + num_parts += 2; } } @@ -1023,8 +1013,8 @@ int partition_with_ram_budget(const std::string data_file, 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); + 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; @@ -1034,17 +1024,17 @@ int partition_with_ram_budget(const std::string data_file, // Instantations of supported templates template void DISKANN_DLLEXPORT -gen_random_slice(const std::string base_file, + 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, + 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, + 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, @@ -1083,9 +1073,15 @@ 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); +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); template DISKANN_DLLEXPORT int generate_pq_data_from_pivots( const std::string data_file, unsigned num_centers, unsigned num_pq_chunks, diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index a8eae0986f..d3b5219147 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -87,7 +87,7 @@ namespace diskann { template<> PQFlashIndex<_u8>::PQFlashIndex( std::shared_ptr &fileReader, diskann::Metric metric) - : reader(fileReader), metric(metric) { + : reader(fileReader), metric(metric) { diskann::cout << "dist_cmp function for _u8 uses slow implementation." " Please contact gopalsr@microsoft.com if you need an AVX/AVX2" @@ -107,7 +107,9 @@ namespace diskann { this->dist_cmp_float = new SlowDistanceL2Float(); } if (metric != diskann::Metric::L2) { - std::cout<<"Only L2 supported for byte vectors for now. Other distance functions are future work. Falling back to L2 distance." << std::endl; + std::cout << "Only L2 supported for byte vectors for now. Other distance " + "functions are future work. Falling back to L2 distance." + << std::endl; this->metric = diskann::Metric::L2; } } @@ -134,42 +136,44 @@ namespace diskann { this->dist_cmp = new SlowDistanceL2Int(); this->dist_cmp_float = new SlowDistanceL2Float(); } - if (metric != diskann::Metric::L2) { - std::cout<<"Only L2 supported for byte vectors for now. Other distance functions are future work. Falling back to L2 distance." << std::endl; + if (metric != diskann::Metric::L2) { + std::cout << "Only L2 supported for byte vectors for now. Other distance " + "functions are future work. Falling back to L2 distance." + << std::endl; this->metric = diskann::Metric::L2; } - } template<> PQFlashIndex::PQFlashIndex( std::shared_ptr &fileReader, diskann::Metric metric) : reader(fileReader), metric(metric) { - if (metric == diskann::Metric::L2) { - if (Avx2SupportedCPU) { - diskann::cout << "Using AVX2 functions for dist_cmp and dist_cmp_float" - << std::endl; - this->dist_cmp = new DistanceL2(); - this->dist_cmp_float = new DistanceL2(); - } else if (AvxSupportedCPU) { - diskann::cout << "No AVX2 support. Switching to AVX functions for " - "dist_cmp and dist_cmp_float." - << std::endl; - this->dist_cmp = new AVXDistanceL2Float(); - this->dist_cmp_float = new AVXDistanceL2Float(); - } else { - diskann::cout << "No AVX/AVX2 support. Switching to slow implementations " - "for dist_cmp and dist_cmp_float" - << std::endl; - this->dist_cmp = new AVXDistanceL2Float(); - this->dist_cmp_float = new AVXDistanceL2Float(); - } + if (metric == diskann::Metric::L2) { + if (Avx2SupportedCPU) { + diskann::cout << "Using AVX2 functions for dist_cmp and dist_cmp_float" + << std::endl; + this->dist_cmp = new DistanceL2(); + this->dist_cmp_float = new DistanceL2(); + } else if (AvxSupportedCPU) { + diskann::cout << "No AVX2 support. Switching to AVX functions for " + "dist_cmp and dist_cmp_float." + << std::endl; + this->dist_cmp = new AVXDistanceL2Float(); + this->dist_cmp_float = new AVXDistanceL2Float(); + } else { + diskann::cout + << "No AVX/AVX2 support. Switching to slow implementations " + "for dist_cmp and dist_cmp_float" + << std::endl; + this->dist_cmp = new AVXDistanceL2Float(); + this->dist_cmp_float = new AVXDistanceL2Float(); + } } else if (metric == diskann::Metric::INNER_PRODUCT) { - std::cout<<"Using inner product distance function" << std::endl; + std::cout << "Using inner product distance function" << std::endl; this->dist_cmp = new DistanceInnerProduct(); this->dist_cmp_float = new DistanceInnerProduct(); } else { - std::cout<<"Unsupported metric type. Reverting to float." << std::endl; + std::cout << "Unsupported metric type. Reverting to float." << std::endl; this->dist_cmp = new AVXDistanceL2Float(); this->dist_cmp_float = new AVXDistanceL2Float(); this->metric = diskann::Metric::L2; @@ -300,7 +304,7 @@ namespace diskann { for (_u64 block = 0; block < num_blocks; block++) { _u64 start_idx = block * BLOCK_SIZE; _u64 end_idx = (std::min)(num_cached_nodes, (block + 1) * BLOCK_SIZE); - std::vector read_reqs; + std::vector read_reqs; std::vector> nhoods; for (_u64 node_idx = start_idx; node_idx < end_idx; node_idx++) { AlignedRead read; @@ -467,7 +471,7 @@ namespace diskann { size_t start = block * BLOCK_SIZE; size_t end = (std::min)((block + 1) * BLOCK_SIZE, nodes_to_expand.size()); - std::vector read_reqs; + std::vector read_reqs; std::vector> nhoods; for (size_t cur_pt = start; cur_pt < end; cur_pt++) { char *buf = nullptr; @@ -565,10 +569,11 @@ namespace diskann { 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]; + 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((_u8*) medoid_coords, (centroid_data + cur_m*aligned_dim)); + disk_pq_table.inflate_vector((_u8 *) medoid_coords, + (centroid_data + cur_m * aligned_dim)); } aligned_free(medoid_buf); delete[] medoid_coords; @@ -596,7 +601,6 @@ namespace diskann { std::string centroids_file = std::string(disk_index_file) + "_centroids.bin"; - 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); @@ -613,11 +617,14 @@ namespace diskann { } this->data_dim = pq_file_dim; - this->disk_data_dim = this->data_dim; // will reset later if we use PQ on disk - this->disk_bytes_per_point = this->data_dim * sizeof(T); // will change later if we use PQ on disk or if we are using inner product without PQ + this->disk_data_dim = + this->data_dim; // will reset later if we use PQ on disk + this->disk_bytes_per_point = + this->data_dim * + sizeof(T); // will change later if we use PQ on disk or if we are using + // inner product without PQ this->aligned_dim = ROUND_UP(pq_file_dim, 8); - size_t npts_u64, nchunks_u64; #ifdef EXEC_ENV_OLS diskann::load_bin<_u8>(files, pq_compressed_vectors, this->data, npts_u64, @@ -642,19 +649,25 @@ namespace diskann { << " #aligned_dim: " << aligned_dim << " #chunks: " << n_chunks << std::endl; - -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 - 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 + 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 + 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 #else - 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); // giving 0 chunks to make the pq_table infer from the + // chunk_offsets file the correct value #endif - disk_pq_n_chunks = disk_pq_table.get_num_chunks(); - disk_bytes_per_point = disk_pq_n_chunks * sizeof(_u8); - std::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(_u8); + std::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 @@ -805,40 +818,40 @@ if (file_exists(disk_pq_pivots_path)) { template void PQFlashIndex::cached_beam_search(const T *query1, const _u64 k_search, const _u64 l_search, _u64 *indices, - float * distances, - const _u64 beam_width, - QueryStats * stats) { + float * distances, + const _u64 beam_width, + QueryStats *stats) { ThreadData data = this->thread_data.pop(); while (data.scratch.sector_scratch == nullptr) { this->thread_data.wait_for_push_notify(); data = this->thread_data.pop(); } + // copy query to thread specific aligned and allocated memory (for distance + // calculations we need aligned data) -// copy query to thread specific aligned and allocated memory (for distance calculations we need aligned data) - - float query_norm = 0; + float query_norm = 0; const T * query = data.scratch.aligned_query_T; const float *query_float = data.scratch.aligned_query_float; for (uint32_t i = 0; i < this->data_dim; i++) { data.scratch.aligned_query_float[i] = query1[i]; data.scratch.aligned_query_T[i] = query1[i]; - query_norm += query1[i]*query1[i]; + query_norm += query1[i] * query1[i]; } -// 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) { - query_norm = std::sqrt(query_norm); - data.scratch.aligned_query_T[this->data_dim -1] = 0; - data.scratch.aligned_query_float[this->data_dim -1] = 0; - for (uint32_t i = 0; i < this->data_dim - 1; i++) { - data.scratch.aligned_query_T[i] /= query_norm; - data.scratch.aligned_query_float[i] /= query_norm; - } + // 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) { + query_norm = std::sqrt(query_norm); + data.scratch.aligned_query_T[this->data_dim - 1] = 0; + data.scratch.aligned_query_float[this->data_dim - 1] = 0; + for (uint32_t i = 0; i < this->data_dim - 1; i++) { + data.scratch.aligned_query_T[i] /= query_norm; + data.scratch.aligned_query_float[i] /= query_norm; + } } - IOContext &ctx = data.ctx; auto query_scratch = &(data.scratch); @@ -867,8 +880,9 @@ if (file_exists(disk_pq_pivots_path)) { _u8 * pq_coord_scratch = query_scratch->aligned_pq_coord_scratch; // lambda to batch compute query<-> node distances in PQ space - auto compute_dists = [this, pq_coord_scratch, pq_dists]( - const unsigned *ids, const _u64 n_ids, float *dists_out) { + auto compute_dists = [this, pq_coord_scratch, pq_dists](const unsigned *ids, + const _u64 n_ids, + float *dists_out) { ::aggregate_coords(ids, n_ids, this->data, this->n_chunks, pq_coord_scratch); ::pq_dist_lookup(pq_coord_scratch, n_ids, this->n_chunks, pq_dists, @@ -912,16 +926,15 @@ if (file_exists(disk_pq_pivots_path)) { unsigned k = 0; // cleared every iteration - std::vector frontier; + std::vector frontier; std::vector> frontier_nhoods; - std::vector frontier_read_reqs; + std::vector frontier_read_reqs; std::vector>> cached_nhoods; while (k < cur_list_size) { auto nk = cur_list_size; - // clear iteration state frontier.clear(); frontier_nhoods.clear(); @@ -963,7 +976,7 @@ if (file_exists(disk_pq_pivots_path)) { if (stats != nullptr) stats->n_hops++; for (_u64 i = 0; i < frontier.size(); i++) { - auto id = frontier[i]; + auto id = frontier[i]; std::pair<_u32, char *> fnhood; fnhood.first = id; fnhood.second = sector_scratch + sector_scratch_idx * SECTOR_LEN; @@ -996,13 +1009,14 @@ if (file_exists(disk_pq_pivots_path)) { float cur_expanded_dist; if (!use_disk_index_pq) { cur_expanded_dist = dist_cmp->compare(query, node_fp_coords_copy, - (unsigned) aligned_dim); - } - else { + (unsigned) aligned_dim); + } 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, (_u8 *) 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, (_u8 *) node_fp_coords_copy); } full_retset.push_back( Neighbor((unsigned) cached_nhood.first, cur_expanded_dist, true)); @@ -1076,13 +1090,14 @@ if (file_exists(disk_pq_pivots_path)) { float cur_expanded_dist; if (!use_disk_index_pq) { cur_expanded_dist = dist_cmp->compare(query, node_fp_coords_copy, - (unsigned) aligned_dim); - } - else { + (unsigned) aligned_dim); + } 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, (_u8 *) 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, (_u8 *) node_fp_coords_copy); } full_retset.push_back( Neighbor(frontier_nhood.first, cur_expanded_dist, true)); @@ -1144,27 +1159,27 @@ if (file_exists(disk_pq_pivots_path)) { hops++; } - // re-sort by distance std::sort(full_retset.begin(), full_retset.end(), [](const Neighbor &left, const Neighbor &right) { return left.distance < right.distance; }); -/* - std::cout<<"return set: \n"; - for (auto &x : full_retset) - std::cout<& read_reqs, - IOContext& ctx, bool async) { + IOContext& ctx, bool async) { using namespace std::chrono_literals; // execute each request sequentially _u64 n_reqs = read_reqs.size(); diff --git a/tests/build_disk_index.cpp b/tests/build_disk_index.cpp index ed0b71b8e8..a83d43dd9f 100644 --- a/tests/build_disk_index.cpp +++ b/tests/build_disk_index.cpp @@ -12,27 +12,31 @@ template bool build_index(const char* dataFilePath, const char* indexFilePath, const char* indexBuildParameters, diskann::Metric metric) { - return diskann::build_disk_index( - dataFilePath, indexFilePath, indexBuildParameters, metric); + return diskann::build_disk_index(dataFilePath, indexFilePath, + indexBuildParameters, metric); } int main(int argc, char** argv) { if (argc != 11) { std::cout << "Usage: " << argv[0] - << " [data_type] [dist_fn: l2/mips] [data_file.bin] " + << " [data_type] [dist_fn: l2/mips] " + "[data_file.bin] " "[index_prefix_path] " - "[R] [L] [B] [M] [T] [PQ_disk_bytes (for very large dimensionality, use 0 for full vectors)]. See README for more information on " + "[R] [L] [B] [M] [T] [PQ_disk_bytes (for very large " + "dimensionality, use 0 for full vectors)]. See README for " + "more information on " "parameters." << std::endl; } else { diskann::Metric metric = diskann::Metric::L2; if (std::string(argv[2]) == std::string("mips")) - metric = diskann::Metric::INNER_PRODUCT; - + metric = diskann::Metric::INNER_PRODUCT; + std::string params = std::string(argv[5]) + " " + std::string(argv[6]) + " " + std::string(argv[7]) + " " + - std::string(argv[8]) + " " + std::string(argv[9]) + " " + std::string(argv[10]); + std::string(argv[8]) + " " + std::string(argv[9]) + + " " + std::string(argv[10]); if (std::string(argv[1]) == std::string("float")) build_index(argv[3], argv[4], params.c_str(), metric); else if (std::string(argv[1]) == std::string("int8")) diff --git a/tests/build_memory_index.cpp b/tests/build_memory_index.cpp index 8faf39dcc8..bfa50c849d 100644 --- a/tests/build_memory_index.cpp +++ b/tests/build_memory_index.cpp @@ -16,7 +16,8 @@ #include "memory_mapper.h" template -int build_in_memory_index(const std::string& data_path, const diskann::Metric &metric, const unsigned R, +int build_in_memory_index(const std::string& data_path, + const diskann::Metric& metric, const unsigned R, const unsigned L, const float alpha, const std::string& save_path, const unsigned num_threads) { @@ -55,16 +56,18 @@ int main(int argc, char** argv) { _u32 ctr = 2; -diskann::Metric metric; - if (std::string(argv[ctr]) == std::string("mips")) - metric = diskann::Metric::INNER_PRODUCT; + diskann::Metric metric; + if (std::string(argv[ctr]) == std::string("mips")) + metric = diskann::Metric::INNER_PRODUCT; else if (std::string(argv[ctr]) == std::string("l2")) - metric = diskann::Metric::L2; + metric = diskann::Metric::L2; else { - std::cout<<"Unsupported distance function. Currently only L2/ Inner Product support." << std::endl; + std::cout << "Unsupported distance function. Currently only L2/ Inner " + "Product support." + << std::endl; return -1; } -ctr++; + ctr++; const std::string data_path(argv[ctr++]); const std::string save_path(argv[ctr++]); diff --git a/tests/search_disk_index.cpp b/tests/search_disk_index.cpp index 241968a7cf..9dacbc394f 100644 --- a/tests/search_disk_index.cpp +++ b/tests/search_disk_index.cpp @@ -56,20 +56,24 @@ int search_disk_index(int argc, char** argv) { size_t query_num, query_dim, query_aligned_dim, gt_num, gt_dim; std::vector<_u64> Lvec; - _u32 ctr = 2; + _u32 ctr = 2; diskann::Metric metric; - if (std::string(argv[ctr]) == std::string("mips")) - metric = diskann::Metric::INNER_PRODUCT; + if (std::string(argv[ctr]) == std::string("mips")) + metric = diskann::Metric::INNER_PRODUCT; else if (std::string(argv[ctr]) == std::string("l2")) - metric = diskann::Metric::L2; + metric = diskann::Metric::L2; else { - std::cout<<"Unsupported distance function. Currently only L2/ Inner Product support." << std::endl; + std::cout << "Unsupported distance function. Currently only L2/ Inner " + "Product support." + << std::endl; return -1; } - - if ((std::string(argv[1]) != std::string("float")) && (metric == diskann::Metric::INNER_PRODUCT)) { - std::cout<<"Currently support only floating point data for Inner Product." << std::endl; + + if ((std::string(argv[1]) != std::string("float")) && + (metric == diskann::Metric::INNER_PRODUCT)) { + std::cout << "Currently support only floating point data for Inner Product." + << std::endl; return -1; } @@ -145,14 +149,13 @@ int search_disk_index(int argc, char** argv) { std::vector node_list; diskann::cout << "Caching " << num_nodes_to_cache << " BFS nodes around medoid(s)" << std::endl; - _pFlashIndex->cache_bfs_levels(num_nodes_to_cache, node_list); -// _pFlashIndex->generate_cache_list_from_sample_queries( -// warmup_query_file, 15, 6, num_nodes_to_cache, num_threads, node_list); + _pFlashIndex->cache_bfs_levels(num_nodes_to_cache, node_list); + // _pFlashIndex->generate_cache_list_from_sample_queries( + // warmup_query_file, 15, 6, num_nodes_to_cache, num_threads, node_list); _pFlashIndex->load_cache_list(node_list); node_list.clear(); node_list.shrink_to_fit(); - omp_set_num_threads(num_threads); uint64_t warmup_L = 20; @@ -216,7 +219,6 @@ int search_disk_index(int argc, char** argv) { uint32_t optimized_beamwidth = 2; - for (uint32_t test_id = 0; test_id < Lvec.size(); test_id++) { _u64 L = Lvec[test_id]; @@ -233,10 +235,9 @@ int search_disk_index(int argc, char** argv) { diskann::QueryStats* 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) +#pragma omp parallel for schedule(dynamic, 1) for (_s64 i = 0; i < (int64_t) query_num; i++) { _pFlashIndex->cached_beam_search( query + (i * query_aligned_dim), recall_at, L, @@ -309,7 +310,8 @@ int main(int argc, char** argv) { if (argc < 12) { diskann::cout << "Usage: " << argv[0] - << " [index_type] [dist_fn] [index_prefix_path] " + << " [index_type] [dist_fn] " + "[index_prefix_path] " " [num_nodes_to_cache] [num_threads] [beamwidth (use 0 to " "optimize internally)] " " [query_file.bin] [truthset.bin (use \"null\" for none)] " diff --git a/tests/search_memory_index.cpp b/tests/search_memory_index.cpp index 64fbcd6909..e9951eb6de 100644 --- a/tests/search_memory_index.cpp +++ b/tests/search_memory_index.cpp @@ -27,29 +27,32 @@ int search_memory_index(int argc, char** argv) { size_t query_num, query_dim, query_aligned_dim, gt_num, gt_dim; std::vector<_u64> Lvec; - _u32 ctr = 2; + _u32 ctr = 2; diskann::Metric metric; - if (std::string(argv[ctr]) == std::string("mips")) - metric = diskann::Metric::INNER_PRODUCT; + if (std::string(argv[ctr]) == std::string("mips")) + metric = diskann::Metric::INNER_PRODUCT; else if (std::string(argv[ctr]) == std::string("l2")) - metric = diskann::Metric::L2; + metric = diskann::Metric::L2; else if (std::string(argv[ctr]) == std::string("fast_l2")) metric = diskann::Metric::FAST_L2; else { - std::cout<<"Unsupported distance function. Currently only L2/ Inner Product/FAST_L2 support." << std::endl; + std::cout << "Unsupported distance function. Currently only L2/ Inner " + "Product/FAST_L2 support." + << std::endl; return -1; } -ctr++; + ctr++; if ((std::string(argv[1]) != std::string("float")) && - ((metric == diskann::Metric::INNER_PRODUCT) || (metric == diskann::Metric::FAST_L2))) { - std::cout << "Error. Inner product and Fast_L2 search currently only supported for " + ((metric == diskann::Metric::INNER_PRODUCT) || + (metric == diskann::Metric::FAST_L2))) { + std::cout << "Error. Inner product and Fast_L2 search currently only " + "supported for " "floating point datatypes." << std::endl; } - std::string data_file(argv[ctr++]); std::string memory_index_file(argv[ctr++]); _u64 num_threads = std::atoi(argv[ctr++]); @@ -57,8 +60,7 @@ ctr++; std::string truthset_bin(argv[ctr++]); _u64 recall_at = std::atoi(argv[ctr++]); std::string result_output_prefix(argv[ctr++]); -// bool use_optimized_search = std::atoi(argv[ctr++]); - + // bool use_optimized_search = std::atoi(argv[ctr++]); bool calc_recall_flag = false; @@ -89,7 +91,6 @@ ctr++; std::cout.setf(std::ios_base::fixed, std::ios_base::floatfield); std::cout.precision(2); - diskann::Index index(metric, data_file.c_str()); index.load(memory_index_file.c_str()); // to load NSG std::cout << "Index loaded" << std::endl; @@ -174,7 +175,8 @@ int main(int argc, char** argv) { if (argc < 11) { std::cout << "Usage: " << argv[0] - << " [index_type] [dist_fn (l2/mips/fast_l2)] [data_file.bin] " + << " [index_type] [dist_fn (l2/mips/fast_l2)] " + "[data_file.bin] " "[memory_index_path] [num_threads] " "[query_file.bin] [truthset.bin (use \"null\" for none)] " " [K] [result_output_prefix]" diff --git a/tests/test_incremental_index.cpp b/tests/test_incremental_index.cpp index 8b98233b28..533ca3a167 100644 --- a/tests/test_incremental_index.cpp +++ b/tests/test_incremental_index.cpp @@ -49,7 +49,7 @@ int main(int argc, char** argv) { paras.Set("saturate_graph", false); paras.Set("num_rnds", num_rnds); - typedef int TagT; + typedef int TagT; diskann::Index index(diskann::L2, argv[1], num_points, num_points - num_incr, num_frozen, true, true, true); diff --git a/tests/utils/bin_to_tsv.cpp b/tests/utils/bin_to_tsv.cpp index 37874e2437..9a7e180a42 100644 --- a/tests/utils/bin_to_tsv.cpp +++ b/tests/utils/bin_to_tsv.cpp @@ -22,7 +22,8 @@ void block_convert(std::ofstream& writer, std::ifstream& reader, T* read_buf, int main(int argc, char** argv) { if (argc != 4) { - std::cout << argv[0] << " input_bin output_tsv" << std::endl; + std::cout << argv[0] << " input_bin output_tsv" + << std::endl; exit(-1); } std::string type_string(argv[1]); @@ -50,9 +51,10 @@ int main(int argc, char** argv) { for (_u64 i = 0; i < nblks; i++) { _u64 cblk_size = std::min(npts - i * blk_size, blk_size); if (type_string == std::string("float")) - block_convert(writer, reader, (float*)read_buf, cblk_size, ndims); + block_convert(writer, reader, (float*) read_buf, cblk_size, ndims); else if (type_string == std::string("int8")) - block_convert(writer, reader, (int8_t*) read_buf, cblk_size, ndims); + block_convert(writer, reader, (int8_t*) read_buf, cblk_size, + ndims); else if (type_string == std::string("uint8")) block_convert(writer, reader, (uint8_t*) read_buf, cblk_size, ndims); diff --git a/tests/utils/compute_groundtruth.cpp b/tests/utils/compute_groundtruth.cpp index de46cab904..d2e1ed2778 100644 --- a/tests/utils/compute_groundtruth.cpp +++ b/tests/utils/compute_groundtruth.cpp @@ -31,7 +31,8 @@ #define ALIGNMENT 512 void command_line_help() { - std::cerr << "./compute_groundtruth " + " " " " << std::endl; @@ -107,8 +108,8 @@ void distsq_to_points( 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, + 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; @@ -120,7 +121,7 @@ void inner_prod_to_points( 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; } @@ -163,9 +164,8 @@ void exact_knn(const size_t dim, const size_t k, 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); + 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; @@ -340,8 +340,6 @@ int aux_main(char **argc) { delete[] closest_points_part; delete[] dist_closest_points_part; - - diskann::aligned_free(base_data); } diff --git a/tests/utils/create_disk_layout.cpp b/tests/utils/create_disk_layout.cpp index 21c6cacede..d0e8712835 100644 --- a/tests/utils/create_disk_layout.cpp +++ b/tests/utils/create_disk_layout.cpp @@ -22,9 +22,10 @@ int create_disk_layout(char **argv) { } int main(int argc, char **argv) { - if (argc != 5) { - std::cout << argv[0] << " data_type data_bin " - "vamana_index_file output_diskann_index_file" + if (argc != 5) { + std::cout << argv[0] + << " data_type data_bin " + "vamana_index_file output_diskann_index_file" << std::endl; exit(-1); } diff --git a/tests/utils/float_bin_to_int8.cpp b/tests/utils/float_bin_to_int8.cpp index 4f422a2336..0620730a51 100644 --- a/tests/utils/float_bin_to_int8.cpp +++ b/tests/utils/float_bin_to_int8.cpp @@ -4,7 +4,6 @@ #include #include "utils.h" - void block_convert(std::ofstream& writer, int8_t* write_buf, std::ifstream& reader, float* read_buf, _u64 npts, _u64 ndims, float bias, float scale) { diff --git a/tests/utils/gen_random_slice.cpp b/tests/utils/gen_random_slice.cpp index dccb50a13c..c91e9d44e4 100644 --- a/tests/utils/gen_random_slice.cpp +++ b/tests/utils/gen_random_slice.cpp @@ -22,7 +22,6 @@ 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])); @@ -31,10 +30,10 @@ int aux_main(char** argv) { } 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" + if (argc != 5) { + std::cout << argv[0] + << " data_type [float/int8/uint8] base_bin_file " + "sample_output_prefix sampling_probability" << std::endl; exit(-1); } diff --git a/tests/utils/partition_data.cpp b/tests/utils/partition_data.cpp index 63476b15e2..593bcb119f 100644 --- a/tests/utils/partition_data.cpp +++ b/tests/utils/partition_data.cpp @@ -11,9 +11,10 @@ int main(int argc, char** argv) { if (argc != 7) { std::cout << "Usage:\n" - << argv[0] << " datatype " - " " - " " + << argv[0] + << " datatype " + " " + " " << std::endl; exit(-1); } diff --git a/tests/utils/partition_with_ram_budget.cpp b/tests/utils/partition_with_ram_budget.cpp index fee3c84e82..6cd6a9401f 100644 --- a/tests/utils/partition_with_ram_budget.cpp +++ b/tests/utils/partition_with_ram_budget.cpp @@ -11,9 +11,10 @@ int main(int argc, char** argv) { if (argc != 8) { std::cout << "Usage:\n" - << argv[0] << " datatype " - " " - " " + << argv[0] + << " datatype " + " " + " " << std::endl; exit(-1); } diff --git a/tests/utils/tsv_to_bin.cpp b/tests/utils/tsv_to_bin.cpp index 111a6bb55d..776e063436 100644 --- a/tests/utils/tsv_to_bin.cpp +++ b/tests/utils/tsv_to_bin.cpp @@ -26,8 +26,9 @@ void block_convert(std::ifstream& reader, std::ofstream& writer, _u64 npts, int main(int argc, char** argv) { if (argc != 6) { std::cout << argv[0] - << " input_filename.tsv output_filename.bin dim num_pts>" - << std::endl; + << " input_filename.tsv output_filename.bin " + "dim num_pts>" + << std::endl; exit(-1); } @@ -50,7 +51,7 @@ int main(int argc, char** argv) { _u64 nblks = ROUND_UP(npts, blk_size) / blk_size; std::cout << "# blks: " << nblks << std::endl; std::ofstream writer(argv[3], std::ios::binary); - auto npts_s32 = (_u32) npts; + auto npts_s32 = (_u32) npts; auto ndims_s32 = (_u32) ndims; writer.write((char*) &npts_s32, sizeof(_u32)); writer.write((char*) &ndims_s32, sizeof(_u32)); diff --git a/tests/utils/uint8_to_float.cpp b/tests/utils/uint8_to_float.cpp index e383489fd4..6a6b3b2fe0 100644 --- a/tests/utils/uint8_to_float.cpp +++ b/tests/utils/uint8_to_float.cpp @@ -11,7 +11,7 @@ int main(int argc, char** argv) { } uint8_t* input; - size_t npts, nd; + 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); diff --git a/tests/utils/vector_analysis.cpp b/tests/utils/vector_analysis.cpp index e370db2774..6538c6823f 100644 --- a/tests/utils/vector_analysis.cpp +++ b/tests/utils/vector_analysis.cpp @@ -22,83 +22,85 @@ template int analyze_norm(std::string base_file) { - std::cout<<"Analyzing data norms" << std::endl; - T* data; + std::cout << "Analyzing data norms" << std::endl; + T* data; _u64 npts, ndims; diskann::load_bin(base_file, data, npts, ndims); std::vector norms(npts, 0); - #pragma omp parallel for schedule(dynamic) - for (_u32 i = 0; i int normalize_base(std::string base_file, std::string out_file) { - std::cout<<"Normalizing base" << std::endl; - T* data; + std::cout << "Normalizing base" << std::endl; + T* data; _u64 npts, ndims; diskann::load_bin(base_file, data, npts, ndims); -// std::vector norms(npts, 0); - #pragma omp parallel for schedule(dynamic) - for (_u32 i = 0; i norms(npts, 0); +#pragma omp parallel for schedule(dynamic) + for (_u32 i = 0; i < npts; i++) { float pt_norm = 0; - for (_u32 d = 0; d < ndims; d++) - pt_norm += data[i*ndims + d]* data[i* ndims + d]; + for (_u32 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++) - data[i*ndims + d] = data[i* ndims + d]/pt_norm; + for (_u32 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; +int augment_base(std::string base_file, std::string out_file, + bool prep_base = true) { + std::cout << "Analyzing data norms" << std::endl; + T* data; _u64 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 (_u32 i = 0; 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; + // std::sort(norms.begin(), norms.end()); + max_norm = std::sqrt(max_norm); + std::cout << "Max norm: " << max_norm << std::endl; + T* new_data; _u64 newdims = ndims + 1; - new_data = new T[npts*newdims]; - for (_u64 i = 0;i < npts; i++) { + new_data = new T[npts * newdims]; + for (_u64 i = 0; i < npts; i++) { if (prep_base) { - for (_u64 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)); + for (_u64 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< int aux_main(char** argv) { - std::string base_file(argv[2]); - _u32 option = atoi(argv[3]); - if (option == 1) - analyze_norm(base_file); - else if (option == 2) - augment_base(base_file, std::string(argv[4]), true); + _u32 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); + augment_base(base_file, std::string(argv[4]), false); else if (option == 4) - normalize_base(base_file, std::string(argv[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]" - << std::endl; + 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]" + << std::endl; exit(-1); } From f5426d88ca0d52c889a38108b0b849260e311a01 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Thu, 15 Jul 2021 10:41:56 +0000 Subject: [PATCH 68/84] clang introduced a bug in distance.h, fixed itt --- include/distance.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/include/distance.h b/include/distance.h index ddfa9bb3b4..086c4286dc 100644 --- a/include/distance.h +++ b/include/distance.h @@ -542,7 +542,7 @@ namespace diskann { using DistanceInnerProduct::compare; float compare(const T *a, const T *b, float norm, unsigned size) const { // not implement - floatresult = -2 * DistanceInnerProduct::inner_product(a, b, size); + float result = -2 * DistanceInnerProduct::inner_product(a, b, size); result += norm; return result; } From d5b528d2d5be064fae1842f59183bdfbae5a679b Mon Sep 17 00:00:00 2001 From: ravishankar Date: Thu, 15 Jul 2021 11:17:23 +0000 Subject: [PATCH 69/84] added unit tester partially --- tests/utils/compute_groundtruth.cpp | 7 +++++-- unit_tester.sh | 30 +++++++++++++++++++++++++++++ 2 files changed, 35 insertions(+), 2 deletions(-) create mode 100755 unit_tester.sh diff --git a/tests/utils/compute_groundtruth.cpp b/tests/utils/compute_groundtruth.cpp index d2e1ed2778..5992c827ed 100644 --- a/tests/utils/compute_groundtruth.cpp +++ b/tests/utils/compute_groundtruth.cpp @@ -34,7 +34,7 @@ void command_line_help() { std::cerr << "./compute_groundtruth " " " - " " + " " << std::endl; } @@ -306,8 +306,11 @@ int aux_main(char **argc) { std::string base_file(argc[2]); std::string query_file(argc[3]); size_t k = atoi(argc[4]); + bool use_mip = false; std::string gt_file(argc[5]); - bool use_mip = atoi(argc[6]); + if (std::string(argv[6]) == std::string("mips")) + use_mip = true; + float *base_data; float *query_data; diff --git a/unit_tester.sh b/unit_tester.sh new file mode 100755 index 0000000000..e441c5c213 --- /dev/null +++ b/unit_tester.sh @@ -0,0 +1,30 @@ +#!/bin/sh + +if [ "$#" -ne "2" ]; then + echo "usage: ./unit_test.sh [path_build_folder] [working_folder]" +else + +BUILD_FOLDER=${1} +WORK_FOLDER=${2} + +echo Running unit testing on various files, with build folder as ${BUILD_FOLDER} and working folder as ${WORK_FOLDER} +# download all unit test files + +#iterate over them and run the corresponding test +CATALOG1="${WORK_FOLDER}/catalog.txt" +CATALOG="${WORK_FOLDER}/catalog_formatted.txt" +sed -e '/^$/d' ${CATALOG1} > ${CATALOG} + + +while IFS= read -r line; do + BASE=$line + read -r QUERY + read -r TYPE + read -r METRIC + echo "Going to run test on ${BASE} base, ${QUERY} query, ${TYPE} datatype, ${METRIC} metric" + echo "Computing Groundtruth" + ${BUILD_FOLDER}/tests/utils/compute_groundtruth ${TYPE} + +done < "${CATALOG}" + +fi From b2078d6cfd6cda40791325df1842473de0fda101 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Thu, 15 Jul 2021 11:23:23 +0000 Subject: [PATCH 70/84] minor bugfix --- tests/utils/compute_groundtruth.cpp | 2 +- unit_tester.sh | 10 ++++++---- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/tests/utils/compute_groundtruth.cpp b/tests/utils/compute_groundtruth.cpp index 5992c827ed..068e28c546 100644 --- a/tests/utils/compute_groundtruth.cpp +++ b/tests/utils/compute_groundtruth.cpp @@ -308,7 +308,7 @@ int aux_main(char **argc) { size_t k = atoi(argc[4]); bool use_mip = false; std::string gt_file(argc[5]); - if (std::string(argv[6]) == std::string("mips")) + if (std::string(argc[6]) == std::string("mips")) use_mip = true; diff --git a/unit_tester.sh b/unit_tester.sh index e441c5c213..a730bc9496 100755 --- a/unit_tester.sh +++ b/unit_tester.sh @@ -17,13 +17,15 @@ sed -e '/^$/d' ${CATALOG1} > ${CATALOG} while IFS= read -r line; do - BASE=$line - read -r QUERY + BASE="${WORK_FOLDER}/${line}" + read -r line + QUERY="${WORK_FOLDER}/${line}" read -r TYPE read -r METRIC - echo "Going to run test on ${BASE} base, ${QUERY} query, ${TYPE} datatype, ${METRIC} metric" + GT="${WORK_FOLDER}/test_gt100" + echo "Going to run test on ${BASE} base, ${QUERY} query, ${TYPE} datatype, ${METRIC} metric, saving gt at ${GT}" echo "Computing Groundtruth" - ${BUILD_FOLDER}/tests/utils/compute_groundtruth ${TYPE} + ${BUILD_FOLDER}/tests/utils/compute_groundtruth ${TYPE} ${BASE} ${QUERY} 30 ${GT} ${METRIC} done < "${CATALOG}" From 361a4c7ea3bda2a97c497294cd3bbbd728e3ab6c Mon Sep 17 00:00:00 2001 From: ravishankar Date: Thu, 15 Jul 2021 13:49:12 +0000 Subject: [PATCH 71/84] finished unit tester --- unit_tester.sh | 30 ++++++++++++++++++++++++++---- 1 file changed, 26 insertions(+), 4 deletions(-) diff --git a/unit_tester.sh b/unit_tester.sh index a730bc9496..1dd0a895bd 100755 --- a/unit_tester.sh +++ b/unit_tester.sh @@ -17,16 +17,38 @@ sed -e '/^$/d' ${CATALOG1} > ${CATALOG} while IFS= read -r line; do - BASE="${WORK_FOLDER}/${line}" + DATASET=${line} + BASE="${WORK_FOLDER}/${DATASET}" read -r line QUERY="${WORK_FOLDER}/${line}" read -r TYPE read -r METRIC - GT="${WORK_FOLDER}/test_gt100" + GT="${WORK_FOLDER}/${DATASET}_gt30_${METRIC}" + MEM="${WORK_FOLDER}/${DATASET}_mem" + DISK="${WORK_FOLDER}/${DATASET}_disk" + MBLOG="${WORK_FOLDER}/${DATASET}_mb.log" + DBLOG="${WORK_FOLDER}/${DATASET}_db.log" + MSLOG="${WORK_FOLDER}/${DATASET}_ms.log" + DSLOG="${WORK_FOLDER}/${DATASET}_ds.log" echo "Going to run test on ${BASE} base, ${QUERY} query, ${TYPE} datatype, ${METRIC} metric, saving gt at ${GT}" echo "Computing Groundtruth" - ${BUILD_FOLDER}/tests/utils/compute_groundtruth ${TYPE} ${BASE} ${QUERY} 30 ${GT} ${METRIC} - + ${BUILD_FOLDER}/tests/utils/compute_groundtruth ${TYPE} ${BASE} ${QUERY} 30 ${GT} ${METRIC} > /dev/null + echo "Building Mem Index" + ${BUILD_FOLDER}/tests/build_memory_index ${TYPE} ${METRIC} ${BASE} ${MEM} 32 50 1.2 0 > ${MBLOG} + awk '/^Degree/' ${MBLOG} + awk '/^Indexing/' ${MBLOG} + echo "Building Disk Index" + ${BUILD_FOLDER}/tests/build_disk_index ${TYPE} ${METRIC} ${BASE} ${DISK} 32 50 0.003 0.001 32 0 > ${DBLOG} + awk '/^Compressing/' ${DBLOG} + echo "#shards in disk index" + awk '/^bin:/' ${DBLOG} + awk '/^Indexing/' ${DBLOG} + echo "Searching Mem Index" + ${BUILD_FOLDER}/tests/search_memory_index ${TYPE} ${METRIC} ${BASE} ${MEM} 16 ${QUERY} ${GT} 10 /tmp/res 10 20 30 40 50 60 70 80 90 100 > ${MSLOG} + awk '/===/{x=NR+10}(NR<=x){print}' ${MSLOG} + echo "Searching Disk Index" + ${BUILD_FOLDER}/tests/search_disk_index ${TYPE} ${METRIC} ${DISK} 10000 10 4 ${QUERY} ${GT} 10 /tmp/res 10 20 30 40 50 60 70 80 90 100 > ${DSLOG} + awk '/===/{x=NR+10}(NR<=x){print}' ${DSLOG} done < "${CATALOG}" fi From b70ca3c03c612a1d830cc5cbfaa47d985c4a7f35 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Thu, 15 Jul 2021 17:29:34 +0000 Subject: [PATCH 72/84] changed back training size to 100K for now, we can increase to 1M later if necessary --- include/aux_utils.h | 2 +- unit_tester.sh | 17 +++++++++-------- 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/include/aux_utils.h b/include/aux_utils.h index 1698b9ed54..d8af561526 100644 --- a/include/aux_utils.h +++ b/include/aux_utils.h @@ -32,7 +32,7 @@ typedef int FileHandle; #include "gperftools/malloc_extension.h" namespace diskann { - const size_t TRAINING_SET_SIZE = 1000000; + const size_t TRAINING_SET_SIZE = 100000; const double SPACE_FOR_CACHED_NODES_IN_GB = 0.25; const double THRESHOLD_FOR_CACHING_IN_GB = 1.0; const uint32_t NUM_NODES_TO_CACHE = 250000; diff --git a/unit_tester.sh b/unit_tester.sh index 1dd0a895bd..a9fc8b4838 100755 --- a/unit_tester.sh +++ b/unit_tester.sh @@ -15,6 +15,7 @@ CATALOG1="${WORK_FOLDER}/catalog.txt" CATALOG="${WORK_FOLDER}/catalog_formatted.txt" sed -e '/^$/d' ${CATALOG1} > ${CATALOG} +mkdir ${WORK_FOLDER}/indices while IFS= read -r line; do DATASET=${line} @@ -23,13 +24,13 @@ while IFS= read -r line; do QUERY="${WORK_FOLDER}/${line}" read -r TYPE read -r METRIC - GT="${WORK_FOLDER}/${DATASET}_gt30_${METRIC}" - MEM="${WORK_FOLDER}/${DATASET}_mem" - DISK="${WORK_FOLDER}/${DATASET}_disk" - MBLOG="${WORK_FOLDER}/${DATASET}_mb.log" - DBLOG="${WORK_FOLDER}/${DATASET}_db.log" - MSLOG="${WORK_FOLDER}/${DATASET}_ms.log" - DSLOG="${WORK_FOLDER}/${DATASET}_ds.log" + GT="${WORK_FOLDER}/indices/${DATASET}_gt30_${METRIC}" + MEM="${WORK_FOLDER}/indices/${DATASET}_mem" + DISK="${WORK_FOLDER}/indices/${DATASET}_disk" + MBLOG="${WORK_FOLDER}/indices/${DATASET}_mb.log" + DBLOG="${WORK_FOLDER}/indices/${DATASET}_db.log" + MSLOG="${WORK_FOLDER}/indices/${DATASET}_ms.log" + DSLOG="${WORK_FOLDER}/indices/${DATASET}_ds.log" echo "Going to run test on ${BASE} base, ${QUERY} query, ${TYPE} datatype, ${METRIC} metric, saving gt at ${GT}" echo "Computing Groundtruth" ${BUILD_FOLDER}/tests/utils/compute_groundtruth ${TYPE} ${BASE} ${QUERY} 30 ${GT} ${METRIC} > /dev/null @@ -38,7 +39,7 @@ while IFS= read -r line; do awk '/^Degree/' ${MBLOG} awk '/^Indexing/' ${MBLOG} echo "Building Disk Index" - ${BUILD_FOLDER}/tests/build_disk_index ${TYPE} ${METRIC} ${BASE} ${DISK} 32 50 0.003 0.001 32 0 > ${DBLOG} + ${BUILD_FOLDER}/tests/build_disk_index ${TYPE} ${METRIC} ${BASE} ${DISK} 32 50 0.03 0.01 32 0 > ${DBLOG} awk '/^Compressing/' ${DBLOG} echo "#shards in disk index" awk '/^bin:/' ${DBLOG} From 48cc376805344adbdf4ae3696a53fe9cada603ac Mon Sep 17 00:00:00 2001 From: ravishankar Date: Thu, 15 Jul 2021 18:36:13 +0000 Subject: [PATCH 73/84] added comments for unit_tester.sh --- src/aux_utils.cpp | 2 +- unit_tester.sh | 47 ++++++++++++++++++++++++++++------------------- 2 files changed, 29 insertions(+), 20 deletions(-) diff --git a/src/aux_utils.cpp b/src/aux_utils.cpp index e1f8205796..a6930f8866 100644 --- a/src/aux_utils.cpp +++ b/src/aux_utils.cpp @@ -797,7 +797,7 @@ namespace diskann { gen_random_slice(data_file_to_use.c_str(), sample_base_prefix, sample_sampling_rate); - // std::remove(mem_index_path.c_str()); + std::remove(mem_index_path.c_str()); if (use_disk_pq) std::remove(disk_pq_compressed_vectors_path.c_str()); diff --git a/unit_tester.sh b/unit_tester.sh index a9fc8b4838..b6d615418e 100755 --- a/unit_tester.sh +++ b/unit_tester.sh @@ -1,36 +1,46 @@ #!/bin/sh -if [ "$#" -ne "2" ]; then - echo "usage: ./unit_test.sh [path_build_folder] [working_folder]" +# Performs build and search test on disk and memory indices (parameters are tuned for 100K-1M sized datasets) +# All indices and logs will be stored in working_folder after run is complete +# To run, create a catalog text file consisting of the following entries +# For each dataset, specify the following 5 lines, in a line by line format, and then move on to next dataset +# dataset_name[used for save file names] +# /path/to/base.bin +# /path/to/query.bin +# data_type[float/uint8/int8] +# metric[l2/mips] + + +if [ "$#" -ne "3" ]; then + echo "usage: ./unit_test.sh [build_folder_path] [catalog] [working_folder]" else BUILD_FOLDER=${1} -WORK_FOLDER=${2} +CATALOG1=${2} +WORK_FOLDER=${3} +mkdir ${WORK_FOLDER} +CATALOG="${WORK_FOLDER}/catalog_formatted.txt" +sed -e '/^$/d' ${CATALOG1} > ${CATALOG} echo Running unit testing on various files, with build folder as ${BUILD_FOLDER} and working folder as ${WORK_FOLDER} # download all unit test files #iterate over them and run the corresponding test -CATALOG1="${WORK_FOLDER}/catalog.txt" -CATALOG="${WORK_FOLDER}/catalog_formatted.txt" -sed -e '/^$/d' ${CATALOG1} > ${CATALOG} -mkdir ${WORK_FOLDER}/indices while IFS= read -r line; do DATASET=${line} - BASE="${WORK_FOLDER}/${DATASET}" - read -r line - QUERY="${WORK_FOLDER}/${line}" + read -r BASE + read -r QUERY read -r TYPE read -r METRIC - GT="${WORK_FOLDER}/indices/${DATASET}_gt30_${METRIC}" - MEM="${WORK_FOLDER}/indices/${DATASET}_mem" - DISK="${WORK_FOLDER}/indices/${DATASET}_disk" - MBLOG="${WORK_FOLDER}/indices/${DATASET}_mb.log" - DBLOG="${WORK_FOLDER}/indices/${DATASET}_db.log" - MSLOG="${WORK_FOLDER}/indices/${DATASET}_ms.log" - DSLOG="${WORK_FOLDER}/indices/${DATASET}_ds.log" + GT="${WORK_FOLDER}/${DATASET}_gt30_${METRIC}" + MEM="${WORK_FOLDER}/${DATASET}_mem" + DISK="${WORK_FOLDER}/${DATASET}_disk" + MBLOG="${WORK_FOLDER}/${DATASET}_mb.log" + DBLOG="${WORK_FOLDER}/${DATASET}_db.log" + MSLOG="${WORK_FOLDER}/${DATASET}_ms.log" + DSLOG="${WORK_FOLDER}/${DATASET}_ds.log" echo "Going to run test on ${BASE} base, ${QUERY} query, ${TYPE} datatype, ${METRIC} metric, saving gt at ${GT}" echo "Computing Groundtruth" ${BUILD_FOLDER}/tests/utils/compute_groundtruth ${TYPE} ${BASE} ${QUERY} 30 ${GT} ${METRIC} > /dev/null @@ -39,7 +49,7 @@ while IFS= read -r line; do awk '/^Degree/' ${MBLOG} awk '/^Indexing/' ${MBLOG} echo "Building Disk Index" - ${BUILD_FOLDER}/tests/build_disk_index ${TYPE} ${METRIC} ${BASE} ${DISK} 32 50 0.03 0.01 32 0 > ${DBLOG} + ${BUILD_FOLDER}/tests/build_disk_index ${TYPE} ${METRIC} ${BASE} ${DISK} 32 50 0.03 0.03 32 0 > ${DBLOG} awk '/^Compressing/' ${DBLOG} echo "#shards in disk index" awk '/^bin:/' ${DBLOG} @@ -51,5 +61,4 @@ while IFS= read -r line; do ${BUILD_FOLDER}/tests/search_disk_index ${TYPE} ${METRIC} ${DISK} 10000 10 4 ${QUERY} ${GT} 10 /tmp/res 10 20 30 40 50 60 70 80 90 100 > ${DSLOG} awk '/===/{x=NR+10}(NR<=x){print}' ${DSLOG} done < "${CATALOG}" - fi From e16e41180bf260abfdf6caa90f670af10504735f Mon Sep 17 00:00:00 2001 From: ravishankar Date: Fri, 16 Jul 2021 04:54:44 +0000 Subject: [PATCH 74/84] added auto tuning parameters for unit tester --- src/linux_aligned_file_reader.cpp | 2 +- unit_tester.sh | 22 +++++++++++++--------- 2 files changed, 14 insertions(+), 10 deletions(-) diff --git a/src/linux_aligned_file_reader.cpp b/src/linux_aligned_file_reader.cpp index bb95201d41..8388886757 100644 --- a/src/linux_aligned_file_reader.cpp +++ b/src/linux_aligned_file_reader.cpp @@ -141,7 +141,7 @@ void LinuxAlignedFileReader::register_thread() { std::cerr << "io_setup() failed; returned " << ret << ", errno=" << errno << ":" << ::strerror(errno) << std::endl; } else { - std::cerr << "allocating ctx: " << ctx << " to thread-id:" << my_id + diskann::cout << "allocating ctx: " << ctx << " to thread-id:" << my_id << std::endl; ctx_map[my_id] = ctx; } diff --git a/unit_tester.sh b/unit_tester.sh index b6d615418e..80eb37802b 100755 --- a/unit_tester.sh +++ b/unit_tester.sh @@ -1,5 +1,4 @@ -#!/bin/sh - +#!/bin/bash # Performs build and search test on disk and memory indices (parameters are tuned for 100K-1M sized datasets) # All indices and logs will be stored in working_folder after run is complete # To run, create a catalog text file consisting of the following entries @@ -9,8 +8,6 @@ # /path/to/query.bin # data_type[float/uint8/int8] # metric[l2/mips] - - if [ "$#" -ne "3" ]; then echo "usage: ./unit_test.sh [build_folder_path] [catalog] [working_folder]" else @@ -41,22 +38,29 @@ while IFS= read -r line; do DBLOG="${WORK_FOLDER}/${DATASET}_db.log" MSLOG="${WORK_FOLDER}/${DATASET}_ms.log" DSLOG="${WORK_FOLDER}/${DATASET}_ds.log" + + FILESIZE=`wc -c "${BASE}" | awk '{print $1}'` + BUDGETBUILD=`bc <<< "scale=2; ${FILESIZE}/(5*1024*1024*1024)"` + BUDGETSERVE=`bc <<< "scale=2; ${FILESIZE}/(10*1024*1024*1024)"` + echo "Going to build with ${BUDGETBUILD} GiB RAM and serve with ${BUDGETSERVE} GiB RAM" + rm ${DISK}_* + echo "Going to run test on ${BASE} base, ${QUERY} query, ${TYPE} datatype, ${METRIC} metric, saving gt at ${GT}" echo "Computing Groundtruth" ${BUILD_FOLDER}/tests/utils/compute_groundtruth ${TYPE} ${BASE} ${QUERY} 30 ${GT} ${METRIC} > /dev/null echo "Building Mem Index" - ${BUILD_FOLDER}/tests/build_memory_index ${TYPE} ${METRIC} ${BASE} ${MEM} 32 50 1.2 0 > ${MBLOG} + /usr/bin/time ${BUILD_FOLDER}/tests/build_memory_index ${TYPE} ${METRIC} ${BASE} ${MEM} 32 50 1.2 0 > ${MBLOG} awk '/^Degree/' ${MBLOG} awk '/^Indexing/' ${MBLOG} + echo "Searching Mem Index" + ${BUILD_FOLDER}/tests/search_memory_index ${TYPE} ${METRIC} ${BASE} ${MEM} 16 ${QUERY} ${GT} 10 /tmp/res 10 20 30 40 50 60 70 80 90 100 > ${MSLOG} + awk '/===/{x=NR+10}(NR<=x){print}' ${MSLOG} echo "Building Disk Index" - ${BUILD_FOLDER}/tests/build_disk_index ${TYPE} ${METRIC} ${BASE} ${DISK} 32 50 0.03 0.03 32 0 > ${DBLOG} + ${BUILD_FOLDER}/tests/build_disk_index ${TYPE} ${METRIC} ${BASE} ${DISK} 32 50 ${BUDGETSERVE} ${BUDGETBUILD} 32 0 > ${DBLOG} awk '/^Compressing/' ${DBLOG} echo "#shards in disk index" awk '/^bin:/' ${DBLOG} awk '/^Indexing/' ${DBLOG} - echo "Searching Mem Index" - ${BUILD_FOLDER}/tests/search_memory_index ${TYPE} ${METRIC} ${BASE} ${MEM} 16 ${QUERY} ${GT} 10 /tmp/res 10 20 30 40 50 60 70 80 90 100 > ${MSLOG} - awk '/===/{x=NR+10}(NR<=x){print}' ${MSLOG} echo "Searching Disk Index" ${BUILD_FOLDER}/tests/search_disk_index ${TYPE} ${METRIC} ${DISK} 10000 10 4 ${QUERY} ${GT} 10 /tmp/res 10 20 30 40 50 60 70 80 90 100 > ${DSLOG} awk '/===/{x=NR+10}(NR<=x){print}' ${DSLOG} From c169ea36dc82500f0956bf604b2b2007532480aa Mon Sep 17 00:00:00 2001 From: ravishankar Date: Fri, 16 Jul 2021 05:35:52 +0000 Subject: [PATCH 75/84] re-ran clang formatting --- include/pq_flash_index.h | 3 +++ include/utils.h | 5 +++-- src/aux_utils.cpp | 5 ++++- src/linux_aligned_file_reader.cpp | 2 +- src/pq_flash_index.cpp | 23 ++++++++++++++++++++--- tests/utils/compute_groundtruth.cpp | 5 ++--- unit_tester.sh | 5 +++-- 7 files changed, 36 insertions(+), 12 deletions(-) diff --git a/include/pq_flash_index.h b/include/pq_flash_index.h index 4d9cb3a7f7..182b8d8f6d 100644 --- a/include/pq_flash_index.h +++ b/include/pq_flash_index.h @@ -130,6 +130,9 @@ namespace diskann { _u64 max_node_len = 0, nnodes_per_sector = 0, max_degree = 0; diskann::Metric metric = diskann::Metric::L2; + float max_base_norm = + 0; // used only for inner product search to re-scale the result value + // (due to the pre-processing of base during index build) // data info _u64 num_points = 0; _u64 data_dim = 0; diff --git a/include/utils.h b/include/utils.h index 86c5f1fd3c..d1d56ce33d 100644 --- a/include/utils.h +++ b/include/utils.h @@ -425,8 +425,8 @@ namespace diskann { // Product Search" by Neyshabur and Srebro template - void prepare_base_for_inner_products(const std::string in_file, - const std::string out_file) { + 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); @@ -496,6 +496,7 @@ namespace diskann { block_pts * out_dims * sizeof(float)); } out_writer.close(); + return max_norm; } // plain saves data as npts X ndims array into filename diff --git a/src/aux_utils.cpp b/src/aux_utils.cpp index a6930f8866..ddccacf81c 100644 --- a/src/aux_utils.cpp +++ b/src/aux_utils.cpp @@ -686,7 +686,10 @@ namespace diskann { << std::endl; std::string prepped_base = index_prefix_path + "_prepped_base.bin"; data_file_to_use = prepped_base; - diskann::prepare_base_for_inner_products(base_file, 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); } unsigned R = (unsigned) atoi(param_list[0].c_str()); diff --git a/src/linux_aligned_file_reader.cpp b/src/linux_aligned_file_reader.cpp index 8388886757..35c8009bd0 100644 --- a/src/linux_aligned_file_reader.cpp +++ b/src/linux_aligned_file_reader.cpp @@ -142,7 +142,7 @@ void LinuxAlignedFileReader::register_thread() { << ":" << ::strerror(errno) << std::endl; } else { diskann::cout << "allocating ctx: " << ctx << " to thread-id:" << my_id - << std::endl; + << std::endl; ctx_map[my_id] = ctx; } lk.unlock(); diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index d3b5219147..a941045453 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -793,6 +793,17 @@ namespace diskann { 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) { + _u64 dumr, dumc; + float *norm_val; + diskann::load_bin(norm_file, norm_val, dumr, dumc); + this->max_base_norm = norm_val[0]; + std::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; } @@ -1177,9 +1188,15 @@ namespace diskann { indices[i] = full_retset[i].id; if (distances != nullptr) { distances[i] = full_retset[i].distance; - if (metric == diskann::Metric::INNER_PRODUCT) // flip the sign from - // convert min to max - distances[i] = -distances[i]; + if (metric == diskann::Metric::INNER_PRODUCT) { // flip the sign from + // convert min to max + distances[i] = (-distances[i]); + if (max_base_norm != 0) + distances[i] *= (max_base_norm * + query_norm); // rescale to revert back to original + // norms (cancelling the effect of + // base and query pre-processing) + } } } diff --git a/tests/utils/compute_groundtruth.cpp b/tests/utils/compute_groundtruth.cpp index 068e28c546..6e331530af 100644 --- a/tests/utils/compute_groundtruth.cpp +++ b/tests/utils/compute_groundtruth.cpp @@ -306,11 +306,10 @@ int aux_main(char **argc) { std::string base_file(argc[2]); std::string query_file(argc[3]); size_t k = atoi(argc[4]); - bool use_mip = false; + bool use_mip = false; std::string gt_file(argc[5]); if (std::string(argc[6]) == std::string("mips")) - use_mip = true; - + use_mip = true; float *base_data; float *query_data; diff --git a/unit_tester.sh b/unit_tester.sh index 80eb37802b..557883fda4 100755 --- a/unit_tester.sh +++ b/unit_tester.sh @@ -59,10 +59,11 @@ while IFS= read -r line; do ${BUILD_FOLDER}/tests/build_disk_index ${TYPE} ${METRIC} ${BASE} ${DISK} 32 50 ${BUDGETSERVE} ${BUDGETBUILD} 32 0 > ${DBLOG} awk '/^Compressing/' ${DBLOG} echo "#shards in disk index" - awk '/^bin:/' ${DBLOG} awk '/^Indexing/' ${DBLOG} echo "Searching Disk Index" - ${BUILD_FOLDER}/tests/search_disk_index ${TYPE} ${METRIC} ${DISK} 10000 10 4 ${QUERY} ${GT} 10 /tmp/res 10 20 30 40 50 60 70 80 90 100 > ${DSLOG} + ${BUILD_FOLDER}/tests/search_disk_index ${TYPE} ${METRIC} ${DISK} 10000 10 4 ${QUERY} ${GT} 10 /tmp/res 100 > ${DSLOG} + echo "# shards used during index construction:" + awk '/medoids/{x=NR+1}(NR<=x){print}' ${DSLOG} awk '/===/{x=NR+10}(NR<=x){print}' ${DSLOG} done < "${CATALOG}" fi From c42f3a1828cb2b13809f7358eb49aab2d0ce6114 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Fri, 16 Jul 2021 05:37:14 +0000 Subject: [PATCH 76/84] small change to unit tester --- unit_tester.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/unit_tester.sh b/unit_tester.sh index 557883fda4..cf26d1672f 100755 --- a/unit_tester.sh +++ b/unit_tester.sh @@ -61,7 +61,7 @@ while IFS= read -r line; do echo "#shards in disk index" awk '/^Indexing/' ${DBLOG} echo "Searching Disk Index" - ${BUILD_FOLDER}/tests/search_disk_index ${TYPE} ${METRIC} ${DISK} 10000 10 4 ${QUERY} ${GT} 10 /tmp/res 100 > ${DSLOG} + ${BUILD_FOLDER}/tests/search_disk_index ${TYPE} ${METRIC} ${DISK} 10000 10 4 ${QUERY} ${GT} 10 /tmp/res 20 40 60 80 100 > ${DSLOG} echo "# shards used during index construction:" awk '/medoids/{x=NR+1}(NR<=x){print}' ${DSLOG} awk '/===/{x=NR+10}(NR<=x){print}' ${DSLOG} From 3a30d70318557055af37d854341c10d0b9309d15 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Fri, 16 Jul 2021 09:21:16 +0000 Subject: [PATCH 77/84] fixed minor bug in unit tester --- unit_tester.sh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/unit_tester.sh b/unit_tester.sh index cf26d1672f..082247057f 100755 --- a/unit_tester.sh +++ b/unit_tester.sh @@ -40,8 +40,8 @@ while IFS= read -r line; do DSLOG="${WORK_FOLDER}/${DATASET}_ds.log" FILESIZE=`wc -c "${BASE}" | awk '{print $1}'` - BUDGETBUILD=`bc <<< "scale=2; ${FILESIZE}/(5*1024*1024*1024)"` - BUDGETSERVE=`bc <<< "scale=2; ${FILESIZE}/(10*1024*1024*1024)"` + BUDGETBUILD=`bc <<< "scale=4; ${FILESIZE}/(5*1024*1024*1024)"` + BUDGETSERVE=`bc <<< "scale=4; ${FILESIZE}/(10*1024*1024*1024)"` echo "Going to build with ${BUDGETBUILD} GiB RAM and serve with ${BUDGETSERVE} GiB RAM" rm ${DISK}_* From c7f39c99de7c2afcbb6dafcbdf44805df833421a Mon Sep 17 00:00:00 2001 From: ravishankar Date: Fri, 16 Jul 2021 09:29:20 +0000 Subject: [PATCH 78/84] fixed some formatting on unit tester --- unit_tester.sh | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/unit_tester.sh b/unit_tester.sh index 082247057f..c4e987a42d 100755 --- a/unit_tester.sh +++ b/unit_tester.sh @@ -40,12 +40,14 @@ while IFS= read -r line; do DSLOG="${WORK_FOLDER}/${DATASET}_ds.log" FILESIZE=`wc -c "${BASE}" | awk '{print $1}'` - BUDGETBUILD=`bc <<< "scale=4; ${FILESIZE}/(5*1024*1024*1024)"` - BUDGETSERVE=`bc <<< "scale=4; ${FILESIZE}/(10*1024*1024*1024)"` - echo "Going to build with ${BUDGETBUILD} GiB RAM and serve with ${BUDGETSERVE} GiB RAM" + BUDGETBUILD=`bc <<< "scale=4; 0.0001 + ${FILESIZE}/(5*1024*1024*1024)"` + BUDGETSERVE=`bc <<< "scale=4; 0.0001 + ${FILESIZE}/(10*1024*1024*1024)"` + echo "=============================================================================================================================================" + echo "Running tests on ${DATASET} dataset, ${TYPE} datatype, $METRIC metric, ${BUDGETBUILD} GiB and ${BUDGETSERVE} GiB build and serve budget" + echo "=============================================================================================================================================" rm ${DISK}_* - echo "Going to run test on ${BASE} base, ${QUERY} query, ${TYPE} datatype, ${METRIC} metric, saving gt at ${GT}" + #echo "Going to run test on ${BASE} base, ${QUERY} query, ${TYPE} datatype, ${METRIC} metric, saving gt at ${GT}" echo "Computing Groundtruth" ${BUILD_FOLDER}/tests/utils/compute_groundtruth ${TYPE} ${BASE} ${QUERY} 30 ${GT} ${METRIC} > /dev/null echo "Building Mem Index" From 18ad0ee61c84b45b945f75e4c51a4b7e29062c24 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Fri, 23 Jul 2021 10:35:41 +0000 Subject: [PATCH 79/84] started code for range search support in pq_flash_index --- include/pq_flash_index.h | 8 +++++++- src/pq_flash_index.cpp | 36 ++++++++++++++++++++++++++++++++++++ 2 files changed, 43 insertions(+), 1 deletion(-) diff --git a/include/pq_flash_index.h b/include/pq_flash_index.h index 182b8d8f6d..f5276969a6 100644 --- a/include/pq_flash_index.h +++ b/include/pq_flash_index.h @@ -114,8 +114,14 @@ namespace diskann { DISKANN_DLLEXPORT void cached_beam_search( const T *query, const _u64 k_search, const _u64 l_search, _u64 *res_ids, float *res_dists, const _u64 beam_width, QueryStats *stats = nullptr); - std::shared_ptr &reader; + + DISKANN_DLLEXPORT void range_search(const T *query1, const double range, + const _u64 l_search, std::vector<_u64> &results, + const _u64 beam_width, + QueryStats *stats = nullptr); + + std::shared_ptr &reader; protected: DISKANN_DLLEXPORT void use_medoids_data_as_centroids(); DISKANN_DLLEXPORT void setup_thread_data(_u64 nthreads); diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index a941045453..3d22e7ddd5 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -1208,6 +1208,42 @@ namespace diskann { } } + template + void PQFlashIndex::range_search(const T *query1, const double range, + const _u64 l_search, std::vector<_u64> &results, + const _u64 beam_width, + QueryStats *stats) { + +tsl::robin_set<_u32> return_ids; +_u32 cur_l = l_search; +std::vector<_u64> cur_results; +std::vector cur_dists; +bool stop_flag = false; +while(!stop_flag) { + cur_results.clear(); + cur_dists.clear(); +cur_results.resize(cur_l); +cur_dists.resize(cur_l); +for (auto &x : cur_dists) +x = std::numeric_limits::max(); +this->cached_beam_search(query1, cur_l, cur_l, cur_results.data(), cur_dists.data(), beam_width, stats); +for (_u32 i = 0; i < cur_l; i++) { + if (cur_dists[i] <= (float) range) { + return_ids.insert(cur_results[i]); + } +} +if (cur_dists[cur_l -1] > (float) range) { + stop_flag = true; +} else { + cur_l *= 2; +} +} +results.clear(); +results.reserve(return_ids.size()); +for (auto &x: return_ids) +results.emplace_back(x); +} + #ifdef EXEC_ENV_OLS template char *PQFlashIndex::getHeaderBytes() { From 42d0c00abcb08a50cf00fcc2bec043146ddbf299 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Fri, 23 Jul 2021 11:17:08 +0000 Subject: [PATCH 80/84] added more code for range search in disk index --- include/aux_utils.h | 3 + include/pq_flash_index.h | 2 +- include/utils.h | 50 +++++ src/aux_utils.cpp | 25 +++ src/pq_flash_index.cpp | 2 +- tests/CMakeLists.txt | 10 + tests/range_search_disk_index.cpp | 311 ++++++++++++++++++++++++++++++ 7 files changed, 401 insertions(+), 2 deletions(-) create mode 100644 tests/range_search_disk_index.cpp diff --git a/include/aux_utils.h b/include/aux_utils.h index d8af561526..f980b7fc5d 100644 --- a/include/aux_utils.h +++ b/include/aux_utils.h @@ -46,6 +46,9 @@ namespace diskann { 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_range_search_recall(unsigned num_queries, std::vector> &groundtruth, + std::vector> &our_results); + DISKANN_DLLEXPORT void read_idmap(const std::string & fname, std::vector &ivecs); diff --git a/include/pq_flash_index.h b/include/pq_flash_index.h index f5276969a6..3a521173b0 100644 --- a/include/pq_flash_index.h +++ b/include/pq_flash_index.h @@ -117,7 +117,7 @@ namespace diskann { DISKANN_DLLEXPORT void range_search(const T *query1, const double range, - const _u64 l_search, std::vector<_u64> &results, + const _u64 l_search, std::vector<_u32> &results, const _u64 beam_width, QueryStats *stats = nullptr); diff --git a/include/utils.h b/include/utils.h index d1d56ce33d..00000d6a7f 100644 --- a/include/utils.h +++ b/include/utils.h @@ -295,6 +295,56 @@ namespace diskann { } } + inline void load_range_truthset(const std::string& bin_file, std::vector> &groundtruth, _u64 & gt_num) { + _u64 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_u32, total_u32; + reader.read((char*) &npts_u32, sizeof(int)); + reader.read((char*) &total_u32, sizeof(int)); + + gt_num = (_u64) npts_u32; + _u64 total_res = (_u64) total_u32; + + 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); + + 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<_u32> gt_count(gt_num); + + reader.read((char*) gt_count.data(), sizeof(_u32)*gt_num); + + for (_u32 i = 0; i < gt_num; i++) { + groundtruth[i].clear(); + groundtruth[i].resize(gt_count[i]); + reader.read((char*) groundtruth[i].data(), sizeof(_u32)*gt_count[i]); + +// debugging code +/* if (i < 10) { + std::cout< inline void load_bin(MemoryMappedFiles& files, const std::string& bin_file, diff --git a/src/aux_utils.cpp b/src/aux_utils.cpp index ddccacf81c..f143b88ea6 100644 --- a/src/aux_utils.cpp +++ b/src/aux_utils.cpp @@ -70,6 +70,31 @@ namespace diskann { return total_recall / (num_queries) * (100.0 / recall_at); } + double calculate_range_search_recall(unsigned 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()); + unsigned 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; + } + return total_recall / (num_queries); + } + template T *generateRandomWarmup(uint64_t warmup_num, uint64_t warmup_dim, uint64_t warmup_aligned_dim) { diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index 3d22e7ddd5..9f39ed715a 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -1210,7 +1210,7 @@ namespace diskann { template void PQFlashIndex::range_search(const T *query1, const double range, - const _u64 l_search, std::vector<_u64> &results, + const _u64 l_search, std::vector<_u32> &results, const _u64 beam_width, QueryStats *stats) { diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index fa6e8ea045..177d34f70b 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -41,3 +41,13 @@ else() endif() +add_executable(range_search_disk_index range_search_disk_index.cpp + ${PROJECT_SOURCE_DIR}/src/aux_utils.cpp ) +if(MSVC) + target_link_options(range_search_disk_index PRIVATE /MACHINE:x64 /DEBUG:FULL) + target_link_libraries(range_search_disk_index debug ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}/diskann_dll.lib) + target_link_libraries(range_search_disk_index optimized ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}/diskann_dll.lib) +else() + target_link_libraries(range_search_disk_index ${PROJECT_NAME} aio -ltcmalloc) +endif() + diff --git a/tests/range_search_disk_index.cpp b/tests/range_search_disk_index.cpp new file mode 100644 index 0000000000..6f82a6bf12 --- /dev/null +++ b/tests/range_search_disk_index.cpp @@ -0,0 +1,311 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "aux_utils.h" +#include "index.h" +#include "math_utils.h" +#include "memory_mapper.h" +#include "partition_and_pq.h" +#include "timer.h" +#include "utils.h" + +#ifndef _WINDOWS +#include +#include +#include +#include "linux_aligned_file_reader.h" +#else +#ifdef USE_BING_INFRA +#include "bing_aligned_file_reader.h" +#else +#include "windows_aligned_file_reader.h" +#endif +#endif + +#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; +} + +template +int search_disk_index(int argc, char** argv) { + // load query bin + T* query = nullptr; +// unsigned* gt_ids = nullptr; +// float* gt_dists = nullptr; +std::vector> groundtruth_ids; + size_t query_num, query_dim, query_aligned_dim, gt_num; + std::vector<_u64> Lvec; + + _u32 ctr = 2; + diskann::Metric metric; + + if (std::string(argv[ctr]) == std::string("mips")) + metric = diskann::Metric::INNER_PRODUCT; + else if (std::string(argv[ctr]) == std::string("l2")) + metric = diskann::Metric::L2; + else { + std::cout << "Unsupported distance function. Currently only L2/ Inner " + "Product support." + << std::endl; + return -1; + } + + if ((std::string(argv[1]) != std::string("float")) && + (metric == diskann::Metric::INNER_PRODUCT)) { + std::cout << "Currently support only floating point data for Inner Product." + << std::endl; + return -1; + } + + ctr++; + + std::string index_prefix_path(argv[ctr++]); + std::string pq_prefix = index_prefix_path + "_pq"; + std::string disk_index_file = index_prefix_path + "_disk.index"; + std::string warmup_query_file = index_prefix_path + "_sample_data.bin"; + _u64 num_nodes_to_cache = std::atoi(argv[ctr++]); + _u32 num_threads = std::atoi(argv[ctr++]); + _u32 beamwidth = std::atoi(argv[ctr++]); + std::string query_bin(argv[ctr++]); + std::string truthset_bin(argv[ctr++]); + double search_range = std::atof(argv[ctr++]); + std::string result_output_prefix(argv[ctr++]); + + bool calc_recall_flag = false; + + for (; ctr < (_u32) argc; ctr++) { + _u64 curL = std::atoi(argv[ctr]); + Lvec.push_back(curL); + } + + if (Lvec.size() == 0) { + diskann::cout + << "No valid Lsearch found." + << std::endl; + return -1; + } + + 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; + + diskann::load_aligned_bin(query_bin, query, query_num, query_dim, + query_aligned_dim); + + if (file_exists(truthset_bin)) { + diskann::load_range_truthset(truthset_bin, groundtruth_ids, gt_num); + 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; +#ifdef _WINDOWS +#ifndef USE_BING_INFRA + reader.reset(new WindowsAlignedFileReader()); +#else + reader.reset(new diskann::BingAlignedFileReader()); +#endif +#else + reader.reset(new LinuxAlignedFileReader()); +#endif + + std::unique_ptr> _pFlashIndex( + new diskann::PQFlashIndex(reader, metric)); + + int res = _pFlashIndex->load(num_threads, pq_prefix.c_str(), + disk_index_file.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)((_u32) 150000, (_u32) 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 (_s64 i = 0; i < (int64_t) warmup_num; i++) { + _pFlashIndex->cached_beam_search(warmup + (i * warmup_aligned_dim), 1, + warmup_L, + warmup_result_ids_64.data() + (i * 1), + warmup_result_dists.data() + (i * 1), 4); + } + 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()); + + uint32_t optimized_beamwidth = 2; + + for (uint32_t test_id = 0; test_id < Lvec.size(); test_id++) { + _u64 L = Lvec[test_id]; + + 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].clear(); + query_result_ids[test_id].resize(query_num); + + diskann::QueryStats* stats = new diskann::QueryStats[query_num]; + + 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++) { + _pFlashIndex->range_search( + query + (i * query_aligned_dim), search_range, L, + query_result_ids[test_id][i], + optimized_beamwidth, stats + i); + } + auto e = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = e - s; + float qps = (1.0 * query_num) / (1.0 * diff.count()); + + float mean_latency = diskann::get_mean_stats( + stats, query_num, + [](const diskann::QueryStats& stats) { return stats.total_us; }); + + float latency_999 = diskann::get_percentile_stats( + stats, query_num, 0.999, + [](const diskann::QueryStats& stats) { return stats.total_us; }); + + float 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; + if (calc_recall_flag) { + recall = diskann::calculate_range_search_recall(query_num, groundtruth_ids, query_result_ids[test_id]); + } + + 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; + } + + 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) { + if (argc < 12) { + diskann::cout + << "Usage: " << argv[0] + << " [index_type] [dist_fn] " + "[index_prefix_path] " + " [num_nodes_to_cache] [num_threads] [beamwidth (use 0 to " + "optimize internally)] " + " [query_file.bin] [truthset.bin (use \"null\" for none)] " + " [range_threshold] [result_output_prefix] " + " [L1] [L2] etc. See README for more information on parameters." + << std::endl; + exit(-1); + } + if (std::string(argv[1]) == std::string("float")) + search_disk_index(argc, argv); + else if (std::string(argv[1]) == std::string("int8")) + search_disk_index(argc, argv); + else if (std::string(argv[1]) == std::string("uint8")) + search_disk_index(argc, argv); + else + diskann::cout << "Unsupported index type. Use float or int8 or uint8" + << std::endl; +} From f21883d4e9b0a41e42788e917d32476937010126 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Sun, 1 Aug 2021 16:25:39 +0000 Subject: [PATCH 81/84] added range search support --- include/pq_flash_index.h | 4 +-- include/utils.h | 14 ++++++++++ src/pq_flash_index.cpp | 46 +++++++++++-------------------- tests/range_search_disk_index.cpp | 15 ++++++++-- 4 files changed, 45 insertions(+), 34 deletions(-) diff --git a/include/pq_flash_index.h b/include/pq_flash_index.h index 3a521173b0..d119b8197a 100644 --- a/include/pq_flash_index.h +++ b/include/pq_flash_index.h @@ -116,8 +116,8 @@ namespace diskann { float *res_dists, const _u64 beam_width, QueryStats *stats = nullptr); - DISKANN_DLLEXPORT void range_search(const T *query1, const double range, - const _u64 l_search, std::vector<_u32> &results, + DISKANN_DLLEXPORT _u32 range_search(const T *query1, const double range, + const _u64 l_search, _u64* indices, float* distances, const _u64 beam_width, QueryStats *stats = nullptr); diff --git a/include/utils.h b/include/utils.h index 00000d6a7f..94940ef983 100644 --- a/include/utils.h +++ b/include/utils.h @@ -327,12 +327,26 @@ namespace diskann { groundtruth.clear(); groundtruth.resize(gt_num); std::vector<_u32> gt_count(gt_num); + + reader.read((char*) gt_count.data(), sizeof(_u32)*gt_num); + std::vector<_u32> 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) + std::cout << "percentile " << p << ": " + << gt_stats[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++) { 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]); // debugging code diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index 9f39ed715a..89a23064b3 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -838,6 +838,7 @@ namespace diskann { data = this->thread_data.pop(); } +//std::cout< - void PQFlashIndex::range_search(const T *query1, const double range, - const _u64 l_search, std::vector<_u32> &results, + _u32 PQFlashIndex::range_search(const T *query1, const double range, + const _u64 l_search, _u64* indices, float* distances, const _u64 beam_width, QueryStats *stats) { - -tsl::robin_set<_u32> return_ids; -_u32 cur_l = l_search; -std::vector<_u64> cur_results; -std::vector cur_dists; -bool stop_flag = false; -while(!stop_flag) { - cur_results.clear(); - cur_dists.clear(); -cur_results.resize(cur_l); -cur_dists.resize(cur_l); -for (auto &x : cur_dists) -x = std::numeric_limits::max(); -this->cached_beam_search(query1, cur_l, cur_l, cur_results.data(), cur_dists.data(), beam_width, stats); -for (_u32 i = 0; i < cur_l; i++) { - if (cur_dists[i] <= (float) range) { - return_ids.insert(cur_results[i]); +_u32 res_count = 0; +this->cached_beam_search(query1, l_search, l_search, indices, distances, beam_width, stats); +for (_u32 i = 0; i < l_search; i++) { + if (distances[i] > (float) range) { + res_count = i; + break; } } -if (cur_dists[cur_l -1] > (float) range) { - stop_flag = true; -} else { - cur_l *= 2; -} -} -results.clear(); -results.reserve(return_ids.size()); -for (auto &x: return_ids) -results.emplace_back(x); + return res_count; } #ifdef EXEC_ENV_OLS diff --git a/tests/range_search_disk_index.cpp b/tests/range_search_disk_index.cpp index 6f82a6bf12..df091c9983 100644 --- a/tests/range_search_disk_index.cpp +++ b/tests/range_search_disk_index.cpp @@ -215,11 +215,17 @@ std::vector> groundtruth_ids; << std::endl; std::vector>> query_result_ids(Lvec.size()); - + std::vector<_u64> indices; + std::vector distances; + uint32_t optimized_beamwidth = 2; for (uint32_t test_id = 0; test_id < Lvec.size(); test_id++) { _u64 L = Lvec[test_id]; + indices.clear(); + distances.clear(); + indices.resize(L*query_num); + distances.resize(L*query_num); if (beamwidth <= 0) { // diskann::cout<<"Tuning beamwidth.." << std::endl; @@ -237,10 +243,15 @@ std::vector> groundtruth_ids; 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++) { + _u32 res_count = _pFlashIndex->range_search( query + (i * query_aligned_dim), search_range, L, - query_result_ids[test_id][i], + indices.data() + i*L, distances.data() + i *L, 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++) + query_result_ids[test_id][i][idx] = indices[i*L + idx]; } auto e = std::chrono::high_resolution_clock::now(); std::chrono::duration diff = e - s; From 197d271ef226817b10cc6908e07ba5b1e52aacd8 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Sun, 1 Aug 2021 17:56:34 +0000 Subject: [PATCH 82/84] tested range search on small dataset --- include/utils.h | 66 +++++++++++++++++++++++++++++++ src/pq_flash_index.cpp | 6 ++- tests/range_search_disk_index.cpp | 4 +- 3 files changed, 74 insertions(+), 2 deletions(-) diff --git a/include/utils.h b/include/utils.h index 94940ef983..9bf2e14f39 100644 --- a/include/utils.h +++ b/include/utils.h @@ -295,6 +295,72 @@ namespace diskann { } } + inline void prune_truthset_for_range(const std::string& bin_file, float range, std::vector> &groundtruth, + size_t& npts) { + _u64 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; + _u64 dim = (unsigned) dim_i32; + _u32* 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 (_u32 i = 0; i < npts; i++) { + groundtruth[i].clear(); + for (_u32 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, _u64 & gt_num) { _u64 read_blk_size = 64 * 1024 * 1024; cached_ifstream reader(bin_file, read_blk_size); diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index 89a23064b3..9c1bea8ab3 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -1222,11 +1222,15 @@ namespace diskann { _u32 res_count = 0; this->cached_beam_search(query1, l_search, l_search, indices, distances, beam_width, stats); for (_u32 i = 0; i < l_search; i++) { + //std::cout< (float) range) { res_count = i; break; - } + } else if (i == l_search -1) + res_count = l_search; } +//std::cout<<"\n\n"<> groundtruth_ids; query_aligned_dim); if (file_exists(truthset_bin)) { - diskann::load_range_truthset(truthset_bin, groundtruth_ids, gt_num); + diskann::load_range_truthset(truthset_bin, groundtruth_ids, gt_num); // use for range search type of truthset +// diskann::prune_truthset_for_range(truthset_bin, 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" @@ -248,6 +249,7 @@ std::vector> groundtruth_ids; query + (i * query_aligned_dim), search_range, L, indices.data() + i*L, distances.data() + i *L, optimized_beamwidth, stats + i); + // std::cout< Date: Tue, 10 Aug 2021 11:48:56 -0700 Subject: [PATCH 83/84] Update memory_mapper.h --- include/memory_mapper.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/include/memory_mapper.h b/include/memory_mapper.h index a0ec974b98..4ccbf6f286 100644 --- a/include/memory_mapper.h +++ b/include/memory_mapper.h @@ -38,4 +38,4 @@ namespace diskann { ~MemoryMapper(); }; -} // namespace diskann \ No newline at end of file +} // namespace diskann From 90d08be9f1a0e50f1678abd2003f12d13e7078c7 Mon Sep 17 00:00:00 2001 From: Harsha Vardhan Simhadri Date: Wed, 11 Aug 2021 16:34:09 -0700 Subject: [PATCH 84/84] minor edits --- include/pq_table.h | 1 - tests/utils/vector_analysis.cpp | 3 ++- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/include/pq_table.h b/include/pq_table.h index 9525c2f5c2..84fe1501af 100644 --- a/include/pq_table.h +++ b/include/pq_table.h @@ -6,7 +6,6 @@ #include "utils.h" namespace diskann { - // template class FixedChunkPQTable { // data_dim = n_chunks * chunk_size; float* tables = diff --git a/tests/utils/vector_analysis.cpp b/tests/utils/vector_analysis.cpp index 6538c6823f..af75ff7721 100644 --- a/tests/utils/vector_analysis.cpp +++ b/tests/utils/vector_analysis.cpp @@ -31,6 +31,7 @@ int analyze_norm(std::string base_file) { for (_u32 i = 0; i < npts; i++) { for (_u32 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) @@ -130,7 +131,7 @@ int main(int argc, char** argv) { << 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]" + "3-prep_query_for_mip, 4-normalize-vecs] [out_file for options 2/3/4]" << std::endl; exit(-1); }