From b7d1520b10d27794473154b513e9eec577acbc59 Mon Sep 17 00:00:00 2001 From: Suryansh Gupta Date: Fri, 24 Jan 2025 19:29:54 +0530 Subject: [PATCH] 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;