From ce56e0c6ff270507661a0458423fb3860226d30e Mon Sep 17 00:00:00 2001 From: Suryansh Gupta Date: Tue, 21 Jan 2025 17:05:19 +0530 Subject: [PATCH 1/3] Account for initialisation and PQ training in latency calculation for disk search. --- src/pq_flash_index.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index d9ad50617..2b93ebe35 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -1276,7 +1276,7 @@ void PQFlashIndex::cached_beam_search(const T *query1, const uint64_t if (beam_width > num_sector_per_nodes * defaults::MAX_N_SECTOR_READS) throw ANNException("Beamwidth can not be higher than defaults::MAX_N_SECTOR_READS", -1, __FUNCSIG__, __FILE__, __LINE__); - + Timer query_timer; ScratchStoreManager> manager(this->_thread_data); auto data = manager.scratch_space(); IOContext &ctx = data->ctx; @@ -1350,7 +1350,7 @@ void PQFlashIndex::cached_beam_search(const T *query1, const uint64_t diskann::aggregate_coords(ids, n_ids, this->data, this->_n_chunks, pq_coord_scratch); diskann::pq_dist_lookup(pq_coord_scratch, n_ids, this->_n_chunks, pq_dists, dists_out); }; - Timer query_timer, io_timer, cpu_timer; + Timer io_timer, cpu_timer; tsl::robin_set &visited = query_scratch->visited; NeighborPriorityQueue &retset = query_scratch->retset; From b7d1520b10d27794473154b513e9eec577acbc59 Mon Sep 17 00:00:00 2001 From: Suryansh Gupta Date: Fri, 24 Jan 2025 19:29:54 +0530 Subject: [PATCH 2/3] Add PQ Computation in cpu_timer --- src/pq_flash_index.cpp | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index 2b93ebe35..4c6b988e2 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -1276,7 +1276,7 @@ void PQFlashIndex::cached_beam_search(const T *query1, const uint64_t if (beam_width > num_sector_per_nodes * defaults::MAX_N_SECTOR_READS) throw ANNException("Beamwidth can not be higher than defaults::MAX_N_SECTOR_READS", -1, __FUNCSIG__, __FILE__, __LINE__); - Timer query_timer; + Timer query_timer, io_timer, cpu_timer; ScratchStoreManager> manager(this->_thread_data); auto data = manager.scratch_space(); IOContext &ctx = data->ctx; @@ -1334,11 +1334,16 @@ void PQFlashIndex::cached_beam_search(const T *query1, const uint64_t const uint64_t num_sectors_per_node = _nnodes_per_sector > 0 ? 1 : DIV_ROUND_UP(_max_node_len, defaults::SECTOR_LEN); + cpu_timer.reset(); // query <-> PQ chunk centers distances _pq_table.preprocess_query(query_rotated); // center the query and rotate if // we have a rotation matrix float *pq_dists = pq_query_scratch->aligned_pqtable_dist_scratch; _pq_table.populate_chunk_distances(query_rotated, pq_dists); + if (stats != nullptr) + { + stats->cpu_us += (float)cpu_timer.elapsed(); + } // query <-> neighbor list float *dist_scratch = pq_query_scratch->aligned_dist_scratch; @@ -1350,7 +1355,6 @@ void PQFlashIndex::cached_beam_search(const T *query1, const uint64_t diskann::aggregate_coords(ids, n_ids, this->data, this->_n_chunks, pq_coord_scratch); diskann::pq_dist_lookup(pq_coord_scratch, n_ids, this->_n_chunks, pq_dists, dists_out); }; - Timer io_timer, cpu_timer; tsl::robin_set &visited = query_scratch->visited; NeighborPriorityQueue &retset = query_scratch->retset; From 68e7bb65720716392252d7c78dc1a2ce7bd4b62a Mon Sep 17 00:00:00 2001 From: Suryansh Gupta Date: Fri, 24 Jan 2025 15:00:26 +0000 Subject: [PATCH 3/3] Run clang-format on the repo --- python/include/builder.h | 11 +++++------ python/include/dynamic_memory_index.h | 18 ++++++++---------- 2 files changed, 13 insertions(+), 16 deletions(-) diff --git a/python/include/builder.h b/python/include/builder.h index 6b1a5b4f3..06102ac62 100644 --- a/python/include/builder.h +++ b/python/include/builder.h @@ -18,10 +18,9 @@ void build_disk_index(diskann::Metric metric, const std::string &data_file_path, template void build_memory_index(diskann::Metric metric, const std::string &vector_bin_path, - const std::string &index_output_path, uint32_t graph_degree, uint32_t complexity, - float alpha, uint32_t num_threads, bool use_pq_build, - size_t num_pq_bytes, bool use_opq, bool use_tags = false, - const std::string& filter_labels_file = "", const std::string& universal_label = "", - uint32_t filter_complexity = 0); + const std::string &index_output_path, uint32_t graph_degree, uint32_t complexity, float alpha, + uint32_t num_threads, bool use_pq_build, size_t num_pq_bytes, bool use_opq, + bool use_tags = false, const std::string &filter_labels_file = "", + const std::string &universal_label = "", uint32_t filter_complexity = 0); -} +} // namespace diskannpy diff --git a/python/include/dynamic_memory_index.h b/python/include/dynamic_memory_index.h index 02d6b8cce..6cbe15ced 100644 --- a/python/include/dynamic_memory_index.h +++ b/python/include/dynamic_memory_index.h @@ -18,8 +18,7 @@ namespace py = pybind11; namespace diskannpy { -template -class DynamicMemoryIndex +template class DynamicMemoryIndex { public: DynamicMemoryIndex(diskann::Metric m, size_t dimensions, size_t max_vectors, uint32_t complexity, @@ -31,19 +30,18 @@ class DynamicMemoryIndex void load(const std::string &index_path); int insert(const py::array_t &vector, DynamicIdType id); py::array_t batch_insert(py::array_t &vectors, - py::array_t &ids, int32_t num_inserts, - int num_threads = 0); + py::array_t &ids, + int32_t num_inserts, int num_threads = 0); int mark_deleted(DynamicIdType id); void save(const std::string &save_path, bool compact_before_save = false); - NeighborsAndDistances search(py::array_t &query, uint64_t knn, - uint64_t complexity); - NeighborsAndDistances batch_search(py::array_t &queries, - uint64_t num_queries, uint64_t knn, uint64_t complexity, - uint32_t num_threads); + NeighborsAndDistances search(py::array_t &query, + uint64_t knn, uint64_t complexity); + NeighborsAndDistances batch_search( + py::array_t &queries, uint64_t num_queries, uint64_t knn, + uint64_t complexity, uint32_t num_threads); void consolidate_delete(); size_t num_points(); - private: const uint32_t _initial_search_complexity; const diskann::IndexWriteParameters _write_parameters;