From 027f6e77f0868c7a6fb1dd5cfacb55b742a0865b Mon Sep 17 00:00:00 2001 From: Gopal Srinivasa Date: Thu, 15 Jun 2023 00:06:38 +0530 Subject: [PATCH 01/15] Added PQ distance hierarchy Changes to CMakelists PQDataStore version that builds correctly Clang-format --- include/abstract_data_store.h | 17 +- include/abstract_scratch.h | 34 ++++ include/in_mem_data_store.h | 3 +- include/index.h | 9 +- include/pq.h | 75 +++++---- include/pq_common.h | 25 +++ include/pq_data_store.h | 87 +++++++++++ include/pq_distance.h | 20 +++ include/pq_l2_distance.h | 87 +++++++++++ include/pq_scratch.h | 43 +++++ include/quantized_distance.h | 56 +++++++ include/scratch.h | 23 +-- src/CMakeLists.txt | 2 +- src/abstract_data_store.cpp | 13 +- src/disk_utils.cpp | 16 +- src/dll/CMakeLists.txt | 4 +- src/in_mem_data_store.cpp | 8 +- src/index.cpp | 164 ++++++++++++------- src/pq_data_store.cpp | 249 +++++++++++++++++++++++++++++ src/pq_flash_index.cpp | 10 +- src/pq_l2_distance.cpp | 285 ++++++++++++++++++++++++++++++++++ src/scratch.cpp | 25 +-- 22 files changed, 1113 insertions(+), 142 deletions(-) create mode 100644 include/abstract_scratch.h create mode 100644 include/pq_common.h create mode 100644 include/pq_data_store.h create mode 100644 include/pq_distance.h create mode 100644 include/pq_l2_distance.h create mode 100644 include/pq_scratch.h create mode 100644 include/quantized_distance.h create mode 100644 src/pq_data_store.cpp create mode 100644 src/pq_l2_distance.cpp diff --git a/include/abstract_data_store.h b/include/abstract_data_store.h index 2e0266814..88542d974 100644 --- a/include/abstract_data_store.h +++ b/include/abstract_data_store.h @@ -13,6 +13,8 @@ namespace diskann { +template struct AbstractScratch; + template class AbstractDataStore { public: @@ -76,11 +78,20 @@ template class AbstractDataStore // num_points) to zero virtual void copy_vectors(const location_t from_loc, const location_t to_loc, const location_t num_points) = 0; - // metric specific operations - + // Some datastores like PQ stores need a preprocessing step before querying. + // Optionally, a scratch object can be passed in to avoid memory allocations + // Default implementation does nothing. + // REFACTOR TODO: Currently, we take an aligned_query as parameter, but this + // should change and this function should do the necessary alignment. + virtual void preprocess_query(const data_t *aligned_query, AbstractScratch *query_scratch = nullptr) const; + // distance functions. virtual float get_distance(const data_t *query, const location_t loc) const = 0; virtual void get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, - float *distances) const = 0; + float *distances, AbstractScratch *scratch_space = nullptr) const = 0; + // Specific overload for index.cpp. + // REFACTOR TODO: Check if the default implementation is sufficient for most cases. + virtual void get_distance(const data_t *preprocessed_query, const std::vector &ids, + std::vector &distances, AbstractScratch *scratch_space) const; virtual float get_distance(const location_t loc1, const location_t loc2) const = 0; // stats of the data stored in store diff --git a/include/abstract_scratch.h b/include/abstract_scratch.h new file mode 100644 index 000000000..57a2f8e34 --- /dev/null +++ b/include/abstract_scratch.h @@ -0,0 +1,34 @@ +#pragma once +namespace diskann +{ + +template struct PQScratch; + +// By somewhat more than a coincidence, it seems that both InMemQueryScratch +// and SSDQueryScratch have the aligned query and PQScratch objects. So we +// can put them in a neat hierarchy and keep PQScratch as a standalone class. +template struct AbstractScratch +{ + AbstractScratch() = default; + // This class does not take any responsibilty for memory management of + // its members. It is the responsibility of the derived classes to do so. + virtual ~AbstractScratch(){}; + + // Scratch objects should not be copied + AbstractScratch(const AbstractScratch &) = delete; + AbstractScratch &operator=(const AbstractScratch &) = delete; + + data_t *aligned_query_T() + { + return _aligned_query_T; + } + PQScratch *pq_scratch() + { + return _pq_scratch; + } + + protected: + data_t *_aligned_query_T = nullptr; + PQScratch *_pq_scratch = nullptr; +}; +} // namespace diskann diff --git a/include/in_mem_data_store.h b/include/in_mem_data_store.h index e9372c43a..8fe467f65 100644 --- a/include/in_mem_data_store.h +++ b/include/in_mem_data_store.h @@ -45,8 +45,9 @@ template class InMemDataStore : public AbstractDataStore *scratch) const override; virtual location_t calculate_medoid() const override; diff --git a/include/index.h b/include/index.h index b6738df52..33ebc3ae8 100644 --- a/include/index.h +++ b/include/index.h @@ -21,6 +21,10 @@ #include "in_mem_data_store.h" #include "abstract_index.h" +// REFACTOR +#include "quantized_distance.h" +#include "pq_data_store.h" + #define OVERHEAD_FACTOR 1.1 #define EXPAND_IF_FULL 0 #define DEFAULT_MAXC 750 @@ -393,7 +397,10 @@ template clas bool _pq_dist = false; bool _use_opq = false; size_t _num_pq_chunks = 0; - uint8_t *_pq_data = nullptr; + // REFACTOR + // uint8_t *_pq_data = nullptr; + std::shared_ptr> _pq_distance_fn = nullptr; + std::shared_ptr> _pq_data_store = nullptr; bool _pq_generated = false; FixedChunkPQTable _pq_table; diff --git a/include/pq.h b/include/pq.h index acfa1b30a..8d8fafcf8 100644 --- a/include/pq.h +++ b/include/pq.h @@ -4,13 +4,7 @@ #pragma once #include "utils.h" - -#define NUM_PQ_BITS 8 -#define NUM_PQ_CENTROIDS (1 << NUM_PQ_BITS) -#define MAX_OPQ_ITERS 20 -#define NUM_KMEANS_REPS_PQ 12 -#define MAX_PQ_TRAINING_SET_SIZE 256000 -#define MAX_PQ_CHUNKS 512 +#include "pq_common.h" namespace diskann { @@ -53,39 +47,40 @@ class FixedChunkPQTable void populate_chunk_inner_products(const float *query_vec, float *dist_vec); }; -template struct PQScratch -{ - float *aligned_pqtable_dist_scratch = nullptr; // MUST BE AT LEAST [256 * NCHUNKS] - float *aligned_dist_scratch = nullptr; // MUST BE AT LEAST diskann MAX_DEGREE - uint8_t *aligned_pq_coord_scratch = nullptr; // MUST BE AT LEAST [N_CHUNKS * MAX_DEGREE] - float *rotated_query = nullptr; - float *aligned_query_float = nullptr; - - PQScratch(size_t graph_degree, size_t aligned_dim) - { - diskann::alloc_aligned((void **)&aligned_pq_coord_scratch, - (size_t)graph_degree * (size_t)MAX_PQ_CHUNKS * sizeof(uint8_t), 256); - diskann::alloc_aligned((void **)&aligned_pqtable_dist_scratch, 256 * (size_t)MAX_PQ_CHUNKS * sizeof(float), - 256); - diskann::alloc_aligned((void **)&aligned_dist_scratch, (size_t)graph_degree * sizeof(float), 256); - diskann::alloc_aligned((void **)&aligned_query_float, aligned_dim * sizeof(float), 8 * sizeof(float)); - diskann::alloc_aligned((void **)&rotated_query, aligned_dim * sizeof(float), 8 * sizeof(float)); - - memset(aligned_query_float, 0, aligned_dim * sizeof(float)); - memset(rotated_query, 0, aligned_dim * sizeof(float)); - } - - void set(size_t dim, T *query, const float norm = 1.0f) - { - for (size_t d = 0; d < dim; ++d) - { - if (norm != 1.0f) - rotated_query[d] = aligned_query_float[d] = static_cast(query[d]) / norm; - else - rotated_query[d] = aligned_query_float[d] = static_cast(query[d]); - } - } -}; +// REFACTOR +// template struct PQScratch +//{ +// float *aligned_pqtable_dist_scratch = nullptr; // MUST BE AT LEAST [256 * NCHUNKS] +// float *aligned_dist_scratch = nullptr; // MUST BE AT LEAST diskann MAX_DEGREE +// uint8_t *aligned_pq_coord_scratch = nullptr; // MUST BE AT LEAST [N_CHUNKS * MAX_DEGREE] +// float *rotated_query = nullptr; +// float *aligned_query_float = nullptr; +// +// PQScratch(size_t graph_degree, size_t aligned_dim) +// { +// diskann::alloc_aligned((void **)&aligned_pq_coord_scratch, +// (size_t)graph_degree * (size_t)MAX_PQ_CHUNKS * sizeof(uint8_t), 256); +// diskann::alloc_aligned((void **)&aligned_pqtable_dist_scratch, 256 * (size_t)MAX_PQ_CHUNKS * sizeof(float), +// 256); +// diskann::alloc_aligned((void **)&aligned_dist_scratch, (size_t)graph_degree * sizeof(float), 256); +// diskann::alloc_aligned((void **)&aligned_query_float, aligned_dim * sizeof(float), 8 * sizeof(float)); +// diskann::alloc_aligned((void **)&rotated_query, aligned_dim * sizeof(float), 8 * sizeof(float)); +// +// memset(aligned_query_float, 0, aligned_dim * sizeof(float)); +// memset(rotated_query, 0, aligned_dim * sizeof(float)); +// } +// +// void set(size_t dim, T *query, const float norm = 1.0f) +// { +// for (size_t d = 0; d < dim; ++d) +// { +// if (norm != 1.0f) +// rotated_query[d] = aligned_query_float[d] = static_cast(query[d]) / norm; +// else +// rotated_query[d] = aligned_query_float[d] = static_cast(query[d]); +// } +// } +// }; void aggregate_coords(const std::vector &ids, const uint8_t *all_coords, const uint64_t ndims, uint8_t *out); diff --git a/include/pq_common.h b/include/pq_common.h new file mode 100644 index 000000000..4cdb69438 --- /dev/null +++ b/include/pq_common.h @@ -0,0 +1,25 @@ +#pragma once + +#include +#include + +#define NUM_PQ_BITS 8 +#define NUM_PQ_CENTROIDS (1 << NUM_PQ_BITS) +#define MAX_OPQ_ITERS 20 +#define NUM_KMEANS_REPS_PQ 12 +#define MAX_PQ_TRAINING_SET_SIZE 256000 +#define MAX_PQ_CHUNKS 512 + +namespace diskann +{ +inline std::string get_quantized_vectors_filename(const std::string &prefix, bool use_opq, uint32_t num_chunks) +{ + return prefix + (use_opq ? "_opq" : "pq") + std::to_string(num_chunks) + "_compressed.bin"; +} + +inline std::string get_pivot_data_filename(const std::string &prefix, bool use_opq, uint32_t num_chunks) +{ + return prefix + (use_opq ? "_opq" : "pq") + std::to_string(num_chunks) + "_pivots.bin"; +} + +} // namespace diskann diff --git a/include/pq_data_store.h b/include/pq_data_store.h new file mode 100644 index 000000000..f3f0d761e --- /dev/null +++ b/include/pq_data_store.h @@ -0,0 +1,87 @@ +#pragma once +#include +#include "distance.h" +#include "quantized_distance.h" +#include "pq.h" +#include "abstract_data_store.h" + +namespace diskann +{ +template class PQDataStore : public AbstractDataStore +{ + + public: + PQDataStore(uint32_t dim, uint32_t num_points, uint32_t num_pq_chunks, + std::shared_ptr> distance_fn, + std::shared_ptr> pq_distance_fn); + ~PQDataStore(); + + // Load quantized vectors from a set of files. Here filename is treated + // as a prefix and the files are assumed to be named with DiskANN + // conventions. + virtual location_t load(const std::string &file_prefix) override; + + // Save quantized vectors to a set of files whose names start with + // file_prefix. + // Currently, the plan is to save the quantized vectors to the quantized + // vectors file. + virtual size_t save(const std::string &file_prefix, const location_t num_points) override; + + // Since base class function is pure virtual, we need to declare it here, even though alignent concept is not needed for Quantized data stores. + virtual size_t get_aligned_dim() const override; + + // Populate quantized data from unaligned data using PQ functionality + virtual void populate_data(const data_t *vectors, const location_t num_pts) override; + virtual void populate_data(const std::string &filename, const size_t offset) override; + + virtual void extract_data_to_bin(const std::string &filename, const location_t num_pts) override; + + virtual void get_vector(const location_t i, data_t *target) const override; + virtual void set_vector(const location_t i, const data_t *const vector) override; + virtual void prefetch_vector(const location_t loc) override; + + virtual void move_vectors(const location_t old_location_start, const location_t new_location_start, + const location_t num_points) override; + virtual void copy_vectors(const location_t from_loc, const location_t to_loc, const location_t num_points) override; + + virtual void preprocess_query(const data_t *query, AbstractScratch *scratch) const override; + + virtual float get_distance(const data_t *query, const location_t loc) const override; + virtual float get_distance(const location_t loc1, const location_t loc2) const override; + + // NOTE: Caller must invoke "PQDistance->preprocess_query" ONCE before calling + // this function. + virtual void get_distance(const data_t *preprocessed_query, const location_t *locations, + const uint32_t location_count, float *distances, + AbstractScratch *scratch_space) const override; + + // NOTE: Caller must invoke "PQDistance->preprocess_query" ONCE before calling + // this function. + virtual void get_distance(const data_t *preprocessed_query, const std::vector &ids, + std::vector &distances, AbstractScratch *scratch_space) const override; + + virtual location_t calculate_medoid() const override; + + virtual size_t get_alignment_factor() const override; + + protected: + virtual location_t expand(const location_t new_size) override; + virtual location_t shrink(const location_t new_size) override; + + virtual location_t load_impl(const std::string &filename); +#ifdef EXEC_ENV_OLS + virtual location_t load_impl(AlignedFileReader &reader); +#endif + + private: + uint8_t *_quantized_data = nullptr; + size_t _num_chunks = 0; + + // REFACTOR TODO: Doing this temporarily before refactoring OPQ into + // its own class. Remove later. + bool _use_opq = false; + + Metric _distance_metric; + std::shared_ptr> _pq_distance_fn = nullptr; +}; +} // namespace diskann diff --git a/include/pq_distance.h b/include/pq_distance.h new file mode 100644 index 000000000..86054643b --- /dev/null +++ b/include/pq_distance.h @@ -0,0 +1,20 @@ +#pragma once + +namespace diskann +{ +template class PQScratch; + +template class PQDistance +{ + // Should this class manage PQScratch? + public: + PQDistance(FixedChunkPQTable &pq_table) : _pq_table(pq_table) + { + } + virtual void preprocess_query(const data_t *query_vec, const size_t query_dim, + PQScratch *scratch_query) = 0; + + private: + FixedChunkPQTable &_pq_table; +}; +} // namespace diskann \ No newline at end of file diff --git a/include/pq_l2_distance.h b/include/pq_l2_distance.h new file mode 100644 index 000000000..7335be209 --- /dev/null +++ b/include/pq_l2_distance.h @@ -0,0 +1,87 @@ +#pragma once +#include "quantized_distance.h" + +namespace diskann +{ +template class PQL2Distance : public QuantizedDistance +{ + public: + // REFACTOR TODO: We could take a file prefix here and load the + // PQ pivots file, so that the distance object is initialized + // immediately after construction. But this would not work well + // with our data store concept where the store is created first + // and data populated after. + // REFACTOR TODO: Ideally, we should only read the num_chunks from + // the pivots file. However, we read the pivots file only later, but + // clients can call functions like get__filename without calling + // load_pivot_data. Hence this. The TODO is whether we should check + // that the num_chunks from the file is the same as this one. + + PQL2Distance(uint32_t num_chunks); + + virtual ~PQL2Distance() override; + + virtual bool is_opq() const override; + + virtual std::string get_quantized_vectors_filename(const std::string &prefix) const override; + virtual std::string get_pivot_data_filename(const std::string &prefix) const override; + virtual std::string get_rotation_matrix_filename(const std::string &prefix) const override; + +#ifdef EXEC_ENV_OLS + virtual void load_pivot_data(MemoryMappedFiles &files, const std::string &pq_table_file, + size_t num_chunks) override; +#else + virtual void load_pivot_data(const std::string &pq_table_file, size_t num_chunks) override; +#endif + + // Number of chunks in the PQ table. Depends on the compression level used. + // Has to be < ndim + virtual uint32_t get_num_chunks() const override; + + // Preprocess the query by computing chunk distances from the query vector to + // various centroids. Since we don't want this class to do scratch management, + // we will take a PQScratch object which can come either from Index class or + // PQFlashIndex class. + virtual void preprocess_query(const data_t *aligned_query, uint32_t original_dim, + PQScratch &pq_scratch) override; + + // Distance function used for graph traversal. This function must be called + // after + // preprocess_query. The reason we do not call preprocess ourselves is because + // that function has to be called once per query, while this function is + // called at each iteration of the graph walk. NOTE: This function expects + // 1. the query to be preprocessed using preprocess_query() + // 2. the scratch object to contain the quantized vectors corresponding to ids + // in aligned_pq_coord_scratch. Done by calling aggregate_coords() + // + virtual void preprocessed_distance(PQScratch &pq_scratch, const uint32_t id_count, + float *dists_out) override; + + // Same as above, but returns the distances in a vector instead of an array. + // Convenience function for index.cpp. + virtual void preprocessed_distance(PQScratch &pq_scratch, const uint32_t n_ids, + std::vector &dists_out) override; + + // Currently this function is required for DiskPQ. However, it too can be + // subsumed under preprocessed_distance if we add the appropriate scratch + // variables to PQScratch and initialize them in + // pq_flash_index.cpp::disk_iterate_to_fixed_point() + virtual float brute_force_distance(const float *query_vec, uint8_t *base_vec) override; + + protected: + // assumes pre-processed query + virtual void prepopulate_chunkwise_distances(const float *query_vec, float *dist_vec); + + // assumes no rotation is involved + // virtual void inflate_vector(uint8_t *base_vec, float *out_vec); + + float *_tables = nullptr; // pq_tables = float array of size [256 * ndims] + uint64_t _ndims = 0; // ndims = true dimension of vectors + uint64_t _num_chunks = 0; + bool _use_rotation = false; + uint32_t *_chunk_offsets = nullptr; + float *_centroid = nullptr; + float *_tables_tr = nullptr; // same as pq_tables, but col-major + float *_rotmat_tr = nullptr; +}; +} // namespace diskann diff --git a/include/pq_scratch.h b/include/pq_scratch.h new file mode 100644 index 000000000..186d9b84d --- /dev/null +++ b/include/pq_scratch.h @@ -0,0 +1,43 @@ +#pragma once +#include +#include "pq_common.h" +#include "utils.h" + +namespace diskann +{ + +template struct PQScratch +{ + float *aligned_pqtable_dist_scratch = nullptr; // MUST BE AT LEAST [256 * NCHUNKS] + float *aligned_dist_scratch = nullptr; // MUST BE AT LEAST diskann MAX_DEGREE + uint8_t *aligned_pq_coord_scratch = nullptr; // AT LEAST [N_CHUNKS * MAX_DEGREE] + float *rotated_query = nullptr; + float *aligned_query_float = nullptr; + + PQScratch(size_t graph_degree, size_t aligned_dim) + { + diskann::alloc_aligned((void **)&aligned_pq_coord_scratch, + (size_t)graph_degree * (size_t)MAX_PQ_CHUNKS * sizeof(uint8_t), 256); + diskann::alloc_aligned((void **)&aligned_pqtable_dist_scratch, 256 * (size_t)MAX_PQ_CHUNKS * sizeof(float), + 256); + diskann::alloc_aligned((void **)&aligned_dist_scratch, (size_t)graph_degree * sizeof(float), 256); + diskann::alloc_aligned((void **)&aligned_query_float, aligned_dim * sizeof(float), 8 * sizeof(float)); + diskann::alloc_aligned((void **)&rotated_query, aligned_dim * sizeof(float), 8 * sizeof(float)); + + memset(aligned_query_float, 0, aligned_dim * sizeof(float)); + memset(rotated_query, 0, aligned_dim * sizeof(float)); + } + + void set_rotated_query(size_t dim, const T *query, const float norm = 1.0f) + { + for (size_t d = 0; d < dim; ++d) + { + if (norm != 1.0f) + rotated_query[d] = aligned_query_float[d] = static_cast(query[d]) / norm; + else + rotated_query[d] = aligned_query_float[d] = static_cast(query[d]); + } + } +}; + +} // namespace diskann \ No newline at end of file diff --git a/include/quantized_distance.h b/include/quantized_distance.h new file mode 100644 index 000000000..2bb927c9d --- /dev/null +++ b/include/quantized_distance.h @@ -0,0 +1,56 @@ +#pragma once +#include +#include +#include +#include "abstract_scratch.h" + +namespace diskann +{ +template struct PQScratch; + +template class QuantizedDistance +{ + public: + QuantizedDistance() = default; + QuantizedDistance(const QuantizedDistance &) = delete; + QuantizedDistance &operator=(const QuantizedDistance &) = delete; + virtual ~QuantizedDistance(){}; + + virtual bool is_opq() const = 0; + virtual std::string get_quantized_vectors_filename(const std::string &prefix) const = 0; + virtual std::string get_pivot_data_filename(const std::string &prefix) const = 0; + virtual std::string get_rotation_matrix_filename(const std::string &prefix) const = 0; + + // Loading the PQ centroid table need not be part of the abstract class. + // However, we want to indicate that this function will change once we have a + // file reader hierarchy, so leave it here as-is. +#ifdef EXEC_ENV_OLS + virtual void load_pivot_data(MemoryMappedFiles &files, const std::String &pq_table_file, size_t num_chunks) = 0; +#else + virtual void load_pivot_data(const std::string &pq_table_file, size_t num_chunks) = 0; +#endif + + // Number of chunks in the PQ table. Depends on the compression level used. + // Has to be < ndim + virtual uint32_t get_num_chunks() const = 0; + + // Preprocess the query by computing chunk distances from the query vector to + // various centroids. Since we don't want this class to do scratch management, + // we will take a PQScratch object which can come either from Index class or + // PQFlashIndex class. + virtual void preprocess_query(const data_t *query_vec, uint32_t query_dim, PQScratch &pq_scratch) = 0; + + // Workhorse + // This function must be called after preprocess_query + virtual void preprocessed_distance(PQScratch &pq_scratch, const uint32_t id_count, float *dists_out) = 0; + + // Same as above, but convenience function for index.cpp. + virtual void preprocessed_distance(PQScratch &pq_scratch, const uint32_t n_ids, + std::vector &dists_out) = 0; + + // Currently this function is required for DiskPQ. However, it too can be subsumed + // under preprocessed_distance if we add the appropriate scratch variables to + // PQScratch and initialize them in pq_flash_index.cpp::disk_iterate_to_fixed_point() + virtual float brute_force_distance(const float *query_vec, uint8_t *base_vec) = 0; +}; +} // namespace diskann diff --git a/include/scratch.h b/include/scratch.h index 3b44f8f80..d788247c6 100644 --- a/include/scratch.h +++ b/include/scratch.h @@ -11,9 +11,9 @@ #include "tsl/robin_map.h" #include "tsl/sparse_map.h" +#include "abstract_scratch.h" #include "neighbor.h" #include "concurrent_queue.h" -#include "pq.h" #include "aligned_file_reader.h" // In-mem index related limits @@ -26,11 +26,12 @@ namespace diskann { +template struct PQScratch; // -// Scratch space for in-memory index based search +// AbstractScratch space for in-memory index based search // -template class InMemQueryScratch +template class InMemQueryScratch : public AbstractScratch { public: ~InMemQueryScratch(); @@ -54,11 +55,11 @@ template class InMemQueryScratch } inline T *aligned_query() { - return _aligned_query; + return this->_aligned_query_T; } inline PQScratch *pq_scratch() { - return _pq_scratch; + return this->_pq_scratch; } inline std::vector &pool() { @@ -106,10 +107,6 @@ template class InMemQueryScratch uint32_t _R; uint32_t _maxc; - T *_aligned_query = nullptr; - - PQScratch *_pq_scratch = nullptr; - // _pool stores all neighbors explored from best_L_nodes. // Usually around L+R, but could be higher. // Initialized to 3L+R for some slack, expands as needed. @@ -146,10 +143,10 @@ template class InMemQueryScratch }; // -// Scratch space for SSD index based search +// AbstractScratch space for SSD index based search // -template class SSDQueryScratch +template class SSDQueryScratch : public AbstractScratch { public: T *coord_scratch = nullptr; // MUST BE AT LEAST [sizeof(T) * data_dim] @@ -157,10 +154,6 @@ template class SSDQueryScratch char *sector_scratch = nullptr; // MUST BE AT LEAST [MAX_N_SECTOR_READS * SECTOR_LEN] size_t sector_idx = 0; // index of next [SECTOR_LEN] scratch to use - T *aligned_query_T = nullptr; - - PQScratch *_pq_scratch; - tsl::robin_set visited; NeighborPriorityQueue retset; std::vector full_retset; diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 2206a01f7..cbca26440 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -13,7 +13,7 @@ else() linux_aligned_file_reader.cpp math_utils.cpp natural_number_map.cpp in_mem_data_store.cpp in_mem_graph_store.cpp natural_number_set.cpp memory_mapper.cpp partition.cpp pq.cpp - pq_flash_index.cpp scratch.cpp logger.cpp utils.cpp filter_utils.cpp index_factory.cpp abstract_index.cpp) + pq_flash_index.cpp scratch.cpp logger.cpp utils.cpp filter_utils.cpp index_factory.cpp abstract_index.cpp pq_l2_distance.cpp pq_data_store.cpp) if (RESTAPI) list(APPEND CPP_SOURCES restapi/search_wrapper.cpp restapi/server.cpp) endif() diff --git a/src/abstract_data_store.cpp b/src/abstract_data_store.cpp index a980bd545..d81545a37 100644 --- a/src/abstract_data_store.cpp +++ b/src/abstract_data_store.cpp @@ -2,7 +2,6 @@ // Licensed under the MIT license. #include - #include "abstract_data_store.h" namespace diskann @@ -40,6 +39,18 @@ template location_t AbstractDataStore::resize(const lo } } +template +void AbstractDataStore::preprocess_query(const data_t *query, AbstractScratch *query_scratch) const +{ +} + +template +void AbstractDataStore::get_distance(const data_t *preprocessed_query, const std::vector &ids, + std::vector &distances, + AbstractScratch *scratch_space) const +{ +} + template DISKANN_DLLEXPORT class AbstractDataStore; template DISKANN_DLLEXPORT class AbstractDataStore; template DISKANN_DLLEXPORT class AbstractDataStore; diff --git a/src/disk_utils.cpp b/src/disk_utils.cpp index aadeb6dd1..bf46b3e9e 100644 --- a/src/disk_utils.cpp +++ b/src/disk_utils.cpp @@ -559,7 +559,21 @@ void breakup_dense_points(const std::string data_file, const std::string labels_ if (dummy_pt_ids.size() != 0) { diskann::cout << dummy_pt_ids.size() << " is the number of dummy points created" << std::endl; - data = (T *)std::realloc((void *)data, labels_per_point.size() * ndims * sizeof(T)); + + T *ptr = (T *)std::realloc((void *)data, labels_per_point.size() * ndims * sizeof(T)); + if (ptr == nullptr) + { + diskann::cerr << "Realloc failed while creating dummy points" << std::endl; + free(data); + data = nullptr; + throw new diskann::ANNException("Realloc failed while expanding data.", -1, __FUNCTION__, __FILE__, + __LINE__); + } + else + { + data = ptr; + } + std::ofstream dummy_writer(out_metadata_file); assert(dummy_writer.is_open()); for (auto i = dummy_pt_ids.begin(); i != dummy_pt_ids.end(); i++) diff --git a/src/dll/CMakeLists.txt b/src/dll/CMakeLists.txt index d00cfeb95..ffde4f797 100644 --- a/src/dll/CMakeLists.txt +++ b/src/dll/CMakeLists.txt @@ -2,8 +2,8 @@ #Licensed under the MIT license. add_library(${PROJECT_NAME} SHARED dllmain.cpp ../abstract_data_store.cpp ../partition.cpp ../pq.cpp ../pq_flash_index.cpp ../logger.cpp ../utils.cpp - ../windows_aligned_file_reader.cpp ../distance.cpp ../memory_mapper.cpp ../index.cpp - ../in_mem_data_store.cpp ../in_mem_graph_store.cpp ../math_utils.cpp ../disk_utils.cpp ../filter_utils.cpp + ../windows_aligned_file_reader.cpp ../distance.cpp ../pq_l2_distance.cpp ../memory_mapper.cpp ../index.cpp + ../in_mem_data_store.cpp ../pq_data_store.cpp ../in_mem_graph_store.cpp ../math_utils.cpp ../disk_utils.cpp ../filter_utils.cpp ../ann_exception.cpp ../natural_number_set.cpp ../natural_number_map.cpp ../scratch.cpp ../index_factory.cpp ../abstract_index.cpp) set(TARGET_DIR "$<$:${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}>$<$:${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}>") diff --git a/src/in_mem_data_store.cpp b/src/in_mem_data_store.cpp index b2f708263..4bf9e25a6 100644 --- a/src/in_mem_data_store.cpp +++ b/src/in_mem_data_store.cpp @@ -156,6 +156,7 @@ void InMemDataStore::extract_data_to_bin(const std::string &filename, co template void InMemDataStore::get_vector(const location_t i, data_t *dest) const { + // REFACTOR TODO: Should we denormalize and return values? memcpy(dest, _data + i * _aligned_dim, this->_dim * sizeof(data_t)); } @@ -172,7 +173,8 @@ template void InMemDataStore::set_vector(const locatio template void InMemDataStore::prefetch_vector(const location_t loc) { - diskann::prefetch_vector((const char *)_data + _aligned_dim * (size_t)loc, sizeof(data_t) * _aligned_dim); + diskann::prefetch_vector((const char *)_data + _aligned_dim * (size_t)loc * sizeof(data_t), + sizeof(data_t) * _aligned_dim); } template float InMemDataStore::get_distance(const data_t *query, const location_t loc) const @@ -182,8 +184,10 @@ template float InMemDataStore::get_distance(const data template void InMemDataStore::get_distance(const data_t *query, const location_t *locations, - const uint32_t location_count, float *distances) const + const uint32_t location_count, float *distances, + AbstractScratch *scratch_space) const { + assert(scratch_space == nullptr); // Scratch space should only be used by PQ data store. for (location_t i = 0; i < location_count; i++) { distances[i] = _distance_fn->compare(query, _data + locations[i] * _aligned_dim, (uint32_t)this->_aligned_dim); diff --git a/src/index.cpp b/src/index.cpp index c30e44ac2..c73b26b80 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -18,6 +18,9 @@ #ifdef _WINDOWS #include #endif +// REFACTOR TODO: Must move to factory. +#include "pq_scratch.h" +#include "pq_l2_distance.h" #include "index.h" #define MAX_POINTS_FOR_USING_BITSET 10000000 @@ -91,14 +94,6 @@ Index::Index(Metric m, const size_t dim, const size_t max_point } const size_t total_internal_points = _max_points + _num_frozen_pts; - if (_pq_dist) - { - if (_num_pq_chunks > _dim) - throw diskann::ANNException("ERROR: num_pq_chunks > dim", -1, __FUNCSIG__, __FILE__, __LINE__); - alloc_aligned(((void **)&_pq_data), total_internal_points * _num_pq_chunks * sizeof(char), 8 * sizeof(char)); - std::memset(_pq_data, 0, total_internal_points * _num_pq_chunks * sizeof(char)); - } - _start = (uint32_t)_max_points; _final_graph.resize(total_internal_points); @@ -125,6 +120,25 @@ Index::Index(Metric m, const size_t dim, const size_t max_point std::make_unique>((location_t)total_internal_points, _dim, this->_distance); } + if (_pq_dist) + { + if (_num_pq_chunks > _dim) + { + throw diskann::ANNException("ERROR: num_pq_chunks > dim", -1, __FUNCSIG__, __FILE__, __LINE__); + } + + // REFACTOR TODO: This should move to a factory method and support OPQ. + _pq_distance_fn = std::make_shared>(_num_pq_chunks); + // REFACTOR TODO: Unlike Distance and DataStore, where distance object is ready when the data store + // is constructed. Here the distance object will not be fully ready until populate_data() is called + _pq_data_store = std::make_shared>(_dim, total_internal_points, _num_pq_chunks, this->_distance, + _pq_distance_fn); + // REFACTOR + // alloc_aligned( + // ((void **)&_pq_data), total_internal_points * _num_pq_chunks * sizeof(char), 8 * sizeof(char)); + // std::memset(_pq_data, 0, total_internal_points * _num_pq_chunks * sizeof(char)); + } + _locks = std::vector(total_internal_points); if (enable_tags) @@ -866,14 +880,13 @@ template int Index template uint32_t Index::calculate_entry_point() { - // TODO: need to compute medoid with PQ data too, for now sample at random if (_pq_dist) { - size_t r = (size_t)rand() * (size_t)RAND_MAX + (size_t)rand(); - return (uint32_t)(r % (size_t)_nd); + // REFACTOR TODO: This function returns a random point. Must change + // to actually compute the medoid. + return _pq_data_store->calculate_medoid(); } - - // TODO: This function does not support multi-threaded calculation of medoid. + // REFACTOR TODO: This function does not support multi-threaded calculation of medoid. // Must revisit if perf is a concern. return _data_store->calculate_medoid(); } @@ -896,7 +909,8 @@ template std::vector Inde return init_ids; } -// Find common filter between a node's labels and a given set of labels, while taking into account universal label +// Find common filter between a node's labels and a given set of labels, while +// taking into account universal label template bool Index::detect_common_filters(uint32_t point_id, bool search_invocation, const std::vector &incoming_labels) @@ -907,8 +921,8 @@ bool Index::detect_common_filters(uint32_t point_id, bool searc curr_node_labels.end(), std::back_inserter(common_filters)); if (common_filters.size() > 0) { - // This is to reduce the repetitive calls. If common_filters size is > 0 , we dont need to check further for - // universal label + // This is to reduce the repetitive calls. If common_filters size is > 0 , + // we dont need to check further for universal label return true; } if (_use_universal_label) @@ -952,31 +966,31 @@ std::pair Index::iterate_to_fixed_point( T *aligned_query = scratch->aligned_query(); - float *query_float = nullptr; - float *query_rotated = nullptr; float *pq_dists = nullptr; - uint8_t *pq_coord_scratch = nullptr; + // uint8_t *pq_coord_scratch = nullptr; // Intialize PQ related scratch to use PQ based distances if (_pq_dist) { - // Get scratch spaces - PQScratch *pq_query_scratch = scratch->pq_scratch(); - query_float = pq_query_scratch->aligned_query_float; - query_rotated = pq_query_scratch->rotated_query; - pq_dists = pq_query_scratch->aligned_pqtable_dist_scratch; + _pq_data_store->preprocess_query(aligned_query, scratch); - // Copy query vector to float and then to "rotated" query - for (size_t d = 0; d < _dim; d++) - { - query_float[d] = (float)aligned_query[d]; - } - pq_query_scratch->set(_dim, aligned_query); + // REFACTOR: Commented out. + // Get scratch spaces + // PQScratch *pq_query_scratch = scratch->pq_scratch(); + // query_float = pq_query_scratch->aligned_query_float; + // query_rotated = pq_query_scratch->rotated_query; + // pq_dists = pq_query_scratch->aligned_pqtable_dist_scratch; - // center the query and rotate if we have a rotation matrix - _pq_table.preprocess_query(query_rotated); - _pq_table.populate_chunk_distances(query_rotated, pq_dists); + //// Copy query vector to float and then to "rotated" query + // for (size_t d = 0; d < _dim; d++) { + // query_float[d] = (float)aligned_query[d]; + // } + // pq_query_scratch->set(_dim, aligned_query); + + //// center the query and rotate if we have a rotation matrix + //_pq_table.preprocess_query(query_rotated); + //_pq_table.populate_chunk_distances(query_rotated, pq_dists); - pq_coord_scratch = pq_query_scratch->aligned_pq_coord_scratch; + // pq_coord_scratch = scratch->pq_scratch()->aligned_pq_coord_scratch; } if (expanded_nodes.size() > 0 || id_scratch.size() > 0) @@ -1006,10 +1020,14 @@ std::pair Index::iterate_to_fixed_point( }; // Lambda to batch compute query<-> node distances in PQ space - auto compute_dists = [this, pq_coord_scratch, pq_dists](const std::vector &ids, - std::vector &dists_out) { - diskann::aggregate_coords(ids, this->_pq_data, this->_num_pq_chunks, pq_coord_scratch); - diskann::pq_dist_lookup(pq_coord_scratch, ids.size(), this->_num_pq_chunks, pq_dists, dists_out); + auto compute_dists = [this, scratch, pq_dists](const std::vector &ids, std::vector &dists_out) { + // REFACTOR + _pq_data_store->get_distance(scratch->aligned_query(), ids, dists_out, scratch); + + // diskann::aggregate_coords(ids, this->_pq_data, this->_num_pq_chunks, + // pq_coord_scratch); + // diskann::pq_dist_lookup(pq_coord_scratch, ids.size(), this->_num_pq_chunks, + // pq_dists, dists_out); }; // Initialize the candidate pool with starting points @@ -1042,7 +1060,13 @@ std::pair Index::iterate_to_fixed_point( float distance; if (_pq_dist) { - pq_dist_lookup(pq_coord_scratch, 1, this->_num_pq_chunks, pq_dists, &distance); + // REFACTOR + std::vector ids = {id}; + std::vector distances = {std::numeric_limits::max()}; + _pq_data_store->get_distance(aligned_query, ids, distances, scratch); + distance = distances[0]; + // pq_dist_lookup(pq_coord_scratch, 1, this->_num_pq_chunks, pq_dists, + // &distance); } else { @@ -1807,8 +1831,8 @@ void Index::build(const char *filename, const size_t num_points << " points, but " << "index can support only " << _max_points << " points as specified in constructor." << std::endl; - if (_pq_dist) - aligned_free(_pq_data); + // REFACTOR PQDataStore will take care of its memory + // if (_pq_dist) aligned_free(_pq_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } @@ -1818,8 +1842,8 @@ void Index::build(const char *filename, const size_t num_points stream << "ERROR: Driver requests loading " << num_points_to_load << " points and file has only " << file_num_points << " points." << std::endl; - if (_pq_dist) - aligned_free(_pq_data); + // REFACTOR: PQDataStore will take care of its memory + // if (_pq_dist) aligned_free(_pq_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } @@ -1830,31 +1854,51 @@ void Index::build(const char *filename, const size_t num_points << "but file has " << file_dim << " dimension." << std::endl; diskann::cerr << stream.str() << std::endl; - if (_pq_dist) - aligned_free(_pq_data); + // REFACTOR: PQDataStore will take care of its memory + // if (_pq_dist) aligned_free(_pq_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } if (_pq_dist) { - double p_val = std::min(1.0, ((double)MAX_PQ_TRAINING_SET_SIZE / (double)file_num_points)); - - std::string suffix = _use_opq ? "_opq" : "_pq"; - suffix += std::to_string(_num_pq_chunks); - auto pq_pivots_file = std::string(filename) + suffix + "_pivots.bin"; - auto pq_compressed_file = std::string(filename) + suffix + "_compressed.bin"; - generate_quantized_data(std::string(filename), pq_pivots_file, pq_compressed_file, _dist_metric, p_val, - _num_pq_chunks, _use_opq); - - copy_aligned_data_from_file(pq_compressed_file.c_str(), _pq_data, file_num_points, _num_pq_chunks, - _num_pq_chunks); + // REFACTOR #ifdef EXEC_ENV_OLS throw ANNException("load_pq_centroid_bin should not be called when " "EXEC_ENV_OLS is defined.", -1, __FUNCSIG__, __FILE__, __LINE__); #else - _pq_table.load_pq_centroid_bin(pq_pivots_file.c_str(), _num_pq_chunks); + // REFACTOR TODO: Both in the previous code and in the current PQDataStore, + // we are writing the PQ files in the same path as the input file. Now we + // may not have write permissions to that folder, but we will always have + // write permissions to the output folder. So we should write the PQ files + // there. The problem is that the Index class gets the output folder prefix + // only at the time of save(), by which time we are too late. So leaving it + // as-is for now. + _pq_data_store->populate_data(filename, 0U); #endif + // double p_val = std::min( + // 1.0, ((double)MAX_PQ_TRAINING_SET_SIZE / (double)file_num_points)); + // + // std::string suffix = _use_opq ? "_opq" : "_pq"; + // suffix += std::to_string(_num_pq_chunks); + // auto pq_pivots_file = std::string(filename) + suffix + "_pivots.bin"; + // auto pq_compressed_file = + // std::string(filename) + suffix + "_compressed.bin"; + // generate_quantized_data(std::string(filename), pq_pivots_file, + // pq_compressed_file, _dist_metric, p_val, + // _num_pq_chunks, _use_opq); + // + // copy_aligned_data_from_file(pq_compressed_file.c_str(), _pq_data, + // file_num_points, _num_pq_chunks, + // _num_pq_chunks); + //#ifdef EXEC_ENV_OLS + // throw ANNException( + // "load_pq_centroid_bin should not be called when " + // "EXEC_ENV_OLS is defined.", + // -1, __FUNCSIG__, __FILE__, __LINE__); + //#else + // _pq_table.load_pq_centroid_bin(pq_pivots_file.c_str(), _num_pq_chunks); + //#endif } _data_store->populate_data(filename, 0U); @@ -2421,8 +2465,10 @@ template void Indexcopy_vectors((location_t)res, (location_t)_max_points, 1); } else { diff --git a/src/pq_data_store.cpp b/src/pq_data_store.cpp new file mode 100644 index 000000000..ff808e9af --- /dev/null +++ b/src/pq_data_store.cpp @@ -0,0 +1,249 @@ +#include + +#include "pq_data_store.h" +#include "pq.h" +#include "pq_scratch.h" +#include "utils.h" +#include "distance.h" + +namespace diskann +{ + +// REFACTOR TODO: Assuming that num_pq_chunks is known already. Must verify if +// this is true. +template +PQDataStore::PQDataStore(uint32_t dim, uint32_t num_points, uint32_t num_pq_chunks, + std::shared_ptr> distance_fn, + std::shared_ptr> pq_distance_fn) + : AbstractDataStore(num_points, dim), _quantized_data(nullptr), _num_chunks(num_pq_chunks), + _distance_metric(distance_fn->get_metric()), _pq_distance_fn(pq_distance_fn) +{ +} + +template PQDataStore::~PQDataStore() +{ + if (_quantized_data != nullptr) + { + aligned_free(_quantized_data); + _quantized_data = nullptr; + } +} + +template location_t PQDataStore::load(const std::string &filename) +{ + return load_impl(filename); +} +template size_t PQDataStore::save(const std::string &filename, const location_t num_points) +{ + return diskann::save_bin(filename, _quantized_data, this->capacity(), _num_chunks, 0); +} + +template size_t PQDataStore::get_aligned_dim() const +{ + return this->get_dims(); +} + +// Populate quantized data from regular data. +template void PQDataStore::populate_data(const data_t *vectors, const location_t num_pts) +{ + throw std::logic_error("Not implemented yet"); +} + +template void PQDataStore::populate_data(const std::string &filename, const size_t offset) +{ + if (_quantized_data != nullptr) + { + aligned_free(_quantized_data); + } + + uint64_t file_num_points = 0, file_dim = 0; + get_bin_metadata(filename, file_num_points, file_dim, offset); + this->_capacity = (location_t)file_num_points; + this->_dim = file_dim; + + double p_val = std::min(1.0, ((double)MAX_PQ_TRAINING_SET_SIZE / (double)file_num_points)); + + auto pivots_file = _pq_distance_fn->get_pivot_data_filename(filename); + auto compressed_file = _pq_distance_fn->get_quantized_vectors_filename(filename); + + generate_quantized_data(filename, pivots_file, compressed_file, _distance_metric, p_val, _num_chunks, + _pq_distance_fn->is_opq()); + + // REFACTOR TODO: Not sure of the alignment. Just copying from index.cpp + alloc_aligned(((void **)&_quantized_data), file_num_points * _num_chunks * sizeof(uint8_t), 8); + copy_aligned_data_from_file(compressed_file.c_str(), _quantized_data, file_num_points, _num_chunks, + _num_chunks); +#ifdef EXEC_ENV_OLS + throw ANNException("load_pq_centroid_bin should not be called when " + "EXEC_ENV_OLS is defined.", + -1, __FUNCSIG__, __FILE__, __LINE__); +#else + _pq_distance_fn->load_pivot_data(pivots_file.c_str(), _num_chunks); +#endif +} + +template +void PQDataStore::extract_data_to_bin(const std::string &filename, const location_t num_pts) +{ + throw std::logic_error("Not implemented yet"); +} + +template void PQDataStore::get_vector(const location_t i, data_t *target) const +{ + // REFACTOR TODO: Should we inflate the compressed vector here? + if (i < this->capacity()) + { + throw std::logic_error("Not implemented yet."); + } + else + { + std::stringstream ss; + ss << "Requested vector " << i << " but only " << this->capacity() << " vectors are present"; + throw diskann::ANNException(ss.str(), -1); + } +} +template void PQDataStore::set_vector(const location_t i, const data_t *const vector) +{ + // REFACTOR TODO: Should we accept a normal vector and compress here? + // memcpy (_data + i * _num_chunks, vector, _num_chunks * sizeof(data_t)); + throw std::logic_error("Not implemented yet"); +} + +template void PQDataStore::prefetch_vector(const location_t loc) +{ + const uint8_t *ptr = _quantized_data + ((size_t)loc) * _num_chunks * sizeof(data_t); + diskann::prefetch_vector((const char *)ptr, _num_chunks * sizeof(data_t)); +} + +template +void PQDataStore::move_vectors(const location_t old_location_start, const location_t new_location_start, + const location_t num_points) +{ + // REFACTOR TODO: Moving vectors is only for in-mem fresh. + throw std::logic_error("Not implemented yet"); +} + +template +void PQDataStore::copy_vectors(const location_t from_loc, const location_t to_loc, const location_t num_points) +{ + // REFACTOR TODO: Is the number of bytes correct? + memcpy(_quantized_data + to_loc * _num_chunks, _quantized_data + from_loc * _num_chunks, _num_chunks * num_points); +} + +// REFACTOR TODO: Currently, we take aligned_query as parameter, but this +// function should also do the alignment. +template +void PQDataStore::preprocess_query(const data_t *aligned_query, AbstractScratch *scratch) const +{ + if (scratch == nullptr) + { + throw diskann::ANNException("Scratch space is null", -1); + } + + PQScratch *pq_scratch = scratch->pq_scratch(); + + if (pq_scratch == nullptr) + { + throw diskann::ANNException("PQScratch space has not been set in the scratch object.", -1); + } + + _pq_distance_fn->preprocess_query(aligned_query, (location_t)this->get_dims(), *pq_scratch); +} + +template float PQDataStore::get_distance(const data_t *query, const location_t loc) const +{ + throw std::logic_error("Not implemented yet"); +} + +template float PQDataStore::get_distance(const location_t loc1, const location_t loc2) const +{ + throw std::logic_error("Not implemented yet"); +} + +template +void PQDataStore::get_distance(const data_t *preprocessed_query, const location_t *locations, + const uint32_t location_count, float *distances, + AbstractScratch *scratch_space) const +{ + if (scratch_space == nullptr) + { + throw diskann::ANNException("Scratch space is null", -1); + } + PQScratch *pq_scratch = scratch_space->pq_scratch(); + if (pq_scratch == nullptr) + { + throw diskann::ANNException("PQScratch not set in scratch space.", -1); + } + diskann::aggregate_coords(locations, location_count, _quantized_data, this->_num_chunks, + pq_scratch->aligned_pq_coord_scratch); + _pq_distance_fn->preprocessed_distance(*pq_scratch, location_count, distances); +} + +template +void PQDataStore::get_distance(const data_t *preprocessed_query, const std::vector &ids, + std::vector &distances, AbstractScratch *scratch_space) const +{ + if (scratch_space == nullptr) + { + throw diskann::ANNException("Scratch space is null", -1); + } + PQScratch *pq_scratch = scratch_space->pq_scratch(); + if (pq_scratch == nullptr) + { + throw diskann::ANNException("PQScratch not set in scratch space.", -1); + } + diskann::aggregate_coords(ids, _quantized_data, this->_num_chunks, pq_scratch->aligned_pq_coord_scratch); + _pq_distance_fn->preprocessed_distance(*pq_scratch, (location_t)ids.size(), distances); +} + +template location_t PQDataStore::calculate_medoid() const +{ + // REFACTOR TODO: Must calculate this just like we do with data store. + size_t r = (size_t)rand() * (size_t)RAND_MAX + (size_t)rand(); + return (uint32_t)(r % (size_t)this->capacity()); +} + +template size_t PQDataStore::get_alignment_factor() const +{ + return 1; +} + +template location_t PQDataStore::load_impl(const std::string &file_prefix) +{ + if (_quantized_data != nullptr) + { + aligned_free(_quantized_data); + } + auto quantized_vectors_file = _pq_distance_fn->get_quantized_vectors_filename(file_prefix); + + size_t num_points; + load_aligned_bin(quantized_vectors_file, _quantized_data, num_points, _num_chunks, _num_chunks); + this->_capacity = (location_t)num_points; + + auto pivots_file = _pq_distance_fn->get_pivot_data_filename(file_prefix); + _pq_distance_fn->load_pivot_data(pivots_file, _num_chunks); + + return this->_capacity; +} + +template location_t PQDataStore::expand(const location_t new_size) +{ + throw std::logic_error("Not implemented yet"); +} + +template location_t PQDataStore::shrink(const location_t new_size) +{ + throw std::logic_error("Not implemented yet"); +} + +#ifdef EXEC_ENV_OLS +template location_t PQDataStore::load_impl(AlignedFileReader &reader) +{ +} +#endif + +template DISKANN_DLLEXPORT class PQDataStore; +template DISKANN_DLLEXPORT class PQDataStore; +template DISKANN_DLLEXPORT class PQDataStore; + +} // namespace diskann \ No newline at end of file diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index 78e44ba70..158994a47 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -4,6 +4,8 @@ #include "common_includes.h" #include "timer.h" +#include "pq.h" +#include "pq_scratch.h" #include "pq_flash_index.h" #include "cosine_similarity.h" @@ -1157,7 +1159,7 @@ void PQFlashIndex::cached_beam_search(const T *query1, const uint64_t auto data = manager.scratch_space(); IOContext &ctx = data->ctx; auto query_scratch = &(data->scratch); - auto pq_query_scratch = query_scratch->_pq_scratch; + auto pq_query_scratch = query_scratch->pq_scratch(); // reset query scratch query_scratch->reset(); @@ -1165,7 +1167,7 @@ void PQFlashIndex::cached_beam_search(const T *query1, const uint64_t // copy query to thread specific aligned and allocated memory (for distance // calculations we need aligned data) float query_norm = 0; - T *aligned_query_T = query_scratch->aligned_query_T; + T *aligned_query_T = query_scratch->aligned_query_T(); float *query_float = pq_query_scratch->aligned_query_float; float *query_rotated = pq_query_scratch->rotated_query; @@ -1186,7 +1188,7 @@ void PQFlashIndex::cached_beam_search(const T *query1, const uint64_t { aligned_query_T[i] = (T)(aligned_query_T[i] / query_norm); } - pq_query_scratch->set(this->data_dim, aligned_query_T); + pq_query_scratch->set_rotated_query(this->data_dim, aligned_query_T); } else { @@ -1194,7 +1196,7 @@ void PQFlashIndex::cached_beam_search(const T *query1, const uint64_t { aligned_query_T[i] = query1[i]; } - pq_query_scratch->set(this->data_dim, aligned_query_T); + pq_query_scratch->set_rotated_query(this->data_dim, aligned_query_T); } // pointers to buffers for data diff --git a/src/pq_l2_distance.cpp b/src/pq_l2_distance.cpp new file mode 100644 index 000000000..24caa9baf --- /dev/null +++ b/src/pq_l2_distance.cpp @@ -0,0 +1,285 @@ + +#include "pq.h" +#include "pq_l2_distance.h" +#include "pq_scratch.h" + +// block size for reading/processing large files and matrices in blocks +#define BLOCK_SIZE 5000000 + +namespace diskann +{ + +template PQL2Distance::PQL2Distance(uint32_t num_chunks) : _num_chunks(num_chunks) +{ +} + +template PQL2Distance::~PQL2Distance() +{ +#ifndef EXEC_ENV_OLS + if (_tables != nullptr) + delete[] _tables; + if (_chunk_offsets != nullptr) + delete[] _chunk_offsets; + if (_centroid != nullptr) + delete[] _centroid; + if (_rotmat_tr != nullptr) + delete[] _rotmat_tr; +#endif + if (_tables_tr != nullptr) + delete[] _tables_tr; +} + +template bool PQL2Distance::is_opq() const +{ + return false; +} + +template +std::string PQL2Distance::get_quantized_vectors_filename(const std::string &prefix) const +{ + if (_num_chunks == 0) + { + throw diskann::ANNException("Must set num_chunks before calling get_quantized_vectors_filename", -1, + __FUNCSIG__, __FILE__, __LINE__); + } + return diskann::get_quantized_vectors_filename(prefix, false, (uint32_t)_num_chunks); +} +template std::string PQL2Distance::get_pivot_data_filename(const std::string &prefix) const +{ + if (_num_chunks == 0) + { + throw diskann::ANNException("Must set num_chunks before calling get_pivot_data_filename", -1, __FUNCSIG__, + __FILE__, __LINE__); + } + return diskann::get_pivot_data_filename(prefix, false, (uint32_t)_num_chunks); +} +template +std::string PQL2Distance::get_rotation_matrix_filename(const std::string &prefix) const +{ + // REFACTOR TODO: Currently, we are assuming that PQ doesn't have a rotation + // matrix. + return ""; +} + +#ifdef EXEC_ENV_OLS +template +void PQL2Distance::load_pivot_data(MemoryMappedFiles &files, const std::string &pq_table_file, + size_t num_chunks) +{ +#else +template +void PQL2Distance::load_pivot_data(const std::string &pq_table_file, size_t num_chunks) +{ +#endif + uint64_t nr, nc; + // std::string rotmat_file = get_opq_rot_matrix_filename(pq_table_file, + // false); + +#ifdef EXEC_ENV_OLS + size_t *file_offset_data; // since load_bin only sets the pointer, no need + // to delete. + diskann::load_bin(files, pq_table_file, file_offset_data, nr, nc); +#else + std::unique_ptr file_offset_data; + diskann::load_bin(pq_table_file, file_offset_data, nr, nc); +#endif + + bool use_old_filetype = false; + + if (nr != 4 && nr != 5) + { + diskann::cout << "Error reading pq_pivots file " << pq_table_file + << ". Offsets dont contain correct metadata, # offsets = " << nr << ", but expecting " << 4 + << " or " << 5; + throw diskann::ANNException("Error reading pq_pivots file at offsets data.", -1, __FUNCSIG__, __FILE__, + __LINE__); + } + + if (nr == 4) + { + diskann::cout << "Offsets: " << file_offset_data[0] << " " << file_offset_data[1] << " " << file_offset_data[2] + << " " << file_offset_data[3] << std::endl; + } + else if (nr == 5) + { + use_old_filetype = true; + diskann::cout << "Offsets: " << file_offset_data[0] << " " << file_offset_data[1] << " " << file_offset_data[2] + << " " << file_offset_data[3] << file_offset_data[4] << std::endl; + } + else + { + throw diskann::ANNException("Wrong number of offsets in pq_pivots", -1, __FUNCSIG__, __FILE__, __LINE__); + } + +#ifdef EXEC_ENV_OLS + diskann::load_bin(files, pq_table_file, tables, nr, nc, file_offset_data[0]); +#else + diskann::load_bin(pq_table_file, _tables, nr, nc, file_offset_data[0]); +#endif + + if ((nr != NUM_PQ_CENTROIDS)) + { + diskann::cout << "Error reading pq_pivots file " << pq_table_file << ". file_num_centers = " << nr + << " but expecting " << NUM_PQ_CENTROIDS << " centers"; + throw diskann::ANNException("Error reading pq_pivots file at pivots data.", -1, __FUNCSIG__, __FILE__, + __LINE__); + } + + this->_ndims = nc; + +#ifdef EXEC_ENV_OLS + diskann::load_bin(files, pq_table_file, centroid, nr, nc, file_offset_data[1]); +#else + diskann::load_bin(pq_table_file, _centroid, nr, nc, file_offset_data[1]); +#endif + + if ((nr != this->_ndims) || (nc != 1)) + { + diskann::cerr << "Error reading centroids from pq_pivots file " << pq_table_file << ". file_dim = " << nr + << ", file_cols = " << nc << " but expecting " << this->_ndims << " entries in 1 dimension."; + throw diskann::ANNException("Error reading pq_pivots file at centroid data.", -1, __FUNCSIG__, __FILE__, + __LINE__); + } + + int chunk_offsets_index = 2; + if (use_old_filetype) + { + chunk_offsets_index = 3; + } +#ifdef EXEC_ENV_OLS + diskann::load_bin(files, pq_table_file, chunk_offsets, nr, nc, file_offset_data[chunk_offsets_index]); +#else + diskann::load_bin(pq_table_file, _chunk_offsets, nr, nc, file_offset_data[chunk_offsets_index]); +#endif + + if (nc != 1 || (nr != num_chunks + 1 && num_chunks != 0)) + { + diskann::cerr << "Error loading chunk offsets file. numc: " << nc << " (should be 1). numr: " << nr + << " (should be " << num_chunks + 1 << " or 0 if we need to infer)" << std::endl; + throw diskann::ANNException("Error loading chunk offsets file", -1, __FUNCSIG__, __FILE__, __LINE__); + } + + this->_num_chunks = nr - 1; + diskann::cout << "Loaded PQ Pivots: #ctrs: " << NUM_PQ_CENTROIDS << ", #dims: " << this->_ndims + << ", #chunks: " << this->_num_chunks << std::endl; + + // For PQ there will be no rotation matrix. + // if (file_exists(rotmat_file)) { + // #ifdef EXEC_ENV_OLS + // diskann::load_bin(files, rotmat_file, (float *&)rotmat_tr, nr, + // nc); + // #else + // diskann::load_bin(rotmat_file, _rotmat_tr, nr, nc); + // #endif + // if (nr != this->_ndims || nc != this->_ndims) { + // diskann::cerr << "Error loading rotation matrix file" << std::endl; + // throw diskann::ANNException("Error loading rotation matrix file", -1, + // __FUNCSIG__, __FILE__, __LINE__); + // } + // _use_rotation = true; + // } + + // alloc and compute transpose + _tables_tr = new float[256 * this->_ndims]; + for (size_t i = 0; i < 256; i++) + { + for (size_t j = 0; j < this->_ndims; j++) + { + _tables_tr[j * 256 + i] = _tables[i * this->_ndims + j]; + } + } +} + +template uint32_t PQL2Distance::get_num_chunks() const +{ + return static_cast(_num_chunks); +} + +// REFACTOR: Instead of doing half the work in the caller and half in this +// function, we let this function +// do all of the work, making it easier for the caller. +template +void PQL2Distance::preprocess_query(const data_t *aligned_query, uint32_t dim, PQScratch &scratch) +{ + // Copy query vector to float and then to "rotated" query + for (size_t d = 0; d < dim; d++) + { + scratch.aligned_query_float[d] = (float)aligned_query[d]; + } + scratch.set_rotated_query(dim, aligned_query); + + for (uint32_t d = 0; d < _ndims; d++) + { + scratch.rotated_query[d] -= _centroid[d]; + } + std::vector tmp(_ndims, 0); + if (_use_rotation) + { + for (uint32_t d = 0; d < _ndims; d++) + { + for (uint32_t d1 = 0; d1 < _ndims; d1++) + { + tmp[d] += scratch.rotated_query[d1] * _rotmat_tr[d1 * _ndims + d]; + } + } + std::memcpy(scratch.rotated_query, tmp.data(), _ndims * sizeof(float)); + } + this->prepopulate_chunkwise_distances(scratch.rotated_query, scratch.aligned_pqtable_dist_scratch); +} + +template +void PQL2Distance::preprocessed_distance(PQScratch &pq_scratch, const uint32_t n_ids, float *dists_out) +{ + pq_dist_lookup(pq_scratch.aligned_pq_coord_scratch, n_ids, _num_chunks, pq_scratch.aligned_pqtable_dist_scratch, + dists_out); +} + +template +void PQL2Distance::preprocessed_distance(PQScratch &pq_scratch, const uint32_t n_ids, + std::vector &dists_out) +{ + pq_dist_lookup(pq_scratch.aligned_pq_coord_scratch, n_ids, _num_chunks, pq_scratch.aligned_pqtable_dist_scratch, + dists_out); +} + +template float PQL2Distance::brute_force_distance(const float *query_vec, uint8_t *base_vec) +{ + float res = 0; + for (size_t chunk = 0; chunk < _num_chunks; chunk++) + { + for (size_t j = _chunk_offsets[chunk]; j < _chunk_offsets[chunk + 1]; j++) + { + const float *centers_dim_vec = _tables_tr + (256 * j); + float diff = centers_dim_vec[base_vec[chunk]] - (query_vec[j]); + res += diff * diff; + } + } + return res; +} + +template +void PQL2Distance::prepopulate_chunkwise_distances(const float *query_vec, float *dist_vec) +{ + memset(dist_vec, 0, 256 * _num_chunks * sizeof(float)); + // chunk wise distance computation + for (size_t chunk = 0; chunk < _num_chunks; chunk++) + { + // sum (q-c)^2 for the dimensions associated with this chunk + float *chunk_dists = dist_vec + (256 * chunk); + for (size_t j = _chunk_offsets[chunk]; j < _chunk_offsets[chunk + 1]; j++) + { + const float *centers_dim_vec = _tables_tr + (256 * j); + for (size_t idx = 0; idx < 256; idx++) + { + double diff = centers_dim_vec[idx] - (query_vec[j]); + chunk_dists[idx] += (float)(diff * diff); + } + } + } +} + +template DISKANN_DLLEXPORT class PQL2Distance; +template DISKANN_DLLEXPORT class PQL2Distance; +template DISKANN_DLLEXPORT class PQL2Distance; + +} // namespace diskann \ No newline at end of file diff --git a/src/scratch.cpp b/src/scratch.cpp index e6305cd29..a88392f89 100644 --- a/src/scratch.cpp +++ b/src/scratch.cpp @@ -5,6 +5,7 @@ #include #include "scratch.h" +#include "pq_scratch.h" namespace diskann { @@ -24,13 +25,13 @@ InMemQueryScratch::InMemQueryScratch(uint32_t search_l, uint32_t indexing_l, throw diskann::ANNException(ss.str(), -1); } - alloc_aligned(((void **)&_aligned_query), aligned_dim * sizeof(T), alignment_factor * sizeof(T)); - memset(_aligned_query, 0, aligned_dim * sizeof(T)); + alloc_aligned(((void **)&this->_aligned_query_T), aligned_dim * sizeof(T), alignment_factor * sizeof(T)); + memset(this->_aligned_query_T, 0, aligned_dim * sizeof(T)); if (init_pq_scratch) - _pq_scratch = new PQScratch(MAX_GRAPH_DEGREE, aligned_dim); + this->_pq_scratch = new PQScratch(MAX_GRAPH_DEGREE, aligned_dim); else - _pq_scratch = nullptr; + this->_pq_scratch = nullptr; _occlude_factor.reserve(maxc); _inserted_into_pool_bs = new boost::dynamic_bitset<>(); @@ -71,12 +72,12 @@ template void InMemQueryScratch::resize_for_new_L(uint32_t new_l template InMemQueryScratch::~InMemQueryScratch() { - if (_aligned_query != nullptr) + if (this->_aligned_query_T != nullptr) { - aligned_free(_aligned_query); + aligned_free(this->_aligned_query_T); } - delete _pq_scratch; + delete this->_pq_scratch; delete _inserted_into_pool_bs; } @@ -97,12 +98,12 @@ template SSDQueryScratch::SSDQueryScratch(size_t aligned_dim, si diskann::alloc_aligned((void **)&coord_scratch, coord_alloc_size, 256); diskann::alloc_aligned((void **)§or_scratch, (size_t)MAX_N_SECTOR_READS * (size_t)SECTOR_LEN, SECTOR_LEN); - diskann::alloc_aligned((void **)&aligned_query_T, aligned_dim * sizeof(T), 8 * sizeof(T)); + diskann::alloc_aligned((void **)&this->_aligned_query_T, aligned_dim * sizeof(T), 8 * sizeof(T)); - _pq_scratch = new PQScratch(MAX_GRAPH_DEGREE, aligned_dim); + this->_pq_scratch = new PQScratch(MAX_GRAPH_DEGREE, aligned_dim); memset(coord_scratch, 0, coord_alloc_size); - memset(aligned_query_T, 0, aligned_dim * sizeof(T)); + memset(this->_aligned_query_T, 0, aligned_dim * sizeof(T)); visited.reserve(visited_reserve); full_retset.reserve(visited_reserve); @@ -112,9 +113,9 @@ template SSDQueryScratch::~SSDQueryScratch() { diskann::aligned_free((void *)coord_scratch); diskann::aligned_free((void *)sector_scratch); - diskann::aligned_free((void *)aligned_query_T); + diskann::aligned_free((void *)this->_aligned_query_T); - delete[] _pq_scratch; + delete[] this->_pq_scratch; } template From cd4a8e8cd18b5045ddee9c645bb2063eab7e7acb Mon Sep 17 00:00:00 2001 From: Gopal Srinivasa Date: Fri, 23 Jun 2023 22:49:09 +0530 Subject: [PATCH 02/15] Fixing compile issues after rebase to main --- include/abstract_data_store.h | 2 +- include/in_mem_data_store.h | 2 +- include/pq_data_store.h | 8 +++++++- src/in_mem_data_store.cpp | 4 ++-- src/index.cpp | 10 +++++----- src/pq_data_store.cpp | 9 +++++++-- 6 files changed, 23 insertions(+), 12 deletions(-) diff --git a/include/abstract_data_store.h b/include/abstract_data_store.h index 88542d974..ab909653e 100644 --- a/include/abstract_data_store.h +++ b/include/abstract_data_store.h @@ -99,7 +99,7 @@ template class AbstractDataStore // in the dataset virtual location_t calculate_medoid() const = 0; - virtual Distance *get_dist_fn() = 0; + virtual std::shared_ptr> get_dist_fn() const = 0; // search helpers // if the base data is aligned per the request of the metric, this will tell diff --git a/include/in_mem_data_store.h b/include/in_mem_data_store.h index 8fe467f65..9ad54df67 100644 --- a/include/in_mem_data_store.h +++ b/include/in_mem_data_store.h @@ -51,7 +51,7 @@ template class InMemDataStore : public AbstractDataStore *get_dist_fn() override; + virtual std::shared_ptr> get_dist_fn() const override; virtual size_t get_alignment_factor() const override; diff --git a/include/pq_data_store.h b/include/pq_data_store.h index f3f0d761e..e87152f91 100644 --- a/include/pq_data_store.h +++ b/include/pq_data_store.h @@ -11,7 +11,7 @@ template class PQDataStore : public AbstractDataStore { public: - PQDataStore(uint32_t dim, uint32_t num_points, uint32_t num_pq_chunks, + PQDataStore(size_t dim, location_t num_points, size_t num_pq_chunks, std::shared_ptr> distance_fn, std::shared_ptr> pq_distance_fn); ~PQDataStore(); @@ -60,6 +60,11 @@ template class PQDataStore : public AbstractDataStore virtual void get_distance(const data_t *preprocessed_query, const std::vector &ids, std::vector &distances, AbstractScratch *scratch_space) const override; + //We are returning the distance function that is used for full precision + //vectors here, not the PQ distance function. This is because the callers + //all are expecting a Distance not QuantizedDistance. + virtual std::shared_ptr> get_dist_fn() const override; + virtual location_t calculate_medoid() const override; virtual size_t get_alignment_factor() const override; @@ -82,6 +87,7 @@ template class PQDataStore : public AbstractDataStore bool _use_opq = false; Metric _distance_metric; + std::shared_ptr> _distance_fn = nullptr; std::shared_ptr> _pq_distance_fn = nullptr; }; } // namespace diskann diff --git a/src/in_mem_data_store.cpp b/src/in_mem_data_store.cpp index 4bf9e25a6..de4ccd5d8 100644 --- a/src/in_mem_data_store.cpp +++ b/src/in_mem_data_store.cpp @@ -361,9 +361,9 @@ template location_t InMemDataStore::calculate_medoid() return min_idx; } -template Distance *InMemDataStore::get_dist_fn() +template std::shared_ptr> InMemDataStore::get_dist_fn() const { - return this->_distance_fn.get(); + return this->_distance_fn; } template DISKANN_DLLEXPORT class InMemDataStore; diff --git a/src/index.cpp b/src/index.cpp index c73b26b80..4d9eba8e2 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -128,10 +128,10 @@ Index::Index(Metric m, const size_t dim, const size_t max_point } // REFACTOR TODO: This should move to a factory method and support OPQ. - _pq_distance_fn = std::make_shared>(_num_pq_chunks); + _pq_distance_fn = std::make_shared>((uint32_t)_num_pq_chunks); // REFACTOR TODO: Unlike Distance and DataStore, where distance object is ready when the data store // is constructed. Here the distance object will not be fully ready until populate_data() is called - _pq_data_store = std::make_shared>(_dim, total_internal_points, _num_pq_chunks, this->_distance, + _pq_data_store = std::make_shared>(_dim, (location_t)total_internal_points, _num_pq_chunks, this->_distance, _pq_distance_fn); // REFACTOR // alloc_aligned( @@ -156,7 +156,7 @@ Index::Index(const IndexConfig &index_config, std::unique_ptrget_dist_fn()); + _distance = _data_store->get_dist_fn(); // enable delete by default for dynamic index if (_dynamic_index) @@ -3288,7 +3288,7 @@ template void Index *dist_fast = (DistanceFastL2 *)_data_store->get_dist_fn(); + auto dist_fast = (DistanceFastL2*) (_data_store->get_dist_fn().get()); for (uint32_t i = 0; i < _nd; i++) { char *cur_node_offset = _opt_graph + i * _node_size; @@ -3337,7 +3337,7 @@ void Index::_search_with_optimized_layout(const DataType &query template void Index::search_with_optimized_layout(const T *query, size_t K, size_t L, uint32_t *indices) { - DistanceFastL2 *dist_fast = (DistanceFastL2 *)_data_store->get_dist_fn(); + DistanceFastL2 *dist_fast = (DistanceFastL2 *)(_data_store->get_dist_fn().get()); NeighborPriorityQueue retset(L); std::vector init_ids(L); diff --git a/src/pq_data_store.cpp b/src/pq_data_store.cpp index ff808e9af..f7bd1f858 100644 --- a/src/pq_data_store.cpp +++ b/src/pq_data_store.cpp @@ -12,11 +12,11 @@ namespace diskann // REFACTOR TODO: Assuming that num_pq_chunks is known already. Must verify if // this is true. template -PQDataStore::PQDataStore(uint32_t dim, uint32_t num_points, uint32_t num_pq_chunks, +PQDataStore::PQDataStore(size_t dim, location_t num_points, size_t num_pq_chunks, std::shared_ptr> distance_fn, std::shared_ptr> pq_distance_fn) : AbstractDataStore(num_points, dim), _quantized_data(nullptr), _num_chunks(num_pq_chunks), - _distance_metric(distance_fn->get_metric()), _pq_distance_fn(pq_distance_fn) + _distance_metric(distance_fn->get_metric()), _distance_fn(distance_fn), _pq_distance_fn(pq_distance_fn) { } @@ -208,6 +208,11 @@ template size_t PQDataStore::get_alignment_factor() co return 1; } +template std::shared_ptr> PQDataStore::get_dist_fn() const +{ + return _distance_fn; +} + template location_t PQDataStore::load_impl(const std::string &file_prefix) { if (_quantized_data != nullptr) From 3225759789312f4eb95b30894df312c0e8f10f79 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Mon, 26 Jun 2023 17:33:04 +0000 Subject: [PATCH 03/15] minor renaming functions --- CMakeLists.txt | 3 ++- include/pq_common.h | 6 +++++ include/pq_data_store.h | 12 ++++----- include/pq_distance.h | 20 --------------- include/pq_l2_distance.h | 6 ++--- include/pq_scratch.h | 2 +- include/quantized_distance.h | 2 +- src/index.cpp | 8 +++--- src/pq_flash_index.cpp | 4 +-- src/pq_l2_distance.cpp | 48 +++++++++++++++++------------------- 10 files changed, 48 insertions(+), 63 deletions(-) delete mode 100644 include/pq_distance.h diff --git a/CMakeLists.txt b/CMakeLists.txt index d4e3979b9..0854e0cbb 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -281,7 +281,8 @@ if(MSVC) else() set(ENV{TCMALLOC_LARGE_ALLOC_REPORT_THRESHOLD} 500000000000) # set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -g -DDEBUG -O0 -fsanitize=address -fsanitize=leak -fsanitize=undefined") - set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -g -DDEBUG -Wall -Wextra") + # set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -g -DDEBUG -Wall -Wextra") + set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -g -DDEBUG") if (NOT PYBIND) set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} -Ofast -DNDEBUG -march=native -mtune=native -ftree-vectorize") else() diff --git a/include/pq_common.h b/include/pq_common.h index 4cdb69438..fb564a8f2 100644 --- a/include/pq_common.h +++ b/include/pq_common.h @@ -22,4 +22,10 @@ inline std::string get_pivot_data_filename(const std::string &prefix, bool use_o return prefix + (use_opq ? "_opq" : "pq") + std::to_string(num_chunks) + "_pivots.bin"; } +inline std::string get_rotation_matrix_suffix(const std::string &pivot_data_filename) +{ + return pivot_data_filename + "_rotation_matrix.bin"; +} + + } // namespace diskann diff --git a/include/pq_data_store.h b/include/pq_data_store.h index e87152f91..82860531a 100644 --- a/include/pq_data_store.h +++ b/include/pq_data_store.h @@ -11,8 +11,7 @@ template class PQDataStore : public AbstractDataStore { public: - PQDataStore(size_t dim, location_t num_points, size_t num_pq_chunks, - std::shared_ptr> distance_fn, + PQDataStore(size_t dim, location_t num_points, size_t num_pq_chunks, std::shared_ptr> distance_fn, std::shared_ptr> pq_distance_fn); ~PQDataStore(); @@ -27,7 +26,8 @@ template class PQDataStore : public AbstractDataStore // vectors file. virtual size_t save(const std::string &file_prefix, const location_t num_points) override; - // Since base class function is pure virtual, we need to declare it here, even though alignent concept is not needed for Quantized data stores. + // Since base class function is pure virtual, we need to declare it here, even though alignent concept is not needed + // for Quantized data stores. virtual size_t get_aligned_dim() const override; // Populate quantized data from unaligned data using PQ functionality @@ -60,9 +60,9 @@ template class PQDataStore : public AbstractDataStore virtual void get_distance(const data_t *preprocessed_query, const std::vector &ids, std::vector &distances, AbstractScratch *scratch_space) const override; - //We are returning the distance function that is used for full precision - //vectors here, not the PQ distance function. This is because the callers - //all are expecting a Distance not QuantizedDistance. + // We are returning the distance function that is used for full precision + // vectors here, not the PQ distance function. This is because the callers + // all are expecting a Distance not QuantizedDistance. virtual std::shared_ptr> get_dist_fn() const override; virtual location_t calculate_medoid() const override; diff --git a/include/pq_distance.h b/include/pq_distance.h deleted file mode 100644 index 86054643b..000000000 --- a/include/pq_distance.h +++ /dev/null @@ -1,20 +0,0 @@ -#pragma once - -namespace diskann -{ -template class PQScratch; - -template class PQDistance -{ - // Should this class manage PQScratch? - public: - PQDistance(FixedChunkPQTable &pq_table) : _pq_table(pq_table) - { - } - virtual void preprocess_query(const data_t *query_vec, const size_t query_dim, - PQScratch *scratch_query) = 0; - - private: - FixedChunkPQTable &_pq_table; -}; -} // namespace diskann \ No newline at end of file diff --git a/include/pq_l2_distance.h b/include/pq_l2_distance.h index 7335be209..e6fc6e41b 100644 --- a/include/pq_l2_distance.h +++ b/include/pq_l2_distance.h @@ -17,7 +17,7 @@ template class PQL2Distance : public QuantizedDistance // load_pivot_data. Hence this. The TODO is whether we should check // that the num_chunks from the file is the same as this one. - PQL2Distance(uint32_t num_chunks); + PQL2Distance(uint32_t num_chunks, bool use_opq = false); virtual ~PQL2Distance() override; @@ -25,7 +25,7 @@ template class PQL2Distance : public QuantizedDistance virtual std::string get_quantized_vectors_filename(const std::string &prefix) const override; virtual std::string get_pivot_data_filename(const std::string &prefix) const override; - virtual std::string get_rotation_matrix_filename(const std::string &prefix) const override; + virtual std::string get_rotation_matrix_suffix(const std::string &pq_pivots_filename) const override; #ifdef EXEC_ENV_OLS virtual void load_pivot_data(MemoryMappedFiles &files, const std::string &pq_table_file, @@ -78,7 +78,7 @@ template class PQL2Distance : public QuantizedDistance float *_tables = nullptr; // pq_tables = float array of size [256 * ndims] uint64_t _ndims = 0; // ndims = true dimension of vectors uint64_t _num_chunks = 0; - bool _use_rotation = false; + bool _is_opq = false; uint32_t *_chunk_offsets = nullptr; float *_centroid = nullptr; float *_tables_tr = nullptr; // same as pq_tables, but col-major diff --git a/include/pq_scratch.h b/include/pq_scratch.h index 186d9b84d..5b599b5be 100644 --- a/include/pq_scratch.h +++ b/include/pq_scratch.h @@ -28,7 +28,7 @@ template struct PQScratch memset(rotated_query, 0, aligned_dim * sizeof(float)); } - void set_rotated_query(size_t dim, const T *query, const float norm = 1.0f) + void initialize(size_t dim, const T *query, const float norm = 1.0f) { for (size_t d = 0; d < dim; ++d) { diff --git a/include/quantized_distance.h b/include/quantized_distance.h index 2bb927c9d..d97d919a8 100644 --- a/include/quantized_distance.h +++ b/include/quantized_distance.h @@ -19,7 +19,7 @@ template class QuantizedDistance virtual bool is_opq() const = 0; virtual std::string get_quantized_vectors_filename(const std::string &prefix) const = 0; virtual std::string get_pivot_data_filename(const std::string &prefix) const = 0; - virtual std::string get_rotation_matrix_filename(const std::string &prefix) const = 0; + virtual std::string get_rotation_matrix_suffix(const std::string &pq_pivots_filename) const = 0; // Loading the PQ centroid table need not be part of the abstract class. // However, we want to indicate that this function will change once we have a diff --git a/src/index.cpp b/src/index.cpp index 4d9eba8e2..dff3ee26f 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -128,11 +128,11 @@ Index::Index(Metric m, const size_t dim, const size_t max_point } // REFACTOR TODO: This should move to a factory method and support OPQ. - _pq_distance_fn = std::make_shared>((uint32_t)_num_pq_chunks); + _pq_distance_fn = std::make_shared>((uint32_t)_num_pq_chunks, use_opq); // REFACTOR TODO: Unlike Distance and DataStore, where distance object is ready when the data store // is constructed. Here the distance object will not be fully ready until populate_data() is called - _pq_data_store = std::make_shared>(_dim, (location_t)total_internal_points, _num_pq_chunks, this->_distance, - _pq_distance_fn); + _pq_data_store = std::make_shared>(_dim, (location_t)total_internal_points, _num_pq_chunks, + this->_distance, _pq_distance_fn); // REFACTOR // alloc_aligned( // ((void **)&_pq_data), total_internal_points * _num_pq_chunks * sizeof(char), 8 * sizeof(char)); @@ -3288,7 +3288,7 @@ template void Index*) (_data_store->get_dist_fn().get()); + auto dist_fast = (DistanceFastL2 *)(_data_store->get_dist_fn().get()); for (uint32_t i = 0; i < _nd; i++) { char *cur_node_offset = _opt_graph + i * _node_size; diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index 158994a47..cf672ba98 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -1188,7 +1188,7 @@ void PQFlashIndex::cached_beam_search(const T *query1, const uint64_t { aligned_query_T[i] = (T)(aligned_query_T[i] / query_norm); } - pq_query_scratch->set_rotated_query(this->data_dim, aligned_query_T); + pq_query_scratch->initialize(this->data_dim, aligned_query_T); } else { @@ -1196,7 +1196,7 @@ void PQFlashIndex::cached_beam_search(const T *query1, const uint64_t { aligned_query_T[i] = query1[i]; } - pq_query_scratch->set_rotated_query(this->data_dim, aligned_query_T); + pq_query_scratch->initialize(this->data_dim, aligned_query_T); } // pointers to buffers for data diff --git a/src/pq_l2_distance.cpp b/src/pq_l2_distance.cpp index 24caa9baf..2949e2646 100644 --- a/src/pq_l2_distance.cpp +++ b/src/pq_l2_distance.cpp @@ -9,7 +9,7 @@ namespace diskann { -template PQL2Distance::PQL2Distance(uint32_t num_chunks) : _num_chunks(num_chunks) +template PQL2Distance::PQL2Distance(uint32_t num_chunks, bool use_opq) : _num_chunks(num_chunks), _is_opq(use_opq) { } @@ -31,7 +31,7 @@ template PQL2Distance::~PQL2Distance() template bool PQL2Distance::is_opq() const { - return false; + return this->_is_opq; } template @@ -42,7 +42,7 @@ std::string PQL2Distance::get_quantized_vectors_filename(const std::stri throw diskann::ANNException("Must set num_chunks before calling get_quantized_vectors_filename", -1, __FUNCSIG__, __FILE__, __LINE__); } - return diskann::get_quantized_vectors_filename(prefix, false, (uint32_t)_num_chunks); + return diskann::get_quantized_vectors_filename(prefix, _is_opq, (uint32_t)_num_chunks); } template std::string PQL2Distance::get_pivot_data_filename(const std::string &prefix) const { @@ -51,14 +51,12 @@ template std::string PQL2Distance::get_pivot_data_file throw diskann::ANNException("Must set num_chunks before calling get_pivot_data_filename", -1, __FUNCSIG__, __FILE__, __LINE__); } - return diskann::get_pivot_data_filename(prefix, false, (uint32_t)_num_chunks); + return diskann::get_pivot_data_filename(prefix, _is_opq, (uint32_t)_num_chunks); } template -std::string PQL2Distance::get_rotation_matrix_filename(const std::string &prefix) const +std::string PQL2Distance::get_rotation_matrix_suffix(const std::string &pq_pivots_filename) const { - // REFACTOR TODO: Currently, we are assuming that PQ doesn't have a rotation - // matrix. - return ""; + return diskann::get_rotation_matrix_suffix(pq_pivots_filename); } #ifdef EXEC_ENV_OLS @@ -163,21 +161,21 @@ void PQL2Distance::load_pivot_data(const std::string &pq_table_file, siz diskann::cout << "Loaded PQ Pivots: #ctrs: " << NUM_PQ_CENTROIDS << ", #dims: " << this->_ndims << ", #chunks: " << this->_num_chunks << std::endl; - // For PQ there will be no rotation matrix. - // if (file_exists(rotmat_file)) { - // #ifdef EXEC_ENV_OLS - // diskann::load_bin(files, rotmat_file, (float *&)rotmat_tr, nr, - // nc); - // #else - // diskann::load_bin(rotmat_file, _rotmat_tr, nr, nc); - // #endif - // if (nr != this->_ndims || nc != this->_ndims) { - // diskann::cerr << "Error loading rotation matrix file" << std::endl; - // throw diskann::ANNException("Error loading rotation matrix file", -1, - // __FUNCSIG__, __FILE__, __LINE__); - // } - // _use_rotation = true; - // } + // For OPQ there will be a rotation matrix to load. + if (this->_is_opq) { + std::string rotmat_file = get_rotation_matrix_suffix(pq_table_file); + #ifdef EXEC_ENV_OLS + diskann::load_bin(files, rotmat_file, (float *&)rotmat_tr, nr, + nc); + #else + diskann::load_bin(rotmat_file, _rotmat_tr, nr, nc); + #endif + if (nr != this->_ndims || nc != this->_ndims) { + diskann::cerr << "Error loading rotation matrix file" << std::endl; + throw diskann::ANNException("Error loading rotation matrix file", -1, + __FUNCSIG__, __FILE__, __LINE__); + } + } // alloc and compute transpose _tables_tr = new float[256 * this->_ndims]; @@ -206,14 +204,14 @@ void PQL2Distance::preprocess_query(const data_t *aligned_query, uint32_ { scratch.aligned_query_float[d] = (float)aligned_query[d]; } - scratch.set_rotated_query(dim, aligned_query); + scratch.initialize(dim, aligned_query); for (uint32_t d = 0; d < _ndims; d++) { scratch.rotated_query[d] -= _centroid[d]; } std::vector tmp(_ndims, 0); - if (_use_rotation) + if (_is_opq) { for (uint32_t d = 0; d < _ndims; d++) { From bcc8eb3ae47b993927ec01951cb8b2bc5373eb63 Mon Sep 17 00:00:00 2001 From: ravishankar Date: Tue, 27 Jun 2023 17:17:27 +0000 Subject: [PATCH 04/15] fixed small bug post rebasing with index factory --- include/pq_common.h | 1 - src/index.cpp | 26 +++++++++++++++++--------- src/pq_l2_distance.cpp | 29 +++++++++++++++-------------- 3 files changed, 32 insertions(+), 24 deletions(-) diff --git a/include/pq_common.h b/include/pq_common.h index fb564a8f2..c6a3a5739 100644 --- a/include/pq_common.h +++ b/include/pq_common.h @@ -27,5 +27,4 @@ inline std::string get_rotation_matrix_suffix(const std::string &pivot_data_file return pivot_data_filename + "_rotation_matrix.bin"; } - } // namespace diskann diff --git a/src/index.cpp b/src/index.cpp index dff3ee26f..01529fa84 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -118,16 +118,7 @@ Index::Index(Metric m, const size_t dim, const size_t max_point // Note: moved this to factory, keeping this for backward compatibility. _data_store = std::make_unique>((location_t)total_internal_points, _dim, this->_distance); - } - - if (_pq_dist) - { - if (_num_pq_chunks > _dim) - { - throw diskann::ANNException("ERROR: num_pq_chunks > dim", -1, __FUNCSIG__, __FILE__, __LINE__); - } - // REFACTOR TODO: This should move to a factory method and support OPQ. _pq_distance_fn = std::make_shared>((uint32_t)_num_pq_chunks, use_opq); // REFACTOR TODO: Unlike Distance and DataStore, where distance object is ready when the data store // is constructed. Here the distance object will not be fully ready until populate_data() is called @@ -139,6 +130,14 @@ Index::Index(Metric m, const size_t dim, const size_t max_point // std::memset(_pq_data, 0, total_internal_points * _num_pq_chunks * sizeof(char)); } + if (_pq_dist) + { + if (_num_pq_chunks > _dim) + { + throw diskann::ANNException("ERROR: num_pq_chunks > dim", -1, __FUNCSIG__, __FILE__, __LINE__); + } + } + _locks = std::vector(total_internal_points); if (enable_tags) @@ -158,6 +157,15 @@ Index::Index(const IndexConfig &index_config, std::unique_ptrget_dist_fn(); + // REFACTOR TODO: This should also move to factory method. + _pq_distance_fn = std::make_shared>((uint32_t)_num_pq_chunks, _use_opq); + + const size_t total_internal_points = _max_points + _num_frozen_pts; + // REFACTOR TODO: Unlike Distance and DataStore, where distance object is ready when the data store + // is constructed. Here the distance object will not be fully ready until populate_data() is called + _pq_data_store = std::make_shared>(_dim, (location_t)total_internal_points, _num_pq_chunks, + this->_distance, _pq_distance_fn); + // enable delete by default for dynamic index if (_dynamic_index) { diff --git a/src/pq_l2_distance.cpp b/src/pq_l2_distance.cpp index 2949e2646..c08744c35 100644 --- a/src/pq_l2_distance.cpp +++ b/src/pq_l2_distance.cpp @@ -9,7 +9,8 @@ namespace diskann { -template PQL2Distance::PQL2Distance(uint32_t num_chunks, bool use_opq) : _num_chunks(num_chunks), _is_opq(use_opq) +template +PQL2Distance::PQL2Distance(uint32_t num_chunks, bool use_opq) : _num_chunks(num_chunks), _is_opq(use_opq) { } @@ -162,20 +163,20 @@ void PQL2Distance::load_pivot_data(const std::string &pq_table_file, siz << ", #chunks: " << this->_num_chunks << std::endl; // For OPQ there will be a rotation matrix to load. - if (this->_is_opq) { + if (this->_is_opq) + { std::string rotmat_file = get_rotation_matrix_suffix(pq_table_file); - #ifdef EXEC_ENV_OLS - diskann::load_bin(files, rotmat_file, (float *&)rotmat_tr, nr, - nc); - #else - diskann::load_bin(rotmat_file, _rotmat_tr, nr, nc); - #endif - if (nr != this->_ndims || nc != this->_ndims) { - diskann::cerr << "Error loading rotation matrix file" << std::endl; - throw diskann::ANNException("Error loading rotation matrix file", -1, - __FUNCSIG__, __FILE__, __LINE__); - } - } +#ifdef EXEC_ENV_OLS + diskann::load_bin(files, rotmat_file, (float *&)rotmat_tr, nr, nc); +#else + diskann::load_bin(rotmat_file, _rotmat_tr, nr, nc); +#endif + if (nr != this->_ndims || nc != this->_ndims) + { + diskann::cerr << "Error loading rotation matrix file" << std::endl; + throw diskann::ANNException("Error loading rotation matrix file", -1, __FUNCSIG__, __FILE__, __LINE__); + } + } // alloc and compute transpose _tables_tr = new float[256 * this->_ndims]; From c6dd6cc38ec38641362c4b876f3ce4933e48c6c9 Mon Sep 17 00:00:00 2001 From: Gopal Srinivasa Date: Mon, 7 Aug 2023 17:21:59 +0530 Subject: [PATCH 05/15] Changes to index factory to support PQDataStore --- include/abstract_data_store.h | 2 +- include/in_mem_data_store.h | 2 + include/index.h | 14 ++- include/index_factory.h | 17 ++- include/pq_data_store.h | 2 + src/abstract_data_store.cpp | 7 +- src/in_mem_data_store.cpp | 12 +- src/index.cpp | 223 ++++++++++------------------------ src/index_factory.cpp | 64 +++++++--- src/pq_data_store.cpp | 3 + 10 files changed, 158 insertions(+), 188 deletions(-) diff --git a/include/abstract_data_store.h b/include/abstract_data_store.h index ab909653e..509cf878a 100644 --- a/include/abstract_data_store.h +++ b/include/abstract_data_store.h @@ -91,7 +91,7 @@ template class AbstractDataStore // Specific overload for index.cpp. // REFACTOR TODO: Check if the default implementation is sufficient for most cases. virtual void get_distance(const data_t *preprocessed_query, const std::vector &ids, - std::vector &distances, AbstractScratch *scratch_space) const; + std::vector &distances, AbstractScratch *scratch_space) const = 0; virtual float get_distance(const location_t loc1, const location_t loc2) const = 0; // stats of the data stored in store diff --git a/include/in_mem_data_store.h b/include/in_mem_data_store.h index 9ad54df67..ec10fdc81 100644 --- a/include/in_mem_data_store.h +++ b/include/in_mem_data_store.h @@ -48,6 +48,8 @@ template class InMemDataStore : public AbstractDataStore *scratch) const override; + virtual void get_distance(const data_t *preprocessed_query, const std::vector &ids, + std::vector &distances, AbstractScratch *scratch_space) const override; virtual location_t calculate_medoid() const override; diff --git a/include/index.h b/include/index.h index 33ebc3ae8..6246a6ca0 100644 --- a/include/index.h +++ b/include/index.h @@ -68,7 +68,10 @@ template clas const bool concurrent_consolidate = false, const bool pq_dist_build = false, const size_t num_pq_chunks = 0, const bool use_opq = false); - DISKANN_DLLEXPORT Index(const IndexConfig &index_config, std::unique_ptr> data_store + //REFACTOR TODO: Ideally, this should take an AbstractPQDataStore, but for now, all our PQDataStores are in-mem + //so this should be ok. + DISKANN_DLLEXPORT Index(const IndexConfig &index_config, std::shared_ptr> data_store_for_reranking, + std::shared_ptr> data_store_for_candidates /* std::unique_ptr graph_store*/); DISKANN_DLLEXPORT ~Index(); @@ -253,9 +256,10 @@ template clas // with iterate_to_fixed_point. std::vector get_init_ids(); - std::pair iterate_to_fixed_point(const T *node_coords, const uint32_t Lindex, + //The query to use is placed in scratch->aligned_query + std::pair iterate_to_fixed_point(InMemQueryScratch *scratch, const uint32_t Lindex, const std::vector &init_ids, - InMemQueryScratch *scratch, bool use_filter, + bool use_filter, const std::vector &filters, bool search_invocation); void search_for_point_and_prune(int location, uint32_t Lindex, std::vector &pruned_list, @@ -336,7 +340,7 @@ template clas std::shared_ptr> _distance; // Data - std::unique_ptr> _data_store; + std::shared_ptr> _data_store; char *_opt_graph = nullptr; // Graph related data structures @@ -400,7 +404,7 @@ template clas // REFACTOR // uint8_t *_pq_data = nullptr; std::shared_ptr> _pq_distance_fn = nullptr; - std::shared_ptr> _pq_data_store = nullptr; + std::shared_ptr> _pq_data_store = nullptr; bool _pq_generated = false; FixedChunkPQTable _pq_table; diff --git a/include/index_factory.h b/include/index_factory.h index 3d1eb7992..5c2e4386b 100644 --- a/include/index_factory.h +++ b/include/index_factory.h @@ -1,6 +1,9 @@ #include "index.h" #include "abstract_graph_store.h" #include "in_mem_graph_store.h" +#include "pq_data_store.h" + + namespace diskann { @@ -13,11 +16,21 @@ class IndexFactory private: void check_config(); + template + Distance* construct_inmem_distance_fn(); + template - std::unique_ptr> construct_datastore(DataStoreStrategy stratagy, size_t num_points, + std::shared_ptr> construct_datastore(DataStoreStrategy stratagy, size_t num_points, size_t dimension); - std::unique_ptr construct_graphstore(GraphStoreStrategy stratagy, size_t size); + std::shared_ptr construct_graphstore(GraphStoreStrategy stratagy, size_t size); + + //For now PQDataStore incorporates within itself all variants of quantization that we support. In the + //future it may be necessary to introduce an AbstractPQDataStore class to spearate various quantization + //flavours. + template + std::shared_ptr> construct_pq_datastore(DataStoreStrategy strategy, size_t num_points, + size_t dimension); template std::unique_ptr create_instance(); diff --git a/include/pq_data_store.h b/include/pq_data_store.h index 82860531a..c6f731bdf 100644 --- a/include/pq_data_store.h +++ b/include/pq_data_store.h @@ -7,6 +7,8 @@ namespace diskann { + //REFACTOR TODO: By default, the PQDataStore is an in-memory datastore because both Vamana and + //DiskANN treat it the same way. But with DiskPQ, that may need to change. template class PQDataStore : public AbstractDataStore { diff --git a/src/abstract_data_store.cpp b/src/abstract_data_store.cpp index d81545a37..1696a9701 100644 --- a/src/abstract_data_store.cpp +++ b/src/abstract_data_store.cpp @@ -44,12 +44,7 @@ void AbstractDataStore::preprocess_query(const data_t *query, AbstractSc { } -template -void AbstractDataStore::get_distance(const data_t *preprocessed_query, const std::vector &ids, - std::vector &distances, - AbstractScratch *scratch_space) const -{ -} + template DISKANN_DLLEXPORT class AbstractDataStore; template DISKANN_DLLEXPORT class AbstractDataStore; diff --git a/src/in_mem_data_store.cpp b/src/in_mem_data_store.cpp index de4ccd5d8..ec523dc3b 100644 --- a/src/in_mem_data_store.cpp +++ b/src/in_mem_data_store.cpp @@ -187,7 +187,6 @@ void InMemDataStore::get_distance(const data_t *query, const location_t const uint32_t location_count, float *distances, AbstractScratch *scratch_space) const { - assert(scratch_space == nullptr); // Scratch space should only be used by PQ data store. for (location_t i = 0; i < location_count; i++) { distances[i] = _distance_fn->compare(query, _data + locations[i] * _aligned_dim, (uint32_t)this->_aligned_dim); @@ -201,6 +200,17 @@ float InMemDataStore::get_distance(const location_t loc1, const location (uint32_t)this->_aligned_dim); } +template +void InMemDataStore::get_distance(const data_t *preprocessed_query, const std::vector &ids, + std::vector &distances, AbstractScratch *scratch_space) const +{ + for (int i = 0; i < ids.size(); i++) + { + distances[i] = + _distance_fn->compare(preprocessed_query, _data + ids[i] * _aligned_dim, (uint32_t)this->_aligned_dim); + } +} + template location_t InMemDataStore::expand(const location_t new_size) { if (new_size == this->capacity()) diff --git a/src/index.cpp b/src/index.cpp index 01529fa84..848885ac5 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -69,19 +69,6 @@ Index::Index(Metric m, const size_t dim, const size_t max_point throw ANNException("ERROR: Dynamic Indexing must have tags enabled.", -1, __FUNCSIG__, __FILE__, __LINE__); } - if (_pq_dist) - { - if (dynamic_index) - throw ANNException("ERROR: Dynamic Indexing not supported with PQ distance based " - "index construction", - -1, __FUNCSIG__, __FILE__, __LINE__); - if (m == diskann::Metric::INNER_PRODUCT) - throw ANNException("ERROR: Inner product metrics not yet supported " - "with PQ distance " - "base index", - -1, __FUNCSIG__, __FILE__, __LINE__); - } - if (dynamic_index && _num_frozen_pts == 0) { _num_frozen_pts = 1; @@ -117,24 +104,16 @@ Index::Index(Metric m, const size_t dim, const size_t max_point } // Note: moved this to factory, keeping this for backward compatibility. _data_store = - std::make_unique>((location_t)total_internal_points, _dim, this->_distance); - - _pq_distance_fn = std::make_shared>((uint32_t)_num_pq_chunks, use_opq); - // REFACTOR TODO: Unlike Distance and DataStore, where distance object is ready when the data store - // is constructed. Here the distance object will not be fully ready until populate_data() is called - _pq_data_store = std::make_shared>(_dim, (location_t)total_internal_points, _num_pq_chunks, - this->_distance, _pq_distance_fn); - // REFACTOR - // alloc_aligned( - // ((void **)&_pq_data), total_internal_points * _num_pq_chunks * sizeof(char), 8 * sizeof(char)); - // std::memset(_pq_data, 0, total_internal_points * _num_pq_chunks * sizeof(char)); - } + std::make_shared>((location_t)total_internal_points, _dim, this->_distance); - if (_pq_dist) - { - if (_num_pq_chunks > _dim) - { - throw diskann::ANNException("ERROR: num_pq_chunks > dim", -1, __FUNCSIG__, __FILE__, __LINE__); + if (_pq_dist) { + _pq_distance_fn = std::make_shared>((uint32_t)_num_pq_chunks, use_opq); + // REFACTOR TODO: Unlike Distance and DataStore, where distance object is ready when the data store + // is constructed. Here the distance object will not be fully ready until populate_data() is called + _pq_data_store = std::make_shared>(_dim, (location_t)total_internal_points, _num_pq_chunks, + this->_distance, _pq_distance_fn); + } else { + _pq_data_store = _data_store; } } @@ -148,23 +127,20 @@ Index::Index(Metric m, const size_t dim, const size_t max_point } template -Index::Index(const IndexConfig &index_config, std::unique_ptr> data_store) +Index::Index(const IndexConfig &index_config, std::shared_ptr> data_store, + std::shared_ptr> pq_data_store) : Index(index_config.metric, index_config.dimension, index_config.max_points, index_config.dynamic_index, index_config.enable_tags, index_config.concurrent_consolidate, index_config.pq_dist_build, index_config.num_pq_chunks, index_config.use_opq, index_config.num_frozen_pts, false) { - _data_store = std::move(data_store); + _data_store = data_store; _distance = _data_store->get_dist_fn(); - // REFACTOR TODO: This should also move to factory method. - _pq_distance_fn = std::make_shared>((uint32_t)_num_pq_chunks, _use_opq); - const size_t total_internal_points = _max_points + _num_frozen_pts; // REFACTOR TODO: Unlike Distance and DataStore, where distance object is ready when the data store // is constructed. Here the distance object will not be fully ready until populate_data() is called - _pq_data_store = std::make_shared>(_dim, (location_t)total_internal_points, _num_pq_chunks, - this->_distance, _pq_distance_fn); + _pq_data_store = pq_data_store; // enable delete by default for dynamic index if (_dynamic_index) @@ -200,13 +176,6 @@ template Index::~I LockGuard lg(lock); } - // if (this->_distance != nullptr) - //{ - // delete this->_distance; - // this->_distance = nullptr; - // } - // REFACTOR - if (_opt_graph != nullptr) { delete[] _opt_graph; @@ -888,12 +857,6 @@ template int Index template uint32_t Index::calculate_entry_point() { - if (_pq_dist) - { - // REFACTOR TODO: This function returns a random point. Must change - // to actually compute the medoid. - return _pq_data_store->calculate_medoid(); - } // REFACTOR TODO: This function does not support multi-threaded calculation of medoid. // Must revisit if perf is a concern. return _data_store->calculate_medoid(); @@ -952,7 +915,7 @@ bool Index::detect_common_filters(uint32_t point_id, bool searc template std::pair Index::iterate_to_fixed_point( - const T *query, const uint32_t Lsize, const std::vector &init_ids, InMemQueryScratch *scratch, + InMemQueryScratch *scratch, const uint32_t Lsize, const std::vector &init_ids, bool use_filter, const std::vector &filter_label, bool search_invocation) { std::vector &expanded_nodes = scratch->pool(); @@ -964,42 +927,16 @@ std::pair Index::iterate_to_fixed_point( std::vector &dist_scratch = scratch->dist_scratch(); assert(id_scratch.size() == 0); - // REFACTOR - // T *aligned_query = scratch->aligned_query(); - // memcpy(aligned_query, query, _dim * sizeof(T)); - // if (_normalize_vecs) - //{ - // normalize((float *)aligned_query, _dim); - // } - T *aligned_query = scratch->aligned_query(); float *pq_dists = nullptr; - // uint8_t *pq_coord_scratch = nullptr; - // Intialize PQ related scratch to use PQ based distances - if (_pq_dist) - { - _pq_data_store->preprocess_query(aligned_query, scratch); - - // REFACTOR: Commented out. - // Get scratch spaces - // PQScratch *pq_query_scratch = scratch->pq_scratch(); - // query_float = pq_query_scratch->aligned_query_float; - // query_rotated = pq_query_scratch->rotated_query; - // pq_dists = pq_query_scratch->aligned_pqtable_dist_scratch; - //// Copy query vector to float and then to "rotated" query - // for (size_t d = 0; d < _dim; d++) { - // query_float[d] = (float)aligned_query[d]; - // } - // pq_query_scratch->set(_dim, aligned_query); - - //// center the query and rotate if we have a rotation matrix - //_pq_table.preprocess_query(query_rotated); - //_pq_table.populate_chunk_distances(query_rotated, pq_dists); - - // pq_coord_scratch = scratch->pq_scratch()->aligned_pq_coord_scratch; - } + //REFACTOR PQ: Preprocess the query with the appropriate "pq" datastore. It could + // also be the "actual" datastore in which case this is a no-op. + //if (_pq_dist) + //{ + _pq_data_store->preprocess_query(aligned_query, scratch); + //} if (expanded_nodes.size() > 0 || id_scratch.size() > 0) { @@ -1031,11 +968,6 @@ std::pair Index::iterate_to_fixed_point( auto compute_dists = [this, scratch, pq_dists](const std::vector &ids, std::vector &dists_out) { // REFACTOR _pq_data_store->get_distance(scratch->aligned_query(), ids, dists_out, scratch); - - // diskann::aggregate_coords(ids, this->_pq_data, this->_num_pq_chunks, - // pq_coord_scratch); - // diskann::pq_dist_lookup(pq_coord_scratch, ids.size(), this->_num_pq_chunks, - // pq_dists, dists_out); }; // Initialize the candidate pool with starting points @@ -1066,20 +998,20 @@ std::pair Index::iterate_to_fixed_point( } float distance; - if (_pq_dist) - { - // REFACTOR - std::vector ids = {id}; - std::vector distances = {std::numeric_limits::max()}; - _pq_data_store->get_distance(aligned_query, ids, distances, scratch); + //if (_pq_dist) + //{ + // REFACTOR PQ. We pay a small price in efficiency for better code structure. + uint32_t ids[] = {id}; + float distances[] = {std::numeric_limits::max()}; + _pq_data_store->get_distance(aligned_query, ids, 1, distances, scratch); distance = distances[0]; // pq_dist_lookup(pq_coord_scratch, 1, this->_num_pq_chunks, pq_dists, // &distance); - } - else - { - distance = _data_store->get_distance(aligned_query, id); - } + //} + //else + //{ + // distance = _data_store->get_distance(aligned_query, id); + //} Neighbor nn = Neighbor(id, distance); best_L_nodes.insert(nn); } @@ -1152,27 +1084,28 @@ std::pair Index::iterate_to_fixed_point( } // Compute distances to unvisited nodes in the expansion - if (_pq_dist) - { + // REFACTOR PQ + //if (_pq_dist) + //{ assert(dist_scratch.capacity() >= id_scratch.size()); compute_dists(id_scratch, dist_scratch); - } - else - { - assert(dist_scratch.size() == 0); - for (size_t m = 0; m < id_scratch.size(); ++m) - { - uint32_t id = id_scratch[m]; - - if (m + 1 < id_scratch.size()) - { - auto nextn = id_scratch[m + 1]; - _data_store->prefetch_vector(nextn); - } - - dist_scratch.push_back(_data_store->get_distance(aligned_query, id)); - } - } + //} + //else + //{ + // assert(dist_scratch.size() == 0); + // for (size_t m = 0; m < id_scratch.size(); ++m) + // { + // uint32_t id = id_scratch[m]; + + // if (m + 1 < id_scratch.size()) + // { + // auto nextn = id_scratch[m + 1]; + // _data_store->prefetch_vector(nextn); + // } + + // dist_scratch.push_back(_data_store->get_distance(aligned_query, id)); + // } + //} cmps += (uint32_t)id_scratch.size(); // Insert pairs into the pool of candidates @@ -1196,7 +1129,7 @@ void Index::search_for_point_and_prune(int location, uint32_t L if (!use_filter) { _data_store->get_vector(location, scratch->aligned_query()); - iterate_to_fixed_point(scratch->aligned_query(), Lindex, init_ids, scratch, false, unused_filter_label, false); + iterate_to_fixed_point(scratch, Lindex, init_ids, false, unused_filter_label, false); } else { @@ -1205,7 +1138,7 @@ void Index::search_for_point_and_prune(int location, uint32_t L filter_specific_start_nodes.emplace_back(_label_to_medoid_id[x]); _data_store->get_vector(location, scratch->aligned_query()); - iterate_to_fixed_point(scratch->aligned_query(), filteredLindex, filter_specific_start_nodes, scratch, true, + iterate_to_fixed_point(scratch, filteredLindex, filter_specific_start_nodes, true, _pts_to_labels[location], false); } @@ -1346,6 +1279,7 @@ void Index::prune_neighbors(const uint32_t location, std::vecto _max_observed_degree = (std::max)(_max_observed_degree, range); // If using _pq_build, over-write the PQ distances with actual distances + //REFACTOR PQ: TODO: How to get rid of this!? if (_pq_dist) { for (auto &ngh : pool) @@ -1790,16 +1724,6 @@ void Index::build(const T *data, const size_t num_points_to_loa _nd = num_points_to_load; _data_store->populate_data(data, (location_t)num_points_to_load); - - // REFACTOR - // memcpy((char *)_data, (char *)data, _aligned_dim * _nd * sizeof(T)); - // if (_normalize_vecs) - //{ - // for (size_t i = 0; i < num_points_to_load; i++) - // { - // normalize(_data + _aligned_dim * i, _aligned_dim); - // } - // } } build_with_data_populated(parameters, tags); @@ -1867,13 +1791,16 @@ void Index::build(const char *filename, const size_t num_points throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } + //REFACTOR PQ TODO: We can remove this if and add a check in the InMemDataStore + //to not populate_data if it has been called once. if (_pq_dist) { // REFACTOR #ifdef EXEC_ENV_OLS - throw ANNException("load_pq_centroid_bin should not be called when " - "EXEC_ENV_OLS is defined.", - -1, __FUNCSIG__, __FILE__, __LINE__); + std::stringstream ss; + ss << "PQ Build is not supported in DLVS environment (i.e. if EXEC_ENV_OLS is defined)" << std::endl; + diskann::cerr << ss.str() << std::endl; + throw ANNException(ss.str(),-1, __FUNCSIG__, __FILE__, __LINE__); #else // REFACTOR TODO: Both in the previous code and in the current PQDataStore, // we are writing the PQ files in the same path as the input file. Now we @@ -1884,29 +1811,6 @@ void Index::build(const char *filename, const size_t num_points // as-is for now. _pq_data_store->populate_data(filename, 0U); #endif - // double p_val = std::min( - // 1.0, ((double)MAX_PQ_TRAINING_SET_SIZE / (double)file_num_points)); - // - // std::string suffix = _use_opq ? "_opq" : "_pq"; - // suffix += std::to_string(_num_pq_chunks); - // auto pq_pivots_file = std::string(filename) + suffix + "_pivots.bin"; - // auto pq_compressed_file = - // std::string(filename) + suffix + "_compressed.bin"; - // generate_quantized_data(std::string(filename), pq_pivots_file, - // pq_compressed_file, _dist_metric, p_val, - // _num_pq_chunks, _use_opq); - // - // copy_aligned_data_from_file(pq_compressed_file.c_str(), _pq_data, - // file_num_points, _num_pq_chunks, - // _num_pq_chunks); - //#ifdef EXEC_ENV_OLS - // throw ANNException( - // "load_pq_centroid_bin should not be called when " - // "EXEC_ENV_OLS is defined.", - // -1, __FUNCSIG__, __FILE__, __LINE__); - //#else - // _pq_table.load_pq_centroid_bin(pq_pivots_file.c_str(), _num_pq_chunks); - //#endif } _data_store->populate_data(filename, 0U); @@ -2220,7 +2124,7 @@ std::pair Index::search(const T *query, con _distance->preprocess_query(query, _data_store->get_dims(), scratch->aligned_query()); auto retval = - iterate_to_fixed_point(scratch->aligned_query(), L, init_ids, scratch, false, unused_filter_label, true); + iterate_to_fixed_point(scratch, L, init_ids, false, unused_filter_label, true); NeighborPriorityQueue &best_L_nodes = scratch->best_l_nodes(); @@ -2321,7 +2225,7 @@ std::pair Index::search_with_filters(const // T *aligned_query = scratch->aligned_query(); // memcpy(aligned_query, query, _dim * sizeof(T)); _distance->preprocess_query(query, _data_store->get_dims(), scratch->aligned_query()); - auto retval = iterate_to_fixed_point(scratch->aligned_query(), L, init_ids, scratch, true, filter_vec, true); + auto retval = iterate_to_fixed_point(scratch, L, init_ids, true, filter_vec, true); auto best_L_nodes = scratch->best_l_nodes(); @@ -2400,7 +2304,7 @@ size_t Index::search_with_tags(const T *query, const uint64_t K const std::vector unused_filter_label; _distance->preprocess_query(query, _data_store->get_dims(), scratch->aligned_query()); - iterate_to_fixed_point(scratch->aligned_query(), L, init_ids, scratch, false, unused_filter_label, true); + iterate_to_fixed_point(scratch, L, init_ids, false, unused_filter_label, true); NeighborPriorityQueue &best_L_nodes = scratch->best_l_nodes(); assert(best_L_nodes.size() <= L); @@ -2469,6 +2373,7 @@ template void Index Distance* IndexFactory::construct_inmem_distance_fn() +{ + if (_config->metric == diskann::Metric::COSINE && std::is_same::value) { + return (Distance *)new AVXNormalizedCosineDistanceFloat(); + } else { + return (Distance *)get_distance_function(_config->metric); + } +} + template -std::unique_ptr> IndexFactory::construct_datastore(DataStoreStrategy strategy, size_t num_points, +std::shared_ptr> IndexFactory::construct_datastore(DataStoreStrategy strategy, size_t num_points, size_t dimension) { const size_t total_internal_points = num_points + _config->num_frozen_pts; @@ -58,28 +69,40 @@ std::unique_ptr> IndexFactory::construct_datastore(DataStor switch (strategy) { case MEMORY: - if (_config->metric == diskann::Metric::COSINE && std::is_same::value) - { - distance.reset((Distance *)new AVXNormalizedCosineDistanceFloat()); - return std::make_unique>((location_t)total_internal_points, dimension, distance); - } - else - { - distance.reset((Distance *)get_distance_function(_config->metric)); - return std::make_unique>((location_t)total_internal_points, dimension, distance); - } - break; + distance.reset(construct_inmem_distance_fn()); + return std::make_unique>((location_t)total_internal_points, dimension, distance); default: break; } return nullptr; } -std::unique_ptr IndexFactory::construct_graphstore(GraphStoreStrategy, size_t size) +std::shared_ptr IndexFactory::construct_graphstore(GraphStoreStrategy, size_t size) { return std::make_unique(size); } +template +std::shared_ptr> IndexFactory::construct_pq_datastore(DataStoreStrategy strategy, size_t num_points, size_t dimension) +{ + std::shared_ptr> distance_fn; + std::shared_ptr> quantized_distance_fn; + + quantized_distance_fn = std::make_shared>(_config->num_pq_chunks, _config->use_opq); + switch (strategy) + { + case MEMORY: + distance_fn.reset(construct_inmem_distance_fn()); + return std::make_unique>(dimension, (location_t)(num_points + _config->num_frozen_pts), + _config->num_pq_chunks, distance_fn, + quantized_distance_fn); + default: + //REFACTOR TODO: We do support diskPQ - so we may need to add a new class for SSDPQDataStore! + break; + } + return nullptr; +} + template std::unique_ptr IndexFactory::create_instance() { @@ -87,7 +110,20 @@ std::unique_ptr IndexFactory::create_instance() size_t dim = _config->dimension; // auto graph_store = construct_graphstore(_config->graph_strategy, num_points); auto data_store = construct_datastore(_config->data_strategy, num_points, dim); - return std::make_unique>(*_config, std::move(data_store)); + std::shared_ptr> pq_data_store = nullptr; + + if (_config->data_strategy == DataStoreStrategy::MEMORY && _config->pq_dist_build) + { + pq_data_store = construct_pq_datastore(_config->data_strategy, num_points, dim); + } + else + { + pq_data_store = data_store; + } + + //REFACTOR TODO: Must construct in-memory PQDatastore if strategy == ONDISK and must construct + //in-mem and on-disk PQDataStore if strategy == ONDISK and diskPQ is required. + return std::make_unique>(*_config, data_store, pq_data_store); } std::unique_ptr IndexFactory::create_instance(const std::string &data_type, const std::string &tag_type, diff --git a/src/pq_data_store.cpp b/src/pq_data_store.cpp index f7bd1f858..a46d86fae 100644 --- a/src/pq_data_store.cpp +++ b/src/pq_data_store.cpp @@ -18,6 +18,9 @@ PQDataStore::PQDataStore(size_t dim, location_t num_points, size_t num_p : AbstractDataStore(num_points, dim), _quantized_data(nullptr), _num_chunks(num_pq_chunks), _distance_metric(distance_fn->get_metric()), _distance_fn(distance_fn), _pq_distance_fn(pq_distance_fn) { + if (num_pq_chunks > dim) { + throw diskann::ANNException("ERROR: num_pq_chunks > dim", -1, __FUNCSIG__, __FILE__, __LINE__); + } } template PQDataStore::~PQDataStore() From 342e19cbaa05aa5cf33f90873cfbd940962f4168 Mon Sep 17 00:00:00 2001 From: Gopal Srinivasa Date: Thu, 24 Aug 2023 23:52:32 +0530 Subject: [PATCH 06/15] Merged graph_store and pq_data_store --- apps/build_stitched_index.cpp | 2 +- apps/utils/count_bfs_levels.cpp | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/apps/build_stitched_index.cpp b/apps/build_stitched_index.cpp index 7767a4bc6..ae07a8d42 100644 --- a/apps/build_stitched_index.cpp +++ b/apps/build_stitched_index.cpp @@ -286,7 +286,7 @@ void prune_and_save(path final_index_path_prefix, path full_index_path_prefix, p diskann::get_bin_metadata(input_data_path, number_of_label_points, dimension); - diskann::Index index(diskann::Metric::L2, dimension, number_of_label_points, nullptr, nullptr, 0, false, false); + diskann::Index index(diskann::Metric::L2, dimension, number_of_label_points, nullptr, nullptr, 0, false, false, false, false, 0, false); // not searching this index, set search_l to 0 index.load(full_index_path_prefix.c_str(), num_threads, 1); diff --git a/apps/utils/count_bfs_levels.cpp b/apps/utils/count_bfs_levels.cpp index 1ec8225db..9bf84b75d 100644 --- a/apps/utils/count_bfs_levels.cpp +++ b/apps/utils/count_bfs_levels.cpp @@ -27,7 +27,7 @@ template void bfs_count(const std::string &index_path, uint32_t dat { using TagT = uint32_t; using LabelT = uint32_t; - diskann::Index index(diskann::Metric::L2, data_dims, 0, nullptr, nullptr, 0, false, false); + diskann::Index index(diskann::Metric::L2, data_dims, 0, nullptr, nullptr, 0, false, false, false, false, 0, false); std::cout << "Index class instantiated" << std::endl; index.load(index_path.c_str(), 1, 100); std::cout << "Index loaded" << std::endl; From b8aeafd3aaaa7a400ac65fe58de7947e2800d59f Mon Sep 17 00:00:00 2001 From: Gopal Srinivasa Date: Mon, 28 Aug 2023 16:17:56 +0530 Subject: [PATCH 07/15] Implementing preprocessing for inmemdatastore --- include/abstract_data_store.h | 10 ++++------ include/in_mem_data_store.h | 6 ++++-- src/abstract_data_store.cpp | 6 ------ src/in_mem_data_store.cpp | 17 +++++++++++++++++ 4 files changed, 25 insertions(+), 14 deletions(-) diff --git a/include/abstract_data_store.h b/include/abstract_data_store.h index aa3dea24b..1e2206923 100644 --- a/include/abstract_data_store.h +++ b/include/abstract_data_store.h @@ -80,12 +80,10 @@ template class AbstractDataStore // num_points) to zero virtual void copy_vectors(const location_t from_loc, const location_t to_loc, const location_t num_points) = 0; - // Some datastores like PQ stores need a preprocessing step before querying. - // Optionally, a scratch object can be passed in to avoid memory allocations - // Default implementation does nothing. - // REFACTOR TODO: Currently, we take an aligned_query as parameter, but this - // should change and this function should do the necessary alignment. - virtual void preprocess_query(const data_t *aligned_query, AbstractScratch *query_scratch = nullptr) const; + //With the PQ Data Store PR, we have also changed iterate_to_fixed_point to NOT take the query + //from the scratch object. Therefore every data store has to implement preprocess_query which + //at the least will be to copy the query into the scratch object. So making this pure virtual. + virtual void preprocess_query(const data_t *aligned_query, AbstractScratch *query_scratch = nullptr) const = 0; // distance functions. virtual float get_distance(const data_t *query, const location_t loc) const = 0; virtual void get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, diff --git a/include/in_mem_data_store.h b/include/in_mem_data_store.h index 33abc12ff..dd152b343 100644 --- a/include/in_mem_data_store.h +++ b/include/in_mem_data_store.h @@ -44,10 +44,12 @@ template class InMemDataStore : public AbstractDataStore *query_scratch) const override; + + virtual float get_distance(const data_t *preprocessed_query, const location_t loc) const override; virtual float get_distance(const location_t loc1, const location_t loc2) const override; - virtual void get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, + virtual void get_distance(const data_t *preprocessed_query, const location_t *locations, const uint32_t location_count, float *distances, AbstractScratch *scratch) const override; virtual void get_distance(const data_t *preprocessed_query, const std::vector &ids, std::vector &distances, AbstractScratch *scratch_space) const override; diff --git a/src/abstract_data_store.cpp b/src/abstract_data_store.cpp index 1696a9701..791d38618 100644 --- a/src/abstract_data_store.cpp +++ b/src/abstract_data_store.cpp @@ -39,12 +39,6 @@ template location_t AbstractDataStore::resize(const lo } } -template -void AbstractDataStore::preprocess_query(const data_t *query, AbstractScratch *query_scratch) const -{ -} - - template DISKANN_DLLEXPORT class AbstractDataStore; template DISKANN_DLLEXPORT class AbstractDataStore; diff --git a/src/in_mem_data_store.cpp b/src/in_mem_data_store.cpp index 2a67b0ecb..1dccfe056 100644 --- a/src/in_mem_data_store.cpp +++ b/src/in_mem_data_store.cpp @@ -2,6 +2,7 @@ // Licensed under the MIT license. #include +#include "abstract_scratch.h" #include "in_mem_data_store.h" #include "utils.h" @@ -178,6 +179,22 @@ template void InMemDataStore::prefetch_vector(const lo sizeof(data_t) * _aligned_dim); } +template +void InMemDataStore::preprocess_query(const data_t *query, AbstractScratch *query_scratch) const +{ + if (query_scratch != nullptr ) + { + memcpy(query_scratch->aligned_query_T(), query, sizeof(data_t) * this->get_dims()); + } + else + { + std::stringstream ss; + ss << "In InMemDataStore::preprocess_query: Query scratch is null"; + diskann::cerr << ss.str() << std::endl; + throw diskann::ANNException(ss.str(), -1); + } +} + template float InMemDataStore::get_distance(const data_t *query, const location_t loc) const { return _distance_fn->compare(query, _data + _aligned_dim * loc, (uint32_t)_aligned_dim); From 604bd9207c6d261e21b821d57da673d5d8a49982 Mon Sep 17 00:00:00 2001 From: Gopal Srinivasa Date: Wed, 6 Sep 2023 20:53:36 +0530 Subject: [PATCH 08/15] Incorporating code review comments --- include/abstract_data_store.h | 1 - include/abstract_scratch.h | 2 +- include/index.h | 7 ++-- include/pq_data_store.h | 10 +++-- include/quantized_distance.h | 2 +- include/scratch.h | 1 - src/index.cpp | 77 ++++------------------------------- src/index_factory.cpp | 8 ++-- src/pq_data_store.cpp | 8 ++-- 9 files changed, 29 insertions(+), 87 deletions(-) diff --git a/include/abstract_data_store.h b/include/abstract_data_store.h index 1e2206923..498723292 100644 --- a/include/abstract_data_store.h +++ b/include/abstract_data_store.h @@ -89,7 +89,6 @@ template class AbstractDataStore virtual void get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, float *distances, AbstractScratch *scratch_space = nullptr) const = 0; // Specific overload for index.cpp. - // REFACTOR TODO: Check if the default implementation is sufficient for most cases. virtual void get_distance(const data_t *preprocessed_query, const std::vector &ids, std::vector &distances, AbstractScratch *scratch_space) const = 0; virtual float get_distance(const location_t loc1, const location_t loc2) const = 0; diff --git a/include/abstract_scratch.h b/include/abstract_scratch.h index 57a2f8e34..45c240198 100644 --- a/include/abstract_scratch.h +++ b/include/abstract_scratch.h @@ -12,7 +12,7 @@ template struct AbstractScratch AbstractScratch() = default; // This class does not take any responsibilty for memory management of // its members. It is the responsibility of the derived classes to do so. - virtual ~AbstractScratch(){}; + virtual ~AbstractScratch() = default; // Scratch objects should not be copied AbstractScratch(const AbstractScratch &) = delete; diff --git a/include/index.h b/include/index.h index bf41cd76c..c321f5de6 100644 --- a/include/index.h +++ b/include/index.h @@ -22,7 +22,6 @@ #include "in_mem_graph_store.h" #include "abstract_index.h" -// REFACTOR #include "quantized_distance.h" #include "pq_data_store.h" @@ -63,9 +62,9 @@ template clas // Constructor for incremental index DISKANN_DLLEXPORT Index(Metric m, const size_t dim, const size_t max_points, const std::shared_ptr index_parameters, - const std::shared_ptr index_search_params, const size_t num_frozen_pts, - const bool dynamic_index, const bool enable_tags, const bool concurrent_consolidate, - const bool pq_dist_build, const size_t num_pq_chunks, const bool use_opq); + const std::shared_ptr index_search_params, const size_t num_frozen_pts = 0, + const bool dynamic_index = false, const bool enable_tags = false, const bool concurrent_consolidate = false, + const bool pq_dist_build = false, const size_t num_pq_chunks = 0, const bool use_opq = false); DISKANN_DLLEXPORT ~Index(); diff --git a/include/pq_data_store.h b/include/pq_data_store.h index 8c7288fd9..f5df3c80c 100644 --- a/include/pq_data_store.h +++ b/include/pq_data_store.h @@ -13,8 +13,10 @@ template class PQDataStore : public AbstractDataStore { public: - PQDataStore(size_t dim, location_t num_points, size_t num_pq_chunks, std::shared_ptr> distance_fn, - std::shared_ptr> pq_distance_fn); + PQDataStore(size_t dim, location_t num_points, size_t num_pq_chunks, std::unique_ptr> distance_fn, + std::unique_ptr> pq_distance_fn); + PQDataStore(const PQDataStore&) = delete; + PQDataStore &operator=(const PQDataStore&) = delete; ~PQDataStore(); // Load quantized vectors from a set of files. Here filename is treated @@ -89,7 +91,7 @@ template class PQDataStore : public AbstractDataStore bool _use_opq = false; Metric _distance_metric; - std::shared_ptr> _distance_fn = nullptr; - std::shared_ptr> _pq_distance_fn = nullptr; + std::unique_ptr> _distance_fn = nullptr; + std::unique_ptr> _pq_distance_fn = nullptr; }; } // namespace diskann diff --git a/include/quantized_distance.h b/include/quantized_distance.h index d97d919a8..8e60f896e 100644 --- a/include/quantized_distance.h +++ b/include/quantized_distance.h @@ -14,7 +14,7 @@ template class QuantizedDistance QuantizedDistance() = default; QuantizedDistance(const QuantizedDistance &) = delete; QuantizedDistance &operator=(const QuantizedDistance &) = delete; - virtual ~QuantizedDistance(){}; + virtual ~QuantizedDistance() = default; virtual bool is_opq() const = 0; virtual std::string get_quantized_vectors_filename(const std::string &prefix) const = 0; diff --git a/include/scratch.h b/include/scratch.h index 89791d76c..49324159e 100644 --- a/include/scratch.h +++ b/include/scratch.h @@ -28,7 +28,6 @@ template class InMemQueryScratch : public AbstractScratch { public: ~InMemQueryScratch(); - // REFACTOR TODO: move all parameters to a new class. InMemQueryScratch(uint32_t search_l, uint32_t indexing_l, uint32_t r, uint32_t maxc, size_t dim, size_t aligned_dim, size_t alignment_factor, bool init_pq_scratch = false); void resize_for_new_L(uint32_t new_search_l); diff --git a/src/index.cpp b/src/index.cpp index a91f0cd7b..4b527dbc7 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -19,9 +19,7 @@ #ifdef _WINDOWS #include #endif -// REFACTOR TODO: Must move to factory. -#include "pq_scratch.h" -#include "pq_l2_distance.h" + #include "index.h" #define MAX_POINTS_FOR_USING_BITSET 10000000 @@ -763,12 +761,7 @@ std::pair Index::iterate_to_fixed_point( float *pq_dists = nullptr; - //REFACTOR PQ: Preprocess the query with the appropriate "pq" datastore. It could - // also be the "actual" datastore in which case this is a no-op. - //if (_pq_dist) - //{ - _pq_data_store->preprocess_query(aligned_query, scratch); - //} + _pq_data_store->preprocess_query(aligned_query, scratch); if (expanded_nodes.size() > 0 || id_scratch.size() > 0) { @@ -798,7 +791,6 @@ std::pair Index::iterate_to_fixed_point( // Lambda to batch compute query<-> node distances in PQ space auto compute_dists = [this, scratch, pq_dists](const std::vector &ids, std::vector &dists_out) { - // REFACTOR _pq_data_store->get_distance(scratch->aligned_query(), ids, dists_out, scratch); }; @@ -830,20 +822,11 @@ std::pair Index::iterate_to_fixed_point( } float distance; - //if (_pq_dist) - //{ - // REFACTOR PQ. We pay a small price in efficiency for better code structure. - uint32_t ids[] = {id}; - float distances[] = {std::numeric_limits::max()}; - _pq_data_store->get_distance(aligned_query, ids, 1, distances, scratch); - distance = distances[0]; - // pq_dist_lookup(pq_coord_scratch, 1, this->_num_pq_chunks, pq_dists, - // &distance); - //} - //else - //{ - // distance = _data_store->get_distance(aligned_query, id); - //} + uint32_t ids[] = {id}; + float distances[] = {std::numeric_limits::max()}; + _pq_data_store->get_distance(aligned_query, ids, 1, distances, scratch); + distance = distances[0]; + Neighbor nn = Neighbor(id, distance); best_L_nodes.insert(nn); } @@ -915,29 +898,8 @@ std::pair Index::iterate_to_fixed_point( } } - // Compute distances to unvisited nodes in the expansion - // REFACTOR PQ - //if (_pq_dist) - //{ - assert(dist_scratch.capacity() >= id_scratch.size()); - compute_dists(id_scratch, dist_scratch); - //} - //else - //{ - // assert(dist_scratch.size() == 0); - // for (size_t m = 0; m < id_scratch.size(); ++m) - // { - // uint32_t id = id_scratch[m]; - - // if (m + 1 < id_scratch.size()) - // { - // auto nextn = id_scratch[m + 1]; - // _data_store->prefetch_vector(nextn); - // } - - // dist_scratch.push_back(_data_store->get_distance(aligned_query, id)); - // } - //} + assert(dist_scratch.capacity() >= id_scratch.size()); + compute_dists(id_scratch, dist_scratch); cmps += (uint32_t)id_scratch.size(); // Insert pairs into the pool of candidates @@ -1586,8 +1548,6 @@ void Index::build(const char *filename, const size_t num_points << " points, but " << "index can support only " << _max_points << " points as specified in constructor." << std::endl; - // REFACTOR PQDataStore will take care of its memory - // if (_pq_dist) aligned_free(_pq_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } @@ -1597,8 +1557,6 @@ void Index::build(const char *filename, const size_t num_points stream << "ERROR: Driver requests loading " << num_points_to_load << " points and file has only " << file_num_points << " points." << std::endl; - // REFACTOR: PQDataStore will take care of its memory - // if (_pq_dist) aligned_free(_pq_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } @@ -1609,8 +1567,6 @@ void Index::build(const char *filename, const size_t num_points << "but file has " << file_dim << " dimension." << std::endl; diskann::cerr << stream.str() << std::endl; - // REFACTOR: PQDataStore will take care of its memory - // if (_pq_dist) aligned_free(_pq_data); throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } @@ -1618,7 +1574,6 @@ void Index::build(const char *filename, const size_t num_points //to not populate_data if it has been called once. if (_pq_dist) { - // REFACTOR #ifdef EXEC_ENV_OLS std::stringstream ss; ss << "PQ Build is not supported in DLVS environment (i.e. if EXEC_ENV_OLS is defined)" << std::endl; @@ -3023,12 +2978,6 @@ template void Index void -// Index::optimize_index_layout() -//{ // use after build or load -//} - // REFACTOR: This should be an OptimizedDataStore class template void Index::optimize_index_layout() { // use after build or load @@ -3065,14 +3014,6 @@ template void Index -// void Index::search_with_optimized_layout(const T *query, -// size_t K, size_t L, uint32_t *indices) -//{ -//} - template void Index::_search_with_optimized_layout(const DataType &query, size_t K, size_t L, uint32_t *indices) { diff --git a/src/index_factory.cpp b/src/index_factory.cpp index a29939b05..f153130b0 100644 --- a/src/index_factory.cpp +++ b/src/index_factory.cpp @@ -94,16 +94,16 @@ std::shared_ptr> IndexFactory::construct_pq_datastore(DataStoreSt size_t dimension, Metric m, size_t num_pq_chunks, bool use_opq) { - std::shared_ptr> distance_fn; - std::shared_ptr> quantized_distance_fn; + std::unique_ptr> distance_fn; + std::unique_ptr> quantized_distance_fn; - quantized_distance_fn = std::make_shared>((uint32_t)num_pq_chunks, use_opq); + quantized_distance_fn = std::move(std::make_unique>((uint32_t)num_pq_chunks, use_opq)); switch (strategy) { case DataStoreStrategy::MEMORY: distance_fn.reset(construct_inmem_distance_fn(m)); return std::make_shared>(dimension, (location_t)(num_points), num_pq_chunks, - distance_fn, quantized_distance_fn); + std::move(distance_fn), std::move(quantized_distance_fn)); default: // REFACTOR TODO: We do support diskPQ - so we may need to add a new class for SSDPQDataStore! break; diff --git a/src/pq_data_store.cpp b/src/pq_data_store.cpp index 899726418..2618ce899 100644 --- a/src/pq_data_store.cpp +++ b/src/pq_data_store.cpp @@ -13,14 +13,16 @@ namespace diskann // this is true. template PQDataStore::PQDataStore(size_t dim, location_t num_points, size_t num_pq_chunks, - std::shared_ptr> distance_fn, - std::shared_ptr> pq_distance_fn) + std::unique_ptr> distance_fn, + std::unique_ptr> pq_distance_fn) : AbstractDataStore(num_points, dim), _quantized_data(nullptr), _num_chunks(num_pq_chunks), - _distance_metric(distance_fn->get_metric()), _distance_fn(distance_fn), _pq_distance_fn(pq_distance_fn) + _distance_metric(distance_fn->get_metric()) { if (num_pq_chunks > dim) { throw diskann::ANNException("ERROR: num_pq_chunks > dim", -1, __FUNCSIG__, __FILE__, __LINE__); } + _distance_fn = std::move(distance_fn); + _pq_distance_fn = std::move(pq_distance_fn); } template PQDataStore::~PQDataStore() From 3ba680d69482a6255e1ae22b929b87f32137dc1b Mon Sep 17 00:00:00 2001 From: rakri Date: Wed, 27 Sep 2023 04:22:06 +0000 Subject: [PATCH 09/15] minor bugfix for PQ data allocation --- src/pq_data_store.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pq_data_store.cpp b/src/pq_data_store.cpp index 2618ce899..6207b75e2 100644 --- a/src/pq_data_store.cpp +++ b/src/pq_data_store.cpp @@ -75,7 +75,7 @@ template void PQDataStore::populate_data(const std::st _pq_distance_fn->is_opq()); // REFACTOR TODO: Not sure of the alignment. Just copying from index.cpp - alloc_aligned(((void **)&_quantized_data), file_num_points * _num_chunks * sizeof(uint8_t), 8); + alloc_aligned(((void **)&_quantized_data), file_num_points * _num_chunks * sizeof(uint8_t), 1); copy_aligned_data_from_file(compressed_file.c_str(), _quantized_data, file_num_points, _num_chunks, _num_chunks); #ifdef EXEC_ENV_OLS From 84f1ee414c0c33fb51a7eeb103d5dacc808e781a Mon Sep 17 00:00:00 2001 From: rakri Date: Sat, 7 Oct 2023 15:51:29 +0000 Subject: [PATCH 10/15] clang-formatted --- apps/build_stitched_index.cpp | 3 ++- apps/utils/count_bfs_levels.cpp | 3 ++- include/abstract_data_store.h | 17 ++++++------ include/in_mem_data_store.h | 7 ++--- include/index.h | 8 ++---- include/index_factory.h | 8 ++---- include/pq_data_store.h | 10 +++---- src/abstract_data_store.cpp | 1 - src/in_mem_data_store.cpp | 6 ++--- src/index.cpp | 34 ++++++++++++----------- src/index_factory.cpp | 48 +++++++++++++++++++-------------- src/pq_data_store.cpp | 5 ++-- 12 files changed, 78 insertions(+), 72 deletions(-) diff --git a/apps/build_stitched_index.cpp b/apps/build_stitched_index.cpp index ae07a8d42..60e38c1be 100644 --- a/apps/build_stitched_index.cpp +++ b/apps/build_stitched_index.cpp @@ -286,7 +286,8 @@ void prune_and_save(path final_index_path_prefix, path full_index_path_prefix, p diskann::get_bin_metadata(input_data_path, number_of_label_points, dimension); - diskann::Index index(diskann::Metric::L2, dimension, number_of_label_points, nullptr, nullptr, 0, false, false, false, false, 0, false); + diskann::Index index(diskann::Metric::L2, dimension, number_of_label_points, nullptr, nullptr, 0, false, false, + false, false, 0, false); // not searching this index, set search_l to 0 index.load(full_index_path_prefix.c_str(), num_threads, 1); diff --git a/apps/utils/count_bfs_levels.cpp b/apps/utils/count_bfs_levels.cpp index 9bf84b75d..6dd2d6233 100644 --- a/apps/utils/count_bfs_levels.cpp +++ b/apps/utils/count_bfs_levels.cpp @@ -27,7 +27,8 @@ template void bfs_count(const std::string &index_path, uint32_t dat { using TagT = uint32_t; using LabelT = uint32_t; - diskann::Index index(diskann::Metric::L2, data_dims, 0, nullptr, nullptr, 0, false, false, false, false, 0, false); + diskann::Index index(diskann::Metric::L2, data_dims, 0, nullptr, nullptr, 0, false, false, false, + false, 0, false); std::cout << "Index class instantiated" << std::endl; index.load(index_path.c_str(), 1, 100); std::cout << "Index loaded" << std::endl; diff --git a/include/abstract_data_store.h b/include/abstract_data_store.h index 606b9fb7f..c5c5a322e 100644 --- a/include/abstract_data_store.h +++ b/include/abstract_data_store.h @@ -80,10 +80,11 @@ template class AbstractDataStore // num_points) to zero virtual void copy_vectors(const location_t from_loc, const location_t to_loc, const location_t num_points) = 0; - //With the PQ Data Store PR, we have also changed iterate_to_fixed_point to NOT take the query - //from the scratch object. Therefore every data store has to implement preprocess_query which - //at the least will be to copy the query into the scratch object. So making this pure virtual. - virtual void preprocess_query(const data_t *aligned_query, AbstractScratch *query_scratch = nullptr) const = 0; + // With the PQ Data Store PR, we have also changed iterate_to_fixed_point to NOT take the query + // from the scratch object. Therefore every data store has to implement preprocess_query which + // at the least will be to copy the query into the scratch object. So making this pure virtual. + virtual void preprocess_query(const data_t *aligned_query, + AbstractScratch *query_scratch = nullptr) const = 0; // distance functions. virtual float get_distance(const data_t *query, const location_t loc) const = 0; virtual void get_distance(const data_t *query, const location_t *locations, const uint32_t location_count, @@ -98,10 +99,10 @@ template class AbstractDataStore // in the dataset virtual location_t calculate_medoid() const = 0; - //REFACTOR PQ TODO: Each data store knows about its distance function, so this is - //redundant. However, we don't have an OptmizedDataStore yet, and to preserve code - //compability, we are exposing this function. - virtual Distance* get_dist_fn() const = 0; + // REFACTOR PQ TODO: Each data store knows about its distance function, so this is + // redundant. However, we don't have an OptmizedDataStore yet, and to preserve code + // compability, we are exposing this function. + virtual Distance *get_dist_fn() const = 0; // search helpers // if the base data is aligned per the request of the metric, this will tell diff --git a/include/in_mem_data_store.h b/include/in_mem_data_store.h index dd152b343..d1ad795f6 100644 --- a/include/in_mem_data_store.h +++ b/include/in_mem_data_store.h @@ -49,14 +49,15 @@ template class InMemDataStore : public AbstractDataStore *scratch) const override; + virtual void get_distance(const data_t *preprocessed_query, const location_t *locations, + const uint32_t location_count, float *distances, + AbstractScratch *scratch) const override; virtual void get_distance(const data_t *preprocessed_query, const std::vector &ids, std::vector &distances, AbstractScratch *scratch_space) const override; virtual location_t calculate_medoid() const override; - virtual Distance* get_dist_fn() const override; + virtual Distance *get_dist_fn() const override; virtual size_t get_alignment_factor() const override; diff --git a/include/index.h b/include/index.h index b6bf2c300..199171020 100644 --- a/include/index.h +++ b/include/index.h @@ -68,8 +68,6 @@ template clas const bool pq_dist_build = false, const size_t num_pq_chunks = 0, const bool use_opq = false, const bool filtered_index = false); - - DISKANN_DLLEXPORT ~Index(); // Saves graph, data, metadata and associated tags. @@ -255,10 +253,9 @@ template clas // with iterate_to_fixed_point. std::vector get_init_ids(); - //The query to use is placed in scratch->aligned_query + // The query to use is placed in scratch->aligned_query std::pair iterate_to_fixed_point(InMemQueryScratch *scratch, const uint32_t Lindex, - const std::vector &init_ids, - bool use_filter, + const std::vector &init_ids, bool use_filter, const std::vector &filters, bool search_invocation); void search_for_point_and_prune(int location, uint32_t Lindex, std::vector &pruned_list, @@ -340,7 +337,6 @@ template clas // Data std::shared_ptr> _data_store; - // Graph related data structures std::unique_ptr _graph_store; diff --git a/include/index_factory.h b/include/index_factory.h index c1a3dd47f..80bc40dba 100644 --- a/include/index_factory.h +++ b/include/index_factory.h @@ -3,8 +3,6 @@ #include "in_mem_graph_store.h" #include "pq_data_store.h" - - namespace diskann { class IndexFactory @@ -13,15 +11,13 @@ class IndexFactory DISKANN_DLLEXPORT explicit IndexFactory(const IndexConfig &config); DISKANN_DLLEXPORT std::unique_ptr create_instance(); - DISKANN_DLLEXPORT static std::unique_ptr construct_graphstore( const GraphStoreStrategy stratagy, const size_t size, const size_t reserve_graph_degree); template DISKANN_DLLEXPORT static std::shared_ptr> construct_datastore(DataStoreStrategy stratagy, size_t num_points, - size_t dimension, - Metric m); + size_t dimension, Metric m); // For now PQDataStore incorporates within itself all variants of quantization that we support. In the // future it may be necessary to introduce an AbstractPQDataStore class to spearate various quantization // flavours. @@ -33,7 +29,7 @@ class IndexFactory template static Distance *construct_inmem_distance_fn(Metric m); private: - void check_config(); + void check_config(); template std::unique_ptr create_instance(); diff --git a/include/pq_data_store.h b/include/pq_data_store.h index f5df3c80c..7c0cb5fe0 100644 --- a/include/pq_data_store.h +++ b/include/pq_data_store.h @@ -7,16 +7,16 @@ namespace diskann { - //REFACTOR TODO: By default, the PQDataStore is an in-memory datastore because both Vamana and - //DiskANN treat it the same way. But with DiskPQ, that may need to change. +// REFACTOR TODO: By default, the PQDataStore is an in-memory datastore because both Vamana and +// DiskANN treat it the same way. But with DiskPQ, that may need to change. template class PQDataStore : public AbstractDataStore { public: PQDataStore(size_t dim, location_t num_points, size_t num_pq_chunks, std::unique_ptr> distance_fn, std::unique_ptr> pq_distance_fn); - PQDataStore(const PQDataStore&) = delete; - PQDataStore &operator=(const PQDataStore&) = delete; + PQDataStore(const PQDataStore &) = delete; + PQDataStore &operator=(const PQDataStore &) = delete; ~PQDataStore(); // Load quantized vectors from a set of files. Here filename is treated @@ -67,7 +67,7 @@ template class PQDataStore : public AbstractDataStore // We are returning the distance function that is used for full precision // vectors here, not the PQ distance function. This is because the callers // all are expecting a Distance not QuantizedDistance. - virtual Distance* get_dist_fn() const override; + virtual Distance *get_dist_fn() const override; virtual location_t calculate_medoid() const override; diff --git a/src/abstract_data_store.cpp b/src/abstract_data_store.cpp index 791d38618..0cff0152e 100644 --- a/src/abstract_data_store.cpp +++ b/src/abstract_data_store.cpp @@ -39,7 +39,6 @@ template location_t AbstractDataStore::resize(const lo } } - template DISKANN_DLLEXPORT class AbstractDataStore; template DISKANN_DLLEXPORT class AbstractDataStore; template DISKANN_DLLEXPORT class AbstractDataStore; diff --git a/src/in_mem_data_store.cpp b/src/in_mem_data_store.cpp index 1dccfe056..cc7acf615 100644 --- a/src/in_mem_data_store.cpp +++ b/src/in_mem_data_store.cpp @@ -182,7 +182,7 @@ template void InMemDataStore::prefetch_vector(const lo template void InMemDataStore::preprocess_query(const data_t *query, AbstractScratch *query_scratch) const { - if (query_scratch != nullptr ) + if (query_scratch != nullptr) { memcpy(query_scratch->aligned_query_T(), query, sizeof(data_t) * this->get_dims()); } @@ -218,7 +218,7 @@ float InMemDataStore::get_distance(const location_t loc1, const location (uint32_t)this->_aligned_dim); } -template +template void InMemDataStore::get_distance(const data_t *preprocessed_query, const std::vector &ids, std::vector &distances, AbstractScratch *scratch_space) const { @@ -389,7 +389,7 @@ template location_t InMemDataStore::calculate_medoid() return min_idx; } -template Distance* InMemDataStore::get_dist_fn() const +template Distance *InMemDataStore::get_dist_fn() const { return this->_distance_fn.get(); } diff --git a/src/index.cpp b/src/index.cpp index 5db9ebb23..97413de07 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -30,8 +30,8 @@ namespace diskann // (bin), and initialize max_points template Index::Index(const IndexConfig &index_config, std::shared_ptr> data_store, - std::unique_ptr graph_store, - std::shared_ptr> pq_data_store) + std::unique_ptr graph_store, + std::shared_ptr> pq_data_store) : _dist_metric(index_config.metric), _dim(index_config.dimension), _max_points(index_config.max_points), _num_frozen_pts(index_config.num_frozen_pts), _dynamic_index(index_config.dynamic_index), _enable_tags(index_config.enable_tags), _indexingMaxC(DEFAULT_MAXC), _query_scratch(nullptr), @@ -142,10 +142,13 @@ Index::Index(Metric m, const size_t dim, const size_t max_point (size_t)((index_parameters == nullptr ? 0 : index_parameters->max_degree) * defaults::GRAPH_SLACK_FACTOR * 1.05))) { - if (_pq_dist) { - _pq_data_store = - IndexFactory::construct_pq_datastore(DataStoreStrategy::MEMORY, max_points + num_frozen_pts, dim, m, num_pq_chunks, use_opq); - } else { + if (_pq_dist) + { + _pq_data_store = IndexFactory::construct_pq_datastore(DataStoreStrategy::MEMORY, max_points + num_frozen_pts, + dim, m, num_pq_chunks, use_opq); + } + else + { _pq_data_store = _data_store; } } @@ -784,8 +787,8 @@ bool Index::detect_common_filters(uint32_t point_id, bool searc template std::pair Index::iterate_to_fixed_point( - InMemQueryScratch *scratch, const uint32_t Lsize, const std::vector &init_ids, - bool use_filter, const std::vector &filter_labels, bool search_invocation) + InMemQueryScratch *scratch, const uint32_t Lsize, const std::vector &init_ids, bool use_filter, + const std::vector &filter_labels, bool search_invocation) { std::vector &expanded_nodes = scratch->pool(); NeighborPriorityQueue &best_L_nodes = scratch->best_l_nodes(); @@ -1143,7 +1146,7 @@ void Index::prune_neighbors(const uint32_t location, std::vecto } // If using _pq_build, over-write the PQ distances with actual distances - //REFACTOR PQ: TODO: How to get rid of this!? + // REFACTOR PQ: TODO: How to get rid of this!? if (_pq_dist) { for (auto &ngh : pool) @@ -1629,15 +1632,15 @@ void Index::build(const char *filename, const size_t num_points throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } - //REFACTOR PQ TODO: We can remove this if and add a check in the InMemDataStore - //to not populate_data if it has been called once. + // REFACTOR PQ TODO: We can remove this if and add a check in the InMemDataStore + // to not populate_data if it has been called once. if (_pq_dist) { #ifdef EXEC_ENV_OLS - std::stringstream ss; + std::stringstream ss; ss << "PQ Build is not supported in DLVS environment (i.e. if EXEC_ENV_OLS is defined)" << std::endl; diskann::cerr << ss.str() << std::endl; - throw ANNException(ss.str(),-1, __FUNCSIG__, __FILE__, __LINE__); + throw ANNException(ss.str(), -1, __FUNCSIG__, __FILE__, __LINE__); #else // REFACTOR TODO: Both in the previous code and in the current PQDataStore, // we are writing the PQ files in the same path as the input file. Now we @@ -1957,8 +1960,7 @@ std::pair Index::search(const T *query, con _data_store->preprocess_query(query, scratch); - auto retval = - iterate_to_fixed_point(scratch, L, init_ids, false, unused_filter_label, true); + auto retval = iterate_to_fixed_point(scratch, L, init_ids, false, unused_filter_label, true); NeighborPriorityQueue &best_L_nodes = scratch->best_l_nodes(); @@ -2228,7 +2230,7 @@ template void Index Distance* IndexFactory::construct_inmem_distance_fn(Metric metric) +template Distance *IndexFactory::construct_inmem_distance_fn(Metric metric) { - if (metric == diskann::Metric::COSINE && std::is_same::value) { + if (metric == diskann::Metric::COSINE && std::is_same::value) + { return (Distance *)new AVXNormalizedCosineDistanceFloat(); - } else { + } + else + { return (Distance *)get_distance_function(metric); } } template -std::shared_ptr> IndexFactory::construct_datastore(DataStoreStrategy strategy, size_t total_internal_points, - size_t dimension, Metric metric) +std::shared_ptr> IndexFactory::construct_datastore(DataStoreStrategy strategy, + size_t total_internal_points, size_t dimension, + Metric metric) { std::unique_ptr> distance; switch (strategy) { case DataStoreStrategy::MEMORY: distance.reset(construct_inmem_distance_fn(metric)); - return std::make_shared>((location_t)total_internal_points, dimension, std::move(distance)); + return std::make_shared>((location_t)total_internal_points, dimension, + std::move(distance)); default: break; } @@ -120,20 +124,24 @@ std::unique_ptr IndexFactory::create_instance() auto data_store = construct_datastore(_config->data_strategy, num_points, dim, _config->metric); std::shared_ptr> pq_data_store = nullptr; - if (_config->data_strategy == DataStoreStrategy::MEMORY && _config->pq_dist_build) { - pq_data_store = construct_pq_datastore(_config->data_strategy, num_points + _config->num_frozen_pts, dim, _config->metric, - _config->num_pq_chunks, _config->use_opq); - } else { + if (_config->data_strategy == DataStoreStrategy::MEMORY && _config->pq_dist_build) + { + pq_data_store = + construct_pq_datastore(_config->data_strategy, num_points + _config->num_frozen_pts, dim, + _config->metric, _config->num_pq_chunks, _config->use_opq); + } + else + { pq_data_store = data_store; } size_t max_reserve_degree = (size_t)(defaults::GRAPH_SLACK_FACTOR * 1.05 * - (_config->index_write_params == nullptr ? 0 : _config->index_write_params->max_degree)); + (_config->index_write_params == nullptr ? 0 : _config->index_write_params->max_degree)); std::unique_ptr graph_store = construct_graphstore(_config->graph_strategy, num_points + _config->num_frozen_pts, max_reserve_degree); - //REFACTOR TODO: Must construct in-memory PQDatastore if strategy == ONDISK and must construct - //in-mem and on-disk PQDataStore if strategy == ONDISK and diskPQ is required. + // REFACTOR TODO: Must construct in-memory PQDatastore if strategy == ONDISK and must construct + // in-mem and on-disk PQDataStore if strategy == ONDISK and diskPQ is required. return std::make_unique>(*_config, data_store, std::move(graph_store), pq_data_store); } @@ -195,11 +203,11 @@ std::unique_ptr IndexFactory::create_instance(const std::string & throw ANNException("Error: unsupported label_type please choose from [uint/ushort]", -1); } -//template DISKANN_DLLEXPORT std::shared_ptr> IndexFactory::construct_datastore( -// DataStoreStrategy stratagy, size_t num_points, size_t dimension, Metric m); -//template DISKANN_DLLEXPORT std::shared_ptr> IndexFactory::construct_datastore( -// DataStoreStrategy stratagy, size_t num_points, size_t dimension, Metric m); -//template DISKANN_DLLEXPORT std::shared_ptr> IndexFactory::construct_datastore( -// DataStoreStrategy stratagy, size_t num_points, size_t dimension, Metric m); +// template DISKANN_DLLEXPORT std::shared_ptr> IndexFactory::construct_datastore( +// DataStoreStrategy stratagy, size_t num_points, size_t dimension, Metric m); +// template DISKANN_DLLEXPORT std::shared_ptr> IndexFactory::construct_datastore( +// DataStoreStrategy stratagy, size_t num_points, size_t dimension, Metric m); +// template DISKANN_DLLEXPORT std::shared_ptr> IndexFactory::construct_datastore( +// DataStoreStrategy stratagy, size_t num_points, size_t dimension, Metric m); } // namespace diskann diff --git a/src/pq_data_store.cpp b/src/pq_data_store.cpp index 6207b75e2..c47c16705 100644 --- a/src/pq_data_store.cpp +++ b/src/pq_data_store.cpp @@ -18,7 +18,8 @@ PQDataStore::PQDataStore(size_t dim, location_t num_points, size_t num_p : AbstractDataStore(num_points, dim), _quantized_data(nullptr), _num_chunks(num_pq_chunks), _distance_metric(distance_fn->get_metric()) { - if (num_pq_chunks > dim) { + if (num_pq_chunks > dim) + { throw diskann::ANNException("ERROR: num_pq_chunks > dim", -1, __FUNCSIG__, __FILE__, __LINE__); } _distance_fn = std::move(distance_fn); @@ -213,7 +214,7 @@ template size_t PQDataStore::get_alignment_factor() co return 1; } -template Distance* PQDataStore::get_dist_fn() const +template Distance *PQDataStore::get_dist_fn() const { return _distance_fn.get(); } From 74330ca7f88747e010dfe6653dedeab8af97346f Mon Sep 17 00:00:00 2001 From: Gopal Srinivasa Date: Mon, 23 Oct 2023 16:08:05 +0530 Subject: [PATCH 11/15] Incorporating CR comments --- include/pq_scratch.h | 26 ++------------------------ src/scratch.cpp | 27 +++++++++++++++++++++++++++ 2 files changed, 29 insertions(+), 24 deletions(-) diff --git a/include/pq_scratch.h b/include/pq_scratch.h index 5b599b5be..4ee5a17fb 100644 --- a/include/pq_scratch.h +++ b/include/pq_scratch.h @@ -14,30 +14,8 @@ template struct PQScratch float *rotated_query = nullptr; float *aligned_query_float = nullptr; - PQScratch(size_t graph_degree, size_t aligned_dim) - { - diskann::alloc_aligned((void **)&aligned_pq_coord_scratch, - (size_t)graph_degree * (size_t)MAX_PQ_CHUNKS * sizeof(uint8_t), 256); - diskann::alloc_aligned((void **)&aligned_pqtable_dist_scratch, 256 * (size_t)MAX_PQ_CHUNKS * sizeof(float), - 256); - diskann::alloc_aligned((void **)&aligned_dist_scratch, (size_t)graph_degree * sizeof(float), 256); - diskann::alloc_aligned((void **)&aligned_query_float, aligned_dim * sizeof(float), 8 * sizeof(float)); - diskann::alloc_aligned((void **)&rotated_query, aligned_dim * sizeof(float), 8 * sizeof(float)); - - memset(aligned_query_float, 0, aligned_dim * sizeof(float)); - memset(rotated_query, 0, aligned_dim * sizeof(float)); - } - - void initialize(size_t dim, const T *query, const float norm = 1.0f) - { - for (size_t d = 0; d < dim; ++d) - { - if (norm != 1.0f) - rotated_query[d] = aligned_query_float[d] = static_cast(query[d]) / norm; - else - rotated_query[d] = aligned_query_float[d] = static_cast(query[d]); - } - } + PQScratch(size_t graph_degree, size_t aligned_dim); + void initialize(size_t dim, const T *query, const float norm = 1.0f); }; } // namespace diskann \ No newline at end of file diff --git a/src/scratch.cpp b/src/scratch.cpp index 0370fe501..a66ad1dc5 100644 --- a/src/scratch.cpp +++ b/src/scratch.cpp @@ -130,6 +130,33 @@ template void SSDThreadData::clear() scratch.reset(); } +template +PQScratch::PQScratch(size_t graph_degree, size_t aligned_dim) +{ + diskann::alloc_aligned((void **)&aligned_pq_coord_scratch, + (size_t)graph_degree * (size_t)MAX_PQ_CHUNKS * sizeof(uint8_t), 256); + diskann::alloc_aligned((void **)&aligned_pqtable_dist_scratch, 256 * (size_t)MAX_PQ_CHUNKS * sizeof(float), 256); + diskann::alloc_aligned((void **)&aligned_dist_scratch, (size_t)graph_degree * sizeof(float), 256); + diskann::alloc_aligned((void **)&aligned_query_float, aligned_dim * sizeof(float), 8 * sizeof(float)); + diskann::alloc_aligned((void **)&rotated_query, aligned_dim * sizeof(float), 8 * sizeof(float)); + + memset(aligned_query_float, 0, aligned_dim * sizeof(float)); + memset(rotated_query, 0, aligned_dim * sizeof(float)); +} + +template +void PQScratch::initialize(size_t dim, const T *query, const float norm = 1.0f) +{ + for (size_t d = 0; d < dim; ++d) + { + if (norm != 1.0f) + rotated_query[d] = aligned_query_float[d] = static_cast(query[d]) / norm; + else + rotated_query[d] = aligned_query_float[d] = static_cast(query[d]); + } +} + + template DISKANN_DLLEXPORT class InMemQueryScratch; template DISKANN_DLLEXPORT class InMemQueryScratch; template DISKANN_DLLEXPORT class InMemQueryScratch; From d3297faeab9a8f9defa0c7596ddab11f256005d3 Mon Sep 17 00:00:00 2001 From: Gopal Srinivasa Date: Tue, 24 Oct 2023 11:54:14 +0530 Subject: [PATCH 12/15] Fixing compile error --- src/scratch.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/scratch.cpp b/src/scratch.cpp index a66ad1dc5..43a7c68a1 100644 --- a/src/scratch.cpp +++ b/src/scratch.cpp @@ -145,7 +145,7 @@ PQScratch::PQScratch(size_t graph_degree, size_t aligned_dim) } template -void PQScratch::initialize(size_t dim, const T *query, const float norm = 1.0f) +void PQScratch::initialize(size_t dim, const T *query, const float norm) { for (size_t d = 0; d < dim; ++d) { From 91600f40ebe08c7d3b6e7020a4a8ed9a1f0c58d3 Mon Sep 17 00:00:00 2001 From: rakri Date: Tue, 24 Oct 2023 08:20:34 +0000 Subject: [PATCH 13/15] minor bug fix + clang-format --- src/scratch.cpp | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/src/scratch.cpp b/src/scratch.cpp index 43a7c68a1..c3836ccf1 100644 --- a/src/scratch.cpp +++ b/src/scratch.cpp @@ -130,8 +130,7 @@ template void SSDThreadData::clear() scratch.reset(); } -template -PQScratch::PQScratch(size_t graph_degree, size_t aligned_dim) +template PQScratch::PQScratch(size_t graph_degree, size_t aligned_dim) { diskann::alloc_aligned((void **)&aligned_pq_coord_scratch, (size_t)graph_degree * (size_t)MAX_PQ_CHUNKS * sizeof(uint8_t), 256); @@ -144,8 +143,7 @@ PQScratch::PQScratch(size_t graph_degree, size_t aligned_dim) memset(rotated_query, 0, aligned_dim * sizeof(float)); } -template -void PQScratch::initialize(size_t dim, const T *query, const float norm) +template void PQScratch::initialize(size_t dim, const T *query, const float norm) { for (size_t d = 0; d < dim; ++d) { @@ -156,7 +154,6 @@ void PQScratch::initialize(size_t dim, const T *query, const float norm) } } - template DISKANN_DLLEXPORT class InMemQueryScratch; template DISKANN_DLLEXPORT class InMemQueryScratch; template DISKANN_DLLEXPORT class InMemQueryScratch; @@ -165,7 +162,12 @@ template DISKANN_DLLEXPORT class SSDQueryScratch; template DISKANN_DLLEXPORT class SSDQueryScratch; template DISKANN_DLLEXPORT class SSDQueryScratch; +template DISKANN_DLLEXPORT class PQScratch; +template DISKANN_DLLEXPORT class PQScratch; +template DISKANN_DLLEXPORT class PQScratch; + template DISKANN_DLLEXPORT class SSDThreadData; template DISKANN_DLLEXPORT class SSDThreadData; template DISKANN_DLLEXPORT class SSDThreadData; + } // namespace diskann From e9ecab0c83ef0133561d726e9714ad4aa9b3a93e Mon Sep 17 00:00:00 2001 From: gopalrs <33950290+gopalrs@users.noreply.github.com> Date: Fri, 27 Oct 2023 11:24:20 +0530 Subject: [PATCH 14/15] Update pq.h --- include/pq.h | 35 ----------------------------------- 1 file changed, 35 deletions(-) diff --git a/include/pq.h b/include/pq.h index 8d8fafcf8..1bf51d0d0 100644 --- a/include/pq.h +++ b/include/pq.h @@ -47,41 +47,6 @@ class FixedChunkPQTable void populate_chunk_inner_products(const float *query_vec, float *dist_vec); }; -// REFACTOR -// template struct PQScratch -//{ -// float *aligned_pqtable_dist_scratch = nullptr; // MUST BE AT LEAST [256 * NCHUNKS] -// float *aligned_dist_scratch = nullptr; // MUST BE AT LEAST diskann MAX_DEGREE -// uint8_t *aligned_pq_coord_scratch = nullptr; // MUST BE AT LEAST [N_CHUNKS * MAX_DEGREE] -// float *rotated_query = nullptr; -// float *aligned_query_float = nullptr; -// -// PQScratch(size_t graph_degree, size_t aligned_dim) -// { -// diskann::alloc_aligned((void **)&aligned_pq_coord_scratch, -// (size_t)graph_degree * (size_t)MAX_PQ_CHUNKS * sizeof(uint8_t), 256); -// diskann::alloc_aligned((void **)&aligned_pqtable_dist_scratch, 256 * (size_t)MAX_PQ_CHUNKS * sizeof(float), -// 256); -// diskann::alloc_aligned((void **)&aligned_dist_scratch, (size_t)graph_degree * sizeof(float), 256); -// diskann::alloc_aligned((void **)&aligned_query_float, aligned_dim * sizeof(float), 8 * sizeof(float)); -// diskann::alloc_aligned((void **)&rotated_query, aligned_dim * sizeof(float), 8 * sizeof(float)); -// -// memset(aligned_query_float, 0, aligned_dim * sizeof(float)); -// memset(rotated_query, 0, aligned_dim * sizeof(float)); -// } -// -// void set(size_t dim, T *query, const float norm = 1.0f) -// { -// for (size_t d = 0; d < dim; ++d) -// { -// if (norm != 1.0f) -// rotated_query[d] = aligned_query_float[d] = static_cast(query[d]) / norm; -// else -// rotated_query[d] = aligned_query_float[d] = static_cast(query[d]); -// } -// } -// }; - void aggregate_coords(const std::vector &ids, const uint8_t *all_coords, const uint64_t ndims, uint8_t *out); void pq_dist_lookup(const uint8_t *pq_ids, const size_t n_pts, const size_t pq_nchunks, const float *pq_dists, From adb9e600f2d3fd9be7d3dd1c1da9c63c39e133c4 Mon Sep 17 00:00:00 2001 From: Gopal Srinivasa Date: Mon, 4 Dec 2023 20:47:14 +0530 Subject: [PATCH 15/15] Fixing warnings about struct/class incompatibility --- include/abstract_data_store.h | 2 +- include/abstract_scratch.h | 5 +++-- include/pq_scratch.h | 3 ++- include/quantized_distance.h | 2 +- include/scratch.h | 2 +- 5 files changed, 8 insertions(+), 6 deletions(-) diff --git a/include/abstract_data_store.h b/include/abstract_data_store.h index c5c5a322e..89856f1fa 100644 --- a/include/abstract_data_store.h +++ b/include/abstract_data_store.h @@ -13,7 +13,7 @@ namespace diskann { -template struct AbstractScratch; +template class AbstractScratch; template class AbstractDataStore { diff --git a/include/abstract_scratch.h b/include/abstract_scratch.h index 45c240198..b42a836f6 100644 --- a/include/abstract_scratch.h +++ b/include/abstract_scratch.h @@ -2,13 +2,14 @@ namespace diskann { -template struct PQScratch; +template class PQScratch; // By somewhat more than a coincidence, it seems that both InMemQueryScratch // and SSDQueryScratch have the aligned query and PQScratch objects. So we // can put them in a neat hierarchy and keep PQScratch as a standalone class. -template struct AbstractScratch +template class AbstractScratch { + public: AbstractScratch() = default; // This class does not take any responsibilty for memory management of // its members. It is the responsibility of the derived classes to do so. diff --git a/include/pq_scratch.h b/include/pq_scratch.h index 4ee5a17fb..2aa90dbe1 100644 --- a/include/pq_scratch.h +++ b/include/pq_scratch.h @@ -6,8 +6,9 @@ namespace diskann { -template struct PQScratch +template class PQScratch { + public: float *aligned_pqtable_dist_scratch = nullptr; // MUST BE AT LEAST [256 * NCHUNKS] float *aligned_dist_scratch = nullptr; // MUST BE AT LEAST diskann MAX_DEGREE uint8_t *aligned_pq_coord_scratch = nullptr; // AT LEAST [N_CHUNKS * MAX_DEGREE] diff --git a/include/quantized_distance.h b/include/quantized_distance.h index 8e60f896e..cc4aea929 100644 --- a/include/quantized_distance.h +++ b/include/quantized_distance.h @@ -6,7 +6,7 @@ namespace diskann { -template struct PQScratch; +template class PQScratch; template class QuantizedDistance { diff --git a/include/scratch.h b/include/scratch.h index 49324159e..2f43e3365 100644 --- a/include/scratch.h +++ b/include/scratch.h @@ -19,7 +19,7 @@ namespace diskann { -template struct PQScratch; +template class PQScratch; // // AbstractScratch space for in-memory index based search