From 4abddf323ef3913ce747bffbf9a91695d09f352f Mon Sep 17 00:00:00 2001 From: huangliang Date: Fri, 26 Jun 2026 03:37:17 +0000 Subject: [PATCH] perf(diskann): reduce search scratch allocations --- include/index/graph/diskann/beam_search.hpp | 110 ++++----- include/index/graph/diskann/diskann_index.hpp | 28 ++- include/index/graph/diskann/pq_table.hpp | 11 +- .../index/graph/diskann/search_scratch.hpp | 225 ++++++++++++++++-- .../index/graph/diskann/visited_bitset.hpp | 70 ++++++ include/index/graph/vamana/robust_prune.hpp | 24 +- tests/diskann/test_diskann_beam_search.cpp | 95 +++++++- tests/diskann/test_diskann_pq.cpp | 23 +- .../diskann/test_diskann_tombstone_search.cpp | 10 +- 9 files changed, 483 insertions(+), 113 deletions(-) create mode 100644 include/index/graph/diskann/visited_bitset.hpp diff --git a/include/index/graph/diskann/beam_search.hpp b/include/index/graph/diskann/beam_search.hpp index 4d38cbc7..5cc84626 100644 --- a/include/index/graph/diskann/beam_search.hpp +++ b/include/index/graph/diskann/beam_search.hpp @@ -50,8 +50,6 @@ #include #include #include -#include -#include #include #include @@ -109,7 +107,7 @@ struct SearchStats { * Extracted as a free function so the contract is unit-testable in isolation. */ inline void scan_and_insert_neighbors(alaya::vamana::NeighborPriorityQueue &retset, - std::unordered_set &visited, + VisitedBitset &visited, const uint32_t *nbrs, uint32_t n_nbrs, const PQTable &pq, @@ -120,7 +118,7 @@ inline void scan_and_insert_neighbors(alaya::vamana::NeighborPriorityQueue &rets if (m >= num_points) { continue; // defensive: ignore corrupt neighbor ids } - if (!visited.insert(m).second) { + if (!visited.test_and_set(m)) { continue; } retset.insert(alaya::vamana::Neighbor(m, pq.pq_distance(m, pq_table))); @@ -160,17 +158,20 @@ inline std::vector> disk_greedy_search(const SearchCo const size_t list_size = std::max(params.search_list_size, top_k); td.reset_query(list_size); auto &frontier = td.retset; - auto &visited = td.visited; - - // Neighbor lists of read nodes, needed when a node is later expanded. - std::unordered_map> nbrs_cache; + auto &visited = td.visited_bits; // Absorb a freshly-read node: cache its neighbor list and insert it into the // frontier with its exact L2 distance (coords are co-located in the record). auto absorb = [&](uint32_t id, const char *rec) { NodeRecordView view{rec, dim}; - nbrs_cache[id].assign(view.nbrs(), view.nbrs() + view.n_nbrs()); - frontier.insert(alaya::vamana::Neighbor(id, l2(query, view.coords(), dim))); + const float distance = l2(query, view.coords(), dim); + const auto insert_result = frontier.insert_with_result(alaya::vamana::Neighbor(id, distance)); + if (insert_result.evicted) { + td.release_cached_neighbors(insert_result.evicted_id); + } + if (insert_result.inserted) { + td.cache_neighbors(id, view.nbrs(), view.n_nbrs()); + } if (stats != nullptr) { stats->read_order.push_back(id); stats->n_nodes_processed++; @@ -199,7 +200,7 @@ inline std::vector> disk_greedy_search(const SearchCo // IP-DiskANN: tombstoned nodes are skipped (graph repaired at delete time). const TombstoneBitmap *tomb = ctx.tombstone; auto consider = [&](uint32_t m, auto &&emit) { - if (m >= ctx.num_points || !visited.insert(m).second) { + if (m >= ctx.num_points || !visited.test_and_set(m)) { return; } if (tomb != nullptr && tomb->is_deleted(m)) { @@ -208,7 +209,7 @@ inline std::vector> disk_greedy_search(const SearchCo emit(m); }; - visited.insert(ctx.medoid); + visited.set(ctx.medoid); absorb(ctx.medoid, read_seed(ctx.medoid)); std::vector reqs; @@ -225,12 +226,7 @@ inline std::vector> disk_greedy_search(const SearchCo std::vector chunk_recs; while (frontier.has_unexpanded_node()) { const uint32_t x = frontier.closest_unexpanded().id; - const auto it = nbrs_cache.find(x); - if (it == nbrs_cache.end()) { - continue; // defensive: a frontier node is always absorbed (and so cached) - } - // Copy: absorb() below may rehash nbrs_cache and invalidate it->second. - const std::vector nbrs = it->second; + const NeighborScratchView nbrs = td.cached_neighbors(x); todo.clear(); for (const uint32_t m : nbrs) { consider(m, [&](uint32_t id) { @@ -273,7 +269,7 @@ inline std::vector> disk_greedy_search(const SearchCo } else { // Async pipeline (default): keep up to n_slots reads in flight across // expansions. `pending` holds discovered (visited-marked) ids awaiting a slot; - // `inflight` maps an in-flight id to its (slot, record pointer). Processing + // ThreadData's inflight array tracks scratch slots and record pointers. Processing // follows I/O completion order, so the strict-greedy expansion order is // relaxed (results may differ only in tie-ordering of equally-distant nodes). std::vector free_slots; @@ -281,7 +277,6 @@ inline std::vector> disk_greedy_search(const SearchCo for (uint64_t s = 0; s < n_slots; ++s) { free_slots.push_back(s); } - std::unordered_map> inflight; std::deque pending; std::vector evts; @@ -292,11 +287,8 @@ inline std::vector> disk_greedy_search(const SearchCo auto refill_pending = [&]() { while (pending.size() < free_slots.size() && frontier.has_unexpanded_node()) { const uint32_t x = frontier.closest_unexpanded().id; - const auto it = nbrs_cache.find(x); - if (it == nbrs_cache.end()) { - continue; - } - for (const uint32_t m : it->second) { + const NeighborScratchView nbrs = td.cached_neighbors(x); + for (const uint32_t m : nbrs) { consider(m, [&](uint32_t id) { pending.push_back(id); }); @@ -322,7 +314,7 @@ inline std::vector> disk_greedy_search(const SearchCo const uint64_t slot = free_slots.back(); free_slots.pop_back(); char *buf = td.sector_scratch + slot * page_size; - inflight.emplace(m, std::make_pair(slot, buf + geom.offset_to_node(m))); + td.set_inflight(slot, m, buf + geom.offset_to_node(m)); reqs.emplace_back(geom.get_page_offset(m), page_size, m, buf); } if (!reqs.empty()) { @@ -338,23 +330,20 @@ inline std::vector> disk_greedy_search(const SearchCo // Continue while any work remains: a read in flight, a queued read, or an // unexpanded frontier node (the latter matters when cache hits add nodes // without ever entering `inflight`). - while (!inflight.empty() || !pending.empty() || frontier.has_unexpanded_node()) { - if (!inflight.empty()) { + while (td.has_inflight() || !pending.empty() || frontier.has_unexpanded_node()) { + if (td.has_inflight()) { reader.get_events(td.ctx_, 1, evts); // get_events clears + fills `evts` for (const auto &e : evts) { if (e.result != static_cast(page_size)) { throw std::runtime_error("disk_greedy_search: short/failed read, result=" + std::to_string(e.result)); } - const auto it = inflight.find(static_cast(e.id)); - if (it == inflight.end()) { + InFlightSlot completed; + if (!td.remove_inflight(static_cast(e.id), completed)) { continue; // defensive: completion for an id we are not tracking } - const uint64_t slot = it->second.first; - const char *rec = it->second.second; - absorb(static_cast(e.id), rec); - free_slots.push_back(slot); - inflight.erase(it); + absorb(static_cast(e.id), completed.record); + free_slots.push_back(completed.page_slot); } } refill_pending(); @@ -404,14 +393,13 @@ inline std::vector> cached_beam_search(const SearchCo const size_t list_size = std::max(params.search_list_size, top_k); td.reset_query(list_size); auto &retset = td.retset; - auto &visited = td.visited; - auto &exact_by_id = td.exact_by_id; + auto &visited = td.visited_bits; - pq.preprocess_query(query, td.pq_table.data()); + pq.preprocess_query(query, td.pq_table.data(), td.pq_qres.data()); const float *pq_table = td.pq_table.data(); // Seed from the medoid (first node expanded — spec scenario). - visited.insert(ctx.medoid); + visited.set(ctx.medoid); retset.insert(alaya::vamana::Neighbor(ctx.medoid, pq.pq_distance(ctx.medoid, pq_table))); const uint64_t beam = std::max(1, params.beam_width); @@ -421,7 +409,7 @@ inline std::vector> cached_beam_search(const SearchCo // frontier with PQ approximate distances. Shared by both search paths. auto process_node = [&](uint32_t id, const char *rec) { NodeRecordView view{rec, dim}; - exact_by_id[id] = l2(query, view.coords(), dim); + td.set_exact_dist(id, l2(query, view.coords(), dim)); if (stats != nullptr) { stats->read_order.push_back(id); stats->n_nodes_processed++; @@ -446,24 +434,24 @@ inline std::vector> cached_beam_search(const SearchCo // (DiskANNSearchParams::deterministic) — it forgoes the cross-beam I/O/compute // overlap of the default path and so runs ~10-15% slower. std::vector batch; - std::unordered_map rec_of; + std::vector batch_recs; while (retset.has_unexpanded_node()) { batch.clear(); + batch_recs.clear(); reqs.clear(); - rec_of.clear(); uint64_t slot = 0; while (retset.has_unexpanded_node() && reqs.size() < beam) { const uint32_t id = retset.closest_unexpanded().id; batch.push_back(id); const char *crec = cache.lookup(id); if (crec != nullptr) { - rec_of[id] = crec; + batch_recs.push_back(crec); if (stats != nullptr) { stats->n_cache_hits++; } } else { char *buf = td.sector_scratch + (slot++) * page_size; - rec_of[id] = buf + geom.offset_to_node(id); + batch_recs.push_back(buf + geom.offset_to_node(id)); reqs.emplace_back(geom.get_page_offset(id), page_size, id, buf); } } @@ -482,8 +470,8 @@ inline std::vector> cached_beam_search(const SearchCo } } - for (const uint32_t id : batch) { - process_node(id, rec_of[id]); + for (size_t i = 0; i < batch.size(); ++i) { + process_node(batch[i], batch_recs[i]); } } } else { @@ -498,9 +486,6 @@ inline std::vector> cached_beam_search(const SearchCo for (uint64_t s = 0; s < beam; ++s) { free_slots.push_back(s); } - // In-flight reads: node id -> (scratch slot, record pointer within that page). - std::unordered_map> inflight; - // Pop closest-unexpanded candidates, serving cache hits inline and submitting // cache misses into free scratch slots. Stops when the slots are full or the // frontier has no unexpanded node left. @@ -519,7 +504,7 @@ inline std::vector> cached_beam_search(const SearchCo const uint64_t slot = free_slots.back(); free_slots.pop_back(); char *buf = td.sector_scratch + slot * page_size; - inflight.emplace(id, std::make_pair(slot, buf + geom.offset_to_node(id))); + td.set_inflight(slot, id, buf + geom.offset_to_node(id)); reqs.emplace_back(geom.get_page_offset(id), page_size, id, buf); } if (!reqs.empty()) { @@ -531,7 +516,7 @@ inline std::vector> cached_beam_search(const SearchCo }; fill_pipe(); - while (!inflight.empty()) { + while (td.has_inflight()) { // Wait for exactly one completion (get_events clears and fills `evts`). The // other in-flight reads keep the disk busy while we process this one — that // is the I/O/compute overlap. NOTE: poll_events/get_events both clear their @@ -542,15 +527,12 @@ inline std::vector> cached_beam_search(const SearchCo throw std::runtime_error("cached_beam_search: short/failed read, result=" + std::to_string(e.result)); } - const auto it = inflight.find(static_cast(e.id)); - if (it == inflight.end()) { + InFlightSlot completed; + if (!td.remove_inflight(static_cast(e.id), completed)) { continue; // defensive: completion for an id we are not tracking } - const uint64_t slot = it->second.first; - const char *rec = it->second.second; - process_node(static_cast(e.id), rec); - free_slots.push_back(slot); - inflight.erase(it); + process_node(static_cast(e.id), completed.record); + free_slots.push_back(completed.page_slot); } fill_pipe(); } @@ -561,16 +543,16 @@ inline std::vector> cached_beam_search(const SearchCo std::vector rerank_req(1); auto read_exact_sync = [&](uint32_t id) -> float { - const auto it = exact_by_id.find(id); - if (it != exact_by_id.end()) { - return it->second; + const float known = td.exact_dist(id); + if (!ThreadData::is_missing_exact(known)) { + return known; } char *buf = td.sector_scratch; rerank_req[0] = {geom.get_page_offset(id), page_size, id, buf}; reader.read(rerank_req, td.ctx_); NodeRecordView view{buf + geom.offset_to_node(id), dim}; const float ex = l2(query, view.coords(), dim); - exact_by_id[id] = ex; + td.set_exact_dist(id, ex); if (stats != nullptr) { stats->n_rerank_reads++; } @@ -593,8 +575,8 @@ inline std::vector> cached_beam_search(const SearchCo out.reserve(n_cand); for (size_t i = 0; i < n_cand; ++i) { const uint32_t id = retset[i].id; - const auto it = exact_by_id.find(id); - out.emplace_back(id, it != exact_by_id.end() ? it->second : retset[i].distance); + const float exact = td.exact_dist(id); + out.emplace_back(id, !ThreadData::is_missing_exact(exact) ? exact : retset[i].distance); } } diff --git a/include/index/graph/diskann/diskann_index.hpp b/include/index/graph/diskann/diskann_index.hpp index f75c84a3..effac570 100644 --- a/include/index/graph/diskann/diskann_index.hpp +++ b/include/index/graph/diskann/diskann_index.hpp @@ -58,6 +58,8 @@ namespace alaya::diskann { +inline constexpr uint32_t kDefaultDiskANNScratchSearchListSize = 150; + /// Build-time configuration. struct DiskANNBuildParams { uint32_t R = 64; ///< graph degree bound @@ -82,6 +84,9 @@ struct DiskANNLoadParams { ///< deeper pipeline overlaps more I/O. Explicit values are floored ///< at 2*beam_width and capped at the libaio context size (1024). ///< Sizes the sector scratch to that many pages per thread. + uint32_t scratch_search_list_size = kDefaultDiskANNScratchSearchListSize; + ///< No-PQ neighbor scratch capacity, in search-list entries. Set this >= the + ///< largest DiskANNSearchParams::search_list_size used after load. // --- In-place update mode (No-PQ only; see disk-update specs) --- bool updatable = false; ///< open O_RDWR + enable insert/remove/update_node/flush @@ -287,16 +292,27 @@ class DiskANNIndex { params.nopq_io_depth == 0 ? kDefaultNoPQIoDepth : params.nopq_io_depth; const uint64_t scratch_slots = std::min(1024, std::max(2ull * beam_width_, nopq_depth)); + const uint32_t scratch_list_size = std::max({DiskANNSearchParams{}.search_list_size, + params.update_search_l, + params.scratch_search_list_size}); + ThreadDataScratchConfig scratch_config; + scratch_config.n_page_slots = scratch_slots; + scratch_config.page_size = geom_.page_size; + scratch_config.pq_table_entries = pq_table_entries; + scratch_config.max_slot_id = max_slot_id_; + scratch_config.max_degree = max_degree_; + scratch_config.search_list_size = scratch_list_size; + scratch_config.query_dim = dim_; thread_data_storage_.resize(pool); { std::vector regs; regs.reserve(pool); for (uint32_t t = 0; t < pool; ++t) { - regs.emplace_back([this, t, pq_table_entries, scratch_slots]() { + regs.emplace_back([this, t, scratch_config]() { reader_->register_thread(); auto td = std::make_unique(); td->ctx_ = reader_->get_ctx(); - td->alloc_scratch(scratch_slots, geom_.page_size, pq_table_entries); + td->alloc_scratch(scratch_config); thread_data_storage_[t] = std::move(td); }); } @@ -500,6 +516,7 @@ class DiskANNIndex { const uint32_t slot = slot_alloc_.alloc(); update_ctx_.forget_slot(slot); max_slot_id_ = std::max(max_slot_id_, slot_alloc_.next_fresh_id()); + resize_thread_data_slot_capacity(); page_io_->write_node(slot, query, static_cast(pruned.size()), pruned.data()); @@ -654,6 +671,7 @@ class DiskANNIndex { if (std::filesystem::exists(slots_path)) { slot_alloc_.load(slots_path); // restore free list + next id + tombstones max_slot_id_ = std::max(max_slot_id_, slot_alloc_.next_fresh_id()); + resize_thread_data_slot_capacity(); } else { slot_alloc_.reset(static_cast(max_slot_id_)); } @@ -852,6 +870,12 @@ class DiskANNIndex { labels_[slot] = label; } + void resize_thread_data_slot_capacity() { + for (auto &td : thread_data_storage_) { + td->resize_slot_capacity(max_slot_id_); + } + } + /// Order-independent equality of two neighbor id lists. static bool same_neighbor_set(std::vector a, std::vector b) { if (a.size() != b.size()) { diff --git a/include/index/graph/diskann/pq_table.hpp b/include/index/graph/diskann/pq_table.hpp index 00c34be6..cd8574e2 100644 --- a/include/index/graph/diskann/pq_table.hpp +++ b/include/index/graph/diskann/pq_table.hpp @@ -220,14 +220,17 @@ class PQTable { * @param table_out Caller-owned buffer of @c n_chunks*256 float32. Entry * [c*256 + k] = squared L2 between the query's chunk-c * residual and centroid k of chunk c. + * @param scratch Caller-owned buffer of @c dim float32 for the query residual. */ - void preprocess_query(const float *query, float *table_out) const { - std::vector qres(dim_); + void preprocess_query(const float *query, float *table_out, float *scratch) const { + if (scratch == nullptr) { + throw std::invalid_argument("PQTable::preprocess_query: scratch must not be null"); + } for (uint64_t d = 0; d < dim_; ++d) { - qres[d] = query[d] - global_centroid_[d]; + scratch[d] = query[d] - global_centroid_[d]; } for (uint32_t c = 0; c < n_chunks_; ++c) { - const float *qchunk = qres.data() + static_cast(c) * chunk_dim_; + const float *qchunk = scratch + static_cast(c) * chunk_dim_; const float *cent = codebook_.data() + static_cast(c) * kPQNumCentroids * chunk_dim_; float *trow = table_out + static_cast(c) * kPQNumCentroids; for (uint32_t k = 0; k < kPQNumCentroids; ++k) { diff --git a/include/index/graph/diskann/search_scratch.hpp b/include/index/graph/diskann/search_scratch.hpp index a94f6065..d70bfd98 100644 --- a/include/index/graph/diskann/search_scratch.hpp +++ b/include/index/graph/diskann/search_scratch.hpp @@ -8,9 +8,9 @@ * * Each search thread borrows one @c ThreadData from a @c ConcurrentQueue for * the duration of a query (LASER pattern, design D8). It bundles: - * - the visited set (dedup of popped nodes), + * - the visited bitset (dedup of popped nodes), * - the @c retset frontier (NeighborPriorityQueue, reused from Vamana), - * - @c exact_by_id: exact L2 distances of nodes actually read from disk/cache, + * - @c exact_dists: exact L2 distances of nodes actually read from disk/cache, * - @c pq_table: the per-query @c n_chunks x 256 PQ distance table (PQ mode), * - @c sector_scratch: a sector-aligned double buffer for async page reads, * - @c ctx_: the thread's AlignedFileReader I/O context. @@ -21,24 +21,66 @@ #pragma once +#include #include -#include -#include +#include +#include +#include +#include #include #include "index/graph/diskann/disk_layout.hpp" +#include "index/graph/diskann/visited_bitset.hpp" #include "index/graph/laser/utils/aligned_file_reader.hpp" #include "index/graph/laser/utils/memory.hpp" #include "index/graph/vamana/robust_prune.hpp" namespace alaya::diskann { +struct ThreadDataScratchConfig { + uint64_t n_page_slots = 0; + uint64_t page_size = 0; + uint32_t pq_table_entries = 0; + uint64_t max_slot_id = 0; + uint32_t max_degree = 0; + uint32_t search_list_size = 0; + uint64_t query_dim = 0; +}; + +struct NeighborScratchView { + const uint32_t *data_ptr = nullptr; + uint32_t len = 0; + + [[nodiscard]] const uint32_t *data() const { return data_ptr; } + [[nodiscard]] size_t size() const { return len; } + [[nodiscard]] const uint32_t *begin() const { return data_ptr; } + [[nodiscard]] const uint32_t *end() const { + return data_ptr == nullptr ? nullptr : data_ptr + len; + } +}; + +struct InFlightSlot { + bool occupied = false; + uint32_t id = 0; + uint64_t page_slot = 0; + const char *record = nullptr; +}; + struct ThreadData { // --- Per-query mutable search state --- - std::unordered_set visited; ///< ids popped/seeded - alaya::vamana::NeighborPriorityQueue retset; ///< exploration frontier - std::unordered_map exact_by_id; ///< read node -> exact L2 sqr - std::vector pq_table; ///< n_chunks*256 (empty if no PQ) + VisitedBitset visited_bits; ///< ids popped/seeded + alaya::vamana::NeighborPriorityQueue retset; ///< exploration frontier + std::vector exact_dists; ///< node id -> exact L2 sqr or NaN + std::vector exact_dirty; ///< exact_dists entries written this query + std::vector pq_table; ///< n_chunks*256 (empty if no PQ) + std::vector pq_qres; ///< dim floats for PQ query residual + std::vector nbrs_buf; ///< contiguous cached neighbor lists + std::vector> nbrs_offsets; ///< id -> (start, len) + std::vector nbrs_dirty; ///< offsets written this query + uint32_t nbrs_buf_next = 0; + uint32_t nbrs_slot_len = 0; + std::vector nbrs_free_offsets; + std::vector inflight; ///< indexed by scratch page slot // --- I/O scratch (allocated once, reused) --- char *sector_scratch = nullptr; ///< n_page_slots * page_size bytes, sector-aligned. @@ -48,26 +90,145 @@ struct ThreadData { uint64_t sector_scratch_bytes = 0; IOContext ctx_{}; ///< AIO context (owned via reader.register_thread()) + void resize_slot_capacity(uint64_t max_slot_id) { + if (max_slot_id <= visited_bits.size_bits()) { + return; + } + visited_bits.resize(max_slot_id); + exact_dists.resize(max_slot_id, std::numeric_limits::quiet_NaN()); + nbrs_offsets.resize(max_slot_id, {0, 0}); + } + /// Reset only the per-query state; keeps allocated buffers. void reset_query(size_t search_list_size) { - visited.clear(); - exact_by_id.clear(); + visited_bits.clear(); + reset_exact_dists(); + reset_neighbors(); + clear_inflight(); retset.reserve(search_list_size); retset.clear(); } - /// Allocate the sector page buffer (@p n_page_slots pages, one per concurrent - /// read) and (optionally) the PQ table. - void alloc_scratch(uint64_t n_page_slots, uint64_t page_size, uint32_t pq_table_entries) { - sector_scratch_bytes = n_page_slots * page_size; + /// Allocate the sector page buffer, flat hot-path structures, and PQ scratch. + void alloc_scratch(const ThreadDataScratchConfig &config) { + if (config.n_page_slots == 0 || config.page_size == 0 || config.max_slot_id == 0) { + throw std::invalid_argument("ThreadData::alloc_scratch: invalid zero-sized config"); + } + sector_scratch_bytes = config.n_page_slots * config.page_size; sector_scratch = reinterpret_cast( alaya::laser::memory::align_allocate(sector_scratch_bytes)); - if (pq_table_entries > 0) { - pq_table.assign(pq_table_entries, 0.0f); + if (config.pq_table_entries > 0) { + pq_table.assign(config.pq_table_entries, 0.0f); + } + if (config.query_dim > 0) { + pq_qres.assign(config.query_dim, 0.0f); + } + visited_bits.resize(config.max_slot_id); + exact_dists.assign(config.max_slot_id, std::numeric_limits::quiet_NaN()); + exact_dirty.reserve(config.search_list_size); + nbrs_slot_len = config.max_degree; + nbrs_buf.assign(static_cast(config.search_list_size) * nbrs_slot_len, 0); + nbrs_offsets.assign(config.max_slot_id, {0, 0}); + nbrs_dirty.reserve(config.search_list_size); + nbrs_free_offsets.reserve(config.search_list_size); + inflight.assign(config.n_page_slots, {}); + reset_neighbors(); + } + + void set_exact_dist(uint32_t id, float distance) { + if (id >= exact_dists.size()) { + throw std::out_of_range("ThreadData::set_exact_dist: id out of range"); + } + if (is_missing_exact(exact_dists[id])) { + exact_dirty.push_back(id); + } + exact_dists[id] = distance; + } + + [[nodiscard]] float exact_dist(uint32_t id) const { + if (id >= exact_dists.size()) { + throw std::out_of_range("ThreadData::exact_dist: id out of range"); + } + return exact_dists[id]; + } + + [[nodiscard]] static bool is_missing_exact(float value) { + static constexpr uint32_t kExponentMask = 0x7F800000u; + static constexpr uint32_t kMantissaMask = 0x007FFFFFu; + uint32_t bits = 0; + std::memcpy(&bits, &value, sizeof(bits)); + return (bits & kExponentMask) == kExponentMask && (bits & kMantissaMask) != 0; + } + + void cache_neighbors(uint32_t id, const uint32_t *nbrs, uint32_t len) { + if (id >= nbrs_offsets.size()) { + throw std::out_of_range("ThreadData::cache_neighbors: id out of range"); + } + if (len > nbrs_slot_len) { + throw std::runtime_error("ThreadData::cache_neighbors: neighbor list exceeds slot size"); + } + if (len == 0) { + nbrs_offsets[id] = {0, 0}; + return; } + if (nbrs_free_offsets.empty()) { + throw std::runtime_error("ThreadData::cache_neighbors: neighbor scratch exhausted"); + } + const uint32_t start = nbrs_free_offsets.back(); + nbrs_free_offsets.pop_back(); + std::copy_n(nbrs, len, nbrs_buf.begin() + start); + nbrs_offsets[id] = {start, len}; + nbrs_dirty.push_back(id); + nbrs_buf_next += nbrs_slot_len; + } + + [[nodiscard]] NeighborScratchView cached_neighbors(uint32_t id) const { + if (id >= nbrs_offsets.size()) { + throw std::out_of_range("ThreadData::cached_neighbors: id out of range"); + } + const auto [start, len] = nbrs_offsets[id]; + return {len == 0 ? nullptr : nbrs_buf.data() + start, len}; + } + + void release_cached_neighbors(uint32_t id) { + if (id >= nbrs_offsets.size()) { + throw std::out_of_range("ThreadData::release_cached_neighbors: id out of range"); + } + const auto [start, len] = nbrs_offsets[id]; + if (len == 0) { + return; + } + nbrs_offsets[id] = {0, 0}; + nbrs_free_offsets.push_back(start); + nbrs_buf_next -= nbrs_slot_len; + } + + void set_inflight(uint64_t page_slot, uint32_t id, const char *record) { + if (page_slot >= inflight.size()) { + throw std::out_of_range("ThreadData::set_inflight: page slot out of range"); + } + inflight[page_slot] = {true, id, page_slot, record}; } - /// Release the sector buffer (PQ table frees with the vector). + [[nodiscard]] bool has_inflight() const { + return std::any_of(inflight.begin(), inflight.end(), [](const InFlightSlot &slot) { + return slot.occupied; + }); + } + + [[nodiscard]] bool remove_inflight(uint32_t id, InFlightSlot &out) { + for (InFlightSlot &slot : inflight) { + if (!slot.occupied || slot.id != id) { + continue; + } + out = slot; + slot.occupied = false; + return true; + } + return false; + } + + /// Release the sector buffer (vectors free with the ThreadData instance). void free_scratch() { if (sector_scratch != nullptr) { alaya::laser::memory::align_free(sector_scratch); @@ -75,6 +236,36 @@ struct ThreadData { } sector_scratch_bytes = 0; } + + private: + void reset_exact_dists() { + for (const uint32_t id : exact_dirty) { + exact_dists[id] = std::numeric_limits::quiet_NaN(); + } + exact_dirty.clear(); + } + + void reset_neighbors() { + for (const uint32_t id : nbrs_dirty) { + nbrs_offsets[id] = {0, 0}; + } + nbrs_dirty.clear(); + nbrs_buf_next = 0; + nbrs_free_offsets.clear(); + if (nbrs_slot_len == 0) { + return; + } + const uint32_t n_slots = static_cast(nbrs_buf.size() / nbrs_slot_len); + for (uint32_t slot = 0; slot < n_slots; ++slot) { + nbrs_free_offsets.push_back(slot * nbrs_slot_len); + } + } + + void clear_inflight() { + for (InFlightSlot &slot : inflight) { + slot.occupied = false; + } + } }; } // namespace alaya::diskann diff --git a/include/index/graph/diskann/visited_bitset.hpp b/include/index/graph/diskann/visited_bitset.hpp new file mode 100644 index 00000000..b1f34d39 --- /dev/null +++ b/include/index/graph/diskann/visited_bitset.hpp @@ -0,0 +1,70 @@ +// SPDX-FileCopyrightText: 2025 AlayaDB.AI +// +// SPDX-License-Identifier: AGPL-3.0-only + +#pragma once + +#include +#include +#include +#include + +namespace alaya::diskann { + +class VisitedBitset { + public: + void resize(uint64_t max_slot_id) { + n_bits_ = max_slot_id; + words_.assign(word_count_for(max_slot_id), 0); + } + + void clear() { + if (!words_.empty()) { + std::memset(words_.data(), 0, words_.size() * sizeof(uint64_t)); + } + } + + [[nodiscard]] bool test(uint32_t id) const { + ensure_in_range(id); + const uint64_t word = words_[id >> kWordShift]; + return (word & bit_mask(id)) != 0; + } + + void set(uint32_t id) { + ensure_in_range(id); + words_[id >> kWordShift] |= bit_mask(id); + } + + [[nodiscard]] bool test_and_set(uint32_t id) { + ensure_in_range(id); + uint64_t &word = words_[id >> kWordShift]; + const uint64_t mask = bit_mask(id); + const bool was_set = (word & mask) != 0; + word |= mask; + return !was_set; + } + + [[nodiscard]] uint64_t size_bits() const { return n_bits_; } + [[nodiscard]] size_t word_count() const { return words_.size(); } + + private: + static constexpr uint32_t kWordShift = 6; + static constexpr uint32_t kWordMask = 63; + + static size_t word_count_for(uint64_t n_bits) { + return static_cast((n_bits + kWordMask) >> kWordShift); + } + + static uint64_t bit_mask(uint32_t id) { return uint64_t{1} << (id & kWordMask); } + + void ensure_in_range(uint32_t id) const { + if (id >= n_bits_) { + throw std::out_of_range("VisitedBitset: id out of range"); + } + } + + uint64_t n_bits_ = 0; + std::vector words_; +}; + +} // namespace alaya::diskann diff --git a/include/index/graph/vamana/robust_prune.hpp b/include/index/graph/vamana/robust_prune.hpp index 4d8b4d61..3ded4f87 100644 --- a/include/index/graph/vamana/robust_prune.hpp +++ b/include/index/graph/vamana/robust_prune.hpp @@ -31,6 +31,12 @@ struct Neighbor { bool operator==(const Neighbor &other) const { return id == other.id; } }; +struct NeighborQueueInsertResult { + bool inserted = false; + bool evicted = false; + uint32_t evicted_id = 0; +}; + // Bounded-capacity priority queue of Neighbors kept in ascending distance // order. Mirrors DiskANN's `NeighborPriorityQueue` (include/neighbor.h): // supports O(log n) insertion with dedup-by-id and a cursor (`cur_`) that @@ -47,9 +53,15 @@ class NeighborPriorityQueue { capacity_ = capacity; } - void insert(const Neighbor &nbr) { + void insert(const Neighbor &nbr) { (void)insert_with_result(nbr); } + + NeighborQueueInsertResult insert_with_result(const Neighbor &nbr) { + NeighborQueueInsertResult result; + if (capacity_ == 0) { + return result; + } if (size_ == capacity_ && data_[size_ - 1] < nbr) { - return; + return result; } size_t lo = 0; size_t hi = size_; @@ -58,11 +70,13 @@ class NeighborPriorityQueue { if (nbr < data_[mid]) { hi = mid; } else if (data_[mid].id == nbr.id) { - return; // dedup + return result; // dedup } else { lo = mid + 1; } } + const bool full = size_ == capacity_; + const uint32_t evicted_id = full ? data_[capacity_ - 1].id : 0; if (lo < capacity_) { std::memmove(&data_[lo + 1], &data_[lo], (size_ - lo) * sizeof(Neighbor)); } @@ -73,6 +87,10 @@ class NeighborPriorityQueue { if (lo < cur_) { cur_ = lo; } + result.inserted = true; + result.evicted = full; + result.evicted_id = evicted_id; + return result; } Neighbor closest_unexpanded() { diff --git a/tests/diskann/test_diskann_beam_search.cpp b/tests/diskann/test_diskann_beam_search.cpp index 5f6145b6..26b52242 100644 --- a/tests/diskann/test_diskann_beam_search.cpp +++ b/tests/diskann/test_diskann_beam_search.cpp @@ -33,6 +33,8 @@ using alaya::diskann::SearchContext; using alaya::diskann::SearchParams; using alaya::diskann::SearchStats; using alaya::diskann::ThreadData; +using alaya::diskann::ThreadDataScratchConfig; +using alaya::diskann::VisitedBitset; using alaya::diskann::write_disk_layout; using alaya::vamana::Neighbor; using alaya::vamana::NeighborPriorityQueue; @@ -82,8 +84,11 @@ uint32_t medoid_of(const std::vector &v, uint64_t n, uint64_t dim) { return best; } -std::vector> brute_force(const std::vector &v, uint64_t n, - uint64_t dim, const float *q, uint32_t k) { +std::vector> brute_force(const std::vector &v, + uint64_t n, + uint64_t dim, + const float *q, + uint32_t k) { std::vector> all; all.reserve(n); for (uint64_t i = 0; i < n; ++i) { @@ -110,12 +115,14 @@ TEST(BeamScanInsert, SkipsVisitedAndInsertsRest) { pq.train(v.data(), n, dim, n_chunks); pq.encode(v.data(), n); std::vector table(static_cast(n_chunks) * 256); - pq.preprocess_query(v.data(), table.data()); + std::vector qres(dim); + pq.preprocess_query(v.data(), table.data(), qres.data()); NeighborPriorityQueue retset(100); - std::unordered_set visited; + VisitedBitset visited; + visited.resize(n); for (uint32_t id = 0; id < 10; ++id) { - visited.insert(id); // 10 already visited + visited.set(id); // 10 already visited } std::vector nbrs(32); for (uint32_t i = 0; i < 32; ++i) { @@ -134,13 +141,20 @@ TEST(BeamScanInsert, PQPrunesUncompetitiveNeighbors) { pq.encode(v.data(), n); std::vector table(static_cast(n_chunks) * 256); - pq.preprocess_query(v.data() + 0 * dim, table.data()); // query == point 0 + std::vector qres(dim); + pq.preprocess_query(v.data() + 0 * dim, table.data(), qres.data()); // query == point 0 NeighborPriorityQueue retset(4); // capacity 4 - std::unordered_set visited; + VisitedBitset visited; + visited.resize(n); std::vector nbrs = {1, 2, 3, 4, 5, 6, 7, 8, 9}; - scan_and_insert_neighbors(retset, visited, nbrs.data(), - static_cast(nbrs.size()), pq, table.data(), n); + scan_and_insert_neighbors(retset, + visited, + nbrs.data(), + static_cast(nbrs.size()), + pq, + table.data(), + n); EXPECT_EQ(retset.size(), 4u); // bounded // The globally farthest point from point 0 must have been pruned. const auto bf = brute_force(v, n, dim, v.data(), n); // sorted ascending @@ -154,11 +168,55 @@ TEST(BeamScanInsert, PQPrunesUncompetitiveNeighbors) { EXPECT_FALSE(present) << "farthest id=" << farthest << " should be pruned"; } +TEST(SearchScratch, ResetQueryClearsFlatPerQueryState) { + ThreadData td; + ThreadDataScratchConfig cfg; + cfg.n_page_slots = 2; + cfg.page_size = 4096; + cfg.pq_table_entries = 512; + cfg.max_slot_id = 130; + cfg.max_degree = 4; + cfg.search_list_size = 8; + cfg.query_dim = 16; + td.alloc_scratch(cfg); + + EXPECT_EQ(td.visited_bits.word_count(), 3u); + EXPECT_EQ(td.exact_dists.size(), 130u); + EXPECT_EQ(td.nbrs_buf.size(), 32u); + EXPECT_EQ(td.nbrs_offsets.size(), 130u); + EXPECT_EQ(td.pq_qres.size(), 16u); + EXPECT_EQ(td.inflight.size(), 2u); + + td.visited_bits.set(7); + td.set_exact_dist(7, 1.25f); + const std::vector nbrs = {1, 2, 3}; + td.cache_neighbors(7, nbrs.data(), static_cast(nbrs.size())); + + ASSERT_TRUE(td.visited_bits.test(7)); + EXPECT_FLOAT_EQ(td.exact_dists[7], 1.25f); + EXPECT_EQ(td.exact_dirty.size(), 1u); + ASSERT_EQ(td.cached_neighbors(7).size(), nbrs.size()); + + td.reset_query(8); + + EXPECT_FALSE(td.visited_bits.test(7)); + EXPECT_TRUE(ThreadData::is_missing_exact(td.exact_dists[7])); + EXPECT_TRUE(td.exact_dirty.empty()); + EXPECT_EQ(td.cached_neighbors(7).size(), 0u); + EXPECT_EQ(td.nbrs_buf_next, 0u); + + td.free_scratch(); +} + // ---- Beam search over a real (tiny) on-disk index -------------------------- class BeamSearchTest : public ::testing::Test { protected: - void build(uint64_t n, uint64_t dim, uint32_t r, uint32_t n_chunks, double cache_ratio, + void build(uint64_t n, + uint64_t dim, + uint32_t r, + uint32_t n_chunks, + double cache_ratio, uint32_t beam = 4) { n_ = n; dim_ = dim; @@ -183,7 +241,15 @@ class BeamSearchTest : public ::testing::Test { reader_->open(index_path_.string()); reader_->register_thread(); td_.ctx_ = reader_->get_ctx(); - td_.alloc_scratch(2u * beam, geom_.page_size, has_pq_ ? n_chunks * 256u : 0u); + ThreadDataScratchConfig cfg; + cfg.n_page_slots = 2u * beam; + cfg.page_size = geom_.page_size; + cfg.pq_table_entries = has_pq_ ? n_chunks * 256u : 0u; + cfg.max_slot_id = n; + cfg.max_degree = r; + cfg.search_list_size = static_cast(std::max(n, 100)); + cfg.query_dim = dim; + td_.alloc_scratch(cfg); beam_ = beam; } @@ -197,8 +263,11 @@ class BeamSearchTest : public ::testing::Test { std::filesystem::remove(index_path_, ec); } - std::vector> search(const float *q, uint32_t top_k, bool use_pq, - bool rerank, uint32_t L = 50) { + std::vector> search(const float *q, + uint32_t top_k, + bool use_pq, + bool rerank, + uint32_t L = 50) { SearchContext ctx; ctx.reader = reader_.get(); ctx.geom = &geom_; diff --git a/tests/diskann/test_diskann_pq.cpp b/tests/diskann/test_diskann_pq.cpp index 53dccde2..ce426f51 100644 --- a/tests/diskann/test_diskann_pq.cpp +++ b/tests/diskann/test_diskann_pq.cpp @@ -115,8 +115,9 @@ TEST_F(PQTableTest, PQDistanceEqualsTrueL2WhenLossless) { pq.encode(data.data(), n); std::vector table(static_cast(n_chunks) * kPQNumCentroids); + std::vector qres(dim); for (uint64_t qi = 0; qi < n; ++qi) { - pq.preprocess_query(data.data() + qi * dim, table.data()); + pq.preprocess_query(data.data() + qi * dim, table.data(), qres.data()); for (uint64_t pi = 0; pi < n; ++pi) { const float approx = pq.pq_distance(pi, table.data()); const float truth = exact_l2_sqr(data.data() + qi * dim, data.data() + pi * dim, dim); @@ -144,7 +145,8 @@ TEST_F(PQTableTest, PreprocessQueryMatchesManualTable) { const std::vector query = {3.0f, 3.0f, 5.0f, 5.0f}; // chunk0=(3,3) chunk1=(5,5) std::vector table(static_cast(n_chunks) * kPQNumCentroids); - pq.preprocess_query(query.data(), table.data()); + std::vector qres(dim); + pq.preprocess_query(query.data(), table.data(), qres.data()); for (uint32_t k = 0; k < kPQNumCentroids; ++k) { const float e0 = 2.0f * (3.0f - k) * (3.0f - k); @@ -163,7 +165,8 @@ TEST_F(PQTableTest, PQDistanceIsSumOfTableLookups) { pq.encode(data.data(), n); std::vector table(static_cast(n_chunks) * kPQNumCentroids); - pq.preprocess_query(data.data() + 5 * dim, table.data()); + std::vector qres(dim); + pq.preprocess_query(data.data() + 5 * dim, table.data(), qres.data()); const auto &codes = pq.codes(); for (uint64_t pi : {0ull, 1ull, 100ull, 399ull}) { float expected = 0.0f; @@ -231,10 +234,13 @@ TEST_F(PQTableTest, FileRoundtripBitIdentical) { // And the loaded table produces identical distances. std::vector t0(static_cast(n_chunks) * kPQNumCentroids); std::vector t1(static_cast(n_chunks) * kPQNumCentroids); - pq.preprocess_query(data.data(), t0.data()); - loaded.preprocess_query(data.data(), t1.data()); + std::vector qres0(dim); + std::vector qres1(dim); + pq.preprocess_query(data.data(), t0.data(), qres0.data()); + loaded.preprocess_query(data.data(), t1.data(), qres1.data()); for (uint64_t pi : {0ull, 250ull, 499ull}) { - EXPECT_FLOAT_EQ(pq.pq_distance(pi, t0.data()), loaded.pq_distance(pi, t1.data())) << "pi=" << pi; + EXPECT_FLOAT_EQ(pq.pq_distance(pi, t0.data()), loaded.pq_distance(pi, t1.data())) + << "pi=" << pi; } } @@ -261,9 +267,8 @@ TEST_F(PQTableTest, LoadRejectsWrongSize) { TEST_F(PQTableTest, FromCodebookRejectsBadShapes) { EXPECT_THROW(PQTable::from_codebook(4, 3, std::vector(4), std::vector(4)), std::invalid_argument); // 4 % 3 != 0 - EXPECT_THROW( - PQTable::from_codebook(4, 2, std::vector(3), std::vector(2 * 256 * 2)), - std::invalid_argument); // global centroid wrong size + EXPECT_THROW(PQTable::from_codebook(4, 2, std::vector(3), std::vector(2 * 256 * 2)), + std::invalid_argument); // global centroid wrong size EXPECT_THROW(PQTable::from_codebook(4, 2, std::vector(4), std::vector(10)), std::invalid_argument); // codebook wrong size } diff --git a/tests/diskann/test_diskann_tombstone_search.cpp b/tests/diskann/test_diskann_tombstone_search.cpp index 4b3ef7ca..5a29d30e 100644 --- a/tests/diskann/test_diskann_tombstone_search.cpp +++ b/tests/diskann/test_diskann_tombstone_search.cpp @@ -28,6 +28,7 @@ using alaya::diskann::SearchContext; using alaya::diskann::SearchParams; using alaya::diskann::SearchStats; using alaya::diskann::ThreadData; +using alaya::diskann::ThreadDataScratchConfig; using alaya::diskann::TombstoneBitmap; using alaya::diskann::write_disk_layout; @@ -72,7 +73,14 @@ class TombstoneSearchTest : public ::testing::Test { reader_->open(index_path_.string()); reader_->register_thread(); td_.ctx_ = reader_->get_ctx(); - td_.alloc_scratch(/*n_page_slots=*/8, geom_.page_size, /*pq_table_entries=*/0); + ThreadDataScratchConfig cfg; + cfg.n_page_slots = 8; + cfg.page_size = geom_.page_size; + cfg.max_slot_id = n_; + cfg.max_degree = scn_.r; + cfg.search_list_size = 50; + cfg.query_dim = scn_.dim; + td_.alloc_scratch(cfg); } void TearDown() override {