From 69d990a62e47a40bbaa09cab27d7b4d9d26e39e8 Mon Sep 17 00:00:00 2001 From: ethanglaser Date: Thu, 20 Aug 2026 16:51:03 -0700 Subject: [PATCH] separate seqlock fine grain into separate index --- include/svs/concurrent/README.md | 263 +++ include/svs/concurrent/blocked_data.h | 433 ++++ include/svs/concurrent/concurrent.h | 54 + include/svs/concurrent/consolidate.h | 525 +++++ include/svs/concurrent/dynamic_index.h | 2080 +++++++++++++++++ .../svs/concurrent/dynamic_search_buffer.h | 625 +++++ include/svs/concurrent/graph.h | 677 ++++++ include/svs/concurrent/graph_concepts.h | 222 ++ include/svs/concurrent/greedy_search.h | 203 ++ include/svs/concurrent/iterator.h | 358 +++ include/svs/concurrent/multi.h | 967 ++++++++ include/svs/concurrent/prune.h | 338 +++ include/svs/concurrent/reverse_edges.h | 99 + include/svs/concurrent/spinlock.h | 49 + include/svs/concurrent/translation.h | 491 ++++ include/svs/concurrent/vamana_build.h | 605 +++++ include/svs/lib/concurrency/atomic_span.h | 93 + include/svs/lib/concurrency/seqlock.h | 141 ++ include/svs/lib/segmented_vector.h | 330 +++ tests/CMakeLists.txt | 75 + tests/svs/concurrent/concurrency.cpp | 578 +++++ tests/svs/concurrent/consolidate.cpp | 82 + tests/svs/concurrent/dynamic_index.cpp | 548 +++++ tests/svs/concurrent/dynamic_index_2.cpp | 727 ++++++ tests/svs/concurrent/graph.cpp | 129 + tests/svs/concurrent/iterator.cpp | 422 ++++ tests/svs/concurrent/multi.cpp | 475 ++++ tests/svs/concurrent/prune.cpp | 59 + tests/svs/concurrent/translation.cpp | 288 +++ tests/svs/lib/segmented_vector.cpp | 237 ++ 30 files changed, 12173 insertions(+) create mode 100644 include/svs/concurrent/README.md create mode 100644 include/svs/concurrent/blocked_data.h create mode 100644 include/svs/concurrent/concurrent.h create mode 100644 include/svs/concurrent/consolidate.h create mode 100644 include/svs/concurrent/dynamic_index.h create mode 100644 include/svs/concurrent/dynamic_search_buffer.h create mode 100644 include/svs/concurrent/graph.h create mode 100644 include/svs/concurrent/graph_concepts.h create mode 100644 include/svs/concurrent/greedy_search.h create mode 100644 include/svs/concurrent/iterator.h create mode 100644 include/svs/concurrent/multi.h create mode 100644 include/svs/concurrent/prune.h create mode 100644 include/svs/concurrent/reverse_edges.h create mode 100644 include/svs/concurrent/spinlock.h create mode 100644 include/svs/concurrent/translation.h create mode 100644 include/svs/concurrent/vamana_build.h create mode 100644 include/svs/lib/concurrency/atomic_span.h create mode 100644 include/svs/lib/concurrency/seqlock.h create mode 100644 include/svs/lib/segmented_vector.h create mode 100644 tests/svs/concurrent/concurrency.cpp create mode 100644 tests/svs/concurrent/consolidate.cpp create mode 100644 tests/svs/concurrent/dynamic_index.cpp create mode 100644 tests/svs/concurrent/dynamic_index_2.cpp create mode 100644 tests/svs/concurrent/graph.cpp create mode 100644 tests/svs/concurrent/iterator.cpp create mode 100644 tests/svs/concurrent/multi.cpp create mode 100644 tests/svs/concurrent/prune.cpp create mode 100644 tests/svs/concurrent/translation.cpp create mode 100644 tests/svs/lib/segmented_vector.cpp diff --git a/include/svs/concurrent/README.md b/include/svs/concurrent/README.md new file mode 100644 index 000000000..40d17f06b --- /dev/null +++ b/include/svs/concurrent/README.md @@ -0,0 +1,263 @@ +# `svs::index::vamana::concurrent` — concurrent dynamic Vamana index + +A dynamic Vamana index that supports **lock-free search concurrent with mutation**: +searches, `add_points`, `delete_entries`, and `consolidate` may all be in flight at the +same time, from any number of threads. + +This is a **separate index**. The pre-existing `svs::index::vamana::VamanaIndex` and +`svs::index::vamana::MutableVamanaIndex` are not modified, and neither is anything else +under `include/svs/`. The only pre-existing file this stack touches at all is +`tests/CMakeLists.txt`, to register the new tests. + +## Why a separate index + +The functionality being reproduced here originates in +[`razdoburdin:ScalableVectorSearch:seqlock`][seqlock], which implements it by editing +`MutableVamanaIndex` and its collaborators in place. That is invasive: the graph, the ID +translator, the search buffer, the pruning heuristics, `greedy_search`, and the build +driver all change shape, and the static index shares most of them. Reproducing the same +functionality as a parallel stack keeps the existing indexes bit-for-bit unchanged, so the +two can coexist and be compared. + +[seqlock]: https://github.com/intel/ScalableVectorSearch/compare/main...razdoburdin:ScalableVectorSearch:seqlock + +## How the separation works + +Everything lives in `svs::index::vamana::concurrent`, a namespace nested inside the one it +shadows. C++ name lookup then gives a **delta layer** for free: a name declared inside +`concurrent` hides the same name in the enclosing namespace, while a name *not* declared +inside `concurrent` resolves to the enclosing namespace's. Each header here therefore only +needs to carry what actually changed; `VamanaSearchParameters`, +`VamanaBuildParameters`, `SearchScratchspace`, the `extensions` customization points, and +the rest are picked up from upstream unchanged. + +Two consequences are easy to get wrong, and both are load-bearing: + +1. **Do not redeclare an entity this stack does not change.** A redeclaration inside + `concurrent` is a *distinct type* that merely looks identical, and values crossing the + boundary stop converting. Redeclaring `GreedySearchPrefetchParameters`, for instance, + makes `SearchScratchspace::prefetch_parameters` (an upstream type, reused verbatim) + fail to bind. Use a `using`-declaration instead — see the block at the top of + `greedy_search.h`, which aliases `GreedySearchPrefetchParameters`, + `GreedySearchTracker`, `NullTracker`, `EntryPointInitializer`, and `NeighborBuilder`. + +2. **Qualify intra-namespace calls.** Once those using-declarations make the upstream + overloads visible, ADL on an upstream argument type makes an unqualified + `greedy_search(...)` ambiguous. Calls inside this stack are written + `concurrent::greedy_search(...)`. + +## The concurrency design + +### Sequence locks on adjacency lists + +`lib::SeqLockCounter` / `lib::SeqLockArray` (`svs/lib/concurrency/seqlock.h`) give each +node an even/odd version counter. A writer bumps it odd, mutates, bumps it even; a reader +snapshots it (`read_begin`), reads, and re-checks (`read_validate`), retrying on +disagreement. `greedy_search` wraps its per-node neighbor expansion in that retry loop, so +searches take no locks at all. Stale neighbors from an invalidated read are harmless: they +carry valid ids and distances, and the search buffer dedupes by id. + +Every adjacency slot is accessed through `SimpleGraphBase::relaxed_load` / +`relaxed_store`, which are relaxed `std::atomic_ref` operations. The *ordering* comes from +the sequence-lock counters, not from these accesses; relaxed atomics compile to plain +loads and stores on the platforms SVS targets, so the only thing they buy is the removal +of a formal data race — which is exactly what makes the TSan run below meaningful. + +### Grow-stable storage + +Appending must never relocate what a lock-free reader is holding. `lib::SegmentedVector` +(`svs/lib/segmented_vector.h`) is a two-level array whose existing elements keep stable +addresses across growth; it backs the per-node locks, the sequence counters, the slot +metadata, and the reverse-edge lists. + +The dataset needs the same property. Rather than change `svs::data::Blocked`, this stack +adds an allocator tag `SegmentedBlocked` and an **additive** partial specialization +`svs::data::SimpleData>` (`blocked_data.h`), identical +to the `Blocked` specialization except that the outer block directory is a +`SegmentedVector`. Because the result is still a `SimpleData`, every generic facility +written against that template — the dataset concepts, the `extensions` customization +points, `compact_data`, the save/load serializer — applies with no further work. + +### Reverse edges: O(|deleted|) consolidation + +`graphs::ReverseEdges` (`reverse_edges.h`) keeps a per-node in-neighbor list `R(n)`, so +`consolidate()` visits only the in-neighbors of deleted nodes instead of scanning the whole +graph. It is off (null) by default; only the dynamic index enables it, so the static index +and the compaction scratch graphs pay a single null check per mutator. + +The maintained invariant is `R(d) ⊇ in(d)`: `record` runs on every created edge, so `R` +may hold stale or duplicated entries but never misses a live in-edge. See *Divergence 7* +below for why the tempting weaker invariant does not work. + +### Lock discipline + +Three mutexes, with a global acquisition order: + +``` +compact_mutex_ -> slot_alloc_mutex_ +compact_mutex_ -> translator_mutex_ +``` + +`slot_alloc_mutex_` and `translator_mutex_` are never held simultaneously. + +- `compact_mutex_` — shared by readers and by `add_points`; exclusive only by `compact()`, + which *shrinks* storage and so must drain readers. Growth needs no exclusion. +- `slot_alloc_mutex_` — held only for Phase 1 of `add_points` (reserving slots). Because + `add_points` holds `compact_mutex_` *shared* for its whole duration and + `slot_alloc_mutex_` only briefly, **`add_points` may be called concurrently from + multiple threads**. +- `translator_mutex_` — guards the ID translator's hash maps. + +`std::shared_mutex` is not recursive, so every translation operation comes in two +flavours: `foo(...)` takes the shared lock itself, and `unsafe_foo(...)` requires the +caller to already hold it (via `lock_for_translation()`). Batch paths and +`BatchIterator::next` take the lock once and use the `unsafe_` variants; anything else +should use the plain form. Calling a self-locking accessor from a context that already +holds the lock is a latent deadlock, not merely slow: a writer arriving between the two +shared acquisitions blocks the second one. + +### Slot lifecycle + +`SlotMetadata` gains a fourth state, `Pending`: a slot reserved by an in-flight +`add_points` whose vector is copied but whose adjacency list is still being built. Pending +slots are invisible to search, to `consolidate`, and to subsequent `add_points` until +promoted to `Valid`. + +## Deliberate divergences from the source branch + +The source branch was reproduced feature-for-feature. Where this stack differs, it is for +one of two reasons: (a) it must avoid editing a pre-existing file, or (b) the source +branch has a defect. Both kinds are listed. + +**1. Search-path extension (avoids editing `extensions.h`).** The source branch edits +`svs::index::vamana::extensions` directly. Here the equivalent behaviour is an +`svs_invoke` override for `single_search` plus a `supplement_search_buffer` step on the +concurrent index itself. The rewrite also fixes two problems in the original: a recursive +`shared_mutex` acquisition that can deadlock, and an iteration over the ID translator with +no lock held. + +**2. Grow-stable dataset (avoids editing `core/data/simple.h`).** The source branch +changes `SimpleData>` in place, which alters an existing dataset +type. Here it is a new allocator tag plus an additive partial specialization; see +*Grow-stable storage* above. + +**3. Spin lock (avoids editing `lib/spinlock.h`).** The source branch extends +`svs::lib::SpinLock`. Here `concurrent::SpinLock` (`spinlock.h`) subclasses it and adds +what this stack needs. + +**4. `capacity()` on scalar quantization — omitted.** The source branch adds a +`capacity()` accessor to `quantization/scalar/scalar.h`. It is used only for memory +reporting and nothing in this stack calls it, so it is left out rather than editing a +pre-existing header. + +**5. `NullLockGuard` and the Python GIL changes — omitted.** The source branch adds a +`NullLockGuard` to `index/vamana/index.h` and releases the GIL around some Python +bindings. Both are about integrating the *modified in place* index into existing call +paths; a separate index does not need them. The Python bindings continue to expose the +pre-existing dynamic index. + +**6. `PruneState` test expectations — corrected.** The source branch's test asserts +`reenable(Candidate) == Candidate`. The implementation returns `Available`, which is what +the two-round pruning heuristic requires. `tests/svs/concurrent/prune.cpp` asserts the +correct values (`reenable(Pruned) == Pruned`, `reenable(Candidate) == Available`, +`excluded(Candidate) == true`). + +**7. Reverse-edge rebuild — bug fixed.** `rebuild_reverse_edges` in the source branch +records `src` into `R(dst)` only when the reverse edge `dst -> src` does not also exist, +halving the index on the reasoning that `gather_work_set` visits `out(d) ∪ R(d)` and so +already covers symmetric in-neighbors. That invariant — *"for every edge `u -> d`: `u ∈ +R(d)` **or** the edge `d -> u` exists"* — holds immediately after a rebuild but is +**not maintainable**: the second disjunct is falsified the moment consolidation rewires +`d` and drops `d -> u`, at which point `u`'s in-edge is invisible, and a later deletion of +`d` leaves `u` pointing at a retired slot. + +The failure is masked in the source branch because every one of its tests pairs +`consolidate()` with `compact()`, and `compact()` rebuilds the index from scratch. +Repeated `consolidate()` *without* `compact()` — the cheap maintenance path, and the +normal one — corrupts the graph within two rounds, which +`debug_check_invariants()` reports as `Node number N has an invalid (Empty) neighbor`. +This stack records unconditionally, giving the strictly stronger `R(d) ⊇ in(d)`, which no +edge *removal* can break. + +**8. Unsynchronized translator reads — bug fixed.** `translate_external_id`, +`translate_external_id_or`, `has_id`, `translate_internal_id`, and `on_ids` read the +translator's `tsl::robin_map`s with no lock in the source branch, while `add_points` +inserts and `consolidate` erases under `translator_mutex_`. A comment there anticipates +reading a stale value, but the actual hazard is worse: an insert can rehash and free the +bucket array a reader is walking. These are now the two-flavour operations described under +*Lock discipline*, and TSan reports the original as a race on every search that overlaps +an insert. + +**9. Greedy-search scaffolding via `using`-declarations.** See *How the separation works*. +The source branch has no analogue because it edits upstream in place and so never crosses +a namespace boundary. + +## Layout + +| File | Contents | +| --- | --- | +| `spinlock.h` | `concurrent::SpinLock` | +| `blocked_data.h` | `SegmentedBlocked` tag + `SimpleData` specialization | +| `graph_concepts.h` | graph concepts for the concurrent graph API | +| `reverse_edges.h` | `graphs::ReverseEdges` in-neighbor index | +| `graph.h` | `SimpleGraphBase`, `SimpleBlockedGraph`, `AddEdgeResult` | +| `translation.h` | `IDTranslator` | +| `dynamic_search_buffer.h` | `MutableBuffer`, `PredicatedSearchNeighbor` | +| `prune.h` | pruning heuristics and `PruneStrategy` | +| `greedy_search.h` | `greedy_search` with the SeqLock retry loop | +| `vamana_build.h` | `VamanaBuilder` | +| `consolidate.h` | reverse-edge-driven and full-scan consolidation | +| `dynamic_index.h` | `MutableVamanaIndex`, `auto_dynamic_assemble` | +| `iterator.h` | `BatchIterator` | +| `multi.h` | `MultiMutableVamanaIndex`, `MultiBatchIterator` | + +Shared building blocks that are not Vamana-specific live under `svs/lib/`: +`lib/segmented_vector.h`, `lib/concurrency/seqlock.h`, `lib/concurrency/atomic_span.h`. + +## Tests + +`tests/svs/concurrent/` mirrors the upstream Vamana tests against this stack — +`translation.cpp`, `graph.cpp`, `prune.cpp`, `consolidate.cpp`, `dynamic_index.cpp`, +`dynamic_index_2.cpp`, `iterator.cpp`, `multi.cpp` — plus `concurrency.cpp`, which is new: +it runs searches against the index while writers insert, delete, and consolidate, and +checks id round-trips, result-set consistency, and post-mutation recall. The source branch +has no multi-threaded test. + +`tests/svs/lib/segmented_vector.cpp` covers `lib::SegmentedVector` directly, including +address stability under concurrent growth. + +Every test case here is tagged `[concurrent]`, so the whole set runs as + +```sh +ctest -L "concurrent|segmented_vector" # 24 tests +./tests/tests "[concurrent],[segmented_vector]" +``` + +`tests/svs/concurrent/dynamic_index_2.cpp` diverges from the upstream file it was ported +from in one respect: the upstream logging tests push a capturing sink onto the +*process-global* logger and never remove it, so the sink outlives by reference the vector it +captures and every later global-logger statement is a use-after-free. The port scopes the +push with a `ScopedGlobalSink` guard. The same leak exists at six sites in pre-existing test +files (`svs/index/flat/flat.cpp`, `svs/index/inverted/{memory_based,clustering}.cpp`, +`svs/index/vamana/{index,dynamic_index_2}.cpp`) and is left alone there; it is why a full +`tests` run can abort with a SIGSEGV inside an unrelated test that happens to log a warning +(commonly `Vamana Index Parameters`). That crash reproduces in a binary built without any of +this stack's sources. + +### ThreadSanitizer + +The correctness of this stack rests almost entirely on memory ordering, which an +uninstrumented test can only fail to disprove. TSan targets are opt-in because they cost +roughly an order of magnitude in time and memory: + +```sh +cmake -DSVS_EXPERIMENTAL_ENABLE_CONCURRENT_TSAN=YES ... +ctest -L tsan +``` + +This builds two targets. `concurrent_tsan` must come out clean. `concurrent_tsan_negative` +is a **negative control**: it defines `SVS_CONCURRENT_UNSAFE_PLAIN_GRAPH_ACCESS`, which +degrades the adjacency-slot accessors to plain loads and stores, and is registered with +`WILL_FAIL TRUE`. A clean positive run only means something if the same run reports races +once the atomics are taken away — otherwise it is equally consistent with TSan watching +the wrong memory. (Never define that macro in a real build.) diff --git a/include/svs/concurrent/blocked_data.h b/include/svs/concurrent/blocked_data.h new file mode 100644 index 000000000..72a07a4a0 --- /dev/null +++ b/include/svs/concurrent/blocked_data.h @@ -0,0 +1,433 @@ +/* + * Copyright 2026 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +/// +/// @file +/// @brief Grow-stable blocked storage for the concurrent Vamana index. +/// +/// ``svs::data::SimpleData>`` keeps its fixed-size blocks in a +/// ``std::vector``, so appending a block may reallocate the outer directory. A lock-free +/// reader evaluating ``blocks_[block_id]`` while a writer grows the dataset would then read +/// through a freed pointer. +/// +/// This header supplies the same storage with a grow-stable outer directory +/// (``lib::SegmentedVector``), selected by a distinct allocator tag so that **no existing +/// dataset type changes**: ``SegmentedBlocked`` picks up a partial specialization of +/// ``svs::data::SimpleData`` that is otherwise a line-for-line copy of the ``Blocked`` +/// specialization. Element addressing ``(block_id, data_id)`` is unchanged; only the outer +/// block directory becomes a two-level lock-free array. +/// +/// Because the result is still a ``svs::data::SimpleData``, every generic facility written +/// against that template — the dataset concepts, the ``extensions`` customization points, +/// ``compact_data``, and the save/load serializer — applies unmodified. +/// + +#include "svs/core/data/simple.h" +#include "svs/lib/segmented_vector.h" + +#include +#include +#include +#include + +namespace svs::index::vamana::concurrent { + +/// +/// @brief Allocator tag selecting grow-stable blocked storage. +/// +/// Behaves exactly like ``svs::data::Blocked`` — same blocking parameters, same +/// inner allocator — and differs only in which ``SimpleData`` specialization it selects. +/// +template class SegmentedBlocked : public svs::data::Blocked { + public: + using parent_type = svs::data::Blocked; + using allocator_type = Alloc; + using value_type = typename std::allocator_traits::value_type; + + constexpr SegmentedBlocked() = default; + explicit SegmentedBlocked(const allocator_type& alloc) + : parent_type{alloc} {} + explicit SegmentedBlocked(const svs::data::BlockingParameters& parameters) + : parent_type{parameters} {} + explicit SegmentedBlocked( + const svs::data::BlockingParameters& parameters, const allocator_type& alloc + ) + : parent_type{parameters, alloc} {} + + // Enable rebinding of allocators. + template friend class SegmentedBlocked; + template + SegmentedBlocked(const SegmentedBlocked& other) + : parent_type{other.parameters(), other.get_allocator()} {} +}; + +} // namespace svs::index::vamana::concurrent + +namespace svs::data { + +// ``SegmentedBlocked`` is a blocked allocator for the purposes of the library's +// blocked/non-blocked dispatch, exactly as ``Blocked`` is. +template +inline constexpr bool + is_blocked_v> = true; + +} // namespace svs::data + +namespace svs::lib::detail { + +// Allow rebinding of allocators through the SegmentedBlocked wrapper. +template +struct AllocatorRebinder> { + using type = + svs::index::vamana::concurrent::SegmentedBlocked>; +}; + +} // namespace svs::lib::detail + +namespace svs::data { + +/// +/// @brief ``SimpleData`` specialization with a grow-stable block directory. +/// +/// A line-for-line copy of the ``SimpleData>`` specialization +/// with two changes: +/// +/// 1. ``blocks_`` is a ``lib::SegmentedVector`` rather than a ``std::vector``, so appending +/// a block never relocates the existing block wrappers (or the heap buffers they point +/// at). A concurrent lock-free reader subscripting ``blocks_[block_id]`` is therefore +/// safe against a writer growing the dataset. +/// 2. ``add_block`` move-constructs into the new slot (``push_back``) instead of +/// ``emplace_back``, because ``DenseArray``'s move-*assignment* compares allocators. +/// +/// Shrinking still frees storage: ``drop_block`` destroys the trailing block, releasing its +/// heap buffer. A reader inside that block would dangle, so the owning index must drain +/// readers (exclusive lock) before shrinking — the same obligation the ``Blocked`` +/// specialization carries. +/// +template +class SimpleData> { + public: + ///// Static Members + static constexpr bool supports_saving = true; + + // Type Aliases + using dim_type = std::tuple>; + using allocator_type = svs::index::vamana::concurrent::SegmentedBlocked; + using inner_allocator_type = Alloc; + using array_type = DenseArray; + + /// Return the underlying allocator. + const allocator_type& get_allocator() const { return allocator_; } + + // value types + using element_type = T; + using value_type = std::span; + using const_value_type = std::span; + + using lib_alloc_data_type = SimpleData< + T, + Extent, + svs::index::vamana::concurrent::SegmentedBlocked>>; + /// Already blocked, so lib_blocked_alloc_data_type is the same as lib_alloc_data_type. + using lib_blocked_alloc_data_type = SimpleData< + T, + Dynamic, + svs::index::vamana::concurrent::SegmentedBlocked>>; + + ///// Constructors + SimpleData(size_t n_elements, size_t n_dimensions, const allocator_type& alloc) + : blocksize_{compute_blocksize(alloc, n_dimensions)} + , blocks_{} + , dimensions_{n_dimensions} + , size_{n_elements} + , allocator_{alloc} { + size_t elements_per_block = blocksize_.value(); + size_t num_blocks = lib::div_round_up(n_elements, elements_per_block); + for (size_t i = 0; i < num_blocks; ++i) { + add_block(); + } + } + + SimpleData(size_t n_elements, size_t n_dimensions) + : SimpleData{n_elements, n_dimensions, allocator_type()} {} + + /// + /// Convert a linear index into an inner-outer index to access the blocked dataset. + /// Returns a pair `p` where: + /// - `p.first` is the block index. + /// - `p.second` is the index within the block. + /// + std::pair resolve(size_t i) const { + return std::pair{i / blocksize_, i % blocksize_}; + } + + /// + /// Return the blocksize with reference to the stored data vectors. + /// + lib::PowerOfTwo blocksize() const { return blocksize_; } + + /// + /// Return the blocksize with respect to bytes. + /// + lib::PowerOfTwo blocksize_bytes() const { + return allocator_.parameters().blocksize_bytes; + } + + /// + /// Return the number of blocks in the dataset. + /// + size_t num_blocks() const { return blocks_.size(); } + + /// + /// Return the maximum number of data vectors that can be stored before a new block is + /// required. + /// + size_t capacity() const { return num_blocks() * blocksize(); } + + /// + /// Return an iterator over each index in the dataset. + /// + threads::UnitRange eachindex() const { + return threads::UnitRange{0, size()}; + } + + /// + /// Add a new data block to the end of the current collection of blocks. + /// + void add_block() { + blocks_.push_back(array_type( + make_dims(blocksize().value(), lib::forward_extent(dimensions())), + allocator_.get_allocator() + )); + } + + /// + /// Remove a data block from the end of the block list. + /// + void drop_block() { + if (!blocks_.empty()) { + blocks_.pop_back(); + } + } + + /// + /// Resizing + /// + void resize(size_t new_size) { + if (new_size > size()) { + // Add blocks until there is sufficient capacity. + while (new_size > capacity()) { + add_block(); + } + // Release: growth is concurrent with lock-free readers calling `size()`, and + // publishing the new size must not be reordered before the blocks that back it. + std::atomic_ref(size_).store(new_size, std::memory_order_release); + } else if (new_size < size()) { + // Reset size then drop blocks until the new size is within the last block. + // Shrinking frees memory, so it only ever runs with readers excluded (the + // owning index holds `compact_mutex_` exclusive); the atomic store is here for + // consistency, not for safety, which the exclusion provides. + std::atomic_ref(size_).store(new_size, std::memory_order_release); + while (capacity() - blocksize().value() > new_size) { + drop_block(); + } + } + } + + void shrink_to_fit() { + // We already shrink when down-sizing, so ``shink_to_fit`` becomes a no-op. + } + + ///// + ///// Dataset API + ///// + + // Acquire, paired with the release store in `resize`: `size_` grows concurrently with + // lock-free readers, so a plain read would be a data race. + size_t size() const { + return std::atomic_ref(const_cast(size_)) + .load(std::memory_order_acquire); + } + constexpr size_t dimensions() const { + if constexpr (Extent != Dynamic) { + return Extent; + } else { + return dimensions_; + } + } + + size_t element_size() const { return sizeof(element_type) * dimensions(); } + + const_value_type get_datum(size_t i) const { + auto [block_id, data_id] = resolve(i); + return getindex(blocks_, block_id).slice(data_id); + } + + value_type get_datum(size_t i) { + auto [block_id, data_id] = resolve(i); + return getindex(blocks_, block_id).slice(data_id); + } + + void prefetch(size_t i) const { lib::prefetch(get_datum(i)); } + + template + void set_datum(size_t i, std::span datum) { + if constexpr (checkbounds_v) { + if (datum.size() != dimensions()) { + throw ANNEXCEPTION( + "Datum with dimensions {} is not equal to internal dimensions {}!", + datum.size(), + dimensions_ + ); + } + } + + if constexpr (std::is_same_v>) { + std::copy(datum.begin(), datum.end(), get_datum(i).begin()); + } else { + std::transform( + datum.begin(), + datum.end(), + get_datum(i).begin(), + [](const U& u) { return lib::relaxed_narrow(u); } + ); + } + } + + template void set_datum(size_t i, const std::vector& v) { + set_datum(i, lib::as_span(v)); + } + + /// + /// Construct an identical copy of the dataset. + /// Not implemented as a copy constructor to avoid unintentional copies. + /// + SimpleData copy() const { + SimpleData other{size(), dimensions(), allocator_}; + for (const auto& i : eachindex()) { + other.set_datum(i, get_datum(i)); + } + return other; + } + + ///// Compaction + template + void + compact(std::span new_to_old, Pool& threadpool, size_t batchsize = 1'000'000) { + // Allocate scratch space. + batchsize = std::min(batchsize, size()); + auto buffer = data::SimpleData(batchsize, dimensions()); + compact_data(*this, buffer, new_to_old, threadpool); + } + + template + void compact(std::span new_to_old, size_t batchsize = 1'000'000) { + auto pool = threads::SequentialThreadPool(); + compact(new_to_old, pool, batchsize); + } + + ///// Saving + lib::SaveTable save(const lib::SaveContext& ctx) const { + return GenericSerializer::save(*this, ctx); + } + + void save(std::ostream& os) const { return GenericSerializer::save(*this, os); } + + lib::SaveTable metadata() const { return GenericSerializer::metadata(*this); } + + static bool check_load_compatibility(std::string_view schema, lib::Version version) { + return GenericSerializer::check_compatibility(schema, version); + } + + static SimpleData + load(const lib::LoadTable& table, const allocator_type& allocator = {}) { + return GenericSerializer::load( + table, lib::Lazy([&allocator](size_t n_elements, size_t n_dimensions) { + return SimpleData(n_elements, n_dimensions, allocator); + }) + ); + } + + static SimpleData load( + const lib::ContextFreeLoadTable& table, + std::istream& is, + const allocator_type& allocator = {} + ) { + return GenericSerializer::load( + table, is, lib::Lazy([&allocator](size_t n_elements, size_t n_dimensions) { + return SimpleData(n_elements, n_dimensions, allocator); + }) + ); + } + + static SimpleData + load(const std::filesystem::path& path, const allocator_type& allocator = {}) { + if (detail::is_likely_reload(path)) { + return lib::load_from_disk(path, allocator); + } + // Try loading directly. + return io::auto_load( + path, lib::Lazy([&allocator](size_t n_elements, size_t n_dimensions) { + return SimpleData(n_elements, n_dimensions, allocator); + }) + ); + } + + private: + // Helper static function to compute blocksize value. + // If blocking parameters have defined blocksize_elements, use it + // directly. Otherwise, compute blocksize based on blocksize_bytes. + static lib::PowerOfTwo compute_blocksize(const allocator_type& alloc, size_t dim) { + if (alloc.parameters().blocksize_elements.has_value()) { + return alloc.parameters().blocksize_elements.value(); + } else { + return lib::prevpow2( + alloc.parameters().blocksize_bytes.value() / (sizeof(T) * dim) + ); + } + } + + private: + // The blocksize in terms of number of vectors. + lib::PowerOfTwo blocksize_; + // Grow-stable directory of fixed-size blocks: appending a block never relocates the + // existing block wrappers (or the heap buffers they point to), so a concurrent + // lock-free reader subscripting blocks_[block_id] is safe against a writer growing the + // dataset. Element addressing (block_id, data_id) is unchanged; only the outer block + // directory is grow-stable (2-level lock-free array). See svs/lib/segmented_vector.h. + lib::SegmentedVector blocks_; + size_t dimensions_; + size_t size_; + allocator_type allocator_; +}; + +} // namespace svs::data + +namespace svs::index::vamana::concurrent { + +/// +/// @brief Grow-stable analogue of ``svs::data::BlockedData``. +/// +/// The dataset type the concurrent index is built on. Interchangeable with +/// ``svs::data::BlockedData`` at every call site; the difference is that growing it does +/// not invalidate a concurrent reader. +/// +template > +using SegmentedBlockedData = svs::data::SimpleData>; + +} // namespace svs::index::vamana::concurrent diff --git a/include/svs/concurrent/concurrent.h b/include/svs/concurrent/concurrent.h new file mode 100644 index 000000000..67c61b4c3 --- /dev/null +++ b/include/svs/concurrent/concurrent.h @@ -0,0 +1,54 @@ +/* + * Copyright 2026 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +/// +/// @file +/// @brief Umbrella header for the concurrent dynamic Vamana index. +/// +/// Pulls in the whole stack and introduces the short alias ``svs::concurrent``. See +/// ``include/svs/concurrent/README.md`` for the design. +/// +/// Including the individual headers directly works equally well; this one exists so that a +/// caller who just wants the index does not have to know which of them declares it. +/// + +#include "svs/concurrent/blocked_data.h" +#include "svs/concurrent/consolidate.h" +#include "svs/concurrent/dynamic_index.h" +#include "svs/concurrent/dynamic_search_buffer.h" +#include "svs/concurrent/graph.h" +#include "svs/concurrent/graph_concepts.h" +#include "svs/concurrent/greedy_search.h" +#include "svs/concurrent/iterator.h" +#include "svs/concurrent/multi.h" +#include "svs/concurrent/prune.h" +#include "svs/concurrent/reverse_edges.h" +#include "svs/concurrent/spinlock.h" +#include "svs/concurrent/translation.h" +#include "svs/concurrent/vamana_build.h" + +namespace svs { + +/// @brief Short alias for the concurrent dynamic Vamana index namespace. +/// +/// The implementation namespace is nested inside the one it extends so that unchanged +/// entities resolve to their pre-existing declarations (see the README); that makes the +/// fully-qualified name long. This alias is purely for convenience at call sites. +namespace concurrent = index::vamana::concurrent; + +} // namespace svs diff --git a/include/svs/concurrent/consolidate.h b/include/svs/concurrent/consolidate.h new file mode 100644 index 000000000..53bb7ba4a --- /dev/null +++ b/include/svs/concurrent/consolidate.h @@ -0,0 +1,525 @@ +/* + * Copyright 2023 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +// local +#include "svs/core/distance.h" +#include "svs/concurrent/graph.h" +#include "svs/index/vamana/extensions.h" +#include "svs/concurrent/prune.h" +#include "svs/lib/array.h" +#include "svs/lib/threads.h" +#include "svs/lib/timing.h" + +// external +#include "tsl/robin_set.h" + +// stdlib +#include +#include + +namespace svs::index::vamana::concurrent { + +/// +/// Parameters controlling aspects of the graph consolidation process. +/// +/// * `update_batch_size`: The algorithm for graph consolidation is a two-phase algorithm +/// over batches of the dataset to facalitate parallelism. +/// +/// The first phase is a read-only phase where updates for the graph are prepared in an +/// auxiliary data structure. The second phase commits these updates to the graph. +/// +/// This multi-phase approach allows for parallelism with both phases without worrying +/// about mutating the graph while reading from it. +/// +/// This parameter controls how large of a batch is processed during each phase. +/// +/// * `prune_to`: The number of candidates to prune to. +/// +/// * `alpha`: The pruning parameter to use when constructing a new set of neighbors for +/// vertices with at least one deleted neighbor. +/// +struct ConsolidationParameters { + size_t update_batch_size; + size_t prune_to; + size_t max_candidate_pool_size; + float alpha; +}; + +/// +/// A helper struct to store pending updates for the consolidated graph. +/// Represents pending updates using several large allocations rather than many small ones. +/// +template class BulkUpdate { + public: + // Constructor + BulkUpdate(size_t max_batch_size, size_t prune_to) + : neighbors_{make_dims(max_batch_size, prune_to)} + , lengths_{max_batch_size} + , needs_update_{max_batch_size} {} + + /// + /// Pre-conditions: + /// * 0 <= src < max_batch_size + /// * neighbors.size() <= prune_to + /// * May be called concurrently from multiple threads as long as `src` is + /// unique for each thread. + /// + template + void insert(size_t src, const std::vector& neighbors) { + assert(neighbors.size() <= getsize<1>(neighbors_)); + assert(src < getsize<0>(neighbors_)); + + needs_update_.at(src) = true; + // N.B.: There is no way we get even close to crossing what can be + // expressed as a 32-bit unsigned number. + lengths_.at(src) = lib::narrow_cast(neighbors.size()); + std::copy(neighbors.begin(), neighbors.end(), neighbors_.slice(src).begin()); + } + + /// + /// Return `true` if index `src` has a pending update. + /// + bool needs_update(size_t src) const { return needs_update_.at(src); } + + /// + /// Return the pending update for `src`. + /// Pre-conditions: + /// * `needs_update(src) == true` + /// + std::span get_update(size_t src) const { + assert(needs_update(src)); + return neighbors_.slice(src).first(lengths_.at(src)); + } + + /// + /// Prepare the data structure for another batch of processing. + /// **Preconditions:** + /// + /// * Must be called from a single thread. + /// + void prepare() { std::fill(needs_update_.begin(), needs_update_.end(), false); } + + private: + Matrix neighbors_; + Vector lengths_; + // N.B.: Use a `Vector` instead of a `std::vector` because `std::vector` + // specialized on `bool` and will thus cannot be guarantee coherent updates from + // multiple threads. + Vector needs_update_; +}; + +template struct ConsolidateThreadLocal { + template using allocator_type = threads::CacheAlignedAllocator; + + // Type Aliases + using set_type = tsl::robin_set, std::equal_to, allocator_type>; + using neighbor_vector_type = std::vector, allocator_type>>; + + // Members + set_type all_candidates{}; + neighbor_vector_type valid_candidates{}; + std::vector> final_candidates{}; +}; + +template < + graphs::MemoryGraph Graph, + data::ImmutableMemoryDataset Data, + threads::ThreadPool Pool, + typename Distance> +class GraphConsolidator { + public: + // Type Aliases + using I = typename Graph::index_type; + using graph_neighbor_container = typename Graph::const_value_type; + using datum_type = typename Data::const_value_type; + + using Compare = typename Distance::compare; + using scratch_type = ConsolidateThreadLocal; + + using set_type = typename scratch_type::set_type; + using neighbor_vector_type = typename scratch_type::neighbor_vector_type; + + // Members + private: + Graph& graph_; + const Data& data_; + Pool& threadpool_; + const Distance& distance_; + ConsolidationParameters params_; + + public: + // Constructor + GraphConsolidator( + Graph& graph, + const Data& data, + Pool& threadpool, + const Distance& distance, + const ConsolidationParameters& params + ) + : graph_{graph} + , data_{data} + , threadpool_{threadpool} + , distance_{distance} + , params_{params} { + // consolidate can observe an in-flight add_points that has + // already grown `data` but not yet `graph` + // assert(graph.n_nodes() <= data.size()); + } + + /// + /// Add all neighbors and neighbor-of-deleted-neighbors to the set `all_candidates`. + /// + /// @param all_candidates In-out parameter. After this function call, `all_candidates` + /// will contain the full set of neighbor candidate indices. + /// @param neighbors The current neighbors of the vertex being processed. + /// @param is_deleted Callable functor returning `true` of a vertex is deleted. + /// + template + void populate_candidates( + set_type& all_candidates, const Neighbors& neighbors, const Deleted& is_deleted + ) const { + all_candidates.clear(); + for (auto dst : neighbors) { + if (is_deleted(dst)) { + // SeqLock retry: a concurrent consolidate may be writing dst's + // neighbors if dst is not deleted in the other consolidate's view. + for (;;) { + auto maybe_seq = graph_.seq_counters()[dst].read_begin(); + if (!maybe_seq) { + svs::detail::pause(); + continue; + } + const auto& others = graph_.get_node(dst); + all_candidates.insert(others.begin(), others.end()); + if (graph_.seq_counters()[dst].read_validate(*maybe_seq)) { + break; + } + svs::detail::pause(); + } + } else { + all_candidates.insert(dst); + } + } + } + + template Accessor, typename SelfDistance, typename Deleted> + void filter_candidates( + neighbor_vector_type& valid_candidates, + const set_type& all_candidates, + const datum_type& src_data, + const Accessor& accessor, + SelfDistance& distance, + const Deleted& is_deleted + ) const { + distance::maybe_fix_argument(distance, src_data); + valid_candidates.clear(); + for (auto dst : all_candidates) { + if (is_deleted(dst)) { + continue; + } + + valid_candidates.push_back( + {dst, distance::compute(distance, src_data, accessor(data_, dst))} + ); + } + + std::sort(valid_candidates.begin(), valid_candidates.end(), Compare{}); + } + + template + void generate_updates( + const GlobalIds& global_ids, + const threads::UnitRange& local_ids, + BulkUpdate& update_buffer, + ConsolidateThreadLocal& tls, + const Deleted& is_deleted + ) const { + auto& [all_candidates, valid_candidates, final_candidates] = tls; + + auto build_adaptor = extensions::build_adaptor(data_, distance_); + + auto accessor = build_adaptor.general_accessor(); + auto&& general_distance = build_adaptor.general_distance(); + + for (auto i : local_ids) { + size_t src = global_ids[i]; + + if (is_deleted(src)) { + continue; + } + + // SeqLock retry: a concurrent consolidate's apply_updates may be + // writing src's neighbors while we read them. + for (;;) { + auto maybe_seq = graph_.seq_counters()[src].read_begin(); + if (!maybe_seq) { + svs::detail::pause(); + continue; + } + + // Determine if any of the neighbors of this node are deleted. + const auto& neighbors = graph_.get_node(src); + if (std::none_of(neighbors.begin(), neighbors.end(), is_deleted)) { + if (graph_.seq_counters()[src].read_validate(*maybe_seq)) { + break; + } + svs::detail::pause(); + continue; + } + + // Add all neighbors and neighbors-of-deleted-neighbors. + populate_candidates(all_candidates, neighbors, is_deleted); + + // Insert non-deleted candidates into the vector to prepare for + // pruning. + filter_candidates( + valid_candidates, + all_candidates, + accessor(data_, src), + accessor, + general_distance, + is_deleted + ); + + size_t new_candidate_size = + std::min(valid_candidates.size(), params_.max_candidate_pool_size); + valid_candidates.resize(new_candidate_size); + heuristic_prune_neighbors( + prune_strategy(distance_), + params_.prune_to, + params_.alpha, + data_, + accessor, + general_distance, + src, + lib::as_const_span(valid_candidates), + final_candidates + ); + + if (graph_.seq_counters()[src].read_validate(*maybe_seq)) { + // Consistent read — commit the results. + update_buffer.insert(i, final_candidates); + break; + } + svs::detail::pause(); + // Retry: discard stale candidates, recompute on next iteration. + } + } + } + + /// + /// Write pending updates to the graph. + /// + template + void apply_updates( + BulkUpdate& update_buffer, + const GlobalIds& global_ids, + const threads::UnitRange& local_ids + ) { + for (auto i : local_ids) { + if (update_buffer.needs_update(i)) { + graph_.replace_node(global_ids[i], update_buffer.get_update(i)); + } + } + } + + /// + /// Gather a superset of the live nodes with an edge to some deleted node in `D`. + /// + /// The in-neighbors of a deleted `d` are contained in `out(d) union R(d)`, so we hand + /// the driver every non-deleted node from both. + /// + template + std::vector + gather_work_set(const DeletedSet& deleted_ids, const Deleted& is_deleted) const { + auto* reverse_edges = graph_.reverse_edges(); + assert(reverse_edges != nullptr); + + tsl::robin_set work{}; + for (auto d : deleted_ids) { + const auto& neighbors = graph_.get_node(lib::narrow_cast(d)); + for (auto a : neighbors) { + reverse_edges->remove(lib::narrow_cast(d), a); + if (!is_deleted(a)) { + work.insert(a); + } + } + reverse_edges->collect(lib::narrow_cast(d), work, is_deleted); + reverse_edges->reset_node(lib::narrow_cast(d)); + } + return std::vector(work.begin(), work.end()); + } + + /// + /// Run the generate/apply driver over an explicit list of node ids `work_ids`. + /// + template + void run_driver(const std::vector& work_ids, const Deleted& is_deleted) { + const size_t num_work = work_ids.size(); + const size_t update_batch_size = std::min(params_.update_batch_size, num_work); + const size_t thread_batch_size = 500; + + // Size the update buffer to the work-set, not the 200k batch constant. + BulkUpdate update_buffer{update_batch_size, params_.prune_to}; + threads::SequentialTLS> tls{threadpool_.size()}; + + size_t start = 0; + while (start < num_work) { + size_t stop = std::min(num_work, start + update_batch_size); + auto global_ids = + std::span(work_ids).subspan(start, stop - start); + threads::UnitRange local_range{0, global_ids.size()}; + + update_buffer.prepare(); + threads::parallel_for( + threadpool_, + threads::DynamicPartition{local_range, thread_batch_size}, + [&](const auto& local_ids, uint64_t tid) { + auto& thread_local_scratch = tls.at(tid); + generate_updates( + global_ids, + threads::UnitRange(local_ids), + update_buffer, + thread_local_scratch, + is_deleted + ); + } + ); + + threads::parallel_for( + threadpool_, + threads::DynamicPartition{local_range, thread_batch_size}, + [&](const auto& local_ids, uint64_t /*tid*/) { + apply_updates(update_buffer, global_ids, threads::UnitRange(local_ids)); + } + ); + + start = stop; + } + } + + /// + /// Reverse-edge-driven consolidation: process only the in-neighbors of the deleted + /// nodes `deleted_ids`, discovered via the graph's reverse-edge index. + /// + template + void operator()(const DeletedSet& deleted_ids, const Deleted& is_deleted) { + auto work_ids = gather_work_set(deleted_ids, is_deleted); + assert(int(work_ids.size()) > -1); + run_driver(work_ids, is_deleted); + } + + /// + /// Full-scan consolidation: examine every node. Used when the reverse-edge index is + /// not available. + /// + template void operator()(const Deleted& is_deleted) { + // Allocate necessary scratch space. + BulkUpdate update_buffer{params_.update_batch_size, params_.prune_to}; + threads::SequentialTLS> tls{threadpool_.size()}; + + const size_t num_nodes = graph_.n_nodes(); + const size_t update_batch_size = std::min(params_.update_batch_size, num_nodes); + const size_t thread_batch_size = 500; + + size_t start = 0; + while (start < num_nodes) { + size_t stop = std::min(num_nodes, start + update_batch_size); + + // Generate updates. + update_buffer.prepare(); + threads::UnitRange global_ids{start, stop}; + threads::parallel_for( + threadpool_, + threads::DynamicPartition{global_ids.eachindex(), thread_batch_size}, + [&](const auto& local_ids, uint64_t tid) { + auto& thread_local_scratch = tls.at(tid); + generate_updates( + global_ids, + threads::UnitRange(local_ids), + update_buffer, + thread_local_scratch, + is_deleted + ); + } + ); + + // Write back results. + threads::parallel_for( + threadpool_, + threads::DynamicPartition{global_ids.eachindex(), thread_batch_size}, + [&](const auto& local_ids, uint64_t /*tid*/) { + apply_updates(update_buffer, global_ids, threads::UnitRange(local_ids)); + } + ); + + // Prepare for the next iteration. + start = stop; + } + } +}; + +template < + graphs::MemoryGraph Graph, + data::ImmutableMemoryDataset Data, + threads::ThreadPool Pool, + typename Distance, + typename Deleted> +void consolidate( + Graph& graph, + const Data& data, + Pool& threadpool, + size_t prune_to, + size_t max_candidate_pool_size, + float alpha, + const Distance& distance, + Deleted&& is_deleted +) { + ConsolidationParameters params{200'000, prune_to, max_candidate_pool_size, alpha}; + auto consolidator = GraphConsolidator{graph, data, threadpool, distance, params}; + consolidator(is_deleted); +} + +/// +/// Reverse-edge-driven consolidation over an explicit deleted set `deleted_ids`. Requires +/// `graph.reverse_edges() != nullptr`; only the in-neighbors of `deleted_ids` are visited. +/// +template < + graphs::MemoryGraph Graph, + data::ImmutableMemoryDataset Data, + threads::ThreadPool Pool, + typename Distance, + typename DeletedSet, + typename Deleted> +void consolidate( + Graph& graph, + const Data& data, + Pool& threadpool, + size_t prune_to, + size_t max_candidate_pool_size, + float alpha, + const Distance& distance, + const DeletedSet& deleted_ids, + Deleted&& is_deleted +) { + ConsolidationParameters params{200'000, prune_to, max_candidate_pool_size, alpha}; + auto consolidator = GraphConsolidator{graph, data, threadpool, distance, params}; + consolidator(deleted_ids, is_deleted); +} + +} // namespace svs::index::vamana::concurrent diff --git a/include/svs/concurrent/dynamic_index.h b/include/svs/concurrent/dynamic_index.h new file mode 100644 index 000000000..d5c381ce4 --- /dev/null +++ b/include/svs/concurrent/dynamic_index.h @@ -0,0 +1,2080 @@ +/* + * Copyright 2023 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +// stdlib +#include +#include +#include +#include +#include + +// Include the flat index to spin-up exhaustive searches on demand. +#include "svs/index/flat/flat.h" + +// svs +#include "svs/concepts/distance.h" +#include "svs/core/data.h" +#include "svs/core/distance.h" +#include "svs/concurrent/graph.h" +#include "svs/core/loading.h" +#include "svs/core/logging.h" +#include "svs/core/medioid.h" +#include "svs/core/query_result.h" +#include "svs/core/recall.h" +#include "svs/concurrent/translation.h" +#include "svs/concurrent/consolidate.h" +#include "svs/concurrent/dynamic_search_buffer.h" +#include "svs/concurrent/greedy_search.h" +#include "svs/index/vamana/index.h" +#include "svs/concurrent/vamana_build.h" +#include "svs/lib/boundscheck.h" +#include "svs/lib/preprocessor.h" +#include "svs/lib/segmented_vector.h" +#include "svs/lib/threads.h" + +namespace svs::index::vamana::concurrent { + +namespace detail { +// Monotonically pull an atomic down to `value` if it is currently larger. +inline void atomic_min(std::atomic& target, size_t value) { + size_t prev = target.load(std::memory_order_relaxed); + while (value < prev && + !target.compare_exchange_weak( + prev, value, std::memory_order_acq_rel, std::memory_order_relaxed + )) { + // prev reloaded by compare_exchange_weak on failure. + } +} + +// Monotonically push an atomic up to `value` if it is currently smaller. +inline void atomic_max(std::atomic& target, size_t value) { + size_t prev = target.load(std::memory_order_relaxed); + while (value > prev && + !target.compare_exchange_weak( + prev, value, std::memory_order_acq_rel, std::memory_order_relaxed + )) { + // prev reloaded by compare_exchange_weak on failure. + } +} +} // namespace detail + +// Forward declaration +template class BatchIterator; +template +class MultiMutableVamanaIndex; + +///// +///// MutableVamanaIndex +///// + +/// +/// Metadata tracking the state of a particular data index. +/// The following states have the given meaning for their corresponding slot: +/// +/// * Valid: Valid and present in the associated dataset. +/// * Deleted: Exists in the associated dataset, but should be considered as "deleted" +/// and not returned from any search algorithms. +/// * Empty: Non-existent and unreachable from standard entry points. +/// +/// Only used for `MutableVamanaIndex`. +/// +enum class SlotMetadata : uint8_t { + Empty = 0x00, + Valid = 0x01, + Deleted = 0x02, + // Reserved by an in-flight add_points: slot owned by the adder, vector + // copied, adjacency list being built. Invisible to search, consolidate, + // and subsequent add_points until promoted to Valid. + Pending = 0x04, +}; + +template inline constexpr std::string_view name(); +template <> inline constexpr std::string_view name() { + return "Empty"; +} +template <> inline constexpr std::string_view name() { + return "Valid"; +} +template <> inline constexpr std::string_view name() { + return "Deleted"; +} +template <> inline constexpr std::string_view name() { + return "Pending"; +} + +// clang-format off +inline constexpr std::string_view name(SlotMetadata metadata) { + #define SVS_SWITCH_RETURN(x) case x: { return name(); } + switch (metadata) { + SVS_SWITCH_RETURN(SlotMetadata::Empty) + SVS_SWITCH_RETURN(SlotMetadata::Valid) + SVS_SWITCH_RETURN(SlotMetadata::Deleted) + SVS_SWITCH_RETURN(SlotMetadata::Pending) + } + #undef SVS_SWITCH_RETURN + throw ANNEXCEPTION("Unreachable!"); +} +// clang-format on + +class ValidBuilder { + public: + ValidBuilder(const lib::SegmentedVector& status) + : status_{status} {} + + template + constexpr PredicatedSearchNeighbor operator()(I i, float distance) const { + // A neighbor is returnable only if its slot is Valid. Deleted slots + // must be skipped; Pending slots are reserved by an in-flight add and + // their vectors/edges are not yet fully published. Empty slots should + // never be reached via a valid edge, but we defend anyway. + bool invalid = + std::atomic_ref(const_cast(getindex(status_, i))) + .load(std::memory_order_acquire) != SlotMetadata::Valid; + // This neighbor should be skipped if the metadata corresponding to the given index + // marks this slot as deleted. + return PredicatedSearchNeighbor(i, distance, !invalid); + } + + private: + const lib::SegmentedVector& status_; +}; + +template +class MutableVamanaIndex { + friend class MultiMutableVamanaIndex; + + public: + // Traits + static constexpr bool supports_insertions = true; + static constexpr bool supports_deletions = true; + static constexpr bool supports_saving = true; + static constexpr bool needs_id_translation = true; + + // Type Aliases + using Idx = typename Graph::index_type; + using internal_id_type = Idx; + using external_id_type = size_t; + using value_type = typename Data::value_type; + using const_value_type = typename Data::const_value_type; + static constexpr size_t extent = Data::extent; + + using distance_type = Dist; + using search_buffer_type = MutableBuffer>; + + using graph_type = Graph; + using data_type = Data; + using entry_point_type = std::vector; + /// The type of the configurable search parameters. + using search_parameters_type = VamanaSearchParameters; + using inner_scratch_type = + svs::tag_t::result_t; + using scratchspace_type = SearchScratchspace; + + // Members + private: + // Invariants: + // * The ID translator should track only valid IDs. + // TODO: + // * Maybe merge some of the `status` metadata tracker with the IDTranslator to reduce + // memory requirements. There are probably some bits we can reclaim there to + // facilitate that. + + graph_type graph_; + data_type data_; + entry_point_type entry_point_; + // Grow-stable per-slot metadata: search reads status_[i] lock-free via ValidBuilder + // while a concurrent add_points grows it. See svs/lib/segmented_vector.h. + lib::SegmentedVector status_; + // a lower bound such that the append path + // (reuse_empty == false) scans for Empty slots starting here. + std::unique_ptr> first_empty_{ + std::make_unique>(0)}; + // lower bound such that no slot with index + // < *first_reusable_ is Empty. + std::unique_ptr> first_reusable_{ + std::make_unique>(0)}; + IDTranslator translator_; + // Count of Valid slots. Maintained atomically in add_points/delete_entry. + // Wrapped in unique_ptr because std::atomic is not movable. + std::unique_ptr> num_valid_{ + std::make_unique>(0)}; + // Protects translator access: exclusive for writes (add/consolidate/compact), + // shared for reads (delete/search). Wrapped in unique_ptr for movability. + std::unique_ptr translator_mutex_{ + std::make_unique()}; + // Reserves slot ownership against compact(). Search and the other readers + // (get_distance/reconstruct_at/batch-iterator) hold this shared so that + // compact()'s shrink — which frees trailing segments of the grow-stable + // containers — drains in-flight readers before destroying storage. + // Writers (add_points, delete_entries, consolidate) also hold it shared; + // compact() holds it exclusive. + // + // Lock acquisition order across the codebase: + // compact_mutex_ -> slot_alloc_mutex_ (never reversed) + // compact_mutex_ -> translator_mutex_ (never reversed) + // slot_alloc_mutex_ and translator_mutex_ are never held simultaneously + // (add_points takes them sequentially), so they have no relative order. + std::unique_ptr compact_mutex_{ + std::make_unique()}; + // Writer-only mutex serializing slot allocation in add_points + std::unique_ptr slot_alloc_mutex_{std::make_unique()}; + + // Thread local data structures. + distance_type distance_; + threads::ThreadPoolHandle threadpool_; + lib::ReadWriteProtected search_parameters_; + + // Configurations + size_t construction_window_size_; + size_t max_candidates_; + size_t prune_to_; + float alpha_ = 1.2; + bool use_full_search_history_ = true; + + // Construction parameters + VamanaBuildParameters build_parameters_{}; + + // SVS logger for per index logging + svs::logging::logger_ptr logger_; + + // Methods + public: + // Constructors + template + MutableVamanaIndex( + Graph graph, + Data data, + Idx entry_point, + Dist distance_function, + const ExternalIds& external_ids, + ThreadPoolProto threadpool_proto, + // Optional logger parameter + svs::logging::logger_ptr logger = svs::logging::get() + ) + : graph_{std::move(graph)} + , data_{std::move(data)} + , entry_point_{entry_point} + , status_(data_.size(), SlotMetadata::Valid) + , first_empty_{std::make_unique>(data_.size())} + , first_reusable_{std::make_unique>(data_.size())} + , translator_() + , num_valid_{std::make_unique>(data_.size())} + , distance_{std::move(distance_function)} + , threadpool_{threads::as_threadpool(std::move(threadpool_proto))} + , search_parameters_{vamana::construct_default_search_parameters(data_)} + , construction_window_size_{2 * graph.max_degree()} + // Ctor accept logger in parameter + , logger_{std::move(logger)} { + translator_.insert(external_ids, threads::UnitRange(0, external_ids.size())); + graph_.enable_reverse_edges(); + graph_.rebuild_reverse_edges(threadpool_); + } + + /// + /// Build a graph from scratch. + /// + template + MutableVamanaIndex( + const VamanaBuildParameters& parameters, + Data data, + const ExternalIds& external_ids, + Dist distance_function, + ThreadPoolProto threadpool_proto, + svs::logging::logger_ptr logger = svs::logging::get() + ) + : graph_(Graph{data.size(), parameters.graph_max_degree}) + , data_(std::move(data)) + , entry_point_{} + , status_(data_.size(), SlotMetadata::Valid) + , first_empty_{std::make_unique>(data_.size())} + , first_reusable_{std::make_unique>(data_.size())} + , translator_() + , num_valid_{std::make_unique>(data_.size())} + , distance_(std::move(distance_function)) + , threadpool_(threads::as_threadpool(std::move(threadpool_proto))) + , search_parameters_(vamana::construct_default_search_parameters(data_)) + , build_parameters_(parameters) + , logger_{std::move(logger)} { + // Verify and set defaults directly on the input parameters + verify_and_set_default_index_parameters(build_parameters_, distance_function); + + // Set graph again as verify function might change graph_max_degree parameter + graph_ = Graph{data_.size(), build_parameters_.graph_max_degree}; + construction_window_size_ = build_parameters_.window_size; + max_candidates_ = build_parameters_.max_candidate_pool_size; + prune_to_ = build_parameters_.prune_to; + alpha_ = build_parameters_.alpha; + use_full_search_history_ = build_parameters_.use_full_search_history; + + // Setup the initial translation of external to internal ids. + translator_.insert(external_ids, threads::UnitRange(0, external_ids.size())); + + // Compute the entry point. + entry_point_.push_back(extensions::compute_entry_point(data_, threadpool_)); + + // Perform graph construction. + auto sp = get_search_parameters(); + auto prefetch_parameters = + GreedySearchPrefetchParameters{sp.prefetch_lookahead_, sp.prefetch_step_}; + auto builder = VamanaBuilder( + graph_, + data_, + distance_, + build_parameters_, + threadpool_, + prefetch_parameters, + logger_ + ); + builder.construct(1.0f, entry_point_[0], logging::Level::Trace, logger_); + builder.construct( + build_parameters_.alpha, entry_point_[0], logging::Level::Trace, logger_ + ); + + graph_.enable_reverse_edges(); + graph_.rebuild_reverse_edges(threadpool_); + } + + /// @brief Post re-load constructor. + /// + /// Preconditions + /// + /// * data.size() == graph.n_nodes(): The graph and the data have the same number of + /// entries. + /// * The data and graph were saved with no "holes". In otherwords, the index was + /// consolidated and compacted prior to saving. + /// * The span of internal ID's in translator covers exactly ``[0, data.size())``. + template + MutableVamanaIndex( + const VamanaIndexParameters& config, + data_type data, + graph_type graph, + const Dist& distance_function, + IDTranslator translator, + Pool threadpool, + svs::logging::logger_ptr logger = svs::logging::get() + ) + : graph_{std::move(graph)} + , data_{std::move(data)} + , entry_point_{lib::narrow(config.entry_point)} + , status_{data_.size(), SlotMetadata::Valid} + , first_empty_{std::make_unique>(data_.size())} + , first_reusable_{std::make_unique>(data_.size())} + , translator_{std::move(translator)} + , num_valid_{std::make_unique>(data_.size())} + , distance_{distance_function} + , threadpool_{std::move(threadpool)} + , search_parameters_{config.search_parameters} + , construction_window_size_{config.build_parameters.window_size} + , max_candidates_{config.build_parameters.max_candidate_pool_size} + , prune_to_{config.build_parameters.prune_to} + , alpha_{config.build_parameters.alpha} + , use_full_search_history_{config.build_parameters.use_full_search_history} + , logger_{std::move(logger)} { + graph_.enable_reverse_edges(); + graph_.rebuild_reverse_edges(threadpool_); + } + + ///// Scratchspace + scratchspace_type scratchspace(const search_parameters_type& sp) const { + return scratchspace_type{ + search_buffer_type( + sp.buffer_config_, + distance::comparator(distance_), + sp.search_buffer_visited_set_ + ), + extensions::single_search_setup(data_, distance_), + {sp.prefetch_lookahead_, sp.prefetch_step_}}; + } + + scratchspace_type scratchspace() const { return scratchspace(get_search_parameters()); } + + ///// Accessors + /// @brief Getter method for logger + svs::logging::logger_ptr get_logger() const { return logger_; } + + /// @brief Get the alpha value used for pruning while mutating the graph. + float get_alpha() const { return alpha_; } + /// @brief Set the alpha value used for pruning while mutating the graph. + void set_alpha(float alpha) { alpha_ = alpha; } + + /// @brief Get the ``graph_max_degree`` used while mutating the graph. + size_t get_graph_max_degree() const { return graph_.max_degree(); } + + /// @brief Return the bytes allocated by each index component. + /// + /// Reports the capacity-based bytes reserved by the graph adjacency lists, the vector + /// data, and the dynamic metadata (per-slot status, entry-point list, and the + /// external/internal ID translation maps). Capacity-based accounting includes the + /// block over-allocation so integrators can report the true memory footprint. + MemoryBreakdown get_memory_breakdown() const { + MemoryBreakdown usage{}; + usage.graph_bytes = svs::data::detail::dataset_allocated_bytes(graph_.get_data()); + usage.data_bytes = svs::data::detail::dataset_allocated_bytes(data_); + + size_t metadata_bytes = status_.capacity() * sizeof(SlotMetadata); + metadata_bytes += + entry_point_.capacity() * sizeof(typename entry_point_type::value_type); + // The IDTranslator holds two tsl::robin_map instances (external->internal and + // internal->external), neither of which exposes its allocated byte count. We + // approximate the storage as the id pair held in each of the two directions. This + // ignores the maps' load-factor slack and control bytes, so it is an estimate of + // the hash-map overhead that is accurate to within a few percent. + metadata_bytes += 2 * translator_.size() * + (sizeof(IDTranslator::external_id_type) + + sizeof(IDTranslator::internal_id_type)); + usage.metadata_bytes = metadata_bytes; + return usage; + } + + /// @brief Get the max candidate pool size used while mutating the graph. + size_t get_max_candidates() const { return max_candidates_; } + /// @brief Set the max candidate pool size to be used while mutating the graph. + void set_max_candidates(size_t max_candidates) { max_candidates_ = max_candidates; } + /// @brief Get the prune_to value used while mutating the graph. + size_t get_prune_to() const { return prune_to_; } + /// @brief Set the prune_to value to be used while mutating the graph. + void set_prune_to(size_t prune_to) { prune_to_ = prune_to; } + + /// @brief Get the window size used while mutating the graph. + size_t get_construction_window_size() const { return construction_window_size_; } + /// @brief Set the window size to be used while mutating the graph. + void set_construction_window_size(size_t window_size) { + construction_window_size_ = window_size; + } + + /// @brief Return whether the full search history is being used while mutating + /// the graph. + bool get_full_search_history() const { return use_full_search_history_; } + /// @brief Enable using the full search history for candidate generation while + /// mutating the graph. + void set_full_search_history(bool enable) { use_full_search_history_ = enable; } + + ///// Index translation. + // + // The translator is a pair of hash maps mutated by `add_points` (insert), + // `consolidate` (erase), and `compact` (remap) under `translator_mutex_` exclusive. A + // hash-map read concurrent with an insert is a data race in the strict sense and a + // crash in practice: an insert can rehash and free the bucket array a reader is walking. + // + // So every read of `translator_` happens under `translator_mutex_` shared. Each + // translation operation comes in two flavours: + // + // * `foo(...)` -- takes the shared lock itself. Use this by default. + // * `unsafe_foo(...)` -- assumes the caller already holds the shared lock (via + // `lock_for_translation()`). Use this to translate a *batch* + // of ids under one lock acquisition, and inside code that + // already holds the lock. + // + // The split is not merely an optimization: `std::shared_mutex` is not recursive, so a + // self-locking accessor called from a context that already holds the lock is a latent + // deadlock (a writer arriving between the two shared acquisitions blocks the second). + // That is why the batch paths below, and `BatchIterator::next`, use the `unsafe_` + // variants under an explicit lock rather than paying per-id locking. + + /// + /// @brief Get the internal ID mapped to be `e`. + /// + /// @param e The external ID to translate to an internal ID. + /// + /// Requires that mapping for `e` exists. Otherwise, all bets are off. + /// + /// @see has_id, translate_internal_id + /// + Idx translate_external_id(size_t e) const { + std::shared_lock lock{*translator_mutex_}; + return unsafe_translate_external_id(e); + } + + /// @copydoc translate_external_id + /// Requires the caller to hold `lock_for_translation()`. + Idx unsafe_translate_external_id(size_t e) const { return translator_.get_internal(e); } + + /// @brief Translate external ID, returning `default_val` if not mapped. + /// + /// Unlike `translate_external_id`, this does not throw on a missing key. + /// Intended for best-effort readers (e.g. search buffer top-up) that may + /// race with `consolidate()` erasing translator entries. + Idx translate_external_id_or(size_t e, Idx default_val) const { + std::shared_lock lock{*translator_mutex_}; + return unsafe_translate_external_id_or(e, default_val); + } + + /// @copydoc translate_external_id_or + /// Requires the caller to hold `lock_for_translation()`. + Idx unsafe_translate_external_id_or(size_t e, Idx default_val) const { + return translator_.get_internal_or(e, default_val); + } + + /// + /// @brief Check whether the external ID `e` exists in the index. + /// + bool has_id(size_t e) const { + std::shared_lock lock{*translator_mutex_}; + return unsafe_has_id(e); + } + + /// @copydoc has_id + /// Requires the caller to hold `lock_for_translation()`. + bool unsafe_has_id(size_t e) const { + if (!translator_.has_external(e)) { + return false; + } + // Check slot is not Deleted (deferred translator cleanup). + auto internal = translator_.get_internal(e); + return std::atomic_ref(const_cast(status_[internal])) + .load(std::memory_order_acquire) == SlotMetadata::Valid; + } + + /// + /// @brief Get the external ID mapped to be `i`. + /// + /// @param i The internal ID to translate to an external ID. + /// + /// Requires that mapping for `i` exists. Otherwise, all bets are off. + /// + size_t translate_internal_id(Idx i) const { + std::shared_lock lock{*translator_mutex_}; + return unsafe_translate_internal_id(i); + } + + /// @copydoc translate_internal_id + /// Requires the caller to hold `lock_for_translation()`. + size_t unsafe_translate_internal_id(Idx i) const { + // Use get_external_or to handle concurrent consolidate erasing entries. + // If the entry was erased, return the internal ID as-is (stale result). + return translator_.get_external_or(i, static_cast(i)); + } + + /// + /// @brief Call the functor with all external IDs in the index. + /// + /// @param f A functor with an overloaded ``operator()(size_t)`` method. Called on + /// each external ID in the index. + /// + /// The translator lock is held for the whole traversal, so `f` must not call back into + /// a translation method that takes the lock itself, and must not mutate the index. + /// + template void on_ids(F&& f) const { + std::shared_lock lock{*translator_mutex_}; + unsafe_on_ids(SVS_FWD(f)); + } + + /// @copydoc on_ids + /// Requires the caller to hold `lock_for_translation()`. + template void unsafe_on_ids(F&& f) const { + // Skip entries whose slot is Deleted (deferred translator cleanup). + for (auto pair : translator_) { + auto internal = pair.second; + if (std::atomic_ref(const_cast(status_[internal])) + .load(std::memory_order_acquire) == SlotMetadata::Valid) { + f(pair.first); + } + } + } + + /// + /// @brief Return a vector of all valid external IDs present in the index. + /// + std::vector external_ids() const { + std::vector ids{}; + on_ids([&ids](size_t id) { ids.push_back(id); }); + return ids; + } + + /// @brief Return the number of **valid** (non-deleted) entries in the index. + size_t size() const { return num_valid_->load(std::memory_order_acquire); } + + /// + /// @brief Translate in-place a collection of internal IDs to external IDs. + /// + /// @param ids The ``DenseArray`` of internal IDs to modify. + /// + /// Modifies each entry in `ids` in place, assumes that entry is an internal ID and + /// remaps it to its external ID. + /// + /// This is used as a post-processing step following search to return the correct + /// external neighbors to the caller, allowing inner search routines to simply return + /// local IDs. + /// + /// Several implementation notes: + /// (1) This is definitely not safe to call multiple times on the same array for obvious + /// reasons. + /// + /// (2) All entries in `ids` should have valid translations. Otherwise, this function's + /// behavior is undefined. + /// + template + requires(std::tuple_size_v == 2) + void translate_to_external(DenseArray& ids) { + std::shared_lock lock{*translator_mutex_}; + threads::parallel_for( + threadpool_, + threads::StaticPartition{getsize<0>(ids)}, + [&](const auto is, uint64_t /*tid*/) { + for (auto i : is) { + for (size_t j = 0, jmax = getsize<1>(ids); j < jmax; ++j) { + auto internal = lib::narrow_cast(ids.at(i, j)); + ids.at(i, j) = unsafe_translate_internal_id(internal); + } + } + } + ); + } + + /// + /// @brief Get the raw data for external id `e`. + /// + auto get_datum(size_t e) const { + // Lock order: compact_mutex_ then translator_mutex_ (global order). + std::shared_lock compact_lock{*compact_mutex_}; + std::shared_lock lock{*translator_mutex_}; + if (!translator_.has_external(e)) { + throw ANNEXCEPTION("External ID {} not found in index!", e); + } + return data_.get_datum(translator_.get_internal(e)); + } + + /// + /// @brief Return the dimensionality of the stored dataset. + /// + /// TODO (MH): This somewhat limits us to using only R^n type datasets. I'd like to see + /// this generalized somewhat. + /// + size_t dimensions() const { return data_.dimensions(); } + + // Return a `greedy_search` compatible builder for this index. + // This is an internal method, mostly used to help implement the batch iterator. + ValidBuilder internal_search_builder() const { return ValidBuilder{status_}; } + + /// @brief RAII reader lock guarding data_/graph_ against compact()'s shrink + /// (which frees segments). Used by BatchIterator::next() to protect the + /// greedy traversal — mirrors the shared lock taken by search(). Growth by + /// add_points needs no lock (grow-stable SegmentedVector storage). + /// + /// Acquire this only around graph traversal, and release it before + /// acquiring lock_for_translation(): the two must never be held nested in + /// the compact->translator order reversed, which would invert the global + /// lock order (compact -> translator) and deadlock against compact. + [[nodiscard]] std::shared_lock lock_for_search() const { + return std::shared_lock(*compact_mutex_); + } + + /// @brief RAII reader lock guarding translator_ against erase/remap by + /// consolidate/compact. Used by BatchIterator::next() to protect + /// internal->external ID translation. + [[nodiscard]] std::shared_lock lock_for_translation() const { + return std::shared_lock(*translator_mutex_); + } + + auto greedy_search_closure( + GreedySearchPrefetchParameters prefetch_parameters, + const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) + ) const { + return [&, prefetch_parameters]( + const auto& query, auto& accessor, auto& distance, auto& buffer + ) { + // Perform the greedy search using the provided resources. + concurrent::greedy_search( + graph_, + data_, + accessor, + query, + distance, + buffer, + vamana::EntryPointInitializer{lib::as_const_span(entry_point_)}, + internal_search_builder(), + prefetch_parameters, + cancel + ); + // Take a pass over the search buffer to remove any deleted elements that + // might remain. + buffer.cleanup(); + }; + } + + // Single Search + template + void search( + const Query& query, + scratchspace_type& scratch, + const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) + ) const { + // Hold compact_mutex_ shared so compact()'s shrink can't free segments + // mid-traversal. add_points growth is lock-free (grow-stable storage). + std::shared_lock compact_lock{*compact_mutex_}; + extensions::single_search( + data_, + scratch.buffer, + scratch.scratch, + query, + greedy_search_closure(scratch.prefetch_parameters, cancel), + *this + ); + } + + template + void search( + QueryResultView results, + const Queries& queries, + const search_parameters_type& sp, + const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) + ) { + { + // compact_mutex_ shared: blocks compact()'s segment-freeing shrink + // during the traversal. Released before translate_to_external() takes + // translator_mutex_ to keep the compact->translator lock order. + std::shared_lock compact_lock{*compact_mutex_}; + threads::parallel_for( + threadpool_, + threads::StaticPartition{queries.size()}, + [&](const auto is, uint64_t SVS_UNUSED(tid)) { + size_t num_neighbors = results.n_neighbors(); + auto buffer = search_buffer_type{ + sp.buffer_config_, distance::comparator(distance_)}; + + auto prefetch_parameters = GreedySearchPrefetchParameters{ + sp.prefetch_lookahead_, sp.prefetch_step_}; + + // Legalize search buffer for this search. + if (buffer.target_capacity() < num_neighbors) { + buffer.change_maxsize(num_neighbors); + } + auto scratch = + extensions::per_thread_batch_search_setup(data_, distance_); + + extensions::per_thread_batch_search( + data_, + buffer, + scratch, + queries, + results, + threads::UnitRange{is}, + greedy_search_closure(prefetch_parameters, cancel), + *this, + cancel + ); + } + ); + } + + // Check if request to cancel the search + if (cancel()) { + return; + } + + // After the search procedure, the indices in `results` are internal. + // Perform one more pass to convert these to external ids. + translate_to_external(results.indices()); + } + + /// + /// @brief Return a unique instance of the distance function. + /// + Dist distance_function() const { return threads::shallow_copy(distance_); } + + /// + /// Perform an exhaustive search on the current state of the index. + /// Useful to understand how well the graph search is doing after index mutation. + /// + template + void exhaustive_search( + const data::ConstSimpleDataView& queries, + size_t num_neighbors, + QueryResultView result + ) { + auto temp_index = temporary_flat_index( + data_, distance_, threads::ThreadPoolReferenceWrapper(threadpool_) + ); + temp_index.search(queries, num_neighbors, result, [&](size_t i) { + return getindex(status_, i) == SlotMetadata::Valid; + }); + + // After the search procedure, the indices in `results` are internal. + // Perform one more pass to convert these to external ids. + translate_to_external(result.indices()); + } + + /// + /// Descriptive Name + /// + // TODO (Mark): Make descriptions better. + constexpr std::string_view name() const { return "dynamic vamana index"; } + + ///// Mutable Interface + + template + void copy_points(const Points& points, const std::vector& slots) { + assert(points.size() == slots.size()); + threads::parallel_for( + threadpool_, + threads::StaticPartition{slots.size()}, + [&](auto is, auto SVS_UNUSED(tid)) { + for (auto i : is) { + data_.set_datum(slots[i], points.get_datum(i)); + } + } + ); + } + + /// + /// @brief Clear the adjacency lists for the given local ids. + /// + /// This ensures that during the rebuild-phase, we don't get any zombie (previously + /// deleted nodes) occurring in the new adjacency lists. + /// + template void clear_lists(const std::vector& local_ids) { + threads::parallel_for( + threadpool_, + threads::StaticPartition(local_ids), + [&](const auto& thread_local_ids, uint64_t /*tid*/) { + for (auto id : thread_local_ids) { + graph_.clear_node(id); + } + } + ); + } + + /// + /// @brief Add the points with the given external IDs to the dataset. + // + /// When `delete_entries` is called, a soft deletion is performed, marking the entries + /// as `deleted`. When `consolidate` is called, the state of these deleted entries + /// becomes `empty`. When `add_points` is called with the `reuse_empty` flag enabled, + /// the memory is scanned from the beginning to locate and fill these empty entries with + /// new points. + /// + /// @param points Dataset of points to add. + /// @param external_ids The external IDs of the corresponding points. Must be a + /// container implementing forward iteration. + /// @param reuse_empty A flag that determines whether to reuse empty entries that may + /// exist after deletion and consolidation. When enabled, scan from the beginning to + /// find and fill these empty entries when adding new points. + /// + template + std::vector add_points( + const Points& points, const ExternalIds& external_ids, bool reuse_empty = false + ) { + const size_t num_points = points.size(); + const size_t num_ids = external_ids.size(); + if (num_points != num_ids) { + throw ANNEXCEPTION( + "Number of points ({}) not equal to the number of external ids ({})!", + num_points, + num_ids + ); + } + + // Reserve slot ownership against compact(). Held for the entire call, + // including the lock-free Phase 2-4 below; compact() takes this + // exclusive and so will block until every in-flight add finishes. + std::shared_lock compact_lock{*compact_mutex_}; + + // Phase 1: reserve slots (Empty->Pending). Pending means + // "reserved by an in-flight add" + std::vector slots{}; + slots.reserve(num_points); + + // Pre-allocation tail mark, restored on rollback so tail slots we consumed + // are not stranded above first_empty_. + size_t first_empty_before = first_empty_->load(std::memory_order_acquire); + + if (reuse_empty) { + size_t s = first_reusable_->load(std::memory_order_acquire); + for (; s < first_empty_before && slots.size() < num_points; ++s) { + SlotMetadata expected = SlotMetadata::Empty; + if (std::atomic_ref(status_[s]) + .compare_exchange_strong( + expected, + SlotMetadata::Pending, + std::memory_order_acq_rel, + std::memory_order_relaxed + )) { + slots.push_back(s); + } + } + detail::atomic_max(*first_reusable_, s); + } + + if (slots.size() < num_points) { + std::lock_guard lock{*slot_alloc_mutex_}; + + size_t s = first_empty_->load(std::memory_order_relaxed); + size_t smax = status_.size(); + for (; s < smax && slots.size() < num_points; ++s) { + if (status_[s] == SlotMetadata::Empty) { + std::atomic_ref(status_[s]) + .store(SlotMetadata::Pending, std::memory_order_release); + slots.push_back(s); + } + } + + if (slots.size() < num_points) { + size_t needed = num_points - slots.size(); + size_t current_size = data_.size(); + size_t new_size = current_size + needed; + data_.resize(new_size); + graph_.unsafe_resize(new_size); + status_.resize(new_size, SlotMetadata::Empty); + for (size_t i = current_size; i < new_size; ++i) { + std::atomic_ref(status_[i]) + .store(SlotMetadata::Pending, std::memory_order_release); + slots.push_back(i); + } + s = new_size; + } + + detail::atomic_max(*first_empty_, s); + } + assert(slots.size() == num_points); + + // Phase 2: Publish the id translation under translator_mutex_ exclusive + // A Pending slot belongs to an in-flight adder and must + // not be treated as stale — that would clobber the other adder's mapping. + // + // replace_stale_and_insert throws if any external id already + // maps to a live slot. + try { + std::lock_guard lock{*translator_mutex_}; + translator_ + .replace_stale_and_insert(external_ids, slots, [this](auto internal) { + return std::atomic_ref( + const_cast(status_[internal]) + ) + .load(std::memory_order_acquire) == SlotMetadata::Deleted; + }); + } catch (...) { + // Release the reserved slots back to Empty. + for (auto s : slots) { + std::atomic_ref(status_[s]) + .store(SlotMetadata::Empty, std::memory_order_release); + } + detail::atomic_min(*first_empty_, first_empty_before); + if (!slots.empty()) { + detail::atomic_min(*first_reusable_, slots.front()); + } + throw; + } + + // Phase 3: Lock-free data copy and adjacency clearing. + // Slots are Pending: invisible to search (ValidBuilder filters), + // reserved against other writers (Empty-slot scan skips Pending). + copy_points(points, slots); + clear_lists(slots); + + // Phase 4: Graph construction — runs without lock. + // VamanaBuilder::construct() is thread-safe via per-node spinlock+seqlock. + // note: VamanaBuilder constructor asserts graph_.n_nodes() == data_.size(). + // Both are grown together under the lock above, so this is always consistent. + auto parameters = VamanaBuildParameters{ + alpha_, + graph_.max_degree(), + construction_window_size_, + max_candidates_, + prune_to_, + use_full_search_history_}; + + auto sp = get_search_parameters(); + auto prefetch_parameters = + GreedySearchPrefetchParameters{sp.prefetch_lookahead_, sp.prefetch_step_}; + VamanaBuilder builder{ + graph_, + data_, + distance_, + parameters, + threadpool_, + prefetch_parameters, + logger_, + logging::Level::Trace}; + builder.construct(alpha_, entry_point(), slots, logging::Level::Trace, logger_); + + // Mark added entries as valid (unique slots per thread, no lock needed). + for (const auto& i : slots) { + std::atomic_ref(status_[i]) + .store(SlotMetadata::Valid, std::memory_order_release); + } + num_valid_->fetch_add(slots.size(), std::memory_order_acq_rel); + + return slots; + } + + /// + /// Delete all IDs stored in the random-access container `ids`. + /// + /// Pre-conditions: + /// * All indices present in `ids` belong to valid slots. + /// + /// Post-conditions: + /// * Deleted slots will not be returned in future calls `search`. + /// + /// Implementation Nodes: + /// * The deletion that happens is a "soft" deletion. This means that the corresponding + /// entries are still present in both the dataset and the graph, and will be navigated + /// through during searched. + /// + /// However, entries marked as `deleted` will not be returned from searches. + /// + /// * Delete consolidation should happen once a large enough percentage of slots have + /// been soft deleted. + /// + /// Delete consolidation performs the actual removal of deleted entries from the + /// graph. + /// + template size_t delete_entries(const T& ids) { + std::shared_lock compact_lock{*compact_mutex_}; + std::shared_lock lock{*translator_mutex_}; + size_t deleted = 0; + for (auto i : ids) { + if (!translator_.has_external(i)) { + continue; // Already deleted + consolidated, or never existed. + } + auto internal = translator_.get_internal(i); + if (is_deleted(internal)) { + continue; // Already soft-deleted, translator entry not yet consolidated. + } + delete_entry(internal); + ++deleted; + } + // Don't erase translator entries here — concurrent search may still + // need them for translate_to_external(). Cleanup happens in + // consolidate()/compact() when deleted slots become empty. + return deleted; + } + + void delete_entry(size_t i) { + auto& meta = getindex(status_, i); + auto ref = std::atomic_ref(meta); + // CAS Valid → Deleted. If the slot is Pending (concurrent adder still + // in phase 2), wait for the adder to promote it to Valid before we + // can soft-delete; otherwise the delete would be silently lost. Only + // the thread that successfully transitions decrements num_valid_; + // double-deletes silently no-op. + for (;;) { + SlotMetadata expected = SlotMetadata::Valid; + if (ref.compare_exchange_strong( + expected, + SlotMetadata::Deleted, + std::memory_order_acq_rel, + std::memory_order_relaxed + )) { + num_valid_->fetch_sub(1, std::memory_order_acq_rel); + return; + } + if (expected != SlotMetadata::Pending) { + // Already Deleted or Empty — no-op. + return; + } + // Pending: adder's Pending → Valid store is imminent; spin. + svs::detail::pause(); + } + } + + bool is_deleted(size_t i) const { + // True only for slots that have been soft-deleted. Pending (in-flight + // add) and Empty are NOT deleted: consolidate must not prune them out + // of other nodes' adjacency lists, and search already filters + // non-Valid slots via ValidBuilder. + return std::atomic_ref(const_cast(status_[i])) + .load(std::memory_order_acquire) == SlotMetadata::Deleted; + } + + Idx entry_point() const { + assert(entry_point_.size() == 1); + return entry_point_[0]; + } + + /// + /// @brief Return all internal IDs whose slot is Valid (live). + /// + /// Used by compact() to pick the surviving set. Pending slots (in-flight + /// adds) are excluded — compact is only safe to run when the caller has + /// ensured no Pending slots exist (compact holds translator_mutex_ + /// exclusive, which prevents a new add from entering phase 1, but an add + /// that reached phase 2 before compact grabbed the lock may still be + /// publishing status Pending → Valid; the compact caller must quiesce + /// these adds first). + std::vector nonmissing_indices() const { + auto indices = std::vector(); + indices.reserve(size()); + for (size_t i = 0, imax = status_.size(); i < imax; ++i) { + if (std::atomic_ref(const_cast(status_[i])) + .load(std::memory_order_acquire) == SlotMetadata::Valid) { + indices.push_back(i); + } + } + return indices; + } + + /// + /// @brief Compact the data and the graph. + /// + /// @param batch_size Granularity at which points are shuffled. Setting this higher can + /// improve performance but requires more working memory. + /// + void compact(Idx batch_size = 1'000) { + std::lock_guard compact_lock{*compact_mutex_}; + + // Consolidate first, under the same exclusive lock. This folds any + // outstanding soft-deletes into the graph. + consolidate_locked(); + compact_locked(batch_size); + } + + // Body of compact() with no compact_mutex_ locking. The caller MUST hold + // compact_mutex_ exclusive + void compact_locked(Idx batch_size = 1'000) { + // The graph is fully remapped below; every reverse-edge entry would be stale. + // Suppress recording during the remap and rebuild from the final graph at the end. + if (auto* re = graph_.reverse_edges()) { + re->set_recording(false); + } + + // Step 1: Compute a prefix-sum matching each valid internal index to its new + // internal index. + // + // In the returned data structure, an entry `j` at index `i` means that the + // data at index `j` is to be moved to index `i`. + auto new_to_old_id_map = nonmissing_indices(); + + // Construct an associative data structure to facilitate graph adjacency list + // remapping. + auto old_to_new_id_map = tsl::robin_map{}; + for (Idx new_id = 0, imax = new_to_old_id_map.size(); new_id < imax; ++new_id) { + Idx old_id = new_to_old_id_map.at(new_id); + old_to_new_id_map.insert({old_id, new_id}); + } + + // Compact the data. + data_.compact(lib::as_const_span(new_to_old_id_map), threadpool_, batch_size); + + // Manually compact the graph. + auto temp_graph = graphs::SimpleGraph(batch_size, graph_.max_degree()); + + // TODO: Write helper classes to do this partitioning. + Idx start = 0; + Idx max_index = new_to_old_id_map.size(); + while (start < max_index) { + Idx stop = std::min(start + batch_size, max_index); + // Remapping of start index to stop index. + auto batch_to_new_id_map = threads::UnitRange{start, stop}; + auto this_batch = batch_to_new_id_map.eachindex(); + + // Copy the graph into the temporary buffer and remap the IDs. + // Edges to non-Valid (Deleted) slots are dropped — those slots + // do not survive compaction, so the edge would dangle. + threads::parallel_for( + threadpool_, + threads::StaticPartition(this_batch), + [&](const auto& batch_ids, uint64_t /*tid*/) { + std::vector buffer{}; + for (auto batch_id : batch_ids) { + auto new_id = batch_to_new_id_map[batch_id]; + auto old_id = new_to_old_id_map[new_id]; + + const auto& list = graph_.get_node(old_id); + buffer.clear(); + buffer.reserve(list.size()); + + for (auto neighbor_old : list) { + auto it = old_to_new_id_map.find(neighbor_old); + if (it != old_to_new_id_map.end()) { + buffer.push_back(it->second); + } + } + + temp_graph.replace_node(batch_id, buffer); + } + } + ); + + // Copy the entries in the temporary graph to the original graph. + threads::parallel_for( + threadpool_, + threads::StaticPartition(this_batch), + [&](const auto& batch_ids, uint64_t /*tid*/) { + for (auto batch_id : batch_ids) { + auto new_id = batch_to_new_id_map[batch_id]; + graph_.replace_node(new_id, temp_graph.get_node(batch_id)); + } + } + ); + start = stop; + } + + ///// Finishing steps. + { + std::lock_guard lock{*translator_mutex_}; + // Shrink the graph and data. compact_mutex_ is held exclusive for the + // whole compact(), so all in-flight readers have drained — freeing + // trailing segments here cannot dangle a concurrent search. + graph_.unsafe_resize(max_index); + data_.resize(max_index); + first_empty_->store(max_index, std::memory_order_release); + first_reusable_->store(max_index, std::memory_order_release); + + // Compact metadata and ID remapping. + for (size_t new_id = 0; new_id < max_index; ++new_id) { + auto old_id = getindex(new_to_old_id_map, new_id); + if (new_id == old_id) { + continue; + } + + auto status = getindex(status_, old_id); + status_[new_id] = status; + if (status == SlotMetadata::Valid) { + translator_.remap_internal_id(old_id, new_id); + } + } + status_.resize(max_index); + + // Update entry points. If an entry point is no longer present + // (e.g. it was Deleted prior to compact), fall back to internal + // ID 0 — by construction max_index > 0 implies a survivor. + for (auto& ep : entry_point_) { + auto it = old_to_new_id_map.find(ep); + if (it != old_to_new_id_map.end()) { + ep = it->second; + } else { + assert(max_index > 0); + ep = 0; + } + } + } + + // Re-derive the reverse-edge index from the fully remapped graph. + if (auto* re = graph_.reverse_edges()) { + re->set_recording(true); + } + graph_.rebuild_reverse_edges(threadpool_); + } + + ///// Threading Interface + + /// @brief Return the current number of threads used for search. + /// + /// @sa set_num_threads + size_t get_num_threads() const { return threadpool_.size(); } + + void set_threadpool(threads::ThreadPoolHandle threadpool) { + threadpool_ = std::move(threadpool); + } + + /// + /// @brief Destroy the original thread pool and set to the provided one. + /// + /// @param threadpool An acceptable thread pool. + /// + /// @copydoc threadpool_requirements + /// + template + void set_threadpool(Pool threadpool) + requires(!std::is_same_v) + { + set_threadpool(threads::ThreadPoolHandle(std::move(threadpool))); + } + + /// + /// @brief Return the current thread pool handle. + /// + threads::ThreadPoolHandle& get_threadpool_handle() { return threadpool_; } + + ///// Window Interface + VamanaSearchParameters get_search_parameters() const { + return search_parameters_.get(); + } + + void populate_search_parameters(VamanaSearchParameters& parameters) const { + parameters = get_search_parameters(); + } + + void set_search_parameters(const VamanaSearchParameters& parameters) { + search_parameters_.set(parameters); + } + + /// + /// @brief Reset performance parameters to their default values for this index. + /// + /// Parameters affected are only those that modify throughput on a given architecture. + /// Accuracy results should not change as a side-effect of calling this function. + /// + void reset_performance_parameters() { + auto sp = get_search_parameters(); + auto prefetch_parameters = extensions::estimate_prefetch_parameters(data_); + sp.prefetch_lookahead_ = prefetch_parameters.lookahead; + sp.prefetch_step_ = prefetch_parameters.step; + set_search_parameters(sp); + } + + ///// Mutation + void consolidate() { + std::shared_lock compact_lock{*compact_mutex_}; + consolidate_locked(); + } + + // Full consolidation: gather the Deleted set (one cheap byte scan of status_) then + // delegate. Concurrent deletes arriving after the scan are picked up next round. + void consolidate_locked() { + tsl::robin_set deleted{}; + for (size_t i = 0, imax = status_.size(); i < imax; ++i) { + if (std::atomic_ref(status_[i]).load(std::memory_order_acquire) == + SlotMetadata::Deleted) { + deleted.insert(lib::narrow_cast(i)); + } + } + consolidate_locked(deleted); + } + + // Body of consolidate()/consolidate(ids) with no compact_mutex_ locking. The caller + // holds compact_mutex_. `deleted` holds the internal slots to prune out of the graph + // and reclaim; those slots must currently be SlotMetadata::Deleted. Consolidation and + // cleanup both range over exactly `deleted`, so both are O(|deleted|), not O(N). + void consolidate_locked(const tsl::robin_set& deleted) { + auto should_remove = [&](size_t i) { + return deleted.contains(lib::narrow_cast(i)); + }; + + // Entry-point candidacy: a replacement must be live (not soft-deleted) and + // not itself about to be removed. + std::function valid = [&](size_t i) { + return !should_remove(i) && !this->is_deleted(i); + }; + + // Determine if the entry point is being removed. + // If so - we need to pick a new one. + assert(entry_point_.size() == 1); + auto entry_point = entry_point_[0]; + if (should_remove(entry_point)) { + svs::logging::debug(logger_, "Replacing entry point."); + auto new_entry_point = + extensions::compute_entry_point(data_, threadpool_, valid); + svs::logging::debug(logger_, "New point: {}", new_entry_point); + assert(valid(new_entry_point)); + entry_point_[0] = new_entry_point; + } + + // Perform graph consolidation over the in-neighbors of `deleted`, discovered via + // the reverse-edge index. + concurrent::consolidate( + graph_, + data_, + threadpool_, + prune_to_, + max_candidates_, + alpha_, + distance_, + deleted, + should_remove + ); + + // After consolidation - clean up the removed slots under lock. O(|deleted|). + { + std::lock_guard lock{*translator_mutex_}; + // Erase translator entries for removed slots (deferred from delete_entries). + // Skip entries already absent — add_points with replace_stale_and_insert + // may have reassigned the external ID and erased the stale reverse entry. + std::vector deleted_internal_ids; + for (auto i : deleted) { + if (status_[i] == SlotMetadata::Deleted && translator_.has_internal(i)) { + deleted_internal_ids.push_back(i); + } + } + if (!deleted_internal_ids.empty()) { + translator_.delete_internal(deleted_internal_ids, false); + } + // Set removed `Deleted` slots to `Empty` + size_t min_freed = std::numeric_limits::max(); + for (auto i : deleted) { + if (status_[i] == SlotMetadata::Deleted) { + std::atomic_ref(status_[i]) + .store(SlotMetadata::Empty, std::memory_order_release); + min_freed = std::min(min_freed, static_cast(i)); + } + } + if (min_freed != std::numeric_limits::max()) { + detail::atomic_min(*first_reusable_, min_freed); + } + } + } + + /// + /// @brief Consolidate only the soft-deleted entries listed in `ids`. + /// + /// * Listed IDs to be removed must have been previously soft-deleted via + /// `delete_entries`. IDs that are not currently soft-deleted (never existed, + /// already consolidated, or still Valid/Pending) are skipped. + /// + /// * For each consolidated ID, no live node retains an edge to it, its + /// translator entry is erased, and its slot is set to `Empty`. + /// + /// @returns The number of listed IDs that were consolidated. + /// + template size_t consolidate(const T& ids) { + std::shared_lock compact_lock{*compact_mutex_}; + + // Collect the internal slots of the listed, already-soft-deleted IDs. + tsl::robin_set targets{}; + { + std::shared_lock lock{*translator_mutex_}; + for (auto i : ids) { + if (!translator_.has_external(i)) { + continue; // Already consolidated, or never existed. + } + auto internal = translator_.get_internal(i); + if (!is_deleted(internal)) { + continue; // Not soft-deleted — nothing to consolidate. + } + targets.insert(lib::narrow_cast(internal)); + } + } + + if (targets.empty()) { + return 0; + } + consolidate_locked(targets); + return targets.size(); + } + + ///// Saving + + VamanaIndexParameters parameters() const { + return { + entry_point_.front(), + {alpha_, + graph_.max_degree(), + get_construction_window_size(), + get_max_candidates(), + prune_to_, + get_full_search_history()}, + get_search_parameters()}; + } + + static constexpr lib::Version save_version = lib::Version(0, 0, 0); + void save( + const std::filesystem::path& config_directory, + const std::filesystem::path& graph_directory, + const std::filesystem::path& data_directory + ) { + // Post-consolidation, all entries should be "valid". + // Therefore, we don't need to save the slot metadata. + consolidate(); + compact(); + + // Save auxiliary data structures. + lib::save_to_disk( + lib::SaveOverride([&](const lib::SaveContext& ctx) { + // Save the construction parameters. + return lib::SaveTable( + "vamana_dynamic_auxiliary_parameters", + save_version, + { + {"name", lib::save(name())}, + {"parameters", lib::save(parameters(), ctx)}, + {"translation", lib::save(translator_, ctx)}, + } + ); + }), + config_directory + ); + + // Save the dataset. + lib::save_to_disk(data_, data_directory); + // Save the graph. + lib::save_to_disk(graph_, graph_directory); + } + + void save(std::ostream& os) { + // Post-consolidation, all entries should be "valid". + // Therefore, we don't need to save the slot metadata. + consolidate(); + compact(); + + lib::begin_serialization(os); + auto save_table = lib::SaveTable( + "vamana_dynamic_auxiliary_parameters", + save_version, + { + {"name", lib::save(name())}, + {"parameters", lib::save(parameters())}, + {"translation", lib::detail::exit_hook(translator_.metadata())}, + } + ); + lib::save_to_stream(save_table, os); + translator_.save(os); + + // Save the dataset. + lib::save_to_stream(data_, os); + // Save the graph. + lib::save_to_stream(graph_, os); + } + + ///// + ///// Calibrate + ///// + + // Return the maximum degree of the graph. + size_t max_degree() const { return graph_.max_degree(); } + + // Experimental algorithm. + // + // Optimize search_window_size and capacity. + // See calibrate.h for more details. + template < + data::ImmutableMemoryDataset Queries, + data::ImmutableMemoryDataset GroundTruth> + VamanaSearchParameters calibrate( + const Queries& queries, + const GroundTruth& groundtruth, + size_t num_neighbors, + double target_recall, + const CalibrationParameters& calibration_parameters = {} + ) { + // Preallocate the destination for search. + // Further, reference the search lambda in the recall lambda. + auto results = svs::QueryResult{queries.size(), num_neighbors}; + + auto do_search = [&](const search_parameters_type& p) { + this->search(results.view(), queries, p); + }; + + auto compute_recall = [&](const search_parameters_type& p) { + // Calling `do_search` will mutate `results`. + do_search(p); + return svs::k_recall_at_n(results, groundtruth, num_neighbors, num_neighbors); + }; + + auto p = vamana::calibrate( + calibration_parameters, + *this, + num_neighbors, + target_recall, + compute_recall, + do_search, + logger_ + ); + + set_search_parameters(p); + return p; + } + + /// @brief Reconstruct vectors. + /// + /// Reconstruct each vector indexed by an external ID and store the results into + /// ``dst``. + /// + /// Preconditions: + /// - ``ids.size() == svs::getsize<0>(dst)``: Each ID has a corresponding entry in the + /// destination array. + /// - All indices are valid external IDs for this index. + /// - ``svs::getsize<1>(dst) == dimensions()``: The space allocated for each vector in + /// ``dst`` is correct. + /// + /// An exception will be thrown if any of these pre-conditions does not hold. + /// If such an exception is thrown, the argument ``dst`` will be left unmodified. + template + void reconstruct_at(data::SimpleDataView dst, std::span ids) { + const size_t ids_size = ids.size(); + const size_t dst_size = dst.size(); + const size_t dst_dims = dst.dimensions(); + + if (ids_size != dst_size) { + throw ANNEXCEPTION( + "IDs span has size {} but destination has {} vectors!", ids_size, dst_size + ); + } + + if (dst_dims != dimensions()) { + throw ANNEXCEPTION( + "Destination has dimensions {} but index is {}!", dst_dims, dimensions() + ); + } + + // Lock order: compact_mutex_ then translator_mutex_ (global order). + // compact_mutex_ shared guards data_/graph_ against compact()'s shrink; + // translator_mutex_ shared guards the ID translation reads below. + std::shared_lock compact_lock{*compact_mutex_}; + std::shared_lock lock{*translator_mutex_}; + + // Bounds checking. + for (size_t i = 0; i < ids_size; ++i) { + I id = ids[i]; // inbounds by loop bounds. + if (!unsafe_has_id(id)) { + throw ANNEXCEPTION("ID {} with value {} is out of bounds!", i, id); + } + } + + // Prerequisites checked - proceed with the operation. + // TODO: Communicate the requested decompression type to the backend dataset to + // allow more fine-grained specialization? + auto threaded_function = [&](auto is, uint64_t SVS_UNUSED(tid)) { + auto accessor = extensions::reconstruct_accessor(data_); + for (auto i : is) { + auto id = unsafe_translate_external_id(ids[i]); + dst.set_datum(i, accessor(data_, id)); + } + }; + threads::parallel_for( + threadpool_, threads::StaticPartition{ids_size}, threaded_function + ); + } + + /// Invoke the provided callable with constant references to the contained graph, data, + /// and entry points. + /// + /// This function is meant to provide a means for implementing experimental algorithms + /// on the contained data structures. + template void experimental_escape_hatch(F&& f) const { + std::invoke(SVS_FWD(f), graph_, data_, distance_, lib::as_const_span(entry_point_)); + } + + ///// + ///// Debug + ///// + + const Data& view_data() const { return data_; } + const Graph& view_graph() const { return graph_; } + + /// + /// @brief Verify the invariants of this data structure. + /// + /// @param allow_deleted Enable or disable deleted entries. + /// + void debug_check_invariants(bool allow_deleted) const { + debug_check_size(); + debug_check_graph_consistency(allow_deleted); + } + + /// + /// Make sure that the capacities of the main data structures (graph, data, metadata) + /// agree. + /// + void debug_check_size() const { + size_t data_size = data_.size(); + auto throw_size_error = [=](const std::string& name, size_t other_size) { + throw ANNEXCEPTION( + "SIZE INVARIANT: Data size is {} but {} is {}.", data_size, name, other_size + ); + }; + + size_t graph_size = graph_.n_nodes(); + if (data_size != graph_size) { + throw_size_error("graph", graph_size); + } + + size_t status_size = status_.size(); + if (data_size != status_size) { + throw_size_error("metadata", status_size); + } + } + + /// + /// @brief Ensure the graph is in a consistent state. + /// + /// @param allow_deleted Flag to indicate if nodes marked as `Deleted` are okay + /// for consideration. Following a consolidation, this should be ``false``. + /// Otherwise, this should be ``true``. + /// + /// In this case, consistency means the that the adjacency lists for all non-deleted + /// vertices contain only non-deleted vertices. + /// + /// This operation should be run after ``debug_check_size()`` to ensure that + /// the sizes of the underlying data structures are consistent. + /// + void debug_check_graph_consistency(bool allow_deleted = false) const { + auto is_valid = [&, allow_deleted = allow_deleted](size_t i) { + const auto& metadata = status_[i]; + // Use a switch to get a compiler error is we add states to `SlotMetadata`. + switch (metadata) { + case SlotMetadata::Valid: { + return true; + } + case SlotMetadata::Deleted: { + return allow_deleted; + } + case SlotMetadata::Empty: { + return false; + } + case SlotMetadata::Pending: { + // In-flight add: edges may be only partially published. + // Treat as not-yet-live for consistency checking. + return false; + } + } + // Make GCC happy. + return false; + }; + + for (size_t i = 0, imax = graph_.n_nodes(); i < imax; ++i) { + if (!is_valid(i)) { + continue; + } + + size_t count = 0; + for (auto j : graph_.get_node(i)) { + if (!is_valid(j)) { + const auto& metadata = status_[j]; + throw ANNEXCEPTION( + "Node number {} has an invalid ({}) neighbor ({}) at position {}!", + i, + concurrent::name(metadata), + j, + count + ); + } + count++; + } + } + } + + ///// Distance + + /// @brief Compute the distance between an external vector and a vector in the index. + template + double get_distance(const ExternalId& external_id, const Query& query) const { + // Lock order: compact_mutex_ then translator_mutex_ (global order). + // compact_mutex_ shared guards data_ against compact()'s shrink; + // translator_mutex_ shared guards the ID translation read. + std::shared_lock compact_lock{*compact_mutex_}; + std::shared_lock lock{*translator_mutex_}; + + // Check if the external ID exists + if (!unsafe_has_id(external_id)) { + throw ANNEXCEPTION( + "ID {} is out of bounds for index of size {}!", external_id, size() + ); + } + // Verify dimensions match + const size_t query_size = query.size(); + const size_t index_vector_size = dimensions(); + if (query_size != index_vector_size) { + throw ANNEXCEPTION( + "Incompatible dimensions. Query has {} while the index expects {}.", + query_size, + index_vector_size + ); + } + + // Translate external ID to internal ID + auto internal_id = unsafe_translate_external_id(external_id); + + // Call extension for distance computation + return extensions::get_distance_ext(data_, distance_, internal_id, query); + } + + /// + /// @brief Top up an under-filled search buffer with a linear scan. + /// + /// When many vectors have been deleted the graph can become sparsely connected and + /// the greedy search may return fewer than ``target_window()`` valid neighbors. This + /// supplements the buffer by scanning the ID translation. + /// + /// **Precondition:** the caller holds ``compact_mutex_`` shared (``search()`` does). + /// This is why the generic ``extensions::check_and_supplement_search_buffer`` is not + /// used here: it reaches back through ``index.get_distance()``, which takes + /// ``compact_mutex_`` shared a *second* time on the same thread. Recursive shared + /// acquisition of a ``std::shared_mutex`` is undefined behavior and deadlocks against + /// a queued writer. It also iterates a snapshot from ``external_ids()``, which reads + /// ``translator_`` unlocked. Both are avoided by doing the scan here, under + /// ``translator_mutex_`` shared (preserving the global compact -> translator order) + /// and skipping slots that are no longer ``Valid``. + /// + /// Called from the ``extensions::single_search`` override below. + /// + template + void supplement_search_buffer(SearchBuffer& search_buffer, const Query& query) const { + if (search_buffer.valid() >= search_buffer.target_window() || + search_buffer.valid() >= size()) { + return; + } + search_buffer.sort(); + + // translator_mutex_ shared: guards both the traversal of ``translator_`` and the + // internal IDs it yields against erasure by consolidate()/compact(). + std::shared_lock translator_lock{*translator_mutex_}; + auto builder = internal_search_builder(); + for (auto pair : translator_) { + auto internal_id = pair.second; + // Skip slots pending deferred translator cleanup, reserved by an in-flight + // add_points, or otherwise not live. + if (std::atomic_ref(const_cast(status_[internal_id] + )) + .load(std::memory_order_acquire) != SlotMetadata::Valid) { + continue; + } + auto dist = extensions::get_distance_ext(data_, distance_, internal_id, query); + search_buffer.insert(builder(internal_id, dist)); + if (search_buffer.valid() >= search_buffer.target_window()) { + break; + } + } + } + + template + auto make_batch_iterator( + std::span query, + size_t extra_search_buffer_capacity = svs::UNSIGNED_INTEGER_PLACEHOLDER + ) const { + return BatchIterator(*this, query, extra_search_buffer_capacity); + } +}; + +/// +/// @brief ``extensions::single_search`` for the concurrent mutable index. +/// +/// Identical to the generic implementation in ``svs/index/vamana/extensions.h`` except +/// that the search-buffer top-up is delegated to the index (see +/// ``MutableVamanaIndex::supplement_search_buffer``). More specialized than the generic +/// overload on the index parameter, so it wins partial ordering for this index type. +/// +template < + typename Data, + typename SearchBuffer, + typename Distance, + typename Query, + typename Search, + typename Graph, + typename IndexData, + typename Dist> +SVS_FORCE_INLINE void svs_invoke( + svs::tag_t, + const Data& SVS_UNUSED(dataset), + SearchBuffer& search_buffer, + Distance& distance, + const Query& query, + const Search& search, + const MutableVamanaIndex& index, + const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) +) { + // Check if request to cancel the search + if (cancel()) { + return; + } + // Perform graph search. + auto accessor = data::GetDatumAccessor(); + search(query, accessor, distance, search_buffer); + + index.supplement_search_buffer(search_buffer, query); +} + +///// Deduction Guides. +// Guide for building. +template +MutableVamanaIndex(const VamanaBuildParameters&, Data, const ExternalIds&, Dist, size_t) + -> MutableVamanaIndex, Data, Dist>; + +template +MutableVamanaIndex(const VamanaBuildParameters&, Data, const ExternalIds&, Dist, Pool) + -> MutableVamanaIndex, Data, Dist>; + +// Guide with logging +template +MutableVamanaIndex( + const VamanaBuildParameters&, + Data, + const ExternalIds&, + Dist, + Pool, + svs::logging::logger_ptr +) -> MutableVamanaIndex, Data, Dist>; + +template +MutableVamanaIndex( + const VamanaBuildParameters&, + Data, + const ExternalIds&, + Dist, + size_t, + svs::logging::logger_ptr +) -> MutableVamanaIndex, Data, Dist>; +namespace detail { + +struct VamanaStateLoader { + ///// Loading + static bool + check_load_compatibility(std::string_view schema, const lib::Version& version) { + // We provide the option to load from a static index. + return VamanaIndexParameters::check_load_compatibility(schema, version) || + (schema == "vamana_dynamic_auxiliary_parameters" && + version == lib::Version(0, 0, 0)); + } + + // Provide a compatibility path for loading static datasets. + static VamanaStateLoader + load(const lib::LoadTable& table, bool debug_load_from_static, size_t assume_datasize) { + if (debug_load_from_static) { + return VamanaStateLoader{ + lib::load(table), + IDTranslator::Identity(assume_datasize)}; + } + + return VamanaStateLoader{ + SVS_LOAD_MEMBER_AT_(table, parameters), + svs::lib::load_at(table, "translation"), + }; + } + + ///// Members + VamanaIndexParameters parameters_; + IDTranslator translator_; +}; + +} // namespace detail + +// Assembly +template < + typename GraphLoader, + typename DataLoader, + typename Distance, + typename ThreadPoolProto> +auto auto_dynamic_assemble( + const std::filesystem::path& config_path, + GraphLoader&& graph_loader, + DataLoader&& data_loader, + Distance distance, + ThreadPoolProto threadpool_proto, + // Set this to `true` to use the identity map for ID translation. + // This allows us to read files generated by the static index construction routines + // to easily benchmark the static versus dynamic implementation. + // + // This is an internal API and should not be considered officially supported nor stable. + bool debug_load_from_static = false, + svs::logging::logger_ptr logger = svs::logging::get() +) { + // Load the dataset + auto threadpool = threads::as_threadpool(std::move(threadpool_proto)); + auto data = svs::detail::dispatch_load(SVS_FWD(data_loader), threadpool); + + // Load the graph. + auto graph = svs::detail::dispatch_load(SVS_FWD(graph_loader), threadpool); + + // Make sure the data and the graph have the same size. + auto datasize = data.size(); + auto graphsize = graph.n_nodes(); + if (datasize != graphsize) { + throw ANNEXCEPTION( + "Reloaded data has {} nodes while the graph has {} nodes!", datasize, graphsize + ); + } + + // // Unload the ID translator and config parameters. + // auto reloader = lib::LoadOverride{[&](const lib::LoadTable& table) { + // // If loading from the static index, then the table we receive is itself the + // // parameters table. + // // + // // There will also be no index translation, so we use the identity translation + // // since the internal and external IDs for the static index are the samen. + // if (debug_load_from_static) { + // return std::make_tuple( + // // TODO: Provide a better method of loading directly from a load-table + // // and correctly handling contexts. + // lib::load(table.unwrap(), ctx), + // IDTranslator(IDTranslator::Identity(datasize)) + // ); + // } else { + // return std::make_tuple( + // lib::load_at(table, "parameters", ctx), + // lib::load_at(table, "translation", ctx) + // ); + // } + // }}; + // auto [parameters, translator] = lib::load_from_disk(reloader, config_path); + auto [parameters, translator] = lib::load_from_disk( + config_path, debug_load_from_static, datasize + ); + + // Make sure that the translator covers all the IDs in the graph and data. + auto translator_size = translator.size(); + if (translator_size != datasize) { + throw ANNEXCEPTION( + "Translator has {} IDs but should have {}", translator_size, datasize + ); + } + + for (size_t i = 0; i < datasize; ++i) { + if (!translator.has_internal(i)) { + throw ANNEXCEPTION("Translator is missing internal id {}", i); + } + } + + // At this point, we should be completely validated. + // Construct the index! + return MutableVamanaIndex{ + parameters, + std::move(data), + std::move(graph), + std::move(distance), + std::move(translator), + std::move(threadpool), + std::move(logger)}; +} + +template < + typename LazyGraphLoader, + typename LazyDataLoader, + typename Distance, + typename ThreadPoolProto> +auto auto_dynamic_assemble( + std::istream& is, + LazyGraphLoader graph_loader, + LazyDataLoader data_loader, + Distance distance, + ThreadPoolProto threadpool_proto, + bool SVS_UNUSED(debug_load_from_static) = false, + svs::logging::logger_ptr logger = svs::logging::get() +) { + // Read the combined TOML (parameters + translation) + // and the translator binary data. + auto table = lib::detail::read_metadata(is); + + auto parameters = lib::load( + table.template cast().at("parameters").template cast() + ); + + auto translation = + table.template cast().at("translation").template cast(); + + auto translator = IDTranslator::load(translation, is); + + auto data = data_loader(); + auto graph = graph_loader(); + + auto datasize = data.size(); + auto graphsize = graph.n_nodes(); + if (datasize != graphsize) { + throw ANNEXCEPTION( + "Reloaded data has {} nodes while the graph has {} nodes!", datasize, graphsize + ); + } + + auto translator_size = translator.size(); + if (translator_size != datasize) { + throw ANNEXCEPTION( + "Translator has {} IDs but should have {}", translator_size, datasize + ); + } + + auto threadpool = threads::as_threadpool(std::move(threadpool_proto)); + return MutableVamanaIndex{ + parameters, + std::move(data), + std::move(graph), + std::move(distance), + std::move(translator), + std::move(threadpool), + std::move(logger)}; +} + +} // namespace svs::index::vamana::concurrent diff --git a/include/svs/concurrent/dynamic_search_buffer.h b/include/svs/concurrent/dynamic_search_buffer.h new file mode 100644 index 000000000..d6811b433 --- /dev/null +++ b/include/svs/concurrent/dynamic_search_buffer.h @@ -0,0 +1,625 @@ +/* + * Copyright 2023 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "svs/index/vamana/search_buffer.h" +#include "svs/lib/boundscheck.h" +#include "svs/lib/datatype.h" +#include "svs/lib/neighbor.h" +#include "svs/lib/threads/threadlocal.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace svs::index::vamana::concurrent { + +/// +/// A search buffer that allows entries to be predicated out. +/// The search buffer will still navigate through those entries, but it won't include +/// them in its final result. +/// +template > class MutableBuffer { + public: + // Type Aliases + using index_type = Idx; + using value_type = PredicatedSearchNeighbor; + using reference = value_type&; + using const_reference = const value_type&; + using vector_type = std::vector>; + using iterator = typename vector_type::iterator; + using const_iterator = typename vector_type::const_iterator; + using compare_type = Cmp; + using filter_type = VisitedFilter; + + private: + ///// Invariants: + // + // 1. 1 <= target_valid_ <= valid_capacity_ + // + // The second inequality is enforced by the `BufferConfig` class. + // TODO: Enforce the first in `BufferConfig` as well. + // + // 2. best_unvisited_ <= size(). + // 3. roi_end_ <= size(). + // 4. valid_ <= valid_capacity_; + // + // 5. WHEN (5A): `valid_ < target_valid_` (haven't yet found enough valid neighbors) + // THEN: `roi_end_ == size()`: Region of interest points to the end of the buffer. + // ELSE (5B): `valid_ >= target_valid_` + // THEN: The number of valid between the start of the buffer and `roi_end_ - 1` is + // exactly `valid_`. + // + // 6. WHEN: `valid_ == valid_capacity_` + // THEN: `back()` must be valid. + // + // 7. The number of valid elements in `[begin(), end())` is equal to `valid_`. + + [[no_unique_address]] Cmp compare_{}; + // Equivalent of the `search_window_size_` in the traditional search buffer. + size_t target_valid_ = 0; + // Number of valid elements can are configured to contain. + // Equivalent to the `search_buffer_capacity_` in the traditional search buffer. + size_t valid_capacity_ = 0; + // Index of the best unvisited candidate. + size_t best_unvisited_ = 0; + // One past the "target_valid_"th entry. + size_t roi_end_ = 0; + // number of valid neighbors. + size_t valid_ = 0; + // A buffer of candidates. + // Unlike the static buffer, this container *does* dynamically change size and does not + // reserve one-past-the-end for copying neighbors. + vector_type candidates_{}; + // An optional visited filter. + std::optional visited_{std::nullopt}; + + public: + MutableBuffer() = default; + + /// Construct a new buffer with the given buffer configuration. + explicit MutableBuffer( + SearchBufferConfig config, Cmp compare = Cmp{}, bool enable_visited = false + ) + : compare_{std::move(compare)} + , target_valid_{config.get_search_window_size()} + , valid_capacity_{config.get_total_capacity()} + , candidates_{valid_capacity_} { + candidates_.clear(); + if (enable_visited) { + enable_visited_set(); + } + } + + /// Construct a new buffer with the given size and capacity. + explicit MutableBuffer(size_t size, Cmp compare = Cmp{}, bool enable_visited = false) + : MutableBuffer{SearchBufferConfig{size}, std::move(compare), enable_visited} {} + + /// Copy the portions of the MutableBuffer that matter for the purposes of scratch + /// space. + /// + /// Preserves the sizes of various containers but not necessarily the values. + MutableBuffer shallow_copy() const { + // We don't care about the contents of the buffer - just its size. + // Therefore, we can construct a new buffer from scratch. + return MutableBuffer{config(), compare_, visited_set_enabled()}; + } + + // TODO: Allow this construction to be noexcept. + SearchBufferConfig config() const { + return SearchBufferConfig{target_valid_, valid_capacity_}; + } + + /// + /// @brief Change the target number of elements to return after search. + /// + /// @param config The new configuration for the buffer. + /// + /// Post conditions: + /// - The target valid capacity for the search buffer will be set to the new capacity. + /// - The actual size (number of contained elements both valid and invalid) will be the + /// minimum of the current size and the new capacity. + /// + void change_maxsize(SearchBufferConfig config) { + // Use temporary variables to ensure the given configuration is valid before + // committing. + size_t target_valid_temp = config.get_search_window_size(); + size_t valid_capacity_temp = config.get_total_capacity(); + + // If the new capacity is lower then the current size, shrink the buffer to fit. + if (valid_capacity_temp < candidates_.size()) { + candidates_.resize(valid_capacity_temp); + } + + // Commit the new sizes. Integer assignment is `noexcept`, so we don't need to worry + // about an exception breaking class invariants. + target_valid_ = target_valid_temp; + valid_capacity_ = valid_capacity_temp; + } + + // Change the maximum number of elements that can be in the search buffer. + void change_maxsize(size_t new_size) { change_maxsize(SearchBufferConfig{new_size}); } + + /// @brief Prepare the buffer for a new search operation. + void clear() { + candidates_.clear(); + best_unvisited_ = 0; + roi_end_ = 0; + valid_ = 0; + if (visited_set_enabled()) { + visited_->reset(); + } + } + + void soft_clear() { + bool use_visited_set = visited_set_enabled(); + if (use_visited_set) { + visited_->reset(); + } + + for (auto& neighbor : candidates_) { + neighbor.clear_visited(); + if (use_visited_set) { + visited_->emplace(neighbor.id()); + } + } + + best_unvisited_ = 0; + } + + size_t capacity() const { return candidates_.capacity(); } + + /// @brief Return the number of valid elements currently contained in the buffer. + size_t valid() const { return valid_; } + + /// @brief Return the target valid capacity as equivalent to buffer capacity + size_t target_capacity() const { return valid_capacity_; } + + /// @brief Return the target valid candidates as equivalent to the search window + size_t target_window() const { return target_valid_; } + + /// @brief Return whether or not the buffer contains its target number of candidates. + bool full() const { return valid() == target_capacity(); } + + /// @brief Return the candidate at index `i`. + /// + /// Element will only be valid if: + /// (A) cleanup() has been invoked. + /// (B) 0 <= i < valid(); + reference operator[](size_t i) { return candidates_[i]; } + + /// @brief Return the candidate at index `i`. + /// + /// Element will only be valid if: + /// (A) cleanup() has been invoked. + /// (B) 0 <= i < valid(); + const_reference operator[](size_t i) const { return candidates_[i]; } + + /// @brief Return the last candidate, whether or not it is valid. + /// + /// It is undefined behavior to call this on an empty buffer. + reference back() { return candidates_.back(); } + + /// @brief Return the last candidate, whether or not it is valid. + /// + /// It is undefined behavior to call this on an empty buffer. + const_reference back() const { return candidates_.back(); } + + /// @brief Return the index of the best unvisited candidate. + size_t best_unvisited() const { return best_unvisited_; } + + /// @brief Return `true` if the search buffer has reached its terminating condition. + /// + /// N.B.: If `done()` evaluates to `true`, do not try to extract further candidates + /// from it using `next()`. + bool done() const { + // Until we've reached the target number of valid elements, we have to ignore + // the state of `roi_end_`. + return best_unvisited_ == ((valid() < target_valid_) ? size() : roi_end_); + } + + /// @brief Return the best unvisited neighbor in the buffer. + /// The returned result will be convertible to `Neighbor`. + /// + /// Pre-conditions: + /// * `search_buffer.done()` must evaluate to `false`, otherwise an out of bounds + /// access will occur. + /// + /// Post-conditions: + /// * The returned neighbor will be marked as visited. + const_reference next() { + // Get the best unvisited node + reference node = getindex(candidates_, best_unvisited_); + node.set_visited(); + + // Increment `best_unvisited_` until it's equal to the size OR until we encounter + // an unvisited node. + while (++best_unvisited_ != roi_end_ && candidates_[best_unvisited_].visited()) {} + return node; + } + + /// Place the neighbor at the end of the search buffer if `full() != true`. + /// Otherwise, do nothing. + /// + /// NOTE: `push_back` does not necessarily maintain the required invariants by this + /// class. + /// + /// These invariants are restored upon calling `sort()`. + /// + /// Therefore, sequences of `push_back` must always be followed by a call to `sort`. + void push_back(value_type neighbor) { + // Follow the contract of not appending more valid elements than the capacity + // allows. + bool valid = neighbor.valid(); + if (full() && valid) { + return; + } + candidates_.push_back(neighbor); + if (valid) { + ++valid_; + } + roi_end_ = candidates_.size(); + } + + // Define iterators. + constexpr const_iterator begin() const noexcept { return candidates_.begin(); } + constexpr const_iterator end() const noexcept { return begin() + size(); } + constexpr iterator begin() noexcept { return candidates_.begin(); } + constexpr iterator end() noexcept { return begin() + size(); } + + /// + /// Insert `neighbor` into the candidates buffer at position `pos`. + /// Afterwards, try to shrink the candidates buffer + /// + /// May invalidate any iterators for `candidates_`. + /// + void unsafe_insert(value_type neighbor, iterator pos, size_t offset) { + // Insert the neighbor into the buffer. + // The rest of this function fixes the invariants that got broken as a result + // of this operation. + candidates_.insert(pos, neighbor); + + // Enough candidates for ROI to kick in (5B instead of 5A). + bool at_threshold = valid() == target_valid_; + bool this_valid = neighbor.valid(); + // Is this candidate begin inserted before the end of the ROI. + bool inserted_below = offset < roi_end_; + if (valid() < target_valid_) { + // Maintain invariant 5A + roi_end_ = size(); + } else if (at_threshold) { + // Several options: + // (1) This is the entry that causes us to cross the valid threshold. If so, + // then we need to move `roi_end_` to the last valid entry; + // (2) We've previously crossed the threshold and this entry is skipped. + // If this is the case, then we must bump `roi_end_` if this neighbor + // is inserted below. + if (this_valid) { + assert(roi_end_ == size() - 1); + // Case 1: This neighbor caused us to cross the threshold. + // + // In this case, the `roi_end_` USED to point to one past the end. + // After the insertion, it points to the end exactly. + // + // To maintain 5B, move `roi_end_` backwards until a valid neighbor is + // reached, then increment by 1 to point to one-past the end. + while (!svs::getindex(candidates_, roi_end_).valid()) { + --roi_end_; + } + ++roi_end_; + } else if (inserted_below) { + // Case 2: We've previously crossed the threshold and this neighbor is + // skipped. + // + // In this case, `roi_end_` used to point to one-past the last valid element + // (according to 5B). + // + // Since we inserted below (moving everything above by 1), we can restore + // 5B by incrementing `roi_end_`. + ++roi_end_; + } + } else { + // At this point - we are in split-buffer territory. + // We can assume that before the insertion, invariante 5B held. + // + // Therefore, we only need to fix 5B if we inserted below the previous + // `roi_end_`. + if (inserted_below) { + // Assert that 5B held prior to this invocation of `unsafe_insert`. + assert(candidates_.at(roi_end_).valid()); + + // If this neighbor is not valid, then the number of valid neighbors in + // `[0, roi_end_]` (with right inclusion) has not changed. + // We can restore 5B by incrementing `roi_end_`. + if (!this_valid) { + ++roi_end_; + } else { + // In this case, the number of neighbors in `[0, roi_end_`] (with + // right inclusion) has been incremented by 1. + // + // We restore 5B by walking back until the next previous valid neighbor + // is found, then step forward by 1 to point to one-past the end. + while (!svs::getindex(candidates_, --roi_end_).valid()) {} + ++roi_end_; + } + } + } + + // Restore Invariant 6. + bool no_split_buffer = target_valid_ == valid_capacity_; + if (no_split_buffer) { + // When using 5A: the ROI already points to the end - so resizing does nothing. + // Furthermore, when using 5A, `valid() > target_valid_` is always false, so + // the branch is never taken. + // + // When using 5B: `roi_end_` points to one-past the `target_valid_`th entry + // Shrinking `candidates_` to `roi_end_` is valid + // + // Prior to this call to, invariants 4 and 7 must have held. + // Therefore, the side-effect of adding this neighbor increased the number of + // valid elements by at most 1 (i.e., `slack() <= 1`). + // + // From this, we deduce that the extra valid neighbor used to be located at or + // above `roi_end_` and was therefore implicitly dropped when shrinking + // `candidates_`. + candidates_.resize(roi_end_); + if (slack() > 0) { + assert(slack() == 1); + valid_ -= 1; + } + } else if (slack() == 0) { + // From invariant 7, we know `candidates_` contains `valid_` number of elements. + // We can restore invariant 6 simply by finding the last valid neighbor and + // resizing appropriately. + candidates_.resize(walk_back(candidates_.size() - 1) + 1); + } else if (slack() >= 0) { + // We know that invariant 6 must have held prior to this invocation of + // `unsafe_insert`. + // + // Therefore the last neighbor in `candidates_` is valid and must be dropped. + // We can restore invariant 6 by starting at this neighbor and walking back + // until the previous valid neighbor is found. + assert(slack() == 1); + assert(back().valid()); + candidates_.resize(walk_back(candidates_.size() - 2) + 1); + // ^ + // | + // Not a typo. It should be "two". + + // We dropped a valid element off the back when shrinking. + valid_ -= 1; + } + } + + bool can_skip(float distance) const { + // If not full - we cannot skip appending this item. + // If we are full, then rely on invariant 6 to compare with the last valid element. + return full() && compare_(back().distance(), distance); + } + + /// Insert the given neighbor into the search buffer. + size_t insert(value_type neighbor) { + if (can_skip(neighbor.distance())) { + return size(); + } + return insert_inner(neighbor); + } + + size_t insert_inner(value_type neighbor) { + const auto start = begin(); + + // Binary search to the first location where `distance` is less than the stored + // neighbor. + auto pos = std::lower_bound( + start, + end(), + neighbor.distance(), + [&](const value_type& other, const float& d) { + return !compare_(d, other.distance()); + } + ); + + // Because repeat ids can exist, we have to search until we're sure no repeat will + // be found. + // + // To do that, we start one before the insertion position. + // Because each instance of repeat ids should have the same distance, we only need + // to look at `ids` until the buffer elements have a distance less than the + // distance of the current node. + if (pos != start) { + auto back = pos; + do { + --back; + const auto& candidate = *back; + if (compare_(candidate.distance(), neighbor.distance())) { + break; + } + if (candidate.id_ == neighbor.id()) { + return size() + 1; + } + } while (back != start); + } + + // Note: Iterators *may* be invalidated as a result of `unsafe_insert`. + // Hoist out the insertion position as an index before inserting. + size_t i = pos - start; + valid_ += static_cast(neighbor.valid()); + unsafe_insert(neighbor, pos, i); + best_unvisited_ = std::min(best_unvisited_, i); + return i; + } + + // Sort all stored elements in the buffer. + void sort() { + assert(valid() <= valid_capacity_); + + // Put the elements in ascending order. + std::sort(begin(), end(), compare_); + + // Now - fix our invariants. + if (valid_ < target_valid_) { + // Maintain 5A + // Invariant 6 has not activated. + roi_end_ = size(); + + // Maintain best_unvisited_ invariant. + // Recompute it from scratch, mirroring the 5B branch below. + for (best_unvisited_ = 0; best_unvisited_ < roi_end_; ++best_unvisited_) { + if (!candidates_[best_unvisited_].visited()) { + break; + } + } + return; + } + + // Maintain 5B + size_t valid_count = 0; + for (size_t i = 0; i < size(); ++i) { + if (candidates_[i].valid()) { + ++valid_count; + if (valid_count == target_valid_) { + // One past the last valid element. + roi_end_ = i + 1; + break; + } + } + } + + // TODO: Switch over to using iterators for the return values to avoid this. + // Maintain best_unvisited_ invariant + // cleanup() and std::sort() may invalidate best_unvisited_. + // so, have to lookup the next unvisited from scratch. + for (best_unvisited_ = 0; best_unvisited_ < roi_end_; ++best_unvisited_) { + if (!candidates_[best_unvisited_].visited()) { + break; + } + } + + // Check if invariant 6 is active. + // If so, drop invalid elements off the end until a valid element is found. + if (slack() == 0) { + assert(!candidates_.empty()); + while (!back().valid()) { + candidates_.pop_back(); + assert(!candidates_.empty()); + } + } + } + + // TODO: Switch over to using iterators for the return values to avoid this. + void cleanup() { + auto new_end = + std::remove_if(begin(), end(), [](const auto& x) { return !x.valid(); }); + candidates_.resize(new_end - begin()); + } + + /// @brief Return the size of the underlying vector. + /// + /// N.B.: THIS FUNCTION IS EASY TO CALL INCORRECTLY. + /// + /// At any given point, the size of the underlying vector and the number of valid + /// candidates in the buffer can be very different. + /// + /// These two are only the same after invoking `clean()`. + /// That is the only context in which non-internal interaction with the search buffer + /// should operate. + size_t size() const { return candidates_.size(); } + + ///// Visited API + bool visited_set_enabled() const { return visited_.has_value(); } + void enable_visited_set() { + if (!visited_set_enabled()) { + visited_.emplace(); + } + } + + /// @brief Enable or disable the visited set based on the argument. + void configure_visited_set(bool enable) { + if (enable) { + enable_visited_set(); + } else { + disable_visited_set(); + } + } + + void disable_visited_set() { + if (visited_set_enabled()) { + visited_.reset(); + } + } + + bool is_visited(Idx i) const { return visited_set_enabled() && unsafe_is_visited(i); } + void prefetch_visited(Idx i) const { + if (visited_set_enabled()) { + unsafe_prefetch_visited(i); + } + } + bool emplace_visited(Idx i) { + return visited_set_enabled() && unsafe_emplace_visited(i); + } + + // Unsafe API + bool unsafe_is_visited(Idx i) const { + assert(visited_); + return visited_->contains(i); + } + void unsafe_prefetch_visited(Idx i) const { + assert(visited_); + return visited_->prefetch(i); + } + bool unsafe_emplace_visited(Idx i) { + assert(visited_); + return visited_->emplace(i); + } + + private: + /// Return how many more valid candidates exist than required. + /// If the number of valid candidates is *less* than the target, a negative number + /// is returned. + int64_t slack() const { + return lib::narrow_cast(valid()) - + lib::narrow_cast(target_capacity()); + } + + /// Return the index of the first preceding valid candidate beginning at the provided + /// index. + /// + /// Requires: + /// * `i` is in the range `[0, size())`. + /// * There exists at least one valid candidate in `[0, i)`. + size_t walk_back(size_t i) const { + while (!candidates_[i].valid()) { + --i; + } + return i; + } +}; + +template +std::ostream& operator<<(std::ostream& io, const MutableBuffer& buffer) { + return io << "MutableBuffer<" << datatype_v << ">(" + << "target_valid = " << buffer.target_capacity() + << ", best_unvisited = " << buffer.best_unvisited() + << ", valid = " << buffer.valid() << ", size = " << buffer.size() << ")"; +} + +} // namespace svs::index::vamana::concurrent diff --git a/include/svs/concurrent/graph.h b/include/svs/concurrent/graph.h new file mode 100644 index 000000000..fb2e2594b --- /dev/null +++ b/include/svs/concurrent/graph.h @@ -0,0 +1,677 @@ +/* + * Copyright 2023 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "svs/concurrent/blocked_data.h" +#include "svs/concurrent/graph_concepts.h" +#include "svs/concurrent/reverse_edges.h" +#include "svs/core/data/simple.h" +#include "svs/lib/algorithms.h" +#include "svs/lib/boundscheck.h" +#include "svs/lib/concurrency/atomic_span.h" +#include "svs/lib/concurrency/seqlock.h" +#include "svs/concurrent/spinlock.h" +#include "svs/lib/saveload.h" +#include "svs/lib/segmented_vector.h" +#include "svs/lib/threads.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace svs::index::vamana::concurrent::graphs { + +// +// We rely on an implicit layout for the graphs where length is stored inline with the +// adjacency list like: +// +// Node 0 : Len N0 N1 N2 .... Nm +// Node 1 : Len N0 N1 N2 .... Nm +// Node 2 : Len N0 N1 N2 .... Nm +// ... +// Node K : Len N0 N1 N2 .... Nm +// +// Note that the the length variable `Len` is the same type as the adjacency list entries. +// +// In general, C++'s support for type-punning, even for trivially constructible and +// copyable types leaves quite a bit of head-scratching. +// +// The utilities developed here are meant to help with dealing with the implicit layout +// described above. +// +// Base class for packed graphs. +// Should not be used directly. Rather, one of it's derived classes should be used instead. +// +template class SimpleGraphBase { + public: + using data_type = Data; + + /// The integer representation used to represent vertices in this graph. + using index_type = Idx; + using value_type = std::span; + using const_value_type = AtomicSpan; + + /// Type used to represent mutable adjacency lists externally. + using reference = std::span; + /// Type used to represent constant adjacency lists externally. + using const_reference = AtomicSpan; + + /// + /// @brief Construct an empty graph of the desired size. + /// + /// @param num_nodes The number of nodes in the graph. + /// @param max_degree The maximum degree of the graph. + /// + /// Implementation notes: Requires that the memory backing the dataset for this graph + /// is default constructible. + /// + explicit SimpleGraphBase(size_t num_nodes, size_t max_degree) + : data_{num_nodes, max_degree + 1} + , max_degree_{lib::narrow(max_degree)} + , seq_counters_(num_nodes) + , node_locks_(num_nodes) { + reset(); + } + + // TODO: Constrain template approparitely. + template + explicit SimpleGraphBase( + size_t num_nodes, size_t max_degree, const Allocator& allocator + ) + : data_{num_nodes, max_degree + 1, allocator} + , max_degree_{lib::narrow(max_degree)} + , seq_counters_(num_nodes) + , node_locks_(num_nodes) { + reset(); + } + + explicit SimpleGraphBase(data_type data) + : data_{std::move(data)} + , max_degree_{lib::narrow(data_.dimensions() - 1)} + , seq_counters_(data_.size()) + , node_locks_(data_.size()) {} + + std::span raw_row(Idx i) const { return data_.get_datum(i); } + + /// + /// @brief Return the outward adjacency list for vertex ``i``. + /// + /// @param i The vertex to get the ID for. + /// + const_reference get_node(Idx i) const { + // Get the raw data. + std::span raw_data = data_.get_datum(i); + Idx num_neighbors = relaxed_load(raw_data.front()); + // Clamp to max_degree to safely handle torn reads of the length field. + num_neighbors = std::min(num_neighbors, max_degree_); + + // Maybe prefetch the rest of the adjacency list. + size_t bytes = (1 + num_neighbors) * sizeof(Idx); + if (bytes > lib::CACHELINE_BYTES) { + lib::prefetch(std::as_bytes(raw_data).subspan(lib::CACHELINE_BYTES)); + } + return AtomicSpan(raw_data.data() + 1, num_neighbors); + } + + /// + /// @brief Return whether or not the adjacency list has an edge from ``src`` to ``dst``. + /// + /// Complexity: Linear in the maximum degree. + /// + bool has_edge(Idx src, Idx dst) const { + for (;;) { + auto maybe_seq = seq_counters_[src].read_begin(); + if (!maybe_seq) { + detail::pause(); + continue; + } + const auto& list = get_node(src); + bool found = (std::find(list.begin(), list.end(), dst) != list.end()); + if (seq_counters_[src].read_validate(*maybe_seq)) { + return found; + } + detail::pause(); + } + } + + /// + /// @brief Return the current out degree of vertex ``i``. + /// + size_t get_node_degree(Idx i) const { + return relaxed_load(data_.get_datum(i).front()); + } + + /// + /// @brief Prefetch the adjacency list for node ``i`` into the L1 cache. + /// + void prefetch_node(Idx i) const { data_.prefetch(i); } + + /// + /// @brief Remove all outgoing neighbors from node ``i``. + /// + /// *Note*: As an implementation detail, this method doesn't mutate the actual adjacency + /// list. Instead, it simply sets the number of neighbors to zero. + /// + /// The complexity of this operation is `O(1)`. + /// + void clear_node(Idx i) { + std::lock_guard lock{node_locks_[i]}; + auto seq = seq_counters_[i].begin_write(); + relaxed_store(data_.get_datum(i).front(), 0); + seq_counters_[i].end_write(seq); + } + + /// + /// @brief Remove all edges from the graph. + /// + void reset() { + for (size_t i = 0; i < n_nodes(); ++i) { + clear_node(i); + } + } + + /// + /// @brief Replace the adjacency list for vertex ``i``. + /// + /// @param i The vertex whose adjacency list is being modified. + /// @param new_neighbors The new adjacency list for vertex ``i``. + /// + /// Takes at most ``max_degree()`` elements from ``new_neighbors``. May silently drop + /// any excess neighbors. + /// + /// **Preconditions:** + /// + /// * All elements of ``new_neighbors`` must be between 0 and ``n_nodes()`` + /// * All elements of ``new_neighbors`` must be unique. + /// + void replace_node(Idx i, const std::vector& new_neighbors) { + replace_node(i, std::span{new_neighbors.data(), new_neighbors.size()}); + } + + /// @copydoc replace_node(Idx,const std::vector&) + void replace_node(Idx i, std::span new_neighbors) { + replace_node_impl(i, new_neighbors); + } + + /// @copydoc replace_node(Idx,const std::vector&) + void replace_node(Idx i, AtomicSpan new_neighbors) { + replace_node_impl(i, new_neighbors); + } + + /// + /// @brief Add an edge from vertex ``src`` to vertex ``dst``. + /// + /// @param src The source vertex. + /// @param dst The destination vertex. + /// + /// @returns The number of out neighbors of ``src`` after ``dst`` is inserted. + /// + /// The adjacency list of ``src`` will be left unchanged if: + /// * ``src == dst`` (no self assignment) + /// * ``get_node_degree(src) == max_degree()`` (adjacency list is already full) + /// * ``dst`` is already an out-neighbor of ``src``. + /// + AddEdgeResult add_edge(Idx src, Idx dst) { + // Don't assign a node as its own neighbor. + if (src == dst) { + return AddEdgeResult::AlreadyExists; + } + + if constexpr (checkbounds_v) { + if (dst >= n_nodes()) { + throw ANNEXCEPTION( + "Trying to assign an edge to node {} when the number of nodes in the " + "graph is {}!", + dst, + n_nodes() + ); + } + } + + // Acquire lock — all reads and writes under the lock to prevent + // concurrent writers from seeing stale state. + std::lock_guard lock{node_locks_[src]}; + + // Check if there's room for the new node. + std::span raw_data = data_.get_datum(src); + Idx current_size = raw_data.front(); + if (current_size == max_degree_) { + return AddEdgeResult::Full; + } + + // At this point, we know there is room. + // Next, we need to find the position where we will insert the new edge. + // We fuse this with redundant edge insertion detection since the insertion + // position will also tell us where the edge would already exist. + Idx new_size = current_size + 1; + value_type adjacency_list = raw_data.subspan(1, new_size); + + auto begin = adjacency_list.begin(); + auto end = adjacency_list.end(); + + // TODO: Replace with binary search eventually. + // Blocking issue: legacy loaded graphs need validation of the sorted adjacency + // lists. + auto it = std::find(begin, end - 1, dst); + // auto it = std::lower_bound(begin, end - 1, dst); + if (it != end - 1 && (*it == dst)) { + return AddEdgeResult::AlreadyExists; + } + + auto seq = seq_counters_[src].begin_write(); + + // Insert at the new location using atomic stores. + for (auto dst_it = end - 1, src_it = end - 2; dst_it != it; --dst_it, --src_it) { + relaxed_store(*dst_it, *src_it); + } + relaxed_store(*it, dst); + + // Update the number of neighbors. + relaxed_store(raw_data.front(), new_size); + + seq_counters_[src].end_write(seq); + if (reverse_edges_) { + reverse_edges_->record(src, dst); + } + return AddEdgeResult::Added; + } + + /// Return the maximum out-degree this graph is capable of containing. + size_t max_degree() const { return max_degree_; } + /// Return the number of vertices currently in the graph. + size_t n_nodes() const { return data_.size(); } + + /// Return the maximum number of vertices this graph can hold without + /// reallocating any of its underlying storage. + size_t capacity() const { + return std::min({data_.capacity(), seq_counters_.capacity(), node_locks_.capacity()} + ); + } + + const data_type& get_data() const { return data_; } + data_type& get_data() { return data_; } + + // Resizeable API + void unsafe_resize(size_t new_size) { + data_.resize(new_size); + seq_counters_.resize(new_size); + node_locks_.resize(new_size); + if (reverse_edges_) { + reverse_edges_->resize(new_size); + } + } + void add_node() { unsafe_resize(n_nodes() + 1); } + + /// @brief Access the per-node sequence lock counters for concurrent read validation. + const SeqLockArray& seq_counters() const { return seq_counters_; } + + /// + /// @brief Enable maintenance of the per-node reverse-edge (in-neighbor) index. + /// + /// Off by default (null): every graph mutator hook is then a single null check, and + /// the static index / compaction scratch graphs pay nothing. Enabled only by the + /// dynamic index. + /// + void enable_reverse_edges() { + reverse_edges_ = std::make_unique>(n_nodes()); + } + + ReverseEdges* reverse_edges() { return reverse_edges_.get(); } + const ReverseEdges* reverse_edges() const { return reverse_edges_.get(); } + + /// + /// @brief Rebuild the reverse-edge index from the current (quiescent) graph. + /// + /// Records ``src`` into ``R(dst)`` for *every* edge ``src -> dst``. Must be called with + /// no concurrent graph mutation (load, post-build, post-compact). + /// + /// It is tempting to skip edges whose reverse ``dst -> src`` also exists, halving the + /// index: `gather_work_set` visits `out(d) union R(d)`, so a symmetric in-neighbor is + /// already covered by `out(d)`. That weaker invariant is *not maintainable*, though. + /// It reads "for every edge `u -> d`: `u` is in `R(d)` **or** the edge `d -> u` + /// exists", and the second disjunct is falsified whenever consolidation rewires `d` + /// and drops `d -> u` -- at which point `u`'s in-edge becomes invisible and a later + /// deletion of `d` leaves `u` pointing at a retired slot. Recording unconditionally + /// gives the strictly stronger `R(d) contains in(d)`, which no edge *removal* can + /// break and which every mutator already preserves on edge creation. + /// + template void rebuild_reverse_edges(Pool& threadpool) { + if (!reverse_edges_) { + return; + } + reverse_edges_->reset(); + threads::parallel_for( + threadpool, + threads::StaticPartition{threads::UnitRange{0, n_nodes()}}, + [&](const auto& is, uint64_t /*tid*/) { + for (auto src : is) { + for (auto dst : get_node(lib::narrow_cast(src))) { + reverse_edges_->record( + lib::narrow_cast(src), lib::narrow_cast(dst) + ); + } + } + } + ); + } + + ///// Saving + static constexpr lib::Version save_version = lib::Version(0, 0, 0); + static constexpr std::string_view serialization_schema = "default_graph"; + + lib::SaveTable metadata() const { + auto table = lib::SaveTable( + serialization_schema, + save_version, + {{"name", "graph"}, + {"max_degree", lib::save(max_degree())}, + {"num_vertices", lib::save(n_nodes())}, + {"eltype", lib::save(datatype_v)}} + ); + return table; + } + + template + lib::SaveTable metadata(const FileName& filename, const lib::UUID& uuid) const { + auto table = metadata(); + table.insert("binary_file", filename); + table.insert("uuid", uuid.str()); + return table; + } + + lib::SaveTable save(const lib::SaveContext& ctx) const { + auto uuid = lib::UUID{}; + auto filename = ctx.generate_name("graph"); + io::save(data_, io::NativeFile(filename), uuid); + return metadata(lib::save(filename.filename()), uuid); + } + + void save(std::ostream& os) const { io::save(data_, os); } + + protected: + template F, typename... Args> + static lib::lazy_result_t + load(const lib::LoadTable& table, const F& lazy, Args&&... args) { + // Perform a sanity check on the element type. + // Make sure we're loading the correct kind. + auto eltype = lib::load_at(table, "eltype"); + if (eltype != datatype_v) { + throw ANNEXCEPTION( + "Trying to load a graph with adjacency list types {} to a graph with " + "adjacency list types {}.", + name(eltype), + name>() + ); + } + + // Now that this is out of the way, resolve the file and load the data. + auto uuid = lib::load_at(table, "uuid"); + auto binaryfile = io::find_uuid(table.context().get_directory(), uuid); + if (!binaryfile.has_value()) { + throw ANNEXCEPTION("Could not open file with uuid {}!", uuid.str()); + } + return lazy(data_type::load(binaryfile.value(), std::forward(args)...)); + } + + template F, typename... AllocArgs> + static lib::lazy_result_t load( + const lib::ContextFreeLoadTable& table, + const F& lazy, + std::istream& is, + AllocArgs&&... alloc_args + ) { + // Perform a sanity check on the element type. + // Make sure we're loading the correct kind. + auto eltype = lib::load_at(table, "eltype"); + if (eltype != datatype_v) { + throw ANNEXCEPTION( + "Trying to load a graph with adjacency list types {} to a graph with " + "adjacency list types {}.", + name(eltype), + name>() + ); + } + + size_t num_vertices = lib::load_at(table, "num_vertices"); + size_t max_degree = lib::load_at(table, "max_degree"); + + // Build a table compatible with GenericSerializer + auto data_table = toml::table{ + {lib::config_schema_key, data::GenericSerializer::serialization_schema}, + {lib::config_version_key, data::GenericSerializer::save_version.str()}, + {"eltype", lib::save(datatype_v)}, + {"num_vectors", lib::save(num_vertices)}, + {"dims", lib::save(max_degree + 1)}, + }; + + return lazy( + data_type::load(lib::ContextFreeLoadTable(data_table), is, alloc_args...) + ); + } + + private: + // Adjacency-slot element access. Every slot is read by lock-free searches while writers + // mutate it, so all accesses go through `std::atomic_ref` with relaxed ordering: the + // per-node sequence-lock counters, not these individual accesses, establish the + // ordering that makes a read consistent. Relaxed atomics compile to plain loads and + // stores on the platforms SVS targets, so this costs nothing at runtime -- it only + // removes the data race that would otherwise make the program ill-formed. + // + // `SVS_CONCURRENT_UNSAFE_PLAIN_GRAPH_ACCESS` degrades these to plain accesses. It + // exists solely as a ThreadSanitizer negative control (see `tests/CMakeLists.txt`): a + // clean TSan run over this graph only means something if the *same* run reports races + // once the atomics are taken away. Never define it in a real build. + static Idx relaxed_load(const Idx& slot) { +#if defined(SVS_CONCURRENT_UNSAFE_PLAIN_GRAPH_ACCESS) + return slot; +#else + return std::atomic_ref(const_cast(slot)).load(std::memory_order_relaxed); +#endif + } + + static void relaxed_store(Idx& slot, Idx value) { +#if defined(SVS_CONCURRENT_UNSAFE_PLAIN_GRAPH_ACCESS) + slot = value; +#else + std::atomic_ref(slot).store(value, std::memory_order_relaxed); +#endif + } + + template void replace_node_impl(Idx i, const Span& new_neighbors) { + std::span old_snapshot{}; + std::array old_buffer; + Idx old_size = 0; + + std::lock_guard lock{node_locks_[i]}; + std::span raw_data = data_.get_datum(i); + + if (reverse_edges_) { + old_size = relaxed_load(raw_data[0]); + old_size = std::min({old_size, max_degree_, Idx{MAX_STACK_DEGREE}}); + for (Idx j = 0; j < old_size; ++j) { + old_buffer[j] = relaxed_load(raw_data[1 + j]); + } + old_snapshot = std::span(old_buffer.data(), old_size); + } + + // Clamp the number of elements to copy to the maximum out degree to correctly + // handle the case where the caller passes in too many neighbors. + Idx elements_to_copy = + std::min(max_degree_, lib::narrow_cast(new_neighbors.size())); + + auto seq = seq_counters_[i].begin_write(); + for (Idx j = 0; j < elements_to_copy; ++j) { + relaxed_store(raw_data[1 + j], new_neighbors[j]); + } + relaxed_store(raw_data[0], elements_to_copy); + seq_counters_[i].end_write(seq); + + if (reverse_edges_) { + for (Idx j = 0; j < elements_to_copy; ++j) { + Idx dst = new_neighbors[j]; + bool existed = std::find(old_snapshot.begin(), old_snapshot.end(), dst) != + old_snapshot.end(); + if (!existed) { + reverse_edges_->record(i, dst); + } + } + } + } + + // Upper bound on adjacency degree we snapshot on the stack for the reverse-edge diff. + // Graphs with larger degree still work; the diff is simply skipped past this bound + // (over-recording, never under-recording, preserving completeness). + static constexpr Idx MAX_STACK_DEGREE = 256; + + protected: + data_type data_; + Idx max_degree_; + SeqLockArray seq_counters_; + // Grow-stable: a concurrent add_points Phase 3 backprop locks node_locks_[other] + // lock-free while another add grows the array; segmented storage keeps existing + // locks at stable addresses. See svs/lib/segmented_vector.h. + lib::SegmentedVector node_locks_; + // Per-node in-neighbor index. Null (disabled) unless the owning index enables it. + std::unique_ptr> reverse_edges_ = nullptr; +}; + +///// +///// Concrete implementations. +///// + +/// +/// @brief Simple graph representation. +/// +/// @tparam Idx The integer type used to encode vertices in this graph. +/// +/// This data structure represents a graph using a single large allocation and a set maximum +/// degree. Accessing adjacency lists takes `O(1)` time. Only out-bound edges are stored. +/// +template > +class SimpleGraph : public SimpleGraphBase> { + public: + using parent_type = SimpleGraphBase>; + using data_type = typename parent_type::data_type; + using parent_type::get_data; + + /// + /// @brief Construct a new empty graph. + /// + /// @param num_nodes The number of nodes in the graph. + /// @param max_degree The maximum allowable degree in the graph. + /// + explicit SimpleGraph(size_t num_nodes, size_t max_degree) + : parent_type{num_nodes, max_degree} {} + + explicit SimpleGraph(size_t num_nodes, size_t max_degree, const Alloc& allocator) + : parent_type{num_nodes, max_degree, allocator} {} + + explicit SimpleGraph(data_type data) + : parent_type{std::move(data)} {} + + explicit SimpleGraph(parent_type&& parent) + : parent_type(std::move(parent)) {} + + ///// Loading + static constexpr SimpleGraph + load(const lib::LoadTable& table, const Alloc& allocator = {}) { + auto lazy = lib::Lazy([](data_type data) { return SimpleGraph(std::move(data)); }); + return parent_type::load(table, lazy, allocator); + } + + template + static constexpr SimpleGraph load( + const lib::ContextFreeLoadTable& table, std::istream& is, AllocArgs&&... alloc_args + ) { + auto lazy = lib::Lazy([](data_type data) { return SimpleGraph(std::move(data)); }); + return parent_type::load(table, lazy, is, std::forward(alloc_args)...); + } + + static constexpr SimpleGraph + load(const std::filesystem::path& path, const Alloc& allocator = {}) { + if (data::detail::is_likely_reload(path)) { + return lib::load_from_disk(path, allocator); + } else { + return SimpleGraph(data_type::load(path, allocator)); + } + } + + template + static constexpr SimpleGraph load(std::istream& is, AllocArgs&&... alloc_args) { + return lib::load_from_stream( + is, std::forward(alloc_args)... + ); + } +}; + +template +bool operator==(const SimpleGraph& x, const SimpleGraph& y) { + return graphs_equal(x, y); +} + +template +class SimpleBlockedGraph + : public SimpleGraphBase>> { + public: + using parent_type = + SimpleGraphBase>>; + using data_type = typename parent_type::data_type; + + // Constructors + SimpleBlockedGraph(size_t num_nodes, size_t max_degree) + : parent_type{num_nodes, max_degree} {} + + explicit SimpleBlockedGraph(data_type data) + : parent_type{std::move(data)} {} + + explicit SimpleBlockedGraph(parent_type&& parent) + : parent_type(std::move(parent)) {} + + ///// Loading + static constexpr SimpleBlockedGraph load(const lib::LoadTable& table) { + auto lazy = + lib::Lazy([](data_type data) { return SimpleBlockedGraph(std::move(data)); }); + return parent_type::load(table, lazy); + } + + static constexpr SimpleBlockedGraph + load(const lib::ContextFreeLoadTable& table, std::istream& is) { + auto lazy = + lib::Lazy([](data_type data) { return SimpleBlockedGraph(std::move(data)); }); + return parent_type::load(table, lazy, is); + } + + static constexpr SimpleBlockedGraph load(const std::filesystem::path& path) { + if (data::detail::is_likely_reload(path)) { + return lib::load_from_disk(path); + } else { + return SimpleBlockedGraph(data_type::load(path)); + } + } + + static constexpr SimpleBlockedGraph load(std::istream& is) { + return lib::load_from_stream(is); + } +}; + +} // namespace svs::index::vamana::concurrent::graphs diff --git a/include/svs/concurrent/graph_concepts.h b/include/svs/concurrent/graph_concepts.h new file mode 100644 index 000000000..12b90e8d0 --- /dev/null +++ b/include/svs/concurrent/graph_concepts.h @@ -0,0 +1,222 @@ +/* + * Copyright 2023 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +// !! NOTICE TO MAINTAINERS !! +// +// Due to limitations in the Doxygen -> Breathe -> Sphinx interface, support for documenting +// C++ 20 Concepts is limited. +// +// For now, we need to put the member-wise documentation in a code block in the main +// concept docstring. +// +// Hopefully this situation changes in future versions of these systems. + +/// +/// @file +/// + +/// +/// @ingroup concepts +/// @defgroup graph_concept_entry Main Graph Concepts. +/// + +/// +/// @ingroup concepts +/// @defgroup graph_concept_public Graph Concept Helpers. +/// + +#include +#include +#include +#include +#include +#include +#include + +namespace svs::index::vamana::concurrent::graphs { + +/// Outcome of `MemoryGraph::add_edge(src, dst)`. Distinguishes three cases so callers +/// can route dropped edges (e.g. to a backedge buffer) without a TOCTOU race between a +/// pre-check and the insert. +enum class AddEdgeResult : uint8_t { + Added, // Edge was inserted. + AlreadyExists, // Edge was already present (or self-loop). Not inserted. + Full, // Node's adjacency list is at max_degree. Edge NOT inserted. +}; + +// clang-format off + +/// +/// @ingroup graph_concept_entry +/// @brief Main concept modeling immutable in-memory graphs. +/// +/// @code{.cpp} +/// template +/// concept ImmutableMemoryGraph = requires(const T& const_g) { +/// // The encoding of vertices in the graph. +/// // At the time of writing, this is expected to be an integer. +/// // This may be relaxed in the future. +/// typename T::index_type; +/// std::integral; +/// +/// // Has `reference` and `const_reference` type aliases. +/// // These types should be at least forward ranges, but preferably random access +/// // ranges. +/// // +/// // Items yieled by the iterators for these ranges should ``index_type``. +/// typename T::reference; +/// typename T::const_reference; +/// +/// // Return the maximum degree that this particular implementation of the graph is +/// // capable of supporting. +/// // +/// // If the graph supports unbounded adjacency lists, may return `Dynamic`. +/// { const_g.max_degree() } -> std::convertible_to; +/// +/// // Return the number of vertices contained in traph. +/// { const_g.n_nodes() } -> std::convertible_to; +/// +/// // Adjacency list operations. +/// requires requires(typename T::index_type i) { +/// // Return a range over the adjacency list for node ``i``. +/// { const_g.get_node(i) } -> std::same_as; +/// +/// // Return the number of out neighbors for node ``i``. +/// { const_g.get_node_degree(i) } -> std::convertible_to; +/// +/// // Prefetch the adjacency list for node ``i``. +/// // This is a performance optimization only and may be implemented as a no-op +/// // without affecting correctness. +/// const_g.prefetch_node(i); +/// }; +/// }; +/// @endcode +/// +template +concept ImmutableMemoryGraph = requires(const T& const_g) { + // The encoding of vertices in the graph. + typename T::index_type; + requires std::integral; + + // Has `reference` and `const_reference` type aliases. + typename T::reference; + typename T::const_reference; + + { const_g.max_degree() } -> std::convertible_to; + { const_g.n_nodes() } -> std::convertible_to; + + // Get an adjacency list. + requires requires(typename T::index_type i) { + { const_g.get_node(i) } -> std::same_as; + { const_g.get_node_degree(i) } -> std::convertible_to; + const_g.prefetch_node(i); + }; +}; +// clang-format on + +/// +/// @ingroup graph_concept_public +/// @brief Obtain the index type used to encode neighbors in the graph type ``G``. +/// +template using index_type_t = typename G::index_type; + +// clang-format off + +/// +/// @ingroup graph_concept_entry +/// @brief Concept modeling mutable in-memory graphs. +/// +/// @code{.cpp} +/// template +/// concept MemoryGraph = requires(T& g, const T& const_g) { +/// // Add an edge to the graph atomically. Returns an AddEdgeResult indicating: +/// // Added - edge was inserted +/// // AlreadyExists - edge was already present (or self-loop); no insert +/// // Full - node is at max_degree; edge NOT inserted (caller should +/// // route to an overflow buffer if needed) +/// requires requires(index_type_t src, index_type_t dst) { +/// { g.add_edge(src, dst) } -> std::convertible_to; +/// }; +/// +/// // Completely clear the adjacency list for vertex ``i``. +/// requires requires(index_type_t i) { +/// g.clear_node(i); +/// }; +/// +/// // Overwrite the adjacency list for `src`. +/// requires requires( +/// index_type_t src, +/// const std::vector>& neighbors_vector, +/// std::span> neighbors_span +/// ) { +/// g.replace_node(src, neighbors_vector); +/// g.replace_node(src, neighbors_span); +/// }; +/// }; +/// @endcode +/// +template +concept MemoryGraph = requires(T& g, const T& const_g) { + // Adding an edge. + requires requires(index_type_t src, index_type_t dst) { + { g.add_edge(src, dst) } -> std::convertible_to; + }; + + // Clear adjacency list. + requires requires(index_type_t i) { + g.clear_node(i); + }; + + // Overwriting an edge. + requires requires( + index_type_t src, + const std::vector>& neighbors_vector, + std::span> neighbors_span + ) { + g.replace_node(src, neighbors_vector); + g.replace_node(src, neighbors_span); + }; +}; +// clang-format on + +/// +/// @ingroup graph_concept_public +/// @brief Compare the equality of two graphs. +/// +/// Two graphs are considered equal if: +/// +/// * The contain the same number of vertices. +/// * The adjacency lists for each vertex compare equal. +/// +template +bool graphs_equal(const Graph1& x, const Graph2& y) { + if (x.num_nodes() != y.num_nodes()) { + return false; + } + + for (size_t i = 0, imax = x.num_nodes(); i < imax; ++i) { + const auto& xa = x.get_node(i); + const auto& ya = y.get_node(i); + if (!std::equal(xa.begin(), xa.end(), ya.begin())) { + return false; + } + } + return true; +} + +} // namespace svs::index::vamana::concurrent::graphs diff --git a/include/svs/concurrent/greedy_search.h b/include/svs/concurrent/greedy_search.h new file mode 100644 index 000000000..b9978813a --- /dev/null +++ b/include/svs/concurrent/greedy_search.h @@ -0,0 +1,203 @@ +/* + * Copyright 2023 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "svs/concepts/data.h" +#include "svs/concepts/distance.h" +#include "svs/concurrent/graph_concepts.h" +// For `svs::index::vamana::GreedySearchPrefetchParameters`, which this stack reuses +// unchanged -- see the using-declaration below. +#include "svs/index/vamana/greedy_search.h" +#include "svs/index/vamana/search_buffer.h" +#include "svs/lib/concurrency/seqlock.h" +#include "svs/concurrent/spinlock.h" + +#include +#include + +namespace svs::index::vamana::concurrent { + +// The greedy-search *scaffolding* -- the tracker API, the tracker concept, the entry-point +// initializer, the default neighbor builder, and the prefetch parameters -- is unchanged by +// this stack; only `greedy_search` itself gains the SeqLock retry loop. Alias the +// pre-existing entities rather than redeclaring them: a redeclaration inside `concurrent` +// would shadow the enclosing namespace's with a distinct look-alike, and values arriving +// from facilities reused verbatim (e.g. `SearchScratchspace::prefetch_parameters`, or +// `RestartInitializer`'s `NullTracker` parameter) would then fail to convert. +using svs::index::vamana::GreedySearchPrefetchParameters; +using svs::index::vamana::GreedySearchTracker; +using svs::index::vamana::NullTracker; + +///// +///// Initialization Customization. +///// + +using svs::index::vamana::EntryPointInitializer; + +///// +///// Greedy Search +///// + +using svs::index::vamana::NeighborBuilder; + +template < + graphs::ImmutableMemoryGraph Graph, + data::ImmutableMemoryDataset Dataset, + data::AccessorFor Accessor, + typename QueryType, + distance::Distance Dist, + typename Buffer, + typename Initializer, + typename Builder, + GreedySearchTracker Tracker> +void greedy_search( + const Graph& graph, + const Dataset& dataset, + Accessor& accessor, + const QueryType& query, + Dist& distance_function, + Buffer& search_buffer, + const Initializer& initializer, + const Builder& builder, + Tracker& search_tracker, + GreedySearchPrefetchParameters prefetch_parameters = {}, + const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) +) { + using I = typename Graph::index_type; + + // Fix the query if needed by the distance function. + distance::maybe_fix_argument(distance_function, query); + + // Initialize the search buffer. + { + // A lambda that wraps the distance computation to avoid propagating everything + // into the initializer. + auto computer = [&](std::integral auto id) { + return distance::compute(distance_function, query, accessor(dataset, id)); + }; + initializer(search_buffer, computer, graph, builder, search_tracker); + } + + // Main search routine. + while (!search_buffer.done()) { + // Check if request to cancel the search + if (cancel()) { + return; + } + // Get the next unvisited vertex. + const auto& node = search_buffer.next(); + auto node_id = node.id(); + + for (;;) { // SeqLock retry loop + auto maybe_seq = graph.seq_counters()[node_id].read_begin(); + if (!maybe_seq) { + svs::detail::pause(); + continue; + } + + // Get the adjacency list for this vertex and prepare prefetching logic. + auto neighbors = graph.get_node(node_id); + const size_t num_neighbors = neighbors.size(); + search_tracker.visited(Neighbor{node}, num_neighbors); + + auto prefetcher = lib::make_prefetcher( + lib::PrefetchParameters{ + prefetch_parameters.lookahead, prefetch_parameters.step}, + num_neighbors, + [&](size_t i) { accessor.prefetch(dataset, neighbors[i]); }, + [&](size_t i) { + // Perform the visited set enabled check just once. + if (search_buffer.visited_set_enabled()) { + // Prefetch next bucket so it's (hopefully) in the cache when + // we next consult the visited filter. + if (i + 1 < num_neighbors) { + search_buffer.unsafe_prefetch_visited(neighbors[i + 1]); + } + return !search_buffer.unsafe_is_visited(neighbors[i]); + } + + // Otherwise, always prefetch the next data item. + return true; + } + ); + + ///// Neighbor expansion. + prefetcher(); + for (auto id : neighbors) { + if (search_buffer.emplace_visited(id)) { + continue; + } + + // Run the prefetcher. + prefetcher(); + + // Compute distance and update search buffer. + auto dist = + distance::compute(distance_function, query, accessor(dataset, id)); + search_buffer.insert(builder(id, dist)); + } + + // Validate that no concurrent write occurred during the read. + if (graph.seq_counters()[node_id].read_validate(*maybe_seq)) { + break; // Consistent read — proceed to the next node. + } + svs::detail::pause(); + // Retry: stale entries from the invalid read remain in the search buffer. + // They have valid IDs and distances, and insert() deduplicates by ID. + } + } +} + +// Overload to provide a default search tracker because search trackers are taken by +// lvalue reference. +template < + graphs::ImmutableMemoryGraph Graph, + data::ImmutableMemoryDataset Dataset, + data::AccessorFor Accessor, + typename QueryType, + distance::Distance Dist, + typename Buffer, + typename Initializer, + typename Builder = NeighborBuilder> +void greedy_search( + const Graph& graph, + const Dataset& dataset, + Accessor& accessor, + QueryType query, + Dist& distance_function, + Buffer& search_buffer, + const Initializer& initializer, + const Builder& builder = NeighborBuilder(), + GreedySearchPrefetchParameters prefetch_parameters = {}, + const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) +) { + auto null_tracker = NullTracker{}; + concurrent::greedy_search( + graph, + dataset, + accessor, + query, + distance_function, + search_buffer, + initializer, + builder, + null_tracker, + prefetch_parameters, + cancel + ); +} +} // namespace svs::index::vamana::concurrent diff --git a/include/svs/concurrent/iterator.h b/include/svs/concurrent/iterator.h new file mode 100644 index 000000000..b33c32c1c --- /dev/null +++ b/include/svs/concurrent/iterator.h @@ -0,0 +1,358 @@ +/* + * Copyright 2024 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +// svs +#include "svs/index/index.h" +#include "svs/concurrent/dynamic_index.h" +#include "svs/index/vamana/index.h" +#include "svs/lib/scopeguard.h" + +// stl +#include +#include + +namespace svs::index::vamana::concurrent { + +/// @brief A graph search initializer that uses the existing contents of the search buffer +/// to initialize the next round of graph search. +/// +/// If a previous search exited with an exception, this initializer allows restarting +/// the search from scratch using the traditional method. +template struct RestartInitializer { + template < + typename Buffer, + typename Computer, + graphs::ImmutableMemoryGraph Graph, + typename Builder> + void operator()( + Buffer& buffer, + const Computer& computer, + const Graph& graph, + const Builder& builder, + vamana::NullTracker tracker // Compile error for non-NullTracker. + ) const { + // Restart the search from scratch if requested. + if (hard_restart_) { + vamana::EntryPointInitializer{ + entry_points_}(buffer, computer, graph, builder, tracker); + return; + } + + // Otherwise, sort the buffer to prepare for the next search. + buffer.sort(); + } + + // Entry points for restarting the search from scratch. + std::span entry_points_; + bool hard_restart_; +}; + +namespace detail { +constexpr void checkdims(size_t query_size, size_t index_dims) { + if (query_size != index_dims) { + throw ANNEXCEPTION( + "Incompatible dimensions. Query has {} while the index expects {}.", + query_size, + index_dims + ); + } +} +} // namespace detail + +/// @brief A batch iterator for retrieving neighbors from the index in batches. +/// +/// This iterator abstracts the process of retrieving neighbors in fixed-size batches +/// while maintaining internal state for efficient graph traversal. +template class BatchIterator { + public: + static_assert( + std::is_trivially_copyable_v, + "The batch iterator requires a trivial (no-throw copy constructible) type to " + "provide its exception guarantees" + ); + + // Public type aliases + using scratchspace_type = index::scratchspace_t; + using internal_id_type = typename Index::internal_id_type; + using external_id_type = size_t; + + // The value type yielded by the iterator. + using value_type = Neighbor; + + private: + // Private type aliases + using result_buffer_type = std::vector; + + /// @brief Copies results from the scratch buffer to the results buffer. + /// Ensures that only unique neighbors are added to the results buffer. + void copy_from_scratch(size_t batch_size) { + results_.clear(); + const auto& buffer = scratchspace_.buffer; + for (size_t i = 0, imax = buffer.size(); i < imax; ++i) { + auto neighbor = buffer[i]; + auto result = yielded_.insert(neighbor.id()); + if (result.second /* inserted */) { + // Rollback insertion into the yielded set if push_back throws. + auto guard = lib::make_dismissable_scope_guard([&]() noexcept { + yielded_.erase(result.first); + }); + results_.push_back(adapt(neighbor)); + guard.dismiss(); + } + + // Stop if the requested batch size is reached. + if (results_.size() == batch_size) { + break; + } + } + } + + /// @brief Initializes the search buffer with the configured capacity. + void initialize_buffer() { + auto config = SearchBufferConfig{0, extra_search_buffer_capacity_}; + scratchspace_.buffer.change_maxsize(config); + scratchspace_.buffer.clear(); + } + + /// @brief Increments the search window and capacity by `batch_size` for the next + /// iteration. + void increment_buffer(size_t batch_size) { + auto config = scratchspace_.buffer.config(); + config.increment(batch_size); + scratchspace_.buffer.change_maxsize(config); + } + + public: + using size_type = typename result_buffer_type::size_type; + using reference = value_type&; + using const_reference = const value_type&; + + /// Random-access iterator to `value_type` over the current batch of results. + using iterator = typename result_buffer_type::iterator; + /// Random-access iterator to `const value_type` over the current batch of results. + using const_iterator = typename result_buffer_type::const_iterator; + + /// @brief Constructs a batch iterator for the given query over the index. + /// @param parent The index to search. + /// @param query The query data. + /// @param extra_search_buffer_capacity Additional buffer capacity for the search. + /// When not provided, ``svs::ITERATOR_EXTRA_BUFFER_CAPACITY_DEFAULT = 100`` is + /// used. + BatchIterator( + const Index& parent, + std::span query, + size_t extra_search_buffer_capacity = svs::UNSIGNED_INTEGER_PLACEHOLDER + ) + : parent_{&parent} + , query_{query.begin(), query.end()} + , scratchspace_{parent_->scratchspace()} { + detail::checkdims(query.size(), parent.dimensions()); + + extra_search_buffer_capacity_ = + extra_search_buffer_capacity == svs::UNSIGNED_INTEGER_PLACEHOLDER + ? svs::ITERATOR_EXTRA_BUFFER_CAPACITY_DEFAULT + : extra_search_buffer_capacity; + initialize_buffer(); + } + + /// @brief Updates the iterator with a new query. + /// Resets the internal state and restarts the search when `next(...)` is called. + void update(std::span newquery) { + detail::checkdims(newquery.size(), parent_->dimensions()); + assert(newquery.size() == query_.size()); + + std::copy(newquery.begin(), newquery.end(), query_.begin()); + + initialize_buffer(); + restart_search_ = true; + iteration_ = 0; + yielded_.clear(); + results_.clear(); + is_exhausted_ = false; + } + + /// @brief Adapts an internal neighbor to an external neighbor. + template svs::Neighbor adapt(N internal) const { + if constexpr (Index::needs_id_translation) { + // `copy_from_scratch` -- our only caller -- runs under + // `parent_->lock_for_translation()`, so use the non-locking translation: + // `std::shared_mutex` is not recursive, and re-acquiring it per neighbor would + // deadlock against a writer that arrives in between. + return Neighbor{ + parent_->unsafe_translate_internal_id(internal.id()), internal.distance()}; + } else { + return internal; + } + } + + /// @brief Returns an iterator to the beginning of the results. + iterator begin() { return results_.begin(); } + /// @brief Returns an iterator to the end of the results. + iterator end() { return results_.end(); } + /// @copydoc begin() + const_iterator begin() const { return results_.begin(); } + /// @copydoc end() + const_iterator end() const { return results_.end(); } + /// @copydoc begin() + const_iterator cbegin() const { return results_.cbegin(); } + /// @copydoc begin() + const_iterator cend() const { return results_.cend(); } + + /// @brief Returns a span over the current batch of neighbors. + /// The span is invalidated by calls to `next(...)`. + std::span contents() const { return lib::as_const_span(results_); } + + /// @brief Returns the number of buffered results. + size_t size() const { return results_.size(); } + + /// @brief Return the batch number corresponding to the current buffer. + size_t batch_number() const { return iteration_; } + + /// @brief Returns whether iterator can find more neighbors or not for the given query. + /// + /// The iterator is considered done when all the available nodes have been yielded or + /// when the search can not find any more neighbors. The transition from not done to + /// done will be triggered by a call to ``next()``. The contents of ``batch_number()`` + /// and ``parameters_for_current_iteration()`` will then remain unchanged by subsequent + /// invocations of ``next()``. + bool done() const { return (yielded_.size() == parent_->size() || is_exhausted_); } + + /// @brief Forces the next iteration to restart the search from scratch. + void restart_next_search() { restart_search_ = true; } + + /// @brief Returns the search parameters used for the current batch. + vamana::VamanaSearchParameters parameters_for_current_iteration() const { + auto& buffer = scratchspace_.buffer; + auto& prefetch = scratchspace_.prefetch_parameters; + return VamanaSearchParameters{ + buffer.config(), + buffer.visited_set_enabled(), + prefetch.lookahead, + prefetch.step}; + } + + /// @brief Prepares the next batch of neighbors (up to ``batch_size``) from the index. + /// Handles exceptions gracefully and ensures iterator state consistency. + void next( + size_t batch_size, + const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) + ) { + if (done()) { + results_.clear(); + return; + } + + increment_buffer(batch_size); + + bool restart_search_copy = std::exchange(restart_search_, true); + + // Hold the search lock (compact_mutex_ shared for a dynamic index; a + // no-op for a static one) so a concurrent compact() cannot free the + // segments under us. add_points growth is lock-free (grow-stable + // storage). The guard is released at the end of this scope — before + // acquiring the translation lock. The two locks must never be held + // nested in the translator-before-compact order: that would invert the + // global lock order (compact -> translator) and deadlock against compact. + { + [[maybe_unused]] auto search_guard = parent_->lock_for_search(); + parent_->experimental_escape_hatch([&]( + const auto& graph, + const auto& data, + const auto& SVS_UNUSED(distance), + std::span entry_points + ) { + auto search_closure = + [&](const auto& query, const auto& accessor, auto& d, auto& buffer) { + constexpr vamana::extensions::UsesReranking< + std::remove_const_t>> + uses_reranking{}; + if constexpr (uses_reranking()) { + distance::maybe_fix_argument(d, query); + for (size_t j = 0, jmax = buffer.size(); j < jmax; ++j) { + auto& neighbor = buffer[j]; + auto id = neighbor.id(); + auto new_distance = + distance::compute(d, query, data.get_primary(id)); + neighbor.set_distance(new_distance); + } + buffer.sort(); + } + + concurrent::greedy_search( + graph, + data, + accessor, + query, + d, + buffer, + RestartInitializer{entry_points, restart_search_copy}, + parent_->internal_search_builder(), + scratchspace_.prefetch_parameters, + cancel + ); + + if constexpr (Index::needs_id_translation) { + buffer.cleanup(); + buffer.sort(); + } + }; + + extensions::single_search( + data, + scratchspace_.buffer, + scratchspace_.scratch, + lib::as_const_span(query_), + search_closure, + *parent_ + ); + }); + } + + ++iteration_; + restart_search_ = false; + // Hold the translation lock (translator_mutex_) so a + // concurrent consolidate/compact cannot erase or remap entries mid-copy. + { + [[maybe_unused]] auto translation_guard = parent_->lock_for_translation(); + copy_from_scratch(batch_size); + } + // If result is empty after calling next(), mark the iterator as exhausted. + // The iterator will not be able to find any more neighbors. + if (results_.size() == 0 && batch_size > 0) { + is_exhausted_ = true; + } + } + + private: + const Index* parent_; // The index being accessed. + std::vector query_; // Local buffer for the query. + scratchspace_type scratchspace_; // Scratch space for search. + std::vector> results_{}; // Filtered results from search. + std::unordered_set yielded_{}; // Set of yielded neighbors. + size_t iteration_ = 0; // Current iteration number. + bool restart_search_ = true; // Whether the next search should restart from scratch. + size_t extra_search_buffer_capacity_ = + svs::UNSIGNED_INTEGER_PLACEHOLDER; // Extra buffer capacity for the next search. + bool is_exhausted_ = false; // Whether the iterator is exhausted. +}; + +// Deduction Guides +template +BatchIterator(const Index*, std::span) -> BatchIterator; + +} // namespace svs::index::vamana::concurrent diff --git a/include/svs/concurrent/multi.h b/include/svs/concurrent/multi.h new file mode 100644 index 000000000..64a163156 --- /dev/null +++ b/include/svs/concurrent/multi.h @@ -0,0 +1,967 @@ +/* + * Copyright 2025 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once +#include "svs/concurrent/dynamic_index.h" +#include "svs/concurrent/iterator.h" + +#include +#include +#include + +namespace svs::index::vamana::concurrent { + +/// @brief A multi-vector batch iterator for retrieving neighbors with unique labels from +/// the index in batches. +/// +/// This iterator abstracts the process of retrieving neighbors in fixed-size batches +/// while maintaining internal state for efficient graph traversal. +/// In multi-vector scenario, +/// each label can have multiple vectors. +/// This iterator ensures that neighbors are retrieved with unique labels +template class MultiBatchIterator { + using label_type = size_t; + using external_id_type = size_t; + using value_type = Neighbor; + + // Private type aliases + using result_buffer_type = std::vector; + + using ParentIndex = typename Index::ParentIndex; + using compare = typename Index::compare; + + public: + /// Random-access iterator to `value_type` over the current batch of results. + using iterator = typename result_buffer_type::iterator; + /// Random-access iterator to `const value_type` over the current batch of results. + using const_iterator = typename result_buffer_type::const_iterator; + + MultiBatchIterator( + const Index& index, + std::span query, + size_t extra_search_buffer_capacity = svs::UNSIGNED_INTEGER_PLACEHOLDER + ) + : index_{index} + , batch_iterator_{index.get_parent_index(), query, extra_search_buffer_capacity} {} + + void next( + size_t batch_size, + const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) + ) { + auto results_copy = results_; + results_.clear(); + get_results_from_extra(batch_size); + + while (results_.size() < batch_size && !batch_iterator_.done()) { + try { + batch_iterator_.next(batch_size, cancel); + } catch (const ANNException& error) { + results_ = std::move(results_copy); + throw error; + } + for (auto& result : batch_iterator_) { + auto label = index_.external_to_label(result.id()); + auto found_in_returned = returned_.find(label); + auto new_result = Neighbor{label, result.distance()}; + + if (found_in_returned == returned_.end()) { + if (results_.size() < batch_size) { + returned_.insert(label); + results_.push_back(std::move(new_result)); + } else { + extra_results_.push_back(std::move(new_result)); + } + } else { + // results_ should be small enough to use find + auto found_in_results = std::find_if( + results_.begin(), + results_.end(), + [label](const auto& res) { return res.id() == label; } + ); + if (found_in_results != results_.end()) { + *found_in_results = + std::min(*found_in_results, new_result, TotalOrder(compare{})); + } + } + } + } + + ++iteration_; + return; + } + + size_t batch_number() const { return iteration_; } + + void update(std::span newquery) { + iteration_ = 0; + returned_.clear(); + results_.clear(); + extra_results_.clear(); + batch_iterator_.update(newquery); + } + + iterator begin() { return results_.begin(); } + iterator end() { return results_.end(); } + const_iterator begin() const { return results_.begin(); } + const_iterator end() const { return results_.end(); } + const_iterator cbegin() const { return results_.cbegin(); } + const_iterator cend() const { return results_.cend(); } + size_t size() const { return results_.size(); } + + bool done() const { + return (batch_iterator_.done() && extra_results_.empty()) || + (returned_.size() == index_.labelcount()); + } + + std::span contents() const { return lib::as_const_span(results_); } + + private: + void get_results_from_extra(size_t batch_size) { + // sort to get the best candidate from the back + std::sort(extra_results_.rbegin(), extra_results_.rend(), TotalOrder(compare{})); + + while (results_.size() < batch_size && !extra_results_.empty()) { + auto best = extra_results_.back(); + extra_results_.pop_back(); + + if (returned_.find(best.id()) == returned_.end()) { + returned_.insert(best.id()); + results_.push_back(std::move(best)); + } + } + + return; + } + + const Index& index_; + size_t iteration_ = 0; + std::unordered_set returned_; + std::vector> results_; + std::vector> extra_results_; + BatchIterator batch_iterator_; +}; + +template +class MultiMutableVamanaIndex { + public: + static constexpr bool supports_insertions = true; + static constexpr bool supports_deletions = true; + static constexpr bool supports_saving = false; // temporary disable for now + static constexpr bool needs_id_translation = true; + + using ParentIndex = MutableVamanaIndex; + using compare = distance::compare_t; + using Idx = typename ParentIndex::Idx; + using search_parameters_type = typename ParentIndex::search_parameters_type; + using external_id_type = typename ParentIndex::external_id_type; + using scratchspace_type = typename ParentIndex::scratchspace_type; + using distance_type = Dist; + using label_type = size_t; + using graph_type = Graph; + using data_type = Data; + using label_to_external_type = + std::unordered_map>; + using external_to_label_type = std::unordered_map; + + private: + distance_type distance_; + // Wrapped in unique_ptr so this class stays movable (std::atomic and + // std::shared_mutex are neither copyable nor movable); the VecSim serializer + // move-constructs MultiMutableVamanaIndex via + // make_unique(std::move(loaded)). + std::unique_ptr> counter_{ + std::make_unique>(0)}; + std::unique_ptr index_{nullptr}; + label_to_external_type label_to_external_; + external_to_label_type external_to_label_; + // External IDs that have been soft-deleted (removed from the live label maps + // above) but not yet consolidated in the parent index. Keyed by the label they + // belonged to so that consolidate(labels) can recover them; the parent index + // erases its translator entries only during (full or partial) consolidation. + // Guarded by l2e_mutex_. + label_to_external_type pending_deletes_; + // Guard the label<->external id translation maps for concurrent add/delete/search. + std::unique_ptr l2e_mutex_{std::make_unique()}; + std::unique_ptr e2l_mutex_{std::make_unique()}; + + template + void + prepare_added_id_by_label(const Labels& labels, std::vector& adds) { + for (const auto l : labels) { + // Ensure a bucket exists for this label (it may be new). + { + l2e_mutex_->lock_shared(); + bool missing = label_to_external_.find(l) == label_to_external_.end(); + l2e_mutex_->unlock_shared(); + if (missing) { + std::lock_guard lock{*l2e_mutex_}; + label_to_external_.insert({l, std::vector{}}); + } + } + + size_t new_external_id = counter_->fetch_add(1, std::memory_order_relaxed); + { + std::lock_guard lock{*l2e_mutex_}; + label_to_external_[l].push_back(new_external_id); + } + { + std::lock_guard lock{*e2l_mutex_}; + external_to_label_.insert({new_external_id, l}); + } + adds.push_back(new_external_id); + } + } + + public: + template + MultiMutableVamanaIndex( + const VamanaBuildParameters& parameters, + Data data, + const Labels& labels, + Dist distance_function, + ThreadPoolProto threadpool_proto, + svs::logging::logger_ptr logger = svs::logging::get() + ) + : distance_(std::move(distance_function)) { + std::vector adds; + adds.reserve(labels.size()); + prepare_added_id_by_label(labels, adds); + index_ = std::make_unique( + parameters, + std::move(data), + std::move(adds), + distance_, + std::move(threadpool_proto), + std::move(logger) + ); + } + + template + MultiMutableVamanaIndex( + Graph graph, + Data data, + Idx entry_point, + Dist distance_function, + const Labels& labels, + ThreadPoolProto threadpool_proto, + svs::logging::logger_ptr logger = svs::logging::get() + ) + : distance_(std::move(distance_function)) { + std::vector adds; + adds.reserve(labels.size()); + prepare_added_id_by_label(labels, adds); + + index_ = std::make_unique( + std::move(graph), + std::move(data), + entry_point, + distance_, + std::move(adds), + std::move(threadpool_proto), + std::move(logger) + ); + } + + /// @brief Constructor for post re-load multi dynamic vamana index. + template + MultiMutableVamanaIndex( + const VamanaIndexParameters& config, + data_type data, + graph_type graph, + const Dist& distance_function, + const std::vector& labels, + Pool threadpool, + svs::logging::logger_ptr logger = svs::logging::get() + ) + : distance_(std::move(distance_function)) { + std::vector adds; + adds.reserve(labels.size()); + prepare_added_id_by_label(labels, adds); + + // create a remapped translator where external_id == internal_id + IDTranslator remapped_translator; + remapped_translator.insert(adds, threads::UnitRange(0, adds.size())); + + index_ = std::make_unique( + config, + std::move(data), + std::move(graph), + distance_, + remapped_translator, + std::move(threadpool), + std::move(logger) + ); + } + + /// @brief Constructor for post re-load dynamic vamana index. + /// This constructor provides a compatibility path for directly loading dynamic vamana + /// datasets. This constructor takes external IDs in translator as labels. The span of + /// internal ID's in translator should be exactly ``[0, data.size())`. + template + MultiMutableVamanaIndex( + const VamanaIndexParameters& config, + data_type data, + graph_type graph, + const Dist& distance_function, + IDTranslator translator, + Pool threadpool, + svs::logging::logger_ptr logger = svs::logging::get() + ) + : distance_(std::move(distance_function)) { + // Create labels where labels = translator.external_ids + std::vector labels(translator.size()); + std::transform( + translator.begin(), + translator.end(), + labels.begin(), + [](const auto& ext_int) { return ext_int.first; } + ); + + std::vector adds; + adds.reserve(translator.size()); + prepare_added_id_by_label(labels, adds); + + // create a remapped translator where external_id == internal_id + IDTranslator remapped_translator; + remapped_translator.insert(adds, threads::UnitRange(0, adds.size())); + + index_ = std::make_unique( + config, + std::move(data), + std::move(graph), + distance_, + remapped_translator, + std::move(threadpool), + std::move(logger) + ); + } + + const label_to_external_type& get_label_to_external_lookup() const { + return label_to_external_; + } + const external_to_label_type& get_external_to_label_lookup() const { + return external_to_label_; + } + const ParentIndex& get_parent_index() const { return *index_; } + + svs::logging::logger_ptr get_logger() const { return index_->get_logger(); } + + template + double get_distance(label_type label, const Query& query) const { + double best = INVALID_DISTANCE; + std::shared_lock l2e_lock{*l2e_mutex_}; + auto it = label_to_external_.find(label); + + if (it != label_to_external_.end()) { + auto& vectors = (*it).second; + for (auto each : vectors) { + best = std::min( + best, + index_->get_distance(each, query), + [](const double a, const double b) { + if (std::isnan(a)) + return false; + if (std::isnan(b)) + return true; + return compare{}(a, b); + } + ); + } + } + + return best; + } + + template + std::vector + add_points(const Points& points, const Labels& labels, bool reuse_empty = false) { + const size_t num_points = points.size(); + const size_t num_labels = labels.size(); + if (num_points != num_labels) { + throw ANNEXCEPTION( + "Number of points ({}) not equal to the number of external ids ({})!", + num_points, + num_labels + ); + } + + std::vector adds; + adds.reserve(num_labels); + prepare_added_id_by_label(labels, adds); + index_->add_points(points, adds, reuse_empty); + + return adds; + } + + // Return the number of deleted vectors + template size_t delete_entries(const T& labels) { + std::vector deletes; + + { + std::lock_guard l2e_lock{*l2e_mutex_}; + for (auto& label : labels) { + auto it = label_to_external_.find(label); + if (it != label_to_external_.end()) { + auto& externals = (*it).second; + deletes.insert(deletes.end(), externals.begin(), externals.end()); + { + std::lock_guard e2l_lock{*e2l_mutex_}; + for (auto& ext : externals) { + external_to_label_.erase(ext); + } + } + // Remember the soft-deleted externals under their label so a + // later consolidate(labels) can consolidate just these points. + auto& pending = pending_deletes_[label]; + pending.insert(pending.end(), externals.begin(), externals.end()); + label_to_external_.erase(it); + } + } + } + index_->delete_entries(deletes); + return deletes.size(); + } + + template + void search( + const Query& query, + scratchspace_type& scratch, + const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) + ) const { + index_->search(query, scratch, cancel); + } + + template + void search( + QueryResultView results, + const Queries& queries, + const search_parameters_type& sp, + const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) + ) { + auto& borrow_threadpool = index_->get_threadpool_handle(); + + threads::parallel_for( + borrow_threadpool, + threads::StaticPartition{queries.size()}, + [&](const auto is, uint64_t SVS_UNUSED(tid)) { + size_t num_neighbors = results.n_neighbors(); + size_t batch_size = + std::max(num_neighbors, sp.buffer_config_.get_search_window_size()); + + // use batch iterator to search + for (auto i : is) { + auto batch_iterator = make_batch_iterator(queries.get_datum(i), 10); + batch_iterator.next(batch_size, cancel); + size_t j{0}; + for (auto& res : batch_iterator) { + if (j == num_neighbors) { + break; + } + results.set(res, i, j++); + } + + for (; j < num_neighbors; ++j) { + // insert default neighbor if not enough + results.set(Neighbor{}, i, j); + } + } + } + ); + + return; + } + + void compact(Idx batch_size = 1'000) { + index_->compact(batch_size); + // compact() fully consolidates, reclaiming every soft-deleted point. + std::lock_guard l2e_lock{*l2e_mutex_}; + pending_deletes_.clear(); + } + + void consolidate() { + index_->consolidate(); + // Full consolidation reclaims every soft-deleted point. + std::lock_guard l2e_lock{*l2e_mutex_}; + pending_deletes_.clear(); + } + + // Consolidate only the soft-deleted vectors belonging to the listed `labels`. + // Mirrors delete_entries(labels): `labels` is any iterable of label values. + // Vectors soft-deleted under other labels remain in the graph and navigable + // until a later consolidate(). Returns the number of external vectors + // consolidated (a label may map to several vectors). + template size_t consolidate(const T& labels) { + std::vector externals; + { + std::lock_guard l2e_lock{*l2e_mutex_}; + for (auto& label : labels) { + auto it = pending_deletes_.find(label); + if (it == pending_deletes_.end()) { + continue; // No soft-deleted vectors pending for this label. + } + externals.insert(externals.end(), it->second.begin(), it->second.end()); + pending_deletes_.erase(it); + } + } + return index_->consolidate(externals); + } + + template + auto make_batch_iterator( + std::span query, + size_t extra_search_buffer_capacity = svs::UNSIGNED_INTEGER_PLACEHOLDER + ) const { + return MultiBatchIterator(*this, query, extra_search_buffer_capacity); + } + + void set_threadpool(threads::ThreadPoolHandle threadpool) { + index_->set_threadpool(std::move(threadpool)); + } + + /// + /// @brief Destroy the original thread pool and set to the provided one. + /// + /// @param threadpool An acceptable thread pool. + /// + /// @copydoc threadpool_requirements + /// + template + void set_threadpool(Pool threadpool) + requires(!std::is_same_v) + { + set_threadpool(threads::ThreadPoolHandle(std::move(threadpool))); + } + + /// + /// @brief Return the current thread pool handle. + /// + threads::ThreadPoolHandle& get_threadpool_handle() { + return index_->get_threadpool_handle(); + } + + bool has_id(size_t e) const { + std::shared_lock l2e_lock{*l2e_mutex_}; + return label_to_external_.find(e) != label_to_external_.end(); + } + + size_t size() const { return index_->size(); } + + size_t labelcount() const { + std::shared_lock l2e_lock{*l2e_mutex_}; + return label_to_external_.size(); + } + + // scrathspace from parent index + scratchspace_type scratchspace(const search_parameters_type& sp) const { + return index_->scratchspace(sp); + } + + // scrathspace from parent index + scratchspace_type scratchspace() const { return scratchspace(get_search_parameters()); } + + // translate internal id -> external id -> label + label_type translate_internal_id(Idx i) const { + std::shared_lock e2l_lock{*e2l_mutex_}; + return external_to_label_.at(index_->translate_internal_id(i)); + } + + // Thread-safe external id -> label lookup. + label_type external_to_label(external_id_type external_id) const { + std::shared_lock e2l_lock{*e2l_mutex_}; + return external_to_label_.at(external_id); + } + + /// @brief Call the functor with all labels in the index. + /// + /// @param f A functor with an overloaded ``operator()(size_t)`` method. Called on + /// each external ID in the index. + /// + template void on_ids(F&& f) const { + std::shared_lock l2e_lock{*l2e_mutex_}; + for (auto pair : label_to_external_) { + f(pair.first); + } + } + + /// + /// @brief Return a vector of all valid labels present in the index. + /// + std::vector external_ids() const { + std::vector ids{}; + on_ids([&ids](size_t id) { ids.push_back(id); }); + return ids; + } + + const Data& view_data() const { return index_->view_data(); } + const Graph& view_graph() const { return index_->view_graph(); } + + void reset_performance_parameters() { index_->reset_performance_parameters(); } + + size_t dimensions() const { return index_->dimensions(); } + + void set_search_parameters(const VamanaSearchParameters& parameters) { + index_->set_search_parameters(parameters); + } + VamanaSearchParameters get_search_parameters() const { + return index_->get_search_parameters(); + } + + void set_construction_window_size(size_t window_size) { + index_->set_construction_window_size(window_size); + } + size_t get_construction_window_size() const { + return index_->get_construction_window_size(); + } + + void set_max_candidates(size_t max_candidate_pool_size) { + index_->set_max_candidates(max_candidate_pool_size); + } + size_t get_max_candidates() const { return index_->get_max_candidates(); } + + void set_prune_to(size_t prune_to) { index_->set_prune_to(prune_to); } + size_t get_prune_to() const { return index_->get_prune_to(); } + + void set_alpha(float alpha) { index_->set_alpha(alpha); } + float get_alpha() const { return index_->get_alpha(); } + + void set_full_search_history(bool use_full_search_history) { + index_->set_full_search_history(use_full_search_history); + } + bool get_full_search_history() const { return index_->get_full_search_history(); } + + size_t max_degree() const { return index_->max_degree(); } + + constexpr std::string_view name() const { return "multi dynamic vamana index"; } + + static constexpr lib::Version save_version = lib::Version(0, 0, 0); + + auto get_labels() const { + // Since data is in order of external ids, + // convert a map of external ids to label types into a sorted vector of labels based + // on external ids. + std::vector> ext_lab_vec( + external_to_label_.begin(), external_to_label_.end() + ); + std::sort(ext_lab_vec.begin(), ext_lab_vec.end(), [](const auto& a, const auto& b) { + return a.first < b.first; + }); + + size_t num_labels = ext_lab_vec.size(); + std::vector labels(num_labels); + std::transform( + ext_lab_vec.begin(), + ext_lab_vec.end(), + labels.begin(), + [](const auto& ext_lab) { return ext_lab.second; } + ); + + return labels; + } + + VamanaIndexParameters get_parameters() const { + return { + index_->entry_point_.front(), + {get_alpha(), + max_degree(), + get_construction_window_size(), + get_max_candidates(), + get_prune_to(), + get_full_search_history()}, + get_search_parameters()}; + } + + void save( + const std::filesystem::path& config_directory, + const std::filesystem::path& graph_directory, + const std::filesystem::path& data_directory + ) { + // Post-consolidation, all entries should be "valid". + // Therefore, we don't need to save the slot metadata. + consolidate(); + compact(); + + auto labels = get_labels(); + size_t num_labels = labels.size(); + + // Save auxiliary data structures. + lib::save_to_disk( + lib::SaveOverride([&](const lib::SaveContext& ctx) { + // Save labels to a file. + auto filename = ctx.generate_name("labels", "binary"); + auto stream = lib::open_write(filename); + lib::write_binary(stream, labels); + + // Save the construction parameters. + auto parameters = get_parameters(); + return lib::SaveTable( + "multi_vamana_dynamic_auxiliary_parameters", + save_version, + {{"name", lib::save(name())}, + {"parameters", lib::save(parameters, ctx)}, + {"num_labels", lib::save(num_labels, ctx)}, + {"filename", lib::save(filename.filename())}} + ); + }), + config_directory + ); + + // Data + lib::save_to_disk(index_->data_, data_directory); + // Graph + lib::save_to_disk(index_->graph_, graph_directory); + } + + void save(std::ostream& os) { + consolidate(); + compact(); + + auto labels = get_labels(); + size_t num_labels = labels.size(); + + lib::begin_serialization(os); + + auto parameters = get_parameters(); + auto save_table = lib::SaveTable( + "multi_vamana_dynamic_auxiliary_parameters", + save_version, + {{"name", lib::save(name())}, + {"parameters", lib::save(parameters)}, + {"num_labels", lib::save(num_labels)}} + ); + lib::save_to_stream(save_table, os); + lib::write_binary(os, labels); + + // Save the dataset. + lib::save_to_stream(index_->data_, os); + // Save the graph. + lib::save_to_stream(index_->graph_, os); + } +}; + +///// Deduction Guides. +// Guide for building. +template +MultiMutableVamanaIndex( + const VamanaBuildParameters&, Data, const ExternalIds&, Dist, size_t +) -> MultiMutableVamanaIndex, Data, Dist>; + +template +MultiMutableVamanaIndex(const VamanaBuildParameters&, Data, const ExternalIds&, Dist, Pool) + -> MultiMutableVamanaIndex, Data, Dist>; + +// Guide with logging +template +MultiMutableVamanaIndex( + const VamanaBuildParameters&, + Data, + const ExternalIds&, + Dist, + Pool, + svs::logging::logger_ptr +) -> MultiMutableVamanaIndex, Data, Dist>; + +enum class MultiMutableVamanaLoad { FROM_MULTI, FROM_DYNAMIC, FROM_STATIC }; + +namespace detail { + +struct MultiVamanaStateLoader { + using label_type = size_t; + ///// Loading + static bool + check_load_compatibility(std::string_view schema, const lib::Version& version) { + // We provide the option to load from a dynamic index. + return VamanaIndexParameters::check_load_compatibility(schema, version) || + (schema == "multi_vamana_dynamic_auxiliary_parameters" && + version == lib::Version(0, 0, 0)); + } + + // Provide compatibility paths for loading dynamic or static vamana datasets. + static MultiVamanaStateLoader load( + const lib::LoadTable& table, + const MultiMutableVamanaLoad load_from, + const size_t assume_datasize + ) { + switch (load_from) { + case MultiMutableVamanaLoad::FROM_MULTI: { + auto num_labels = lib::load_at(table, "num_labels"); + std::vector labels(num_labels); + auto resolved = table.resolve_at("filename"); + auto stream = lib::open_read(resolved); + lib::read_binary(stream, labels); + return MultiVamanaStateLoader{ + SVS_LOAD_MEMBER_AT_(table, parameters), + IDTranslator{}, + std::move(labels)}; + } + case MultiMutableVamanaLoad::FROM_DYNAMIC: + return MultiVamanaStateLoader{ + SVS_LOAD_MEMBER_AT_(table, parameters), + svs::lib::load_at(table, "translation"), + std::vector{}}; + case MultiMutableVamanaLoad::FROM_STATIC: + return MultiVamanaStateLoader{ + lib::load(table), + IDTranslator::Identity(assume_datasize), + std::vector{}}; + default: + throw ANNEXCEPTION("Invalid multi vamana load type"); + } + } + + ///// Members + VamanaIndexParameters parameters_; + IDTranslator translator_; + std::vector labels_; +}; +} // namespace detail + +template < + typename GraphLoader, + typename DataLoader, + typename Distance, + typename ThreadPoolProto> +auto auto_multi_dynamic_assemble( + const std::filesystem::path& config_path, + GraphLoader&& graph_loader, + DataLoader&& data_loader, + Distance distance, + ThreadPoolProto threadpool_proto, + /// This flag provides compatibility paths for directly loading dynamic vamana or static + /// vamana datasets. + MultiMutableVamanaLoad load_from = MultiMutableVamanaLoad::FROM_MULTI, + svs::logging::logger_ptr logger = svs::logging::get() +) { + // Load the dataset + auto threadpool = threads::as_threadpool(std::move(threadpool_proto)); + auto data = svs::detail::dispatch_load(SVS_FWD(data_loader), threadpool); + + // Load the graph. + auto graph = svs::detail::dispatch_load(SVS_FWD(graph_loader), threadpool); + + // Make sure the data and the graph have the same size. + auto datasize = data.size(); + auto graphsize = graph.n_nodes(); + if (datasize != graphsize) { + throw ANNEXCEPTION( + "Reloaded data has {} nodes while the graph has {} nodes!", datasize, graphsize + ); + } + auto [parameters, translator, labels] = + lib::load_from_disk( + config_path, load_from, datasize + ); + + switch (load_from) { + case MultiMutableVamanaLoad::FROM_MULTI: { + if (labels.size() != datasize) { + throw ANNEXCEPTION( + "Labels has {} IDs but should have {}", labels.size(), datasize + ); + } + return MultiMutableVamanaIndex{ + parameters, + std::move(data), + std::move(graph), + std::move(distance), + labels, + std::move(threadpool), + std::move(logger)}; + } + case MultiMutableVamanaLoad::FROM_DYNAMIC: + case MultiMutableVamanaLoad::FROM_STATIC: { + // Make sure that the translator covers all the IDs in the graph and data. + auto translator_size = translator.size(); + if (translator_size != datasize) { + throw ANNEXCEPTION( + "Translator has {} IDs but should have {}", translator_size, datasize + ); + } + + for (size_t i = 0; i < datasize; ++i) { + if (!translator.has_internal(i)) { + throw ANNEXCEPTION("Translator is missing internal id {}", i); + } + } + + return MultiMutableVamanaIndex{ + parameters, + std::move(data), + std::move(graph), + std::move(distance), + std::move(translator), + std::move(threadpool), + std::move(logger)}; + } + default: + throw ANNEXCEPTION("Invalid multi vamana load type"); + } +} + +template < + typename LazyGraphLoader, + typename LazyDataLoader, + typename Distance, + typename ThreadPoolProto> +auto auto_multi_dynamic_assemble( + std::istream& is, + LazyGraphLoader graph_loader, + LazyDataLoader data_loader, + Distance distance, + ThreadPoolProto threadpool_proto, + svs::logging::logger_ptr logger = svs::logging::get() +) { + using label_type = size_t; + + auto table = lib::detail::read_metadata(is); + + auto parameters = lib::load( + table.template cast().at("parameters").template cast() + ); + + auto num_labels = + lib::load(table.template cast().at("num_labels")); + + // Read labels binary data directly from the stream. + std::vector labels(num_labels); + lib::read_binary(is, labels); + + auto data = data_loader(); + auto graph = graph_loader(); + + auto datasize = data.size(); + auto graphsize = graph.n_nodes(); + if (datasize != graphsize) { + throw ANNEXCEPTION( + "Reloaded data has {} nodes while the graph has {} nodes!", datasize, graphsize + ); + } + + if (labels.size() != datasize) { + throw ANNEXCEPTION("Labels has {} IDs but should have {}", labels.size(), datasize); + } + + auto threadpool = threads::as_threadpool(std::move(threadpool_proto)); + return MultiMutableVamanaIndex{ + parameters, + std::move(data), + std::move(graph), + std::move(distance), + labels, + std::move(threadpool), + std::move(logger)}; +} + +} // namespace svs::index::vamana::concurrent diff --git a/include/svs/concurrent/prune.h b/include/svs/concurrent/prune.h new file mode 100644 index 000000000..1c6e33cad --- /dev/null +++ b/include/svs/concurrent/prune.h @@ -0,0 +1,338 @@ +/* + * Copyright 2023 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "svs/concepts/data.h" +#include "svs/concepts/distance.h" +#include "svs/core/distance.h" +#include "svs/lib/neighbor.h" +#include "svs/lib/type_traits.h" + +#include +#include + +namespace svs::index::vamana::concurrent { + +// Dispatch tags for pruning strategies. +struct IterativePruneStrategy {}; +struct ProgressivePruneStrategy {}; +struct LegacyPruneStrategy {}; + +// Default strategy is the iterative strategy. +template struct PruneStrategy; + +// Strategy for L2 +template <> struct PruneStrategy { + using type = ProgressivePruneStrategy; +}; + +// Specialize IP to use the iterative strategy. +template <> struct PruneStrategy { + using type = IterativePruneStrategy; +}; +template <> struct PruneStrategy { + using type = IterativePruneStrategy; +}; + +template using prune_strategy_t = typename PruneStrategy::type; +template constexpr prune_strategy_t prune_strategy() { + return prune_strategy_t(); +} +template +constexpr prune_strategy_t prune_strategy(const T& SVS_UNUSED(dist)) { + return prune_strategy(); +} + +namespace detail { + +template +concept IntegerOrNeighbor = std::integral || svs::NeighborLike; + +template +As construct_as(lib::Type, const N& n) { + return n.id(); +} + +template +As construct_as(lib::Type, const N& n) { + // N.B.: Be sure to use the copy constructor for `As` to preserve any metadata + // attached to `n`. + return As(n); +} + +} // namespace detail + +///// +///// Iterative Prune Strategy +///// + +enum class PruneState : uint8_t { Available, Added, Pruned, Candidate }; + +inline PruneState reenable(PruneState state) { + return (state == PruneState::Candidate) ? PruneState::Available : state; +} + +inline bool excluded(PruneState state) { return state != PruneState::Available; } + +/// +/// @brief Function to prune neighbors using MRNG rule (extended with alpha as in Vamana). +/// +/// @tparam Data The full type of the given dataset. +/// @tparam Dist The distance functor use when comparing vectors. +/// @tparam Neighbors The full neighbor-type of the candidate pool. +/// @tparam I The type of the resulting index for each neighbor. +/// @tparam Alloc Allocator for the result vector. +/// +template < + data::ImmutableMemoryDataset Data, + data::AccessorFor Accessor, + distance::Distance, data::const_value_type_t> Dist, + NeighborLike Neighbors, + detail::IntegerOrNeighbor I, + typename Alloc> +void heuristic_prune_neighbors( + IterativePruneStrategy SVS_UNUSED(dispatch), + size_t max_result_size, + float alpha, + const Data& dataset, + const Accessor& accessor, + Dist& distance_function, + size_t current_node_id, + const std::span& pool, + std::vector& result +) { + auto cmp = distance::comparator(distance_function); + assert(std::is_sorted(pool.begin(), pool.end(), cmp)); + if (pool.empty()) { + return; + } + + result.clear(); + result.reserve(max_result_size); + size_t poolsize = pool.size(); + if (poolsize == 0) { + return; + } + + auto pruned = std::vector(poolsize, PruneState::Available); + // the first round + size_t start = 0; + while (result.size() < max_result_size && start < poolsize) { + auto id = pool[start].id(); + if (excluded(pruned[start]) || id == current_node_id) { + ++start; + continue; + } + pruned[start] = PruneState::Added; + + // Only once we know this item needs to be processed to we retrieve + // the corresponding data and perform preprocessing. + const auto& query = accessor(dataset, id); + distance::maybe_fix_argument(distance_function, query); + result.push_back(detail::construct_as(lib::Type(), pool[start])); + for (size_t t = start + 1; t < poolsize; ++t) { + if (pruned[t] == PruneState::Pruned) { + continue; + } + + const auto& candidate = pool[t]; + auto djk = distance::compute( + distance_function, query, accessor(dataset, candidate.id()) + ); + + if (cmp(djk, candidate.distance())) { + if (cmp(alpha * djk, candidate.distance())) { + pruned[t] = PruneState::Pruned; + } else { + pruned[t] = PruneState::Candidate; + } + } + } + ++start; + } + + // the second round + start = 0; + while (result.size() < max_result_size && start < poolsize) { + auto id = pool[start].id(); + if (pruned[start] != PruneState::Candidate || id == current_node_id) { + ++start; + continue; + } + + // Only once we know this item needs to be processed to we retrieve + // the corresponding data and perform preprocessing. + const auto& query = accessor(dataset, id); + distance::maybe_fix_argument(distance_function, query); + + const auto& candidate = pool[start]; + for (size_t t = 0; t < start; ++t) { + if (pruned[t] != PruneState::Candidate) { + continue; + } + + auto djk = distance::compute( + distance_function, query, accessor(dataset, pool[t].id()) + ); + + if (cmp(alpha * djk, candidate.distance())) { + pruned[start] = PruneState::Pruned; + break; + } + } + + if (pruned[start] == PruneState::Candidate) { + result.push_back(detail::construct_as(lib::Type(), pool[start])); + } + ++start; + } +} + +template < + data::ImmutableMemoryDataset Data, + data::AccessorFor Accessor, + distance::Distance, data::const_value_type_t> Dist, + NeighborLike Neighbors, + detail::IntegerOrNeighbor I, + typename Alloc> +void heuristic_prune_neighbors( + ProgressivePruneStrategy SVS_UNUSED(dispatch), + size_t max_result_size, + float alpha, + const Data& dataset, + const Accessor& accessor, + Dist& distance_function, + size_t current_node_id, + const std::span& pool, + std::vector& result +) { + auto cmp = distance::comparator(distance_function); + assert(std::is_sorted(pool.begin(), pool.end(), cmp)); + if (pool.empty()) { + return; + } + + result.clear(); + result.reserve(max_result_size); + size_t poolsize = pool.size(); + std::vector pruned(poolsize, type_traits::tombstone_v); + + float current_alpha = 1.0f; + while (result.size() < max_result_size && !cmp(alpha, current_alpha)) { + size_t start = 0; + while (result.size() < max_result_size && start < poolsize) { + auto id = pool[start].id(); + if (cmp(current_alpha, pruned[start]) || id == current_node_id) { + ++start; + continue; + } + pruned[start] = type_traits::sentinel_v; + + // Only once we know this item needs to be processed to we retrieve + // the corresponding data and perform preprocessing. + const auto& query = accessor(dataset, id); + distance::maybe_fix_argument(distance_function, query); + result.push_back(detail::construct_as(lib::Type(), pool[start])); + for (size_t t = start + 1; t < poolsize; ++t) { + if (cmp(current_alpha, pruned[t])) { + continue; + } + + const auto& candidate = pool[t]; + auto djk = distance::compute( + distance_function, query, accessor(dataset, candidate.id()) + ); + pruned[t] = std::max(pruned[t], candidate.distance() / djk, cmp); + } + ++start; + } + if (alpha == 1) { + break; + } + current_alpha *= alpha; + } +} + +/// +/// @brief Function to prune neighbors using MRNG rule (extended with alpha as in Vamana). +/// +/// @tparam Data The full type of the given dataset. +/// @tparam Dist The distance functor use when comparing vectors. +/// @tparam Neighbors The full neighbor-type of the candidate pool. +/// @tparam I The type of the reusting index for each neighbor. +/// @tparam Alloc Allocator for the result vector. +/// +template < + data::ImmutableMemoryDataset Data, + data::AccessorFor Accessor, + distance::Distance, data::const_value_type_t> Dist, + NeighborLike Neighbors, + detail::IntegerOrNeighbor I, + typename Alloc> +void heuristic_prune_neighbors( + LegacyPruneStrategy SVS_UNUSED(dispatch), + size_t max_result_size, + float alpha, + const Data& dataset, + const Accessor& accessor, + Dist& distance_function, + size_t current_node_id, + const std::span& pool, + std::vector& result +) { + auto cmp = distance::comparator(distance_function); + assert(std::is_sorted(pool.begin(), pool.end(), cmp)); + if (pool.empty()) { + return; + } + + result.clear(); + result.reserve(max_result_size); + size_t poolsize = pool.size(); + std::vector pruned(poolsize, false); + size_t start = 0; + + while (result.size() < max_result_size && start < poolsize) { + auto id = pool[start].id(); + if (pruned[start] || id == current_node_id) { + ++start; + continue; + } + pruned[start] = true; + + // Only once we know this item needs to be processed to we retrieve + // the corresponding data and perform preprocessing. + const auto& query = accessor(dataset, id); + distance::maybe_fix_argument(distance_function, query); + result.push_back(detail::construct_as(lib::Type(), pool[start])); + for (size_t t = start + 1; t < poolsize; ++t) { + if (pruned[t]) { + continue; + } + + const auto& candidate = pool[t]; + auto djk = distance::compute( + distance_function, query, accessor(dataset, candidate.id()) + ); + if (cmp(alpha * djk, candidate.distance())) { + pruned[t] = true; + } + } + ++start; + } +} +} // namespace svs::index::vamana::concurrent diff --git a/include/svs/concurrent/reverse_edges.h b/include/svs/concurrent/reverse_edges.h new file mode 100644 index 000000000..24c760183 --- /dev/null +++ b/include/svs/concurrent/reverse_edges.h @@ -0,0 +1,99 @@ +/* + * Copyright 2026 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "svs/concurrent/spinlock.h" +#include "svs/lib/segmented_vector.h" + +// external +#include "tsl/robin_set.h" + +#include +#include +#include +#include +#include + +namespace svs::index::vamana::concurrent::graphs { + +/// +/// @brief Per-node index of in-neighbors: `R(n)` is the list of nodes that point at `n`. +/// +/// Stored as one `std::vector` per node, indexed by node id in a grow-stable +/// `SegmentedVector`, with a per-node `SpinLock`. Every operation touches only the target +/// node's list under its own lock. +/// +/// `R(n)` is a complete superset of `n`'s in-neighbors: `record` is called unconditionally +/// on every created edge, so it may hold stale (edge later dropped) or duplicate entries, +/// but never misses a live in-edge. Consolidation reads it to find who points at a deleted +/// node, and prunes stale entries via `remove`/`reset_node`. +/// +template class ReverseEdges { + public: + explicit ReverseEdges(size_t num_nodes) + : lists_(num_nodes) + , locks_(num_nodes) {} + + void resize(size_t new_size) { + lists_.resize(new_size); + locks_.resize(new_size); + } + + void set_recording(bool on) { recording_.store(on, std::memory_order_release); } + + void record(Idx m, Idx n) { + if (!recording_.load(std::memory_order_acquire)) { + return; + } + std::lock_guard lock{locks_[n]}; + lists_[n].push_back(m); + } + + void remove(Idx m, Idx n) { + std::lock_guard lock{locks_[n]}; + auto& list = lists_[n]; + list.erase(std::remove(list.begin(), list.end(), m), list.end()); + } + + template + void collect(Idx n, tsl::robin_set& out, const Deleted& is_deleted) const { + std::lock_guard lock{locks_[n]}; + for (auto m : lists_[n]) { + if (!is_deleted(m)) { + out.insert(m); + } + } + } + + void reset_node(Idx n) { + std::lock_guard lock{locks_[n]}; + lists_[n].clear(); + } + + void reset() { + for (size_t i = 0, imax = lists_.size(); i < imax; ++i) { + lists_[i].clear(); + } + } + + private: + lib::SegmentedVector> lists_; + mutable lib::SegmentedVector locks_; + std::atomic recording_{true}; +}; + +} // namespace svs::index::vamana::concurrent::graphs diff --git a/include/svs/concurrent/spinlock.h b/include/svs/concurrent/spinlock.h new file mode 100644 index 000000000..6d1085c15 --- /dev/null +++ b/include/svs/concurrent/spinlock.h @@ -0,0 +1,49 @@ +/* + * Copyright 2026 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "svs/lib/spinlock.h" + +#include + +namespace svs::index::vamana::concurrent { + +/// +/// @brief ``svs::SpinLock`` plus value-initializing copy/move operations. +/// +/// ``lib::SegmentedVector`` constructs each segment's elements and, on +/// ``resize``, needs the element type to be constructible from an existing element. A +/// ``std::atomic`` member makes ``svs::SpinLock`` neither copyable nor movable, so +/// this subclass supplies those operations. A copied or moved lock is always born +/// *unlocked*: the operations exist to satisfy container requirements, never to transfer +/// ownership of a held lock, and a container is only ever grown while its existing +/// elements are untouched. +/// +class SpinLock : public svs::SpinLock { + public: + SpinLock() = default; + + SpinLock(const SpinLock& /*unused*/) + : svs::SpinLock{} {} + SpinLock& operator=(const SpinLock& /*unused*/) { return *this; } + SpinLock(SpinLock&& /*unused*/) noexcept + : svs::SpinLock{} {} + SpinLock& operator=(SpinLock&& /*unused*/) noexcept { return *this; } + ~SpinLock() = default; +}; + +} // namespace svs::index::vamana::concurrent diff --git a/include/svs/concurrent/translation.h b/include/svs/concurrent/translation.h new file mode 100644 index 000000000..fbca4f410 --- /dev/null +++ b/include/svs/concurrent/translation.h @@ -0,0 +1,491 @@ +/* + * Copyright 2023 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +// Index translation. + +// svs +#include "svs/lib/algorithms.h" +#include "svs/lib/boundscheck.h" +#include "svs/lib/exception.h" +#include "svs/lib/file.h" +#include "svs/lib/misc.h" +#include "svs/lib/narrow.h" +#include "svs/lib/readwrite.h" +#include "svs/lib/saveload.h" +#include "svs/lib/threads.h" + +#include "svs/third-party/fmt.h" + +// tsl +#include "tsl/robin_map.h" + +// stl +#include +#include +#include + +namespace svs::index::vamana::concurrent { + +class IDTranslator { + public: + using internal_id_type = uint32_t; + using external_id_type = uint64_t; + + using const_iterator = + tsl::robin_map::const_iterator; + using value_type = typename const_iterator::value_type; + + // Construct the identity transformation of size `n`. + struct Identity { + Identity() = delete; + Identity(size_t n) + : n_{n} {} + size_t n_; + }; + + IDTranslator() = default; + IDTranslator(Identity tag) { + auto ids = threads::UnitRange{0, tag.n_}; + insert(ids, ids); + } + + /// + /// @brief Return the number of translations. + /// + size_t size() const { + if constexpr (checkbounds_v) { + const size_t e2i = external_to_internal_.size(); + const size_t i2e = internal_to_external_.size(); + + if (e2i != i2e) { + throw ANNEXCEPTION("Size mismatch! E2I is {} while I2E is {}!", e2i, i2e); + } + } + return external_to_internal_.size(); + } + + /// + /// @brief Insert the two ranges. + /// + /// @param external Container implementing forward iteration of the external IDs to add. + /// @param internal Container implementing forward iteration of the internal IDs to add. + /// @param check Check that none of the external and internal ids have an assigned + /// mapping yet **and** both only contain unique elements. Only safe to set to + /// ``false`` if this holds true. + /// + /// If any of the checks associated with the ``check`` parameter fail, the container + /// is left unmodified. + /// + template + void insert(const External& external, const Internal& internal, bool check = true) { + insert(external.begin(), external.end(), internal.begin(), internal.end(), check); + } + + template + void insert( + const ExtBegin& ext_begin, + const ExtEnd& ext_end, + const IntBegin& int_begin, + const IntEnd& int_end, + bool check = true + ) { + // Make sure the two iterators are the same length. + auto external_count = std::distance(ext_begin, ext_end); + auto internal_count = std::distance(int_begin, int_end); + if (external_count != internal_count) { + throw ANNEXCEPTION( + "Length of external IDs is {} while the length of internal IDs is {}!", + external_count, + internal_count + ); + } + + // Perform checks before actually modifying data structures to prevent a partially + // completed operation. + if (check) { + if (!lib::all_unique(ext_begin, ext_end)) { + throw ANNEXCEPTION("External IDs contain repeat elements!"); + } + if (!lib::all_unique(int_begin, int_end)) { + throw ANNEXCEPTION("Internal IDs contain repeat elements!"); + } + check_external_free(ext_begin, ext_end); + check_internal_free(int_begin, int_end); + } + + // Now, we actually begin the insertion. + auto i = int_begin; + auto e = ext_begin; + while (i != int_end || e != ext_end) { + insert_translation(*e, *i); + ++e; + ++i; + } + } + + template + void insert_translation(Ext external_id, Int internal_id) { + external_to_internal_[external_id] = lib::narrow(internal_id); + internal_to_external_[internal_id] = lib::narrow(external_id); + } + + /// + /// @brief Insert mappings, replacing any existing mapping whose current internal + /// ID is considered stale by the caller (e.g. the associated slot has been + /// marked for deletion but the translator entry was not yet cleaned up). + /// + /// Throws if an external ID already maps to a non-stale internal ID. + /// The batch is validated in full before any mutation, so on throw the + /// translation maps are left unmodified. + /// + /// @param external Container of external IDs to insert. + /// @param internal Container of internal IDs to insert (same length as external). + /// @param is_stale Callable ``bool(internal_id_type)``. Returns true if the given + /// internal ID is stale and its existing mapping may be overwritten. + /// + template + void replace_stale_and_insert( + const External& external, const Internal& internal, IsStale&& is_stale + ) { + auto ext_begin = external.begin(); + auto ext_end = external.end(); + auto int_begin = internal.begin(); + auto int_end = internal.end(); + + if (std::distance(ext_begin, ext_end) != std::distance(int_begin, int_end)) { + throw ANNEXCEPTION( + "Length of external IDs is {} while the length of internal IDs is {}!", + std::distance(ext_begin, ext_end), + std::distance(int_begin, int_end) + ); + } + if (!lib::all_unique(ext_begin, ext_end)) { + throw ANNEXCEPTION("External IDs contain repeat elements!"); + } + if (!lib::all_unique(int_begin, int_end)) { + throw ANNEXCEPTION("Internal IDs contain repeat elements!"); + } + + // Validation pass + for (auto e = ext_begin; e != ext_end; ++e) { + auto found = external_to_internal_.find(*e); + if (found != external_to_internal_.end() && !is_stale(found->second)) { + throw ANNEXCEPTION("Index already contains external ID {}!", *e); + } + } + + // Commit pass + auto e = ext_begin; + auto i = int_begin; + for (; e != ext_end; ++e, ++i) { + auto found = external_to_internal_.find(*e); + if (found != external_to_internal_.end()) { + internal_to_external_.erase(found->second); + } + insert_translation(*e, *i); + } + } + + /// + /// @brief Return whether the external ID exists. + /// + /// @param e The external ID to check. + /// + bool has_external(external_id_type e) const { + return external_to_internal_.contains(e); + } + + /// + /// @brief Return whether the internal ID exists. + /// + /// @param e The internal ID to check. + /// + bool has_internal(internal_id_type e) const { + return internal_to_external_.contains(e); + } + + /// + /// @brief Return the internal ID mapped to by the external ID. + /// + /// @param e The external ID to translate to an internal ID. + /// + internal_id_type get_internal(external_id_type e) const { + return external_to_internal_.at(e); + } + + /// + /// @brief Return the external ID mapped to by the internal ID. + /// + /// @param i The internal ID to translate to an external ID. + /// + external_id_type get_external(internal_id_type i) const { + return internal_to_external_.at(i); + } + + /// @brief Return the external ID, or a default if not found. + external_id_type + get_external_or(internal_id_type i, external_id_type default_val) const { + auto it = internal_to_external_.find(i); + if (it == internal_to_external_.end()) { + return default_val; + } + return it->second; + } + + /// @brief Return the internal ID, or a default if not found. + internal_id_type + get_internal_or(external_id_type e, internal_id_type default_val) const { + auto it = external_to_internal_.find(e); + if (it == external_to_internal_.end()) { + return default_val; + } + return it->second; + } + + /// + /// @brief Return a start forward iterator over the external->internal IDs. + /// + const_iterator begin() const { return external_to_internal_.begin(); } + + /// + /// @brief Return an end forward iterator over the external->internal IDs. + /// + const_iterator end() const { return external_to_internal_.end(); } + + /// + /// @brief Remap the internal ID. + /// + /// Assumptions: + /// * `from` exists. + /// * `to` does not exist. + /// + void remap_internal_id(internal_id_type from, internal_id_type to) { + assert(has_internal(from)); + assert(!has_internal(to)); + + auto itr = internal_to_external_.find(from); + auto external = itr->second; + + // Updating the internal-to-external ID is easy. + external_to_internal_[external] = to; + internal_to_external_.erase(itr); + internal_to_external_.insert({to, external}); + } + + /// + /// @brief Delete entries from internal IDs. + /// + /// @param internal_ids A container with the internal ids to delete. Must implement a + /// forward iterator. Furthermore, all entries must be unique (not checked). + /// @param check Check if all ids actually exist in the translation table. This is only + /// safe to set to ``false`` if it can be guaranteed that all entries in + /// ``internal_ids`` exist in the table. + /// + /// Note, if ``check == true`` and an internal id is found not to exist, the underlying + /// translation tables will not be modified. + /// + template + void delete_internal(const Internal& internal_ids, bool check = true) { + // First pass - ensure that all IDs to remove actually exist. + if (check) { + check_internal_exist(internal_ids.begin(), internal_ids.end()); + } + + for (const auto i : internal_ids) { + auto e = get_external(i); + internal_to_external_.erase(i); + external_to_internal_.erase(e); + } + } + + /// + /// @brief Delete entries from external IDs. + /// + /// @param external_ids A container with the external ids to delete. Must implement a + /// forward iterator. Furthermore, all entries must be unique (not checked). + /// @param check Check if all ids actually exist in the translation table. This is only + /// safe to set to ``false`` if it can be guaranteed that all entries in + /// ``external_ids`` exist in the table. + /// + /// Note, if ``check == true`` and an external id is found not to exist, the underlying + /// translation tables will not be modified. + /// + template + void delete_external(const External& external_ids, bool check = true) { + // First pass - ensure that all IDs to remove actually exist. + if (check) { + check_external_exist(external_ids.begin(), external_ids.end()); + } + + for (const auto e : external_ids) { + auto i = get_internal(e); + internal_to_external_.erase(i); + external_to_internal_.erase(e); + } + } + + /// + /// @brief Ensure that **none** of the external ids in the iterator exist yet. + /// + /// @param begin Start forward iterator to the external IDs to check. + /// @param end End forward iterator to the external IDs to check. + /// + template + void check_external_free(const Begin& begin, const End& end) const { + check(begin, end, external_to_internal_, "Index already contains external"); + } + + /// + /// @brief Ensure that **all** of the external ids in the iterator exist. + /// + /// @param begin Start forward iterator to the external IDs to check. + /// @param end End forward iterator to the external IDs to check. + /// + template + void check_external_exist(const Begin& begin, const End& end) const { + check( + begin, + end, + external_to_internal_, + "Index does not contain external", + std::logical_not() + ); + } + + /// + /// @brief Ensure that **none** of the internal ids in the iterator exist yet. + /// + /// @param begin Start forward iterator to the internal IDs to check. + /// @param end End forward iterator to the internal IDs to check. + /// + template + void check_internal_free(const Begin& begin, const End& end) const { + check(begin, end, internal_to_external_, "Index already contains internal"); + } + + /// + /// @brief Ensure that **all** of the internal ids in the iterator exist. + /// + /// @param begin Start forward iterator to the internal IDs to check. + /// @param end End forward iterator to the internal IDs to check. + /// + template + void check_internal_exist(const Begin& begin, const End& end) const { + check( + begin, + end, + internal_to_external_, + "Index does not contain internal", + std::logical_not() + ); + } + + ///// Saving and Loading + static constexpr std::string_view kind = "external to internal id translation"; + static constexpr std::string_view serialization_schema = + "external_to_internal_translation"; + static constexpr lib::Version save_version = lib::Version(0, 0, 0); + + lib::SaveTable metadata() const { + return lib::SaveTable( + serialization_schema, + save_version, + {{"kind", kind}, + {"num_points", lib::save(size())}, + {"external_id_type", lib::save(datatype_v)}, + {"internal_id_type", lib::save(datatype_v)}} + ); + } + + void save(std::ostream& os) const { + for (auto i = begin(), iend = end(); i != iend; ++i) { + // N.B.: Apparently `std::pair` of integers is not trivially copyable ... + lib::write_binary(os, i->first); + lib::write_binary(os, i->second); + } + } + + lib::SaveTable save(const lib::SaveContext& ctx) const { + auto filename = ctx.generate_name("id_translation", "binary"); + // Save the translations to a file. + auto os = lib::open_write(filename); + save(os); + auto table = metadata(); + table.insert("filename", lib::save(filename.filename())); + return table; + } + + static void validate(const lib::ContextFreeLoadTable& table) { + if (kind != lib::load_at(table, "kind")) { + throw ANNEXCEPTION("Mismatched kind!"); + } + + constexpr std::string_view external_id_name = name>(); + constexpr std::string_view internal_id_name = name>(); + if (external_id_name != lib::load_at(table, "external_id_type")) { + throw ANNEXCEPTION("Mismatched external id types!"); + } + if (internal_id_name != lib::load_at(table, "internal_id_type")) { + throw ANNEXCEPTION("Mismatched internal id types!"); + } + } + + static IDTranslator load(const lib::ContextFreeLoadTable& table, std::istream& is) { + IDTranslator::validate(table); + auto num_points = lib::load_at(table, "num_points"); + + auto translator = IDTranslator{}; + for (size_t i = 0; i < num_points; ++i) { + auto external_id = lib::read_binary(is); + auto internal_id = lib::read_binary(is); + translator.insert_translation(external_id, internal_id); + } + return translator; + } + + static IDTranslator load(const lib::LoadTable& table) { + IDTranslator::validate(table); + + // Now that we've more-or-less validated the metadata, time to start loading + // the points. + auto resolved = table.resolve_at("filename"); + auto is = lib::open_read(resolved); + return IDTranslator::load(table, is); + } + + private: + template + void check( + const Begin& begin, + const End& end, + const Map& map, + const char* message, + Modifier modify = lib::identity() + ) const { + for (auto i = begin; i != end; ++i) { + if (modify(map.contains(*i))) { + throw ANNEXCEPTION("{} ID {}!", message, *i); + } + } + } + + tsl::robin_map external_to_internal_{}; + tsl::robin_map internal_to_external_{}; +}; + +} // namespace svs::index::vamana::concurrent diff --git a/include/svs/concurrent/vamana_build.h b/include/svs/concurrent/vamana_build.h new file mode 100644 index 000000000..1ca380507 --- /dev/null +++ b/include/svs/concurrent/vamana_build.h @@ -0,0 +1,605 @@ +/* + * Copyright 2023 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +// local +#include "svs/concepts/data.h" +#include "svs/concepts/distance.h" +#include "svs/core/logging.h" +#include "svs/index/vamana/build_params.h" +#include "svs/index/vamana/extensions.h" +#include "svs/concurrent/greedy_search.h" +#include "svs/concurrent/prune.h" +#include "svs/index/vamana/search_buffer.h" +#include "svs/index/vamana/search_tracker.h" +#include "svs/lib/boundscheck.h" +#include "svs/lib/exception.h" +#include "svs/lib/narrow.h" +#include "svs/lib/neighbor.h" +#include "svs/concurrent/spinlock.h" +#include "svs/lib/threads/threadlocal.h" +#include "svs/lib/threads/threadpool.h" +#include "svs/lib/timing.h" +#include "svs/third-party/fmt.h" + +// external +#include "tsl/robin_map.h" +#include "tsl/robin_set.h" + +// stdlib +#include +#include +#include +#include +#include +#include + +namespace svs::index::vamana::concurrent { + +// Optional search tracker to get full history of graph search. +template class OptionalTracker { + public: + using set_type = tsl::robin_set, IDHash, IDEqual>; + using const_iterator = typename set_type::const_iterator; + + private: + std::optional neighbors_; + + public: + ///// Constructors + OptionalTracker(bool enable) + : neighbors_{std::nullopt} { + if (enable) { + neighbors_.emplace(); + } + } + + ///// Methods + bool enabled() const { return neighbors_.has_value(); } + size_t size() const { return enabled() ? (*neighbors_).size() : 0; } + + const_iterator begin() const { return neighbors_.value().begin(); } + const_iterator end() const { return neighbors_.value().end(); } + + void clear() { + // This method is safe to call even if the tracker isn't being used. + if (enabled()) { + (*neighbors_).clear(); + } + } + + ///// Search Tracker API + void visited(const Neighbor& neighbor, size_t SVS_UNUSED(distance_computations)) { + if (enabled()) { + (*neighbors_).insert(neighbor); + } + } +}; + +// Define an auxiliary struct to disambiguate constructor calls. +struct BackedgeBufferParameters { + size_t bucket_size_; + size_t num_buckets_; +}; + +/// +/// @brief A helper type for managing synchronization and parallelism of backedges +/// +/// The big idea is to use locking over coarse regions of indices. +/// This still provides synchronized access to individual entries, but allows parallelized +/// access to multiple buckets. +/// +template class BackedgeBuffer { + public: + // Map an vertex to it's expanded adjacency list. + using set_type = tsl::robin_set; + using map_type = tsl::robin_map; + + private: + // The number of elements assigned to each bucket - starting sequentially from zero. + // Used to determine which bucket an index belongs to. + size_t bucket_size_; + std::vector buckets_; + std::vector bucket_locks_; + + public: + ///// Constructors + BackedgeBuffer(BackedgeBufferParameters parameters) + : bucket_size_{parameters.bucket_size_} + , buckets_(parameters.num_buckets_) + , bucket_locks_{parameters.num_buckets_} {} + + BackedgeBuffer(size_t num_elements, size_t bucket_size) + : BackedgeBuffer(BackedgeBufferParameters{ + bucket_size, lib::div_round_up(num_elements, bucket_size)}) {} + + // Add a point. + void add_edge(Idx src, Idx dst) { + // Get the bucket that the source vertex belongs to. + size_t bucket = src / bucket_size_; + // The bucket array is sized once from the graph size snapshotted at + // VamanaBuilder construction. A concurrent add_points may have grown the + // graph past that snapshot, so a greedy-search neighbor `src` can fall + // beyond our bucket range. That node belongs to another in-flight add and + // gets its own back-edges from that add's builder; drop the overflow edge + // here rather than indexing out of bounds — consistent with the + // best-effort dropping already done on AddEdgeResult::Full. + if (bucket >= buckets_.size()) { + return; + } + // Lock the bucket and update the adjacency list. + std::lock_guard lock(bucket_locks_.at(bucket)); + + // The "try_emplace" method will default construct the set if it doesn't exist. + // Whether or not the set existed to begin with, we get an iterator to the set + // which we can then add the destination to. + auto& map = buckets_.at(bucket); + auto [iterator, _] = map.try_emplace(src); + iterator.value().insert(dst); + } + + // Return the underlying buckets directly. + // Buckets can be iterated over to add back edges. + std::vector& buckets() { return buckets_; } + + // Return the number of buckets in the buffer + size_t num_buckets() const { + assert(buckets_.size() == bucket_locks_.size()); + return buckets_.size(); + } + + // Reset the container for another iteration. + void reset() { + for (size_t i = 0, imax = num_buckets(); i < imax; ++i) { + std::lock_guard lock(bucket_locks_.at(i)); + buckets_.at(i).clear(); + } + } +}; + +template < + graphs::MemoryGraph Graph, + data::ImmutableMemoryDataset Data, + typename Dist, + threads::ThreadPool Pool> +class VamanaBuilder { + public: + // Type Aliases + using Idx = typename Graph::index_type; + using search_buffer_type = SearchBuffer>; + + template using set_type = tsl::robin_set; + + using update_type = + threads::SequentialTLS>>>; + + /// Constructor + VamanaBuilder( + Graph& graph, + const Data& data, + Dist distance_function, + const VamanaBuildParameters& params, + Pool& threadpool, + GreedySearchPrefetchParameters prefetch_hint = {}, + svs::logging::logger_ptr logger = svs::logging::get(), + logging::Level level = logging::Level::Debug + ) + : graph_{graph} + , data_{data} + , distance_function_{std::move(distance_function)} + , params_{params} + , prefetch_hint_{prefetch_hint} + , threadpool_{threadpool} + , backedge_buffer_{data.size(), 1000} { + // Print all parameters + svs::logging::log( + logger, + level, + "Vamana Build Parameters: alpha={}, graph_max_degree={}, " + "max_candidate_pool_size={}, prune_to={}, window_size={}, " + "use_full_search_history={}", + params.alpha, + params.graph_max_degree, + params.max_candidate_pool_size, + params.prune_to, + params.window_size, + params.use_full_search_history + ); + // Note: graph/data size invariant (graph_.n_nodes() == data_.size()) is + // maintained under mutation_mutex_ in add_points(). During concurrent + // add_points() calls, sizes may temporarily differ between lock release + // and this point, but both are always >= the slots we will operate on. + } + + void construct( + float alpha, + Idx entry_point, + logging::Level level = logging::Level::Trace, + logging::logger_ptr logger = svs::logging::get() + ) { + construct( + alpha, entry_point, threads::UnitRange{0, data_.size()}, level, logger + ); + } + + template + void construct( + float alpha, + Idx entry_point, + const R& range, + logging::Level level = logging::Level::Trace, + logging::logger_ptr logger = svs::logging::get() + ) { + size_t num_nodes = range.size(); + size_t num_batches = std::max( + size_t{40}, lib::div_round_up(num_nodes, lib::narrow_cast(64 * 64)) + ); + + // If num_batches is greater than num_nodes, set + // num_batches to num_nodes to avoid unnecessary iterations. + if (num_batches > num_nodes) { + num_batches = num_nodes; + } + + size_t batchsize = lib::div_round_up(num_nodes, num_batches); + std::vector entry_points{entry_point}; + + // Runtime variables + double search_time = 0; + double reverse_time = 0; + unsigned progress_counter = 0; + + svs::logging::log(logger, level, "Number of syncs: {}", num_batches); + svs::logging::log(logger, level, "Batch Size: {}", batchsize); + + // The base point for iteration. + auto&& base = range.begin(); + auto timer = lib::Timer(); + for (size_t batch_id = 0; batch_id < num_batches; ++batch_id) { + // Set up batch parameters + auto start = std::min(num_nodes, batchsize * batch_id) + base; + auto stop = std::min(num_nodes, batchsize * (batch_id + 1)) + base; + + // Perform search. + // N.B. - We purposely pass "params_.alpha" instead of the external "alpha" + // because it seems to generally yield better results. + auto x = timer.push_back("generate neighbors"); + generate_neighbors( + threads::IteratorPair{start, stop}, params_.alpha, entry_points, timer + ); + search_time += lib::as_seconds(x.finish()); + + auto y = timer.push_back("reverse edges"); + add_reverse_edges(threads::IteratorPair{start, stop}, alpha, timer); + reverse_time += lib::as_seconds(y.finish()); + + auto this_progress = lib::narrow_cast(batch_id) * 1e2 / + lib::narrow_cast(num_batches); + if (this_progress > progress_counter && batch_id > 0) { + auto total_elapsed_time = lib::as_seconds(timer.elapsed()); + auto num_batches_f = lib::narrow_cast(num_batches); + auto batch_id_f = lib::narrow_cast(batch_id); + + double estimated_remaining_time = + total_elapsed_time * (num_batches_f / batch_id_f - 1); + constexpr std::string_view message = "Completed round {} of {}. " + "Search Time: {:.4}s, " + "Reverse Time: {:.4}s, " + "Total Time: {:.4}s, " + "Estimated Remaining Time: {:.4}s"; + + svs::logging::log( + logger, + level, + message, + batch_id + 1, + num_batches, + search_time, + reverse_time, + total_elapsed_time, + estimated_remaining_time + ); + search_time = 0; + reverse_time = 0; + progress_counter += 1; + } + } + svs::logging::log( + logger, level, "Completed pass using window size {}.", params_.window_size + ); + svs::logging::log(logger, level, "{}", timer); + } + + /// + /// Generate Adjacency lists for new collection of nodes. + /// As far as the algorithm is concerned, this implements the search and heuristic + /// pruning for the vertices. + /// + /// Addition of back edges is saved for another step. + /// + template + void generate_neighbors( + const R& indices, + float alpha, + const std::vector& entry_points, + lib::Timer& timer + ) { + auto range = threads::StaticPartition{indices}; + + update_type updates{threadpool_.size()}; + auto main = timer.push_back("main"); + threads::parallel_for( + threadpool_, + range, + [&](const auto& local_indices, uint64_t tid) { + // Thread local variables + auto& thread_local_updates = updates.at(tid); + + // Scratch space. + std::vector> pool{}; + auto search_buffer = search_buffer_type{params_.window_size}; + + // Enable use of the visited filter of the search buffer. + // It seems to help in high-window-size scenarios. + search_buffer.enable_visited_set(); + set_type visited{}; + auto tracker = OptionalTracker(params_.use_full_search_history); + + // Unpack adaptor. + auto build_adaptor = extensions::build_adaptor(data_, distance_function_); + auto&& graph_search_distance = build_adaptor.graph_search_distance(); + auto&& general_distance = build_adaptor.general_distance(); + auto general_accessor = build_adaptor.general_accessor(); + + for (auto node_id : local_indices) { + pool.clear(); + search_buffer.clear(); + visited.clear(); + tracker.clear(); + + const auto& graph_search_query = + build_adaptor.access_query_for_graph_search(data_, node_id); + + // Perform the greedy search. + // The search tracker will be used if it is enabled. + { + auto accessor = build_adaptor.graph_search_accessor(); + concurrent::greedy_search( + graph_, + data_, + accessor, + graph_search_query, + graph_search_distance, + search_buffer, + vamana::EntryPointInitializer{lib::as_const_span(entry_points)}, + NeighborBuilder(), + tracker, + prefetch_hint_ + ); + } + + const auto& post_search_query = build_adaptor.modify_post_search_query( + data_, node_id, graph_search_query + ); + + // If the query and distance functors are sufficiently different for the + // graph search and the general case, then we *may* need to reapply fix + // argument before we can do any further distance computations. + // + // Decide whether we need to make this call. + if constexpr (decltype(build_adaptor)::refix_argument_after_search) { + distance::maybe_fix_argument(general_distance, post_search_query); + } + + auto modify_distance = [&](NeighborLike auto const& n) { + return build_adaptor.post_search_modify( + data_, general_distance, post_search_query, n + ); + }; + + // If the full search history is to be used, then use the tracker to + // populate the candidate pool. + // + // Otherwise, pull results directly out of the search buffer. + if (tracker.enabled()) { + for (const auto& neighbor : tracker) { + pool.push_back(modify_distance(neighbor)); + visited.insert(neighbor.id()); + } + } else { + for (size_t i = 0, imax = search_buffer.size(); i < imax; ++i) { + const auto& neighbor = search_buffer[i]; + pool.push_back(modify_distance(neighbor)); + visited.insert(neighbor.id()); + } + } + + // Add neighbors of the query that are not part of `visited`. + for (auto id : graph_.get_node(node_id)) { + assert(id != node_id); + // Try to emplace the node id into the visited set. + // If the id was inserted, then it didn't already exist in the + // visited set and we need to add it to the candidate pool. + auto [_, inserted] = visited.emplace(id); + if (inserted) { + pool.emplace_back( + id, + distance::compute( + general_distance, + post_search_query, + general_accessor(data_, id) + ) + ); + } + } + + std::sort( + pool.begin(), + pool.end(), + TotalOrder(distance::comparator(general_distance)) + ); + pool.resize(std::min(pool.size(), params_.max_candidate_pool_size)); + + // Prune and wait for an update. + thread_local_updates.emplace_back(node_id, std::vector{}); + auto& pruned_results = thread_local_updates.back().second; + heuristic_prune_neighbors( + prune_strategy(distance_function_), + params_.graph_max_degree, + alpha, + data_, + general_accessor, + general_distance, + node_id, + lib::as_const_span(pool), + pruned_results + ); + } + } + ); + + main.finish(); + + // Apply updates. + auto update = timer.push_back("updates"); + threads::parallel_for(threadpool_, [&](uint64_t tid) { + const auto& thread_local_updates = updates.at(tid); + for (auto [node_id, update] : thread_local_updates) { + graph_.replace_node(node_id, update); + } + }); + } + + /// + /// Add reverse edges to the graph. + /// + template + void add_reverse_edges(const R& indices, float alpha, lib::Timer& timer) { + // Apply backedges to all new candidate adjacency lists. + // If adding an edge to the graph will cause it to violate the maximum degree + // constraint, save the excess to the backedge buffer. + auto backedge_timer = timer.push_back("backedge generation"); + auto range = threads::StaticPartition{indices}; + backedge_buffer_.reset(); + threads::parallel_for( + threadpool_, + range, + [&](const auto& is, uint64_t SVS_UNUSED(tid)) { + for (auto node_id : is) { + for (auto other_id : graph_.get_node(node_id)) { + // graph_.add_edge is atomic under node_locks_[other_id]. + // If it reports Full, route to the overflow buffer — + // no TOCTOU race between a pre-check and the insert. + if (graph_.add_edge(other_id, node_id) == + graphs::AddEdgeResult::Full) { + backedge_buffer_.add_edge(other_id, node_id); + } + } + } + } + ); + backedge_timer.finish(); + + // For all vertices that now exceed the max degree requirement, run the pruning + // procedure on the union of their current adjacency list as well as any extra edges + // that were recorded in the previous process. + // + // Take care to avoid duplicate entries. + auto prune_timer = timer.push_back("pruning backedges"); + threads::parallel_for( + threadpool_, + threads::DynamicPartition{backedge_buffer_.buckets(), 1}, + [&](auto& buckets, uint64_t SVS_UNUSED(tid)) { + // Thread local auxiliary data structures. + std::vector> candidates{}; + std::vector pruned_results{}; + auto build_adaptor = extensions::build_adaptor(data_, distance_function_); + + auto general_accessor = build_adaptor.general_accessor(); + auto&& general_distance = build_adaptor.general_distance(); + + auto cmp = distance::comparator(general_distance); + for (auto& bucket : buckets) { + for (const auto& kv : bucket) { + // The ``neighbors`` class is a set. + auto src = kv.first; + const auto& neighbors = kv.second; + const auto& src_data = general_accessor(data_, src); + distance::maybe_fix_argument(general_distance, src_data); + + // Helper lambda to make distance computations look a little + // cleaner. + auto make_neighbor = [&](auto i) { + return Neighbor{ + i, + distance::compute( + general_distance, src_data, general_accessor(data_, i) + )}; + }; + + candidates.clear(); + // Add the overflow candidates. + for (auto n : neighbors) { + candidates.push_back(make_neighbor(n)); + } + + // Add the old adjacency list. + for (auto n : graph_.get_node(src)) { + if (!neighbors.contains(n)) { + candidates.push_back(make_neighbor(n)); + } + } + std::sort(candidates.begin(), candidates.end(), TotalOrder(cmp)); + candidates.resize( + std::min(candidates.size(), params_.max_candidate_pool_size) + ); + + heuristic_prune_neighbors( + prune_strategy(distance_function_), + params_.prune_to, + alpha, + data_, + general_accessor, + general_distance, + src, + lib::as_const_span(candidates), + pruned_results + ); + graph_.replace_node(src, pruned_results); + } + } + } + ); + } + + private: + /// The graph being constructed. + Graph& graph_; + /// The dataset we're building the graph over. + const Data& data_; + /// The distance function to use. + Dist distance_function_; + /// Parameters regarding index construction. + VamanaBuildParameters params_; + /// Prefetch parameters to use during the graph search. + GreedySearchPrefetchParameters prefetch_hint_; + /// Worker threadpool. + Pool& threadpool_; + /// Overflow backedge buffer. + BackedgeBuffer backedge_buffer_; +}; +} // namespace svs::index::vamana::concurrent diff --git a/include/svs/lib/concurrency/atomic_span.h b/include/svs/lib/concurrency/atomic_span.h new file mode 100644 index 000000000..4d7d4ae50 --- /dev/null +++ b/include/svs/lib/concurrency/atomic_span.h @@ -0,0 +1,93 @@ +/* + * Copyright 2026 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include + +namespace svs { + +/// +/// @brief A non-owning, zero-copy view over a contiguous range of ``T`` that performs +/// atomic loads on every element access. +/// +/// Each dereference uses ``std::atomic_ref::load(std::memory_order_relaxed)``. +/// On x86, this compiles to a plain MOV instruction — identical to non-atomic access. +/// +/// This type is designed to be used as a drop-in replacement for ``std::span`` +/// when concurrent reads and writes are possible, ensuring no undefined behavior +/// while maintaining zero-copy semantics. +/// +template class AtomicSpan { + public: + using value_type = std::remove_const_t; + + class iterator { + public: + using value_type = AtomicSpan::value_type; + using difference_type = std::ptrdiff_t; + using iterator_category = std::input_iterator_tag; + + explicit iterator(const T* p) + : ptr_(p) {} + + value_type operator*() const { + return std::atomic_ref(const_cast(*ptr_)) + .load(std::memory_order_relaxed); + } + + iterator& operator++() { + ++ptr_; + return *this; + } + + iterator operator++(int) { + auto tmp = *this; + ++ptr_; + return tmp; + } + + bool operator==(const iterator& other) const { return ptr_ == other.ptr_; } + bool operator!=(const iterator& other) const { return ptr_ != other.ptr_; } + + private: + const T* ptr_; + }; + + AtomicSpan(const T* data, size_t size) + : data_(data) + , size_(size) {} + + size_t size() const { return size_; } + bool empty() const { return size_ == 0; } + const T* data() const { return data_; } + + value_type operator[](size_t i) const { + return std::atomic_ref(const_cast(data_[i])) + .load(std::memory_order_relaxed); + } + + iterator begin() const { return iterator{data_}; } + iterator end() const { return iterator{data_ + size_}; } + + private: + const T* data_; + size_t size_; +}; + +} // namespace svs diff --git a/include/svs/lib/concurrency/seqlock.h b/include/svs/lib/concurrency/seqlock.h new file mode 100644 index 000000000..7fe3ea70e --- /dev/null +++ b/include/svs/lib/concurrency/seqlock.h @@ -0,0 +1,141 @@ +/* + * Copyright 2026 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "svs/lib/segmented_vector.h" + +#include +#include +#include +#include +#include + +namespace svs { + +/// +/// @brief Per-element sequence lock counter for reader-writer synchronization. +/// +/// Uses a uint8_t counter: odd values indicate a write in progress, even values indicate +/// a stable state. +/// +/// **Writer-writer serialization is the caller's responsibility.** Only one writer +/// at a time may call ``begin_write``/``end_write`` on a given counter. Use an external +/// lock (e.g., per-node ``SpinLock``) to serialize concurrent writers to the same element. +/// +class SeqLockCounter { + using counter_type = uint8_t; + + public: + SeqLockCounter() = default; + + SeqLockCounter(const SeqLockCounter& other) + : seq_(other.seq_.load(std::memory_order_relaxed)) {} + + SeqLockCounter& operator=(const SeqLockCounter& other) { + seq_.store(other.seq_.load(std::memory_order_relaxed), std::memory_order_relaxed); + return *this; + } + + SeqLockCounter(SeqLockCounter&& other) noexcept + : seq_(other.seq_.load(std::memory_order_relaxed)) {} + + SeqLockCounter& operator=(SeqLockCounter&& other) noexcept { + seq_.store(other.seq_.load(std::memory_order_relaxed), std::memory_order_relaxed); + return *this; + } + + /// + /// @brief Begin a write operation. Returns the pre-write sequence value. + /// + /// Increments the counter to an odd value, signaling to readers that a write is in + /// progress. The returned value must be passed to ``end_write``. + /// + counter_type begin_write() { + auto seq = seq_.load(std::memory_order_relaxed); + seq_.store(seq + 1, std::memory_order_relaxed); + std::atomic_thread_fence(std::memory_order_release); + return seq; + } + + /// + /// @brief End a write operation. + /// + /// @param seq The value returned by the corresponding ``begin_write`` call. + /// + /// Increments the counter to an even value, signaling that the write is complete + /// and data is consistent. + /// + void end_write(counter_type seq) { seq_.store(seq + 2, std::memory_order_release); } + + /// + /// @brief Begin a read operation. + /// + /// @returns The current sequence value if it is even (no write in progress), + /// or ``std::nullopt`` if a write is in progress. + /// + /// The returned value (if present) must be passed to ``read_validate`` after the + /// read is complete. + /// + std::optional read_begin() const { + auto seq = seq_.load(std::memory_order_acquire); + if (seq % 2 > 0) { + return std::nullopt; + } + return seq; + } + + /// + /// @brief Validate that no write occurred during the read. + /// + /// @param seq The value returned by ``read_begin``. + /// + /// @returns ``true`` if the data read between ``read_begin`` and ``read_validate`` + /// is consistent (no concurrent write occurred). + /// + bool read_validate(counter_type seq) const { + std::atomic_thread_fence(std::memory_order_acquire); + return seq_.load(std::memory_order_relaxed) == seq; + } + + private: + std::atomic seq_{0}; +}; + +/// +/// @brief Array of SeqLock counters, one per element (e.g., one per graph node). +/// +class SeqLockArray { + public: + SeqLockArray() = default; + explicit SeqLockArray(size_t n) + : counters_(n) {} + + SeqLockCounter& operator[](size_t i) { return counters_[i]; } + const SeqLockCounter& operator[](size_t i) const { return counters_[i]; } + + void resize(size_t n) { counters_.resize(n); } + size_t size() const { return counters_.size(); } + size_t capacity() const { return counters_.capacity(); } + + private: + // Grow-stable storage: appending counters never relocates existing ones, so a + // concurrent lock-free reader (greedy_search reading seq_counters_[i]) is safe + // against a writer's grow. See svs/lib/segmented_vector.h. + lib::SegmentedVector counters_; +}; + +} // namespace svs diff --git a/include/svs/lib/segmented_vector.h b/include/svs/lib/segmented_vector.h new file mode 100644 index 000000000..7f642ae9c --- /dev/null +++ b/include/svs/lib/segmented_vector.h @@ -0,0 +1,330 @@ +/* + * Copyright 2026 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "svs/lib/boundscheck.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace svs::lib { + +/// +/// @brief An unbounded, grow-stable vector for single-writer/many-reader use. +/// +/// Two-level "lock-free dynamic array" (Dechev et al.). A fixed top-level directory of +/// ``kDirBuckets`` bucket pointers; directory bucket ``k`` is a single contiguous +/// heap array of ``kFirstBucket << k`` elements (the first bucket holds ``kFirstBucket``, +/// each subsequent bucket doubles). Grouping elements into chunks of ``kFirstBucket`` and +/// applying the power-of-two layout to the chunk index gives: ``q = i / kFirstBucket``, +/// ``bucket = floor(log2(q + 1))``, with the bucket's first global index +/// ``kFirstBucket * (2^bucket - 1)``. Sixty-four buckets address far more than any real +/// dataset, so there is no practical size cap. +/// +/// The directory array is a fixed member (never relocates), and each bucket array is +/// allocated once and never moved or reallocated. Therefore the address of any element +/// ``i < size()`` is stable for the lifetime of that element — a concurrent reader +/// indexing element ``i`` is unaffected by appends that grow the structure past ``i``. +/// +/// Concurrency contract: +/// * **One writer at a time.** ``resize`` / ``push_back`` / ``pop_back`` / ``shrink_to`` +/// (the only operations that change the structure) must be serialized by the caller +/// (e.g. under a mutex). +/// * **Many concurrent readers.** ``operator[]`` and ``size`` may run concurrently with a +/// writer's *grow*: a new bucket is allocated and its elements constructed, the bucket +/// pointer is published with a release store, and ``size_`` is bumped last (release). A +/// reader does an acquire load of ``size_`` then an acquire load of the bucket pointer, +/// so for any ``i < size()`` it observes, the bucket and element are fully published. +/// * **Shrink frees storage.** ``shrink_to`` destroys trailing elements and frees buckets +/// that lie entirely above the new size; a reader holding a reference to a freed element +/// would dangle, so the caller must drain readers (e.g. via an exclusive lock) first. +/// +/// This mirrors the std::vector subset used by the dynamic Vamana index: ``operator[]``, +/// ``at``, ``size``, ``empty``, ``capacity``, ``resize(n)``, ``resize(n, fill)``, +/// ``push_back``, ``pop_back``, ``shrink_to(n)``. +/// +template class SegmentedVector { + static constexpr std::size_t kDirBuckets = 64; + // Number of elements in the first directory bucket. Bucket ``k`` then holds + // ``kFirstBucket << k`` elements, so a small first bucket means many tiny + // allocations near the start while a large one front-loads capacity. + static constexpr std::size_t kFirstBucket = 1; + static_assert( + (kFirstBucket & (kFirstBucket - 1)) == 0, "kFirstBucket must be a power of two" + ); + + public: + using value_type = T; + using size_type = std::size_t; + using reference = T&; + using const_reference = const T&; + + SegmentedVector() = default; + explicit SegmentedVector(size_type n) { resize(n); } + SegmentedVector(size_type n, const T& fill) { resize(n, fill); } + + SegmentedVector(const SegmentedVector& other) { copy_from_(other); } + SegmentedVector& operator=(const SegmentedVector& other) { + if (this != &other) { + destroy_all_(); + copy_from_(other); + } + return *this; + } + + SegmentedVector(SegmentedVector&& other) noexcept { steal_from_(other); } + SegmentedVector& operator=(SegmentedVector&& other) noexcept { + if (this != &other) { + destroy_all_(); + steal_from_(other); + } + return *this; + } + + ~SegmentedVector() { destroy_all_(); } + + /// + /// @brief Access element ``i``. Precondition: ``i < size()``. + /// + /// Safe to call concurrently with a writer's grow ``resize``/``push_back``, provided + /// ``i`` was ``< size()`` as observed by the reader. + /// + const_reference operator[](size_type i) const noexcept { + auto [b, off] = locate_(i); + return dir_[b].load(std::memory_order_acquire)[off]; + } + reference operator[](size_type i) noexcept { + const auto& self = *this; + return const_cast(self[i]); + } + + /// @brief Bounds-checked access (used by ``svs::getindex`` when bounds checking is on). + reference at(size_type i) { + if (i >= size()) { + throw std::out_of_range("SegmentedVector::at index out of range"); + } + return (*this)[i]; + } + const_reference at(size_type i) const { + if (i >= size()) { + throw std::out_of_range("SegmentedVector::at index out of range"); + } + return (*this)[i]; + } + + size_type size() const noexcept { return size_.load(std::memory_order_acquire); } + bool empty() const noexcept { return size() == 0; } + + /// @brief Logical capacity: number of elements addressable without allocating a new + /// bucket. With ``m`` buckets this is ``kFirstBucket * (2^m - 1)``. + size_type capacity() const noexcept { return bucket_first_index_(allocated_buckets_); } + + /// @brief Grow or shrink the logical size. New elements are default-constructed. + /// Single-writer; concurrent readers safe on grow (see class contract). + void resize(size_type n) { resize_impl_(n, nullptr); } + + /// @brief Grow or shrink the logical size, filling new elements with ``fill``. + void resize(size_type n, const T& fill) { resize_impl_(n, &fill); } + + /// @brief Append one element, move-*constructing* it into the new slot. + /// + /// Grows logical size by one, allocating a new bucket if needed. The element is + /// constructed in place via T's move constructor, so types whose move-*assignment* is + /// unavailable or expensive (e.g. DenseArray, whose move-assign compares allocators) + /// still work. The element is published (bucket pointer first, then ``size_``) so a + /// concurrent reader that observes the new ``size()`` sees the constructed value. + /// Single-writer. + void push_back(T&& value) { + size_type i = size_.load(std::memory_order_relaxed); + auto [b, off] = locate_(i); + T* bucket = ensure_bucket_(b); + new (&bucket[off]) T(std::move(value)); + size_.store(i + 1, std::memory_order_release); + } + + /// @brief Drop the last element, destroying it (logical only; does not free the + /// bucket). Single-writer; caller must have drained readers if a bucket is later freed. + void pop_back() { + size_type i = size_.load(std::memory_order_relaxed); + if (i > 0) { + auto [b, off] = locate_(i - 1); + dir_[b].load(std::memory_order_relaxed)[off].~T(); + size_.store(i - 1, std::memory_order_release); + } + } + + /// @brief Shrink to ``n`` elements, destroying the dropped elements and freeing buckets + /// that lie entirely above ``n``. Single-writer; caller must have drained readers. + void shrink_to(size_type n) { + size_type old = size_.load(std::memory_order_relaxed); + if (n >= old) { + return; + } + // Stop readers from seeing the elements about to be destroyed. + size_.store(n, std::memory_order_release); + destroy_range_(n, old); + free_buckets_above_(n); + } + + private: + // Fixed top-level directory. Bucket k (when non-null) is a contiguous heap array of + // (kFirstBucket << k) elements; the elements with global index < size_ are constructed. + std::atomic dir_[kDirBuckets] = {}; + std::atomic size_{0}; + size_type allocated_buckets_{0}; + + // Number of elements held by bucket ``b`` (= kFirstBucket << b). + static constexpr size_type bucket_size_(size_type b) noexcept { + return kFirstBucket << b; + } + + // First global element index held by bucket ``b`` (= kFirstBucket * (2^b - 1)). + static constexpr size_type bucket_first_index_(size_type b) noexcept { + return kFirstBucket * ((size_type{1} << b) - 1); + } + + // Map element index ``i`` to (bucket, offset-within-bucket). Group elements into + // chunks of kFirstBucket, then apply the power-of-two bucket layout to the chunk + // index: q = i / kFirstBucket; bucket = floor(log2(q+1)); the bucket's first global + // index is kFirstBucket * (2^bucket - 1). + static constexpr std::pair locate_(size_type i) noexcept { + size_type q = i / kFirstBucket; + size_type bucket = static_cast(std::bit_width(q + 1)) - 1; + return {bucket, i - bucket_first_index_(bucket)}; + } + + // Ensure bucket ``b`` is allocated (single-writer) and return its base pointer. The + // bucket's elements are raw storage until constructed by the caller; the pointer is + // published with release so readers that later observe a matching size see it. + T* ensure_bucket_(size_type b) { + T* bucket = dir_[b].load(std::memory_order_relaxed); + if (bucket == nullptr) { + bucket = static_cast(::operator new[](bucket_size_(b) * sizeof(T))); + dir_[b].store(bucket, std::memory_order_release); + if (b + 1 > allocated_buckets_) { + allocated_buckets_ = b + 1; + } + } + return bucket; + } + + // Construct elements [from, to) in place (single-writer). ``fill`` is nullptr for + // default-construction. Allocates buckets as needed. + void construct_range_(size_type from, size_type to, const T* fill) { + for (size_type i = from; i < to; ++i) { + auto [b, off] = locate_(i); + T* bucket = ensure_bucket_(b); + if (fill == nullptr) { + new (&bucket[off]) T(); + } else { + new (&bucket[off]) T(*fill); + } + } + } + + // Destroy elements [from, to) (single-writer). Does not free buckets. + void destroy_range_(size_type from, size_type to) { + for (size_type i = from; i < to; ++i) { + auto [b, off] = locate_(i); + dir_[b].load(std::memory_order_relaxed)[off].~T(); + } + } + + void resize_impl_(size_type n, const T* fill) { + size_type old = size_.load(std::memory_order_relaxed); + if (n == old) { + return; + } + if (n < old) { + // Logical-only shrink (no bucket freeing — use shrink_to for reclamation), + // but still destroy the dropped elements to run their destructors. + size_.store(n, std::memory_order_release); + destroy_range_(n, old); + return; + } + // Grow: construct the new elements, then publish the new size last so a reader + // that observes it sees fully-constructed elements in published buckets. + construct_range_(old, n, fill); + size_.store(n, std::memory_order_release); + } + + // Free every bucket whose entire index range lies at or above ``n`` (single-writer; + // readers drained). A bucket straddling ``n`` keeps its allocation. + void free_buckets_above_(size_type n) { + for (size_type b = allocated_buckets_; b-- > 0;) { + if (bucket_first_index_(b) < n) { + break; // this and all lower buckets contain live (or kept) elements + } + T* bucket = dir_[b].load(std::memory_order_relaxed); + if (bucket != nullptr) { + ::operator delete[](static_cast(bucket)); + dir_[b].store(nullptr, std::memory_order_relaxed); + } + allocated_buckets_ = b; + } + } + + void destroy_all_() { + size_type n = size_.load(std::memory_order_relaxed); + destroy_range_(0, n); + for (size_type b = 0; b < allocated_buckets_; ++b) { + T* bucket = dir_[b].load(std::memory_order_relaxed); + if (bucket != nullptr) { + ::operator delete[](static_cast(bucket)); + dir_[b].store(nullptr, std::memory_order_relaxed); + } + } + size_.store(0, std::memory_order_relaxed); + allocated_buckets_ = 0; + } + + void copy_from_(const SegmentedVector& other) { + size_type n = other.size_.load(std::memory_order_relaxed); + for (size_type i = 0; i < n; ++i) { + auto [b, off] = locate_(i); + T* bucket = ensure_bucket_(b); + new (&bucket[off]) T(other[i]); + } + size_.store(n, std::memory_order_release); + } + + void steal_from_(SegmentedVector& other) noexcept { + for (size_type b = 0; b < kDirBuckets; ++b) { + dir_[b].store( + other.dir_[b].load(std::memory_order_relaxed), std::memory_order_relaxed + ); + other.dir_[b].store(nullptr, std::memory_order_relaxed); + } + size_.store(other.size_.load(std::memory_order_relaxed), std::memory_order_relaxed); + allocated_buckets_ = other.allocated_buckets_; + other.size_.store(0, std::memory_order_relaxed); + other.allocated_buckets_ = 0; + } +}; + +} // namespace svs::lib + +namespace svs { +// Opt SegmentedVector into svs::getindex's optional bounds checking. +template +inline constexpr bool enable_boundschecking> = true; +} // namespace svs diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 8c812d35a..225854781 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -92,6 +92,7 @@ set(TEST_SOURCES ${TEST_DIR}/svs/lib/version.cpp ${TEST_DIR}/svs/lib/uuid.cpp ${TEST_DIR}/svs/lib/concurrency/readwrite_protected.cpp + ${TEST_DIR}/svs/lib/segmented_vector.cpp ${TEST_DIR}/svs/lib/avx_detection.cpp # Third Party ${TEST_DIR}/svs/third-party/fmt.cpp @@ -150,6 +151,16 @@ set(TEST_SOURCES ${TEST_DIR}/svs/quantization/scalar/scalar.cpp ${TEST_DIR}/svs/index/vamana/dynamic_index.cpp + + # Concurrent (lock-free search / concurrent mutation) Vamana index. + ${TEST_DIR}/svs/concurrent/translation.cpp + ${TEST_DIR}/svs/concurrent/graph.cpp + ${TEST_DIR}/svs/concurrent/prune.cpp + ${TEST_DIR}/svs/concurrent/consolidate.cpp + ${TEST_DIR}/svs/concurrent/iterator.cpp + ${TEST_DIR}/svs/concurrent/multi.cpp + ${TEST_DIR}/svs/concurrent/concurrency.cpp + ${TEST_DIR}/svs/concurrent/dynamic_index.cpp ) ##### @@ -158,6 +169,7 @@ set(TEST_SOURCES SET(INTEGRATION_TESTS ${TEST_DIR}/svs/index/vamana/dynamic_index_2.cpp + ${TEST_DIR}/svs/concurrent/dynamic_index_2.cpp # Higher level constructs ${TEST_DIR}/svs/orchestrators/vamana.cpp ${TEST_DIR}/svs/orchestrators/dynamic_vamana.cpp @@ -237,3 +249,66 @@ list(APPEND CMAKE_MODULE_PATH ${catch2_SOURCE_DIR}/extras) include(CTest) include(Catch) catch_discover_tests(tests ADD_TAGS_AS_LABELS SKIP_IS_FAILURE) + +##### +##### ThreadSanitizer targets for the concurrent index +##### + +# The correctness of `svs::index::vamana::concurrent` rests almost entirely on memory +# ordering, which a non-instrumented test can only fail to disprove. These targets are +# opt-in because TSan costs roughly an order of magnitude in both time and memory: +# +# cmake -DSVS_EXPERIMENTAL_ENABLE_CONCURRENT_TSAN=YES ... +# ctest -L tsan +# +# `svs::svs` and Catch2 are not instrumented here, only the test translation units. That is +# enough: the concurrent index is header-only, so the accesses under test -- the adjacency +# lists, the sequence-lock counters, the slot metadata -- are all inlined into the +# instrumented TUs. +option(SVS_EXPERIMENTAL_ENABLE_CONCURRENT_TSAN + "Build ThreadSanitizer targets for the concurrent Vamana index (slow)" OFF +) + +if (SVS_EXPERIMENTAL_ENABLE_CONCURRENT_TSAN) + message("Enabling ThreadSanitizer targets for the concurrent Vamana index!") + + set(CONCURRENT_TSAN_SOURCES + ${TEST_DIR}/svs/concurrent/graph.cpp + ${TEST_DIR}/svs/concurrent/concurrency.cpp + ) + + add_executable(concurrent_tsan ${CONCURRENT_TSAN_SOURCES}) + target_compile_options(concurrent_tsan PRIVATE -fsanitize=thread -g -O1) + target_link_options(concurrent_tsan PRIVATE -fsanitize=thread) + target_compile_definitions(concurrent_tsan PRIVATE SVS_THREAD_SANITIZER) + target_compile_definitions(concurrent_tsan PRIVATE SVS_TEST_DATA_DIR="${DATA_DIRECTORY}") + target_link_libraries(concurrent_tsan PRIVATE svs::svs Catch2::Catch2WithMain) + target_link_libraries( + concurrent_tsan PRIVATE svs_compile_options svs_x86_options_base + ) + target_include_directories(concurrent_tsan PRIVATE ${PROJECT_SOURCE_DIR}) + add_test(NAME concurrent_tsan COMMAND concurrent_tsan) + set_tests_properties(concurrent_tsan PROPERTIES LABELS "tsan") + + # Negative control. Building the graph's adjacency-slot accessors as plain (non-atomic) + # loads and stores must make TSan complain -- otherwise a clean run above proves nothing + # about whether TSan was watching the right memory at all. + add_executable(concurrent_tsan_negative ${CONCURRENT_TSAN_SOURCES}) + target_compile_options(concurrent_tsan_negative PRIVATE -fsanitize=thread -g -O1) + target_link_options(concurrent_tsan_negative PRIVATE -fsanitize=thread) + target_compile_definitions(concurrent_tsan_negative + PRIVATE SVS_THREAD_SANITIZER SVS_CONCURRENT_UNSAFE_PLAIN_GRAPH_ACCESS + ) + target_compile_definitions(concurrent_tsan_negative + PRIVATE SVS_TEST_DATA_DIR="${DATA_DIRECTORY}" + ) + target_link_libraries( + concurrent_tsan_negative PRIVATE svs::svs Catch2::Catch2WithMain + ) + target_link_libraries( + concurrent_tsan_negative PRIVATE svs_compile_options svs_x86_options_base + ) + target_include_directories(concurrent_tsan_negative PRIVATE ${PROJECT_SOURCE_DIR}) + add_test(NAME concurrent_tsan_negative COMMAND concurrent_tsan_negative) + set_tests_properties(concurrent_tsan_negative PROPERTIES WILL_FAIL TRUE LABELS "tsan") +endif() diff --git a/tests/svs/concurrent/concurrency.cpp b/tests/svs/concurrent/concurrency.cpp new file mode 100644 index 000000000..96a7ddf52 --- /dev/null +++ b/tests/svs/concurrent/concurrency.cpp @@ -0,0 +1,578 @@ +/* + * Copyright 2026 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Concurrency tests for ``svs::index::vamana::concurrent::MutableVamanaIndex``: correctness +// of searches issued while other threads insert, delete, and consolidate. +// +// The rest of the suite (``dynamic_index.cpp``, ``dynamic_index_2.cpp``, ``multi.cpp``, +// ``iterator.cpp``) covers single-threaded functional parity with the pre-existing dynamic +// index. This file covers the property that motivates a separate index in the first place: +// searches and mutations may overlap in time. +// +// Recall thresholds here are deliberately coarse. They exist to catch a graph that has been +// corrupted into uselessness, not to track search quality -- that is the job of the +// benchmark suite. +// +// Assertions inside the hot loops accumulate into counters and are checked once at the end. +// Catch2's assertion bookkeeping is not free, and these loops run for millions of +// iterations. + +// header under test +#include "svs/concurrent/dynamic_index.h" + +// For the definition of the `BatchIterator` that `make_batch_iterator` returns. +#include "svs/concurrent/iterator.h" + +#include "svs/core/data.h" +#include "svs/core/distance.h" +#include "svs/lib/threads.h" + +// catch2 +#include "catch2/catch_test_macros.hpp" + +// stl +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace cc = svs::index::vamana::concurrent; +namespace ccg = svs::index::vamana::concurrent::graphs; + +namespace { + +// ThreadSanitizer costs roughly an order of magnitude in both time and memory, and it is +// looking for *races*, which show up just as readily in a small index. Shrink the problem +// rather than skipping the run. +#if defined(__SANITIZE_THREAD__) || defined(SVS_THREAD_SANITIZER) +constexpr size_t kInitialPoints = 4000; +constexpr size_t kIncrementalPoints = 1000; +constexpr size_t kNumQueries = 50; +#else +constexpr size_t kInitialPoints = 20000; +constexpr size_t kIncrementalPoints = 5000; +constexpr size_t kNumQueries = 200; +#endif + +constexpr size_t kDim = 32; +constexpr size_t kMaxDegree = 32; +constexpr size_t kNumNeighbors = 10; +constexpr size_t kBuildThreads = 8; + +using Idx = uint32_t; +using Distance = svs::distance::DistanceL2; +using ConcurrentData = cc::SegmentedBlockedData; +using ConcurrentGraph = ccg::SimpleBlockedGraph; +using ConcurrentIndex = cc::MutableVamanaIndex; + +std::vector random_vectors(size_t n, size_t dim, uint32_t seed) { + std::mt19937 rng{seed}; + std::normal_distribution dist{0.0f, 1.0f}; + std::vector out(n * dim); + for (auto& v : out) { + v = dist(rng); + } + return out; +} + +svs::data::SimpleData make_dataset(const std::vector& raw, size_t dim) { + const size_t n = raw.size() / dim; + auto data = svs::data::SimpleData(n, dim); + for (size_t i = 0; i < n; ++i) { + data.set_datum(i, std::span(raw.data() + i * dim, dim)); + } + return data; +} + +// Build a concurrent index over ``raw`` through the ordinary build constructor. +std::unique_ptr build_index( + const std::vector& raw, size_t dim, std::span ids, size_t threads +) { + const size_t n = raw.size() / dim; + auto data = ConcurrentData(n, dim); + for (size_t i = 0; i < n; ++i) { + data.set_datum(i, std::span(raw.data() + i * dim, dim)); + } + + auto parameters = svs::index::vamana::VamanaBuildParameters{ + 1.2f, kMaxDegree, 2 * kMaxDegree, 750, kMaxDegree, true}; + return std::make_unique( + parameters, + std::move(data), + std::vector(ids.begin(), ids.end()), + Distance{}, + threads + ); +} + +// Brute-force ground truth over the given set of live external IDs. +std::vector> ground_truth( + const std::vector& base, + const std::unordered_set& live, + const std::vector& queries, + size_t dim, + size_t k +) { + const size_t nq = queries.size() / dim; + std::vector> result(nq); + for (size_t q = 0; q < nq; ++q) { + std::vector> scored; + scored.reserve(live.size()); + for (size_t id : live) { + float d = 0; + for (size_t j = 0; j < dim; ++j) { + float diff = queries[q * dim + j] - base[id * dim + j]; + d += diff * diff; + } + scored.emplace_back(d, id); + } + std::partial_sort( + scored.begin(), + scored.begin() + static_cast(std::min(k, scored.size())), + scored.end() + ); + for (size_t i = 0; i < std::min(k, scored.size()); ++i) { + result[q].push_back(scored[i].second); + } + } + return result; +} + +double recall_at_k( + const svs::QueryResult& got, const std::vector>& expected +) { + size_t hits = 0, total = 0; + for (size_t q = 0; q < expected.size(); ++q) { + std::unordered_set truth{expected[q].begin(), expected[q].end()}; + for (size_t j = 0; j < got.n_neighbors(); ++j) { + if (truth.count(got.index(q, j))) { + ++hits; + } + } + total += truth.size(); + } + return total == 0 ? 1.0 : static_cast(hits) / static_cast(total); +} + +// Insert ``[first, first + n)`` of ``base`` as a single batch. +void add_batch( + ConcurrentIndex& index, const std::vector& base, size_t first, size_t n +) { + auto batch = svs::data::SimpleData(n, kDim); + std::vector batch_ids(n); + for (size_t i = 0; i < n; ++i) { + batch.set_datum(i, std::span(base.data() + (first + i) * kDim, kDim)); + batch_ids[i] = first + i; + } + index.add_points(batch, batch_ids); +} + +} // namespace + +CATCH_TEST_CASE("Concurrent MutableVamanaIndex quiescent recall", "[concurrent][index]") { + auto base = random_vectors(kInitialPoints, kDim, 1234); + std::vector ids(kInitialPoints); + std::iota(ids.begin(), ids.end(), 0); + + auto index = build_index(base, kDim, ids, kBuildThreads); + CATCH_REQUIRE(index->size() == kInitialPoints); + + auto queries_raw = random_vectors(kNumQueries, kDim, 999); + auto queries = make_dataset(queries_raw, kDim); + + auto sp = index->get_search_parameters(); + sp.buffer_config({100}); + index->set_search_parameters(sp); + + auto results = svs::QueryResult{kNumQueries, kNumNeighbors}; + index->search(results.view(), queries, index->get_search_parameters()); + + std::unordered_set live{ids.begin(), ids.end()}; + auto truth = ground_truth(base, live, queries_raw, kDim, kNumNeighbors); + const double recall = recall_at_k(results, truth); + + // A correctly built Vamana graph at this window size should be well above 0.9. + CATCH_INFO("quiescent recall@" << kNumNeighbors << " = " << recall); + CATCH_REQUIRE(recall > 0.90); +} + +// The core test: searches run continuously while writers insert new vectors and delete +// existing ones. Any torn adjacency read, use-after-free from a resize, or missing ID +// translation surfaces as a crash, an exception, or an inconsistent ID round-trip. +CATCH_TEST_CASE( + "Concurrent MutableVamanaIndex search during mutation", "[concurrent][index]" +) { + const size_t total = kInitialPoints + kIncrementalPoints; + auto base = random_vectors(total, kDim, 4321); + + std::vector initial_ids(kInitialPoints); + std::iota(initial_ids.begin(), initial_ids.end(), 0); + // Build over the first ``kInitialPoints`` only; the tail is inserted concurrently + // below. + auto initial_slice = std::vector( + base.begin(), base.begin() + static_cast(kInitialPoints * kDim) + ); + auto index = build_index(initial_slice, kDim, initial_ids, kBuildThreads); + + auto sp = index->get_search_parameters(); + sp.buffer_config({100}); + index->set_search_parameters(sp); + + auto queries_raw = random_vectors(kNumQueries, kDim, 777); + auto queries = make_dataset(queries_raw, kDim); + + std::atomic writers_running{0}; + std::atomic searches_completed{0}; + std::atomic bad_roundtrips{0}; + std::atomic duplicate_ids{0}; + std::atomic exceptions{0}; + + // Two writer threads insert *disjoint* halves of the incremental points concurrently. + // Concurrent `add_points` is the headline capability of this index: the slot allocator + // is the only serialized section, and graph/data growth is lock-free. + constexpr size_t kWriters = 2; + constexpr size_t kBatch = 500; + const size_t per_writer = kIncrementalPoints / kWriters; + + auto writer = [&](size_t w) { + try { + const size_t begin = kInitialPoints + w * per_writer; + for (size_t offset = 0; offset < per_writer; offset += kBatch) { + const size_t n = std::min(kBatch, per_writer - offset); + add_batch(*index, base, begin + offset, n); + + // Delete a fresh, disjoint set of the original IDs each round, *spread* + // across the whole ID range with a stride. A contiguous low-ID slice would + // almost never intersect a query's top-k, so the interesting race -- a + // result slot retired between selection and ID translation -- would go + // unexercised. Each (writer, round) pair takes its own residue class mod + // 40, so the rounds are disjoint and together retire a quarter of the + // original vectors. + const size_t round = w * (per_writer / kBatch) + offset / kBatch; + std::vector to_delete; + for (size_t id = round; id < kInitialPoints; id += 40) { + to_delete.push_back(id); + } + index->delete_entries(to_delete); + } + } catch (const std::exception& e) { + CATCH_WARN("writer threw: " << e.what()); + exceptions.fetch_add(1); + } + writers_running.fetch_sub(1); + }; + + // Searchers: hammer the index with single-query searches throughout. + auto searcher = [&] { + try { + auto scratch = index->scratchspace(); + std::unordered_set seen_ids; + while (writers_running.load(std::memory_order_relaxed) != 0) { + for (size_t q = 0; q < kNumQueries; ++q) { + auto query = + std::span(queries_raw.data() + q * kDim, kDim); + index->search(query, scratch); + // ``[0, valid())`` is the region a caller is allowed to read: skipped + // (deleted) candidates have been compacted out by this point. + const size_t n = + std::min(kNumNeighbors, scratch.buffer.valid()); + seen_ids.clear(); + for (size_t j = 0; j < n; ++j) { + auto internal = scratch.buffer[j].id(); + // A ranked result must never list the same vector twice. Torn + // adjacency reads or a botched retry in the seqlock section would + // show up here, because the search buffer dedupes by ID and can + // only be fooled by inconsistent input. + if (!seen_ids.insert(internal).second) { + duplicate_ids.fetch_add(1, std::memory_order_relaxed); + } + // Every ID a search hands back either still maps to a live external + // ID -- in which case the mapping must round-trip exactly -- or it + // names a slot retired by a concurrent deleter, which is expected + // and unobservable from here (``translate_internal_id`` degrades to + // returning the internal ID when the entry has been erased). + auto external = index->translate_internal_id(internal); + if (index->has_id(external) && + index->translate_external_id(external) != internal) { + bad_roundtrips.fetch_add(1, std::memory_order_relaxed); + } + } + searches_completed.fetch_add(1, std::memory_order_relaxed); + } + } + } catch (const std::exception& e) { + CATCH_WARN("searcher threw: " << e.what()); + exceptions.fetch_add(1); + } + }; + + std::vector threads; + writers_running.store(kWriters); + for (size_t w = 0; w < kWriters; ++w) { + threads.emplace_back(writer, w); + } + for (int i = 0; i < 8; ++i) { + threads.emplace_back(searcher); + } + for (auto& t : threads) { + t.join(); + } + + CATCH_INFO("searches completed: " << searches_completed.load()); + CATCH_REQUIRE(exceptions.load() == 0); + CATCH_REQUIRE(bad_roundtrips.load() == 0); + CATCH_REQUIRE(duplicate_ids.load() == 0); + CATCH_REQUIRE(searches_completed.load() > 0); + + // Post-mutation the index must still be a correct index. + index->debug_check_invariants(true); + std::unordered_set live; + index->on_ids([&live](size_t id) { live.insert(id); }); + CATCH_REQUIRE(live.size() == index->size()); + + auto truth = ground_truth(base, live, queries_raw, kDim, kNumNeighbors); + auto results = svs::QueryResult{kNumQueries, kNumNeighbors}; + index->search(results.view(), queries, index->get_search_parameters()); + const double recall = recall_at_k(results, truth); + CATCH_INFO("post-mutation recall@" << kNumNeighbors << " = " << recall); + CATCH_REQUIRE(recall > 0.85); + + // And consolidation/compaction must leave it correct too. + index->consolidate(); + index->compact(); + index->debug_check_invariants(false); + auto results2 = svs::QueryResult{kNumQueries, kNumNeighbors}; + index->search(results2.view(), queries, index->get_search_parameters()); + const double recall2 = recall_at_k(results2, truth); + CATCH_INFO("post-consolidate/compact recall@" << kNumNeighbors << " = " << recall2); + CATCH_REQUIRE(recall2 > 0.85); + CATCH_REQUIRE(index->size() == live.size()); +} + +// `consolidate()` walks the reverse-edge index and rewires in-neighbors of deleted slots in +// place. Unlike `compact()` it does not shrink storage, so it is allowed to run while +// searches are in flight. That is the property under test here. +CATCH_TEST_CASE( + "Concurrent MutableVamanaIndex consolidate during search", "[concurrent][index]" +) { + auto base = random_vectors(kInitialPoints, kDim, 8642); + std::vector ids(kInitialPoints); + std::iota(ids.begin(), ids.end(), 0); + auto index = build_index(base, kDim, ids, kBuildThreads); + + auto sp = index->get_search_parameters(); + sp.buffer_config({100}); + index->set_search_parameters(sp); + + auto queries_raw = random_vectors(kNumQueries, kDim, 555); + + std::atomic writer_done{false}; + std::atomic searches_completed{0}; + std::atomic exceptions{0}; + + std::thread writer{[&] { + try { + // Five rounds of "delete a stride, then consolidate", retiring 25% overall. + for (size_t round = 0; round < 5; ++round) { + std::vector to_delete; + for (size_t id = round; id < kInitialPoints; id += 20) { + to_delete.push_back(id); + } + index->delete_entries(to_delete); + index->consolidate(); + } + } catch (const std::exception& e) { + CATCH_WARN("writer threw: " << e.what()); + exceptions.fetch_add(1); + } + writer_done.store(true); + }}; + + auto searcher = [&] { + try { + auto scratch = index->scratchspace(); + while (!writer_done.load(std::memory_order_relaxed)) { + for (size_t q = 0; q < kNumQueries; ++q) { + auto query = + std::span(queries_raw.data() + q * kDim, kDim); + index->search(query, scratch); + searches_completed.fetch_add(1, std::memory_order_relaxed); + } + } + } catch (const std::exception& e) { + CATCH_WARN("searcher threw: " << e.what()); + exceptions.fetch_add(1); + } + }; + + std::vector searchers; + for (int i = 0; i < 6; ++i) { + searchers.emplace_back(searcher); + } + writer.join(); + for (auto& t : searchers) { + t.join(); + } + + CATCH_INFO("searches completed: " << searches_completed.load()); + CATCH_REQUIRE(exceptions.load() == 0); + CATCH_REQUIRE(searches_completed.load() > 0); + index->debug_check_invariants(false); + + // The surviving 75% must still be reachable at a reasonable rate. + std::unordered_set live; + index->on_ids([&live](size_t id) { live.insert(id); }); + CATCH_REQUIRE(live.size() == index->size()); + + auto queries = make_dataset(queries_raw, kDim); + auto truth = ground_truth(base, live, queries_raw, kDim, kNumNeighbors); + auto results = svs::QueryResult{kNumQueries, kNumNeighbors}; + index->search(results.view(), queries, index->get_search_parameters()); + const double recall = recall_at_k(results, truth); + CATCH_INFO("post-consolidate recall@" << kNumNeighbors << " = " << recall); + CATCH_REQUIRE(recall > 0.85); +} + +// The batch iterator holds a cursor across calls. This checks it works at all, and that it +// keeps working while a writer mutates the index. +CATCH_TEST_CASE("Concurrent MutableVamanaIndex batch iterator", "[concurrent][index]") { + const size_t total = kInitialPoints + kIncrementalPoints; + auto base = random_vectors(total, kDim, 24680); + + std::vector initial_ids(kInitialPoints); + std::iota(initial_ids.begin(), initial_ids.end(), 0); + auto initial_slice = std::vector( + base.begin(), base.begin() + static_cast(kInitialPoints * kDim) + ); + auto index = build_index(initial_slice, kDim, initial_ids, kBuildThreads); + + auto queries_raw = random_vectors(kNumQueries, kDim, 13579); + + CATCH_SECTION("quiescent") { + // Batches must be non-overlapping, and each batch must be sorted. + // + // Batches are *not* globally monotonic and it would be wrong to assert that: the + // iterator is approximate, so a later batch can surface a vector closer than one + // the earlier batch's window had already returned. Count those inversions and + // report them as a quality signal rather than a correctness one -- this is + // pre-existing behaviour and has nothing to do with concurrency. + auto query = std::span(queries_raw.data(), kDim); + auto it = index->make_batch_iterator(query); + std::unordered_set all; + float previous_worst = -1.0f; + size_t batches = 0; + size_t inversions = 0; + size_t repeats = 0; + size_t unsorted = 0; + for (; batches < 5 && !it.done(); ++batches) { + it.next(10); + float last = -1.0f; + for (const auto& n : it) { + if (!all.insert(n.id()).second) { + ++repeats; + } + if (n.distance() < last) { + ++unsorted; + } + last = n.distance(); + if (n.distance() < previous_worst) { + ++inversions; + } + } + if (it.size() != 0) { + previous_worst = (it.end() - 1)->distance(); + } + } + CATCH_INFO( + batches << " batches, " << all.size() << " distinct vectors, " << inversions + << " cross-batch inversions" + ); + CATCH_REQUIRE(repeats == 0); + CATCH_REQUIRE(unsorted == 0); + CATCH_REQUIRE(all.size() >= 40); + } + + CATCH_SECTION("during mutation") { + // The claim under test is memory safety and per-batch consistency, *not* that a + // long-lived cursor sees a stable snapshot. + std::atomic writer_done{false}; + std::atomic batches_completed{0}; + std::atomic exceptions{0}; + std::atomic bad_ids{0}; + + std::thread writer{[&] { + try { + constexpr size_t kBatch = 500; + for (size_t offset = 0; offset < kIncrementalPoints; offset += kBatch) { + const size_t n = std::min(kBatch, kIncrementalPoints - offset); + add_batch(*index, base, kInitialPoints + offset, n); + } + } catch (const std::exception& e) { + CATCH_WARN("writer threw: " << e.what()); + exceptions.fetch_add(1); + } + writer_done.store(true); + }}; + + auto reader = [&] { + try { + while (!writer_done.load(std::memory_order_relaxed)) { + for (size_t q = 0; q < kNumQueries; ++q) { + auto query = + std::span(queries_raw.data() + q * kDim, kDim); + auto it = index->make_batch_iterator(query); + for (size_t b = 0; b < 3 && !it.done(); ++b) { + it.next(10); + for (const auto& n : it) { + // Any ID the iterator yields must name a slot that either + // is live or was retired mid-flight; a bad graph read shows + // up as an ID no larger than the highest slot ever handed + // out. + if (n.id() >= total) { + bad_ids.fetch_add(1, std::memory_order_relaxed); + } + } + batches_completed.fetch_add(1, std::memory_order_relaxed); + } + } + } + } catch (const std::exception& e) { + CATCH_WARN("reader threw: " << e.what()); + exceptions.fetch_add(1); + } + }; + + std::vector readers; + for (int i = 0; i < 4; ++i) { + readers.emplace_back(reader); + } + writer.join(); + for (auto& t : readers) { + t.join(); + } + + CATCH_INFO("batches completed during mutation: " << batches_completed.load()); + CATCH_REQUIRE(exceptions.load() == 0); + CATCH_REQUIRE(bad_ids.load() == 0); + CATCH_REQUIRE(batches_completed.load() > 0); + } +} diff --git a/tests/svs/concurrent/consolidate.cpp b/tests/svs/concurrent/consolidate.cpp new file mode 100644 index 000000000..4f042a391 --- /dev/null +++ b/tests/svs/concurrent/consolidate.cpp @@ -0,0 +1,82 @@ +/* + * Copyright 2023 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// header under test +#include "svs/concurrent/consolidate.h" + +// svs +#include "svs/core/distance.h" +#include "svs/lib/timing.h" + +// test utilities +#include "tests/utils/test_dataset.h" + +// stl +#include + +// catch2 +#include "catch2/catch_test_macros.hpp" + +namespace graphs = svs::index::vamana::concurrent::graphs; + +namespace { + +template +void check_post_conditions(const Graph& graph, Predicate&& predicate) { + bool contains_deleted = false; + svs::threads::UnitRange node_range{0, graph.n_nodes()}; + for (size_t i : node_range) { + if (predicate(i)) { + contains_deleted = true; + continue; + } + + const auto& neighbors = graph.get_node(i); + CATCH_REQUIRE(std::none_of(neighbors.begin(), neighbors.end(), predicate)); + + // Don't invent nodes out of thin air. + CATCH_REQUIRE(std::all_of(neighbors.begin(), neighbors.end(), [&](const auto& i) { + return node_range.contains(i); + })); + } + CATCH_REQUIRE(contains_deleted); +} + +} // namespace + +CATCH_TEST_CASE("Concurrent Graph Consolidation", "[concurrent][graph_index]") { + // The concurrent consolidate requires the concurrent graph type: its + // `add_edge` returns `AddEdgeResult` rather than a degree, and it carries the + // per-node sequence locks the lock-free readers validate against. + auto graph = graphs::SimpleGraph::load(test_dataset::graph_file()); + auto data = test_dataset::data_f32(); + auto threadpool = svs::threads::DefaultThreadPool(2); + + CATCH_SECTION("Remove Even Nodes") { + auto tic = svs::lib::now(); + auto predicate = [](const auto& i) { return (i % 10) == 0; }; + + svs::distance::DistanceL2 distance{}; + svs::index::vamana::concurrent::consolidate( + graph, data, threadpool, graph.max_degree(), 750, 1.2, distance, predicate + ); + std::cout << "Pruning took " << svs::lib::time_difference(svs::lib::now(), tic) + << std::endl; + + // Ensure that all non-deleted nodes only have non-deleted neighbors. + check_post_conditions(graph, predicate); + } +} diff --git a/tests/svs/concurrent/dynamic_index.cpp b/tests/svs/concurrent/dynamic_index.cpp new file mode 100644 index 000000000..593c4e023 --- /dev/null +++ b/tests/svs/concurrent/dynamic_index.cpp @@ -0,0 +1,548 @@ +/* + * Copyright 2023 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// header under test. +#include "svs/concurrent/dynamic_index.h" +#include "svs/concurrent/consolidate.h" + +// stl +#include +#include +#include +#include +#include +#include +#include +#include + +// svs +#include "svs/core/recall.h" +#include "svs/lib/timing.h" + +// catch2 +#include "catch2/catch_test_macros.hpp" +#include + +// tests +#include "tests/utils/test_dataset.h" +#include "tests/utils/utils.h" + +// The concurrent index is a separate type from `svs::index::vamana::MutableVamanaIndex`; +// the pre-existing static and dynamic indexes are untouched by these tests. +namespace cc = svs::index::vamana::concurrent; +namespace ccg = svs::index::vamana::concurrent::graphs; + +namespace { + +// `test_dataset::data_blocked_f32()` returns `svs::data::BlockedData`. The +// concurrent index is built on the grow-stable variant of the same storage, so load +// the same file into that type instead. +cc::SegmentedBlockedData data_segmented_f32() { + return cc::SegmentedBlockedData::load(test_dataset::data_svs_file()); +} + +using ConcurrentGraph = ccg::SimpleBlockedGraph; +using ConcurrentData = cc::SegmentedBlockedData; + +} // namespace + +// The MutableVamanaIndex "Soft Deletion" test uses outdated API. +#if 0 +namespace { +template auto copy_dataset(const T& data) { + auto copy = svs::data::SimplePolymorphicData{ + data.size(), data.dimensions()}; + for (size_t i = 0; i < data.size(); ++i) { + copy.set_datum(i, data.get_datum(i)); + } + return copy; +} + +template void check_results(const T& results, const U& deleted) { + for (size_t i = 0; i < svs::getsize<0>(results); ++i) { + for (size_t j = 0; j < svs::getsize<1>(results); ++j) { + CATCH_REQUIRE(!deleted.contains(results.at(i, j))); + } + } +} + +template +void check_deleted(const T& index, const U& deleted, size_t imax) { + for (size_t i = 0; i < imax; ++i) { + if (deleted.contains(i)) { + CATCH_REQUIRE(index.is_deleted(i)); + } else { + CATCH_REQUIRE(!index.is_deleted(i)); + } + } +} + +template +void check_equal(const Left& left, const Right& right) { + CATCH_REQUIRE(left.size() == right.size()); + CATCH_REQUIRE(left.dimensions() == right.dimensions()); + + for (size_t i = 0, imax = left.size(); i < imax; ++i) { + const auto& datum_left = left.get_datum(i); + const auto& datum_right = right.get_datum(i); + CATCH_REQUIRE(std::equal(datum_left.begin(), datum_left.end(), datum_right.begin()) + ); + } +} + +} // namespace + +#if defined(NDEBUG) +const double DELETE_PERCENT = 0.3; +#else +const double DELETE_PERCENT = 0.05; +#endif + +CATCH_TEST_CASE("Concurrent MutableVamanaIndex", "[concurrent][graph_index]") { + const size_t num_threads = 2; + const size_t num_neighbors = 10; + + const auto base_data = test_dataset::data_blocked_f32(); + // const auto base_data = test_dataset::data_f32(); + const auto queries = test_dataset::queries(); + const auto groundtruth = test_dataset::groundtruth_euclidean(); + + CATCH_SECTION("Soft Deletion") { + // In this section, we test soft deletion. + // The idea is as follows: + // + // (1) Load the test index. + // (2) Run a round of queries to ensure that everything loading correctly. + // (3) Set a target deletion percentage where all the neighbors returned by + // all results returned by the previous query plus a random collection of extras + // are deleted. + // + // (4) Rerun queries, make sure accuracy is still high and that no deleted indices + // are present in the results. + auto entry_point = svs::index::load_entry_point(test_dataset::metadata_file()); + + auto index = svs::index::MutableVamanaIndex{ + test_dataset::graph_blocked(), + base_data.copy(), + entry_point, + svs::distance::DistanceL2(), + svs::threads::UnitRange(0, base_data.size()), + num_threads}; + + check_equal(base_data, index); + index.debug_check_graph_consistency(false); + + auto results = svs::QueryResult(queries.size(), num_neighbors); + index.set_search_window_size(num_neighbors); + + auto tic = svs::lib::now(); + index.search(queries.view(), num_neighbors, results.view()); + auto original_time = svs::lib::time_difference(svs::lib::now(), tic); + auto original_recall = svs::k_recall_at_n(groundtruth, results); + CATCH_REQUIRE(index.entry_point() == entry_point); + + std::unordered_set ids_to_delete{}; + double delete_percent = DELETE_PERCENT; + for (size_t i = 0; i < groundtruth.size(); ++i) { + auto slice = groundtruth.get_datum(i); + for (size_t j = 0; j < num_neighbors; ++j) { + auto id = slice[j]; + + // For now - don't delete the entry point. + if (id != entry_point) { + ids_to_delete.insert(slice[j]); + } + } + + if (ids_to_delete.size() > delete_percent * base_data.size()) { + break; + } + } + + index.set_threadpool(threads::CppAsyncThreadPool(num_threads)); + + std::cout << "Deleting " << ids_to_delete.size() << " entries!" << std::endl; + index.delete_entries(ids_to_delete); + check_deleted(index, ids_to_delete, base_data.size()); + index.debug_check_graph_consistency(true); + CATCH_REQUIRE_THROWS_AS( + index.debug_check_graph_consistency(false), svs::ANNException + ); + CATCH_REQUIRE(index.entry_point() == entry_point); + // Make sure the correct points were deleted. + tic = svs::lib::now(); + index.search(queries.view(), num_neighbors, results.view()); + auto new_time = svs::lib::time_difference(tic); + + // Make sure none of the returned results are in the deleted list. + check_results(results.indices(), ids_to_delete); + + index.set_threadpool(threads::QueueThreadPoolWrapper(num_threads)); + + auto results_reference = svs::QueryResult(queries.size(), num_neighbors); + index.exhaustive_search(queries.view(), num_neighbors, results_reference.view()); + auto new_recall = svs::k_recall_at_n(results_reference.indices(), results); + + // Perform graph consolidation and see how the results are effected. + index.set_alpha(1.2); + index.consolidate(); + index.debug_check_graph_consistency(false); + tic = svs::lib::now(); + index.search(queries.view(), num_neighbors, results.view()); + auto post_consolidate_time = svs::lib::time_difference(tic); + auto post_consolidate_recall = + svs::k_recall_at_n(results_reference.indices(), results); + + // Check deletion again. + check_deleted(index, ids_to_delete, base_data.size()); + CATCH_REQUIRE(index.entry_point() == entry_point); + + std::cout << "Original recall: " << original_recall + << ", New Recall: " << new_recall + << ", Post Recall: " << post_consolidate_recall << std::endl; + std::cout << "Original Time: " << original_time << " (s), New Time: " << new_time + << " (s) Post Time: " << post_consolidate_time << std::endl; + CATCH_REQUIRE(new_recall > original_recall); + check_results(results.indices(), ids_to_delete); + + // Now - delete the entry point and consolidate. + ids_to_delete.insert(entry_point); + std::vector entry_point_vector{}; + entry_point_vector.push_back(entry_point); + index.delete_entries(entry_point_vector); + index.set_alpha(1.2); + index.consolidate(); + index.debug_check_graph_consistency(false); + + auto& threadpool = + index.get_threadpool_handle().get.get(); + threadpool.resize(3); + CATCH_REQUIRE(index.get_num_threads() == 3); + threadpool.resize(num_threads); + CATCH_REQUIRE(index.get_num_threads() == num_threads); + + CATCH_REQUIRE(index.entry_point() != entry_point); + index.search(queries.view(), num_neighbors, results.view()); + auto post_entrypoint_recall = + svs::k_recall_at_n(results_reference.indices(), results); + std::cout << "Post entry-point deletion recall: " << post_entrypoint_recall + << std::endl; + + // Add the deleted points back in. + auto points = svs::data::SimpleData( + ids_to_delete.size(), base_data.dimensions() + ); + + size_t i = 0; + for (const auto& j : ids_to_delete) { + points.set_datum(i, base_data.get_datum(j)); + ++i; + } + + index.set_threadpool(threads::DefaultThreadPool(num_threads)); + tic = svs::lib::now(); + index.add_points(points, ids_to_delete); + auto insert_time = svs::lib::time_difference(tic); + std::cout << "Insertion took: " << insert_time << " seconds!" << std::endl; + + // Check that the stored dataset and the original dataset are equal. + check_equal(base_data, index); + index.debug_check_graph_consistency(false); + + tic = svs::lib::now(); + index.search(queries.view(), num_neighbors, results.view()); + auto post_add_time = svs::lib::time_difference(tic); + auto post_reinsertion_recall = svs::k_recall_at_n(groundtruth, results); + std::cout << "Post reinsertion recall: " << post_reinsertion_recall << " in " + << post_add_time << " seconds." << std::endl; + } +} +#endif + +namespace { + +// Build a one-row dataset holding a copy of `source[row]`, to feed add_points. +svs::data::SimpleData +one_point_from(const cc::SegmentedBlockedData& source, size_t row) { + auto points = svs::data::SimpleData(1, source.dimensions()); + points.set_datum(0, source.get_datum(row)); + return points; +} + +} // namespace + +// Re-adding an external ID that already exists must throw *without mutating the +// index*. These tests pin the two failure modes the transactional-insert + +// Phase-1-rollback fix addresses: a leaked reserved slot (Defect A) and a +// half-inserted "ghost" ID that would later hang delete (Defect B). +CATCH_TEST_CASE( + "Concurrent MutableVamana add_points duplicate external IDs", + "[concurrent][graph_index][dynamic_index]" +) { + const size_t num_threads = 2; + using Distance = svs::distance::DistanceL2; + + auto base = data_segmented_f32(); + const size_t n = base.size(); + std::vector indices(n); + std::iota(indices.begin(), indices.end(), 0); + + svs::index::vamana::VamanaBuildParameters parameters{1.2, 64, 10, 20, 10, true}; + auto index = cc::MutableVamanaIndex( + parameters, base.copy(), indices, Distance(), num_threads + ); + + CATCH_SECTION("Re-adding a live ID throws and leaks no slot") { + // All n slots are Valid, so a re-add must grow to reserve a slot, then + // throw. With the fix that slot is rolled back to Empty and reused by the + // next add — the physical footprint grows by exactly one across both + // operations. Without the fix the reserved slot leaks (stuck Pending) and + // the footprint grows by two. + const size_t footprint_before = index.view_data().size(); + CATCH_REQUIRE(index.size() == n); + + auto dup = one_point_from(base, 0); + CATCH_REQUIRE_THROWS_AS( + index.add_points(dup, std::vector{0}), svs::ANNException + ); + // Throw left the live mapping intact and added nothing. + CATCH_REQUIRE(index.size() == n); + CATCH_REQUIRE(index.has_id(0)); + + // A subsequent fresh add must reuse the rolled-back slot. + const size_t fresh_id = n + 100; + auto fresh = one_point_from(base, 1); + index.add_points(fresh, std::vector{fresh_id}); + CATCH_REQUIRE(index.has_id(fresh_id)); + CATCH_REQUIRE(index.size() == n + 1); + CATCH_REQUIRE(index.view_data().size() == footprint_before + 1); + index.debug_check_invariants(false); + } + + CATCH_SECTION("Mixed [fresh, live] batch throws without a ghost ID") { + // The batch validates fully before mutating the translator, so the fresh + // ID is never committed. It must not be searchable, and a later delete of + // it must be a prompt no-op (not spin forever on a Pending slot). + const size_t fresh_id = n + 100; + auto points = svs::data::SimpleData(2, base.dimensions()); + points.set_datum(0, base.get_datum(1)); // fresh_id + points.set_datum(1, base.get_datum(0)); // live id 0 -> triggers throw + + CATCH_REQUIRE_THROWS_AS( + index.add_points(points, std::vector{fresh_id, 0}), svs::ANNException + ); + + CATCH_REQUIRE_FALSE(index.has_id(fresh_id)); // no ghost mapping + CATCH_REQUIRE(index.size() == n); + // Deleting the never-added ID resolves nothing and returns immediately. + CATCH_REQUIRE(index.delete_entries(std::vector{fresh_id}) == 0); + index.debug_check_invariants(false); + } + + CATCH_SECTION("Re-adding a soft-deleted (stale) ID still succeeds") { + // The stale-replacement path must be unaffected by the transactional split. + index.delete_entries(std::vector{0}); + CATCH_REQUIRE(index.size() == n - 1); + + auto readd = one_point_from(base, 0); + CATCH_REQUIRE_NOTHROW(index.add_points(readd, std::vector{0})); + CATCH_REQUIRE(index.has_id(0)); + CATCH_REQUIRE(index.size() == n); + index.debug_check_invariants(true); + } +} + +CATCH_TEST_CASE( + "Concurrent MutableVamana Index Save and Load", + "[concurrent][graph_index][dynamic_index][saveload]" +) { + const size_t num_threads = 2; + using Distance = svs::distance::DistanceL2; + + auto data = data_segmented_f32(); + std::vector indices(data.size()); + std::iota(indices.begin(), indices.end(), 0); + + svs::index::vamana::VamanaBuildParameters parameters{1.2, 64, 10, 20, 10, true}; + auto index = cc::MutableVamanaIndex( + parameters, std::move(data), indices, Distance(), num_threads + ); + + const size_t num_neighbors = 10; + auto queries = test_dataset::queries(); + auto search_params = svs::index::vamana::VamanaSearchParameters{}; + search_params.buffer_config_ = svs::index::vamana::SearchBufferConfig{num_neighbors}; + auto results = svs::QueryResult(queries.size(), num_neighbors); + index.search(results.view(), queries.cview(), search_params); + + CATCH_SECTION("Load MutableVamana Index being serialized natively to stream") { + std::stringstream stream; + index.save(stream); + { + // `svs::DynamicVamana::assemble(stream, ...)` peels the archive header off the + // stream before delegating, and `auto_dynamic_assemble` expects a stream + // positioned just past it. Do the same here. + auto deserializer = svs::lib::detail::Deserializer::build(stream); + CATCH_REQUIRE(deserializer.is_native()); + + // Reassembled directly through the concurrent index's own loader rather + // than through `svs::DynamicVamana`: the type-erased orchestrator wraps the + // pre-existing dynamic index, which this stack deliberately leaves alone. + auto loaded = cc::auto_dynamic_assemble( + stream, + [&]() -> ConcurrentGraph { return ConcurrentGraph::load(stream); }, + [&]() -> ConcurrentData { + return svs::lib::load_from_stream(stream); + }, + Distance(), + num_threads + ); + + CATCH_REQUIRE(loaded.size() == index.size()); + CATCH_REQUIRE(loaded.dimensions() == index.dimensions()); + CATCH_REQUIRE(loaded.get_alpha() == index.get_alpha()); + CATCH_REQUIRE(loaded.get_graph_max_degree() == index.get_graph_max_degree()); + CATCH_REQUIRE(loaded.get_max_candidates() == index.get_max_candidates()); + CATCH_REQUIRE( + loaded.get_construction_window_size() == + index.get_construction_window_size() + ); + CATCH_REQUIRE(loaded.get_prune_to() == index.get_prune_to()); + CATCH_REQUIRE( + loaded.get_full_search_history() == index.get_full_search_history() + ); + index.on_ids([&](size_t e) { CATCH_REQUIRE(loaded.has_id(e)); }); + + auto loaded_results = svs::QueryResult(queries.size(), num_neighbors); + loaded.search(loaded_results.view(), queries.cview(), search_params); + for (size_t q = 0; q < queries.size(); ++q) { + for (size_t i = 0; i < num_neighbors; ++i) { + CATCH_REQUIRE(loaded_results.index(q, i) == results.index(q, i)); + CATCH_REQUIRE( + loaded_results.distance(q, i) == + Catch::Approx(results.distance(q, i)).epsilon(1e-5) + ); + } + } + } + } + + CATCH_SECTION("Load MutableVamana Index being serialized with intermediate files") { + std::stringstream stream; + { + svs::lib::UniqueTempDirectory tempdir{"svs_dynvamana_save"}; + const auto config_dir = tempdir.get() / "config"; + const auto graph_dir = tempdir.get() / "graph"; + const auto data_dir = tempdir.get() / "data"; + std::filesystem::create_directories(config_dir); + std::filesystem::create_directories(graph_dir); + std::filesystem::create_directories(data_dir); + index.save(config_dir, graph_dir, data_dir); + svs::lib::DirectoryArchiver::pack(tempdir, stream); + } + { + svs::lib::UniqueTempDirectory tempdir{"svs_concurrent_dynvamana_load"}; + // Read the archive header first: `unpack` consumes the rest of the stream. + auto deserializer = svs::lib::detail::Deserializer::build(stream); + svs::lib::DirectoryArchiver::unpack(stream, tempdir, deserializer.magic()); + + auto loaded = cc::auto_dynamic_assemble( + tempdir.get() / "config", + svs::lib::Lazy([&]() { + return ConcurrentGraph::load(tempdir.get() / "graph"); + }), + svs::lib::Lazy([&]() { + return ConcurrentData::load(tempdir.get() / "data"); + }), + Distance(), + num_threads + ); + + CATCH_REQUIRE(loaded.size() == index.size()); + CATCH_REQUIRE(loaded.dimensions() == index.dimensions()); + CATCH_REQUIRE(loaded.get_alpha() == index.get_alpha()); + CATCH_REQUIRE(loaded.get_graph_max_degree() == index.get_graph_max_degree()); + CATCH_REQUIRE(loaded.get_max_candidates() == index.get_max_candidates()); + CATCH_REQUIRE( + loaded.get_construction_window_size() == + index.get_construction_window_size() + ); + CATCH_REQUIRE(loaded.get_prune_to() == index.get_prune_to()); + CATCH_REQUIRE( + loaded.get_full_search_history() == index.get_full_search_history() + ); + index.on_ids([&](size_t e) { CATCH_REQUIRE(loaded.has_id(e)); }); + + auto loaded_results = svs::QueryResult(queries.size(), num_neighbors); + loaded.search(loaded_results.view(), queries.cview(), search_params); + for (size_t q = 0; q < queries.size(); ++q) { + for (size_t i = 0; i < num_neighbors; ++i) { + CATCH_REQUIRE(loaded_results.index(q, i) == results.index(q, i)); + CATCH_REQUIRE( + loaded_results.distance(q, i) == + Catch::Approx(results.distance(q, i)).epsilon(1e-5) + ); + } + } + } + } +} + +CATCH_TEST_CASE( + "Concurrent MutableVamana Index Memory Usage", "[concurrent][graph_index][dynamic_index]" +) { + const size_t num_threads = 2; + using Distance = svs::distance::DistanceL2; + + auto data = data_segmented_f32(); + const size_t data_size = data.size(); + // Expected data bytes are capacity-based; capture them before the dataset is moved + // into the index so the test can pin the exact value. + const size_t expected_data_bytes = data.capacity() * data.element_size(); + std::vector indices(data_size); + std::iota(indices.begin(), indices.end(), 0); + + svs::index::vamana::VamanaBuildParameters parameters{1.2, 64, 10, 20, 10, true}; + auto index = cc::MutableVamanaIndex( + parameters, std::move(data), indices, Distance(), num_threads + ); + + const size_t expected_graph_bytes = index.view_graph().get_data().capacity() * + index.view_graph().get_data().element_size(); + using Index = decltype(index); + // The per-slot status array is a SegmentedVector, whose capacity rounds up to the + // power-of-two bucket layout rather than tracking size() exactly. + const size_t expected_status_bytes = + svs::lib::SegmentedVector(data_size).capacity() * + sizeof(cc::SlotMetadata); + const size_t expected_metadata_bytes = expected_status_bytes + + sizeof(typename Index::internal_id_type) + + 2 * indices.size() * + (sizeof(typename Index::external_id_type) + + sizeof(typename Index::internal_id_type)); + const size_t expected_total_bytes = + expected_data_bytes + expected_graph_bytes + expected_metadata_bytes; + + // Dynamic get_memory_usage() should exactly match the capacity-based graph and data + // bytes plus the deterministic metadata implied by the input ids. + const auto breakdown = index.get_memory_breakdown(); + CATCH_REQUIRE(breakdown.graph_bytes == expected_graph_bytes); + CATCH_REQUIRE(breakdown.data_bytes == expected_data_bytes); + CATCH_REQUIRE(breakdown.metadata_bytes == expected_metadata_bytes); + CATCH_REQUIRE(breakdown.total() == expected_total_bytes); + const size_t usage = index.get_memory_breakdown().total(); + CATCH_REQUIRE(usage == expected_total_bytes); +} diff --git a/tests/svs/concurrent/dynamic_index_2.cpp b/tests/svs/concurrent/dynamic_index_2.cpp new file mode 100644 index 000000000..ac3f9e505 --- /dev/null +++ b/tests/svs/concurrent/dynamic_index_2.cpp @@ -0,0 +1,727 @@ +/* + * Copyright 2023 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// svs +#include "svs/core/medioid.h" +#include "svs/core/recall.h" +#include "svs/index/flat/flat.h" +#include "svs/concurrent/dynamic_index.h" +#include "svs/lib/float16.h" +#include "svs/lib/preprocessor.h" +#include "svs/lib/timing.h" + +#include "svs/misc/dynamic_helper.h" + +// tests +#include "spdlog/sinks/callback_sink.h" +#include "tests/utils/test_dataset.h" +#include "tests/utils/utils.h" +#include "tests/utils/vamana_reference.h" + +// catch +#include "catch2/catch_test_macros.hpp" + +// stl +#include +#include +#include +#include + +// The concurrent index is a separate type from `cc::MutableVamanaIndex`; +// the pre-existing static and dynamic indexes are untouched by these tests. +namespace cc = svs::index::vamana::concurrent; +namespace ccg = svs::index::vamana::concurrent::graphs; + +// Internal linkage for the whole helper block. The upstream +// `tests/svs/index/vamana/dynamic_index_2.cpp`, which this file is a port of, declares +// several of the same helpers at global scope (`Report`, its stream operator); both files +// link into the same `tests` binary, so leaving these external would be an ODR violation. +namespace { + +// Adds a sink to the process-global logger for the duration of a scope. +// +// The upstream test this file is ported from pushes a capturing sink onto the global logger +// and never removes it. The sink outlives the vector it captures by reference, so every +// later log statement that reaches the global logger is a use-after-free. It is latent +// because whether the freed memory still looks like a `std::vector` is a matter of heap +// luck; when the luck runs out it surfaces as a `std::bad_alloc` or a SIGSEGV inside a +// completely unrelated test that happens to log a warning. Scope the push instead. +class ScopedGlobalSink { + public: + ScopedGlobalSink(svs::logging::logger_ptr logger, spdlog::sink_ptr sink) + : logger_{std::move(logger)} + , sink_{std::move(sink)} { + logger_->sinks().push_back(sink_); + } + + ScopedGlobalSink(const ScopedGlobalSink&) = delete; + ScopedGlobalSink& operator=(const ScopedGlobalSink&) = delete; + ScopedGlobalSink(ScopedGlobalSink&&) = delete; + ScopedGlobalSink& operator=(ScopedGlobalSink&&) = delete; + + ~ScopedGlobalSink() { + auto& sinks = logger_->sinks(); + sinks.erase(std::remove(sinks.begin(), sinks.end(), sink_), sinks.end()); + } + + private: + svs::logging::logger_ptr logger_; + spdlog::sink_ptr sink_; +}; + +using Idx = uint32_t; +using Eltype = float; +using QueryEltype = float; +using Distance = svs::distance::DistanceL2; +const size_t N = 128; + +const size_t NUM_NEIGHBORS = 10; +const double TARGET_RECALL = 0.95; + +/// +/// Utility Methods +/// + +template I div(I i, float fraction) { + return svs::lib::narrow(std::floor(svs::lib::narrow(i) * fraction)); +} + +template std::string stringify(Args&&... args) { + std::ostringstream stream{}; + ((stream << args), ...); + return stream.str(); +} + +/// +/// @brief Compute the window size required to achieve the desired recall. +/// +template +auto find_windowsize( + MutableIndex& index, + const Groundtruth& groundtruth, + const Queries& queries, + double target_recall = TARGET_RECALL, + size_t window_lower = NUM_NEIGHBORS, + size_t window_upper = 1000 +) -> svs::index::vamana::VamanaSearchParameters { + auto range = svs::threads::UnitRange(window_lower, window_upper); + auto parameters = svs::index::vamana::VamanaSearchParameters(); + size_t window_size = *std::lower_bound( + range.begin(), + range.end(), + target_recall, + [&](size_t window_size, double recall) { + parameters.buffer_config(window_size); + auto result = + svs::index::search_batch_with(index, queries, NUM_NEIGHBORS, parameters); + auto this_recall = svs::k_recall_at_n(groundtruth, result); + return this_recall < recall; + } + ); + + parameters.buffer_config(window_size); + return parameters; +} + +/// +/// @brief A report regarding a mutating operation. +/// +struct Report { + template + Report( + double operation_time, + double groundtruth_time, + double search_time, + double recall, + Args&&... args + ) + : operation_time_{operation_time} + , groundtruth_time_{groundtruth_time} + , search_time_{search_time} + , recall_{recall} + , message_{stringify(std::forward(args)...)} {} + + Report( + double operation_time, + double groundtruth_time, + double search_time, + double recall, + std::string message + ) + : operation_time_{operation_time} + , groundtruth_time_{groundtruth_time} + , search_time_{search_time} + , recall_{recall} + , message_{std::move(message)} {} + + ///// Members + double operation_time_; + double groundtruth_time_; + double search_time_; + double recall_; + std::string message_; +}; + +std::ostream& operator<<(std::ostream& stream, const Report& report) { + stream << "[" << report.message_ << "] -- {" + << "operation: " << report.operation_time_ + << ", groundtruth: " << report.groundtruth_time_ + << ", search: " << report.search_time_ << ", recall: " << report.recall_ << "}"; + return stream; +} + +/// +/// Reference Dataset. +/// + +template +void do_check( + MutableIndex& index, + svs::misc::ReferenceDataset& reference, + const Queries& queries, + double operation_time, + std::string message, + bool calibrate = false +) { + // Compute groundtruth + auto tic = svs::lib::now(); + auto gt = reference.groundtruth(); + CATCH_REQUIRE(gt.n_neighbors() == NUM_NEIGHBORS); + CATCH_REQUIRE(gt.n_queries() == queries.size()); + + double groundtruth_time = svs::lib::time_difference(tic); + + if (calibrate) { + auto parameters = find_windowsize(index, gt, queries); + index.set_search_parameters(parameters); + } + + // Run search + tic = svs::lib::now(); + auto result = svs::index::search_batch(index, queries, NUM_NEIGHBORS); + double search_time = svs::lib::time_difference(tic); + + // Extra ID checks + reference.check_ids(result); + reference.check_equal_ids(index); + + // compute recall + double recall = svs::k_recall_at_n(gt, result, NUM_NEIGHBORS, NUM_NEIGHBORS); + + // Report the calibrated search window size if we calibrated this round. + if (calibrate) { + auto search_window_size = + index.get_search_parameters().buffer_config_.get_search_window_size(); + message += stringify(" - Calibrate window size: ", search_window_size); + } + + std::cout + << Report(operation_time, groundtruth_time, search_time, recall, std::move(message)) + << '\n'; +} + +/// +/// Main Loop. +/// + +template +void test_loop( + MutableIndex& index, + svs::misc::ReferenceDataset& reference, + const Queries& queries, + size_t num_points, + size_t consolidate_every, + size_t iterations +) { + size_t consolidate_count = 0; + for (size_t i = 0; i < iterations; ++i) { + // Add Points + { + auto [points, time] = reference.add_points(index, num_points); + CATCH_REQUIRE(points <= num_points); + CATCH_REQUIRE(points > num_points - reference.bucket_size()); + index.debug_check_invariants(true); + do_check(index, reference, queries, time, stringify("add ", points, " points")); + } + + // Delete Points + { + auto [points, time] = reference.delete_points(index, num_points); + CATCH_REQUIRE(points <= num_points); + CATCH_REQUIRE(points > num_points - reference.bucket_size()); + index.debug_check_invariants(true); + do_check( + index, reference, queries, time, stringify("delete ", points, " points") + ); + } + + // Maybe consolidate. + ++consolidate_count; + if (consolidate_count == consolidate_every) { + auto tic = svs::lib::now(); + index.consolidate(); + double diff = svs::lib::time_difference(tic); + index.debug_check_invariants(false); + do_check(index, reference, queries, diff, "consolidate"); + consolidate_count = 0; + + // Compact + tic = svs::lib::now(); + // Use a batchsize smaller than the whole dataset to ensure that the compaction + // algorithm correctly handles this case. + index.compact(reference.valid() / 10); + diff = svs::lib::time_difference(tic); + index.debug_check_invariants(false); + do_check(index, reference, queries, diff, "compact"); + } + } +} + +} // namespace + +CATCH_TEST_CASE( + "Testing Concurrent Graph Index", "[concurrent][graph_index][dynamic_index]" +) { + // Set hyper parameters here + const size_t max_degree = 64; +#if defined(NDEBUG) + const float initial_fraction = 0.25; + const float modify_fraction = 0.05; +#else + const float initial_fraction = 0.05; + const float modify_fraction = 0.005; +#endif + const size_t num_threads = 10; + const float alpha = 1.2; + + // Set up log + std::vector captured_logs; + std::vector captured_levels; + + auto callback_sink = std::make_shared( + [&captured_logs, &captured_levels](const spdlog::details::log_msg& msg) { + captured_logs.emplace_back(msg.payload.data(), msg.payload.size()); + captured_levels.push_back(svs::logging::detail::from_spdlog(msg.level)); + } + ); + callback_sink->set_level(spdlog::level::trace); + auto test_logger = std::make_shared("test_logger", callback_sink); + test_logger->set_level(spdlog::level::trace); + std::vector global_captured_logs; + auto global_callback_sink = std::make_shared( + [&global_captured_logs](const spdlog::details::log_msg& msg) { + global_captured_logs.emplace_back(msg.payload.data(), msg.payload.size()); + } + ); + global_callback_sink->set_level(spdlog::level::trace); + auto global_sink_guard = ScopedGlobalSink{svs::logging::get(), global_callback_sink}; + + // Load the base dataset and queries. + auto data = svs::data::SimpleData::load(test_dataset::data_svs_file()); + auto data_copy = data; + auto num_points = data.size(); + auto queries = test_dataset::queries(); + + auto reference = svs::misc::ReferenceDataset( + std::move(data), + Distance(), + num_threads, + div(num_points, 0.5 * modify_fraction), + NUM_NEIGHBORS, + queries, + 0x12345678, + test_logger + ); + + auto num_indices_to_add = div(reference.size(), initial_fraction); + std::cout << "Initializing with " << num_indices_to_add << " entries!\n"; + + // Construct a blocked dataset consisting of 50% of the base dataset. + auto data_mutable = cc::SegmentedBlockedData(num_indices_to_add, N); + std::vector initial_indices{}; + { + auto [vectors, indices] = reference.generate(num_indices_to_add); + // Copy assign ``initial_indices`` + auto num_points_added = indices.size(); + CATCH_REQUIRE(vectors.size() == num_points_added); + CATCH_REQUIRE(num_points_added <= num_indices_to_add); + CATCH_REQUIRE(num_points_added > num_indices_to_add - reference.bucket_size()); + + initial_indices = indices; + if (vectors.size() != num_indices_to_add || indices.size() != num_indices_to_add) { + throw ANNEXCEPTION("Something when horribly wrong!"); + } + + for (size_t i = 0; i < num_indices_to_add; ++i) { + data_mutable.set_datum(i, vectors.get_datum(i)); + } + } + + svs::index::vamana::VamanaBuildParameters parameters{ + 1.2, max_degree, 2 * max_degree, 1000, max_degree - 4, true}; + + auto tic = svs::lib::now(); + auto index = cc::MutableVamanaIndex( + parameters, + std::move(data_mutable), + initial_indices, + Distance(), + num_threads, + test_logger + ); + double build_time = svs::lib::time_difference(tic); + index.debug_check_invariants(false); + + CATCH_REQUIRE(captured_logs[0].find("Total / % Measured:") != std::string::npos); + CATCH_REQUIRE(captured_levels[0] == svs::logging::Level::Debug); + CATCH_REQUIRE(captured_logs[1].find("Vamana Build Parameters:") != std::string::npos); + CATCH_REQUIRE(captured_levels[1] == svs::logging::Level::Debug); + CATCH_REQUIRE(captured_logs[2].find("Number of syncs:") != std::string::npos); + CATCH_REQUIRE(captured_levels[2] == svs::logging::Level::Trace); + CATCH_REQUIRE(captured_logs[3].find("Batch Size:") != std::string::npos); + CATCH_REQUIRE(captured_levels[3] == svs::logging::Level::Trace); + + // Test get_distance functionality + svs::DistanceDispatcher dispatcher(svs::L2); + dispatcher([&](auto dist) { + svs_test::GetDistanceTester::test(index, dist, data_copy, initial_indices); + }); + + // Verify that we can get and set build parameters. + CATCH_REQUIRE(index.get_alpha() == alpha); + index.set_alpha(1.0); + CATCH_REQUIRE(index.get_alpha() == 1.0); + index.set_alpha(alpha); + CATCH_REQUIRE(index.get_alpha() == alpha); + + CATCH_REQUIRE(index.get_graph_max_degree() == max_degree); + + const size_t expected_construction_window = 2 * max_degree; + CATCH_REQUIRE(index.get_construction_window_size() == expected_construction_window); + index.set_construction_window_size(10); + CATCH_REQUIRE(index.get_construction_window_size() == 10); + index.set_construction_window_size(expected_construction_window); + CATCH_REQUIRE(index.get_construction_window_size() == expected_construction_window); + + CATCH_REQUIRE(index.get_max_candidates() == 1000); + index.set_max_candidates(750); + CATCH_REQUIRE(index.get_max_candidates() == 750); + + CATCH_REQUIRE(index.get_prune_to() == max_degree - 4); + index.set_prune_to(max_degree - 2); + CATCH_REQUIRE(index.get_prune_to() == max_degree - 2); + + CATCH_REQUIRE(index.get_full_search_history() == true); + index.set_full_search_history(false); + CATCH_REQUIRE(index.get_full_search_history() == false); + + reference.configure_extra_checks(true); + CATCH_REQUIRE(reference.extra_checks_enabled()); + + do_check( + index, + reference, + queries, + build_time, + stringify("initial build (", num_indices_to_add, ") points"), + true + ); + + test_loop(index, reference, queries, div(reference.size(), modify_fraction), 2, 6); + + // Try saving the index. + svs_test::prepare_temp_directory(); + auto tmp = svs_test::temp_directory(); + index.save(tmp / "config", tmp / "graph", tmp / "data"); + + auto reloaded = cc::auto_dynamic_assemble( + tmp / "config", + SVS_LAZY(ccg::SimpleBlockedGraph::load(tmp / "graph")), + SVS_LAZY(cc::SegmentedBlockedData::load(tmp / "data")), + svs::DistanceL2(), + 2 + ); + + do_check( + reloaded, + reference, + queries, + build_time, + stringify("initial build (", num_indices_to_add, ") points"), + true + ); + + reloaded = cc::auto_dynamic_assemble( + tmp / "config", + SVS_LAZY(ccg::SimpleBlockedGraph::load(tmp / "graph")), + SVS_LAZY(cc::SegmentedBlockedData::load(tmp / "data")), + svs::DistanceL2(), + svs::threads::CppAsyncThreadPool(2) + ); + + do_check( + reloaded, + reference, + queries, + build_time, + stringify("initial build (", num_indices_to_add, ") points"), + true + ); + + reloaded = cc::auto_dynamic_assemble( + tmp / "config", + SVS_LAZY(ccg::SimpleBlockedGraph::load(tmp / "graph")), + SVS_LAZY(cc::SegmentedBlockedData::load(tmp / "data")), + svs::DistanceL2(), + svs::threads::QueueThreadPoolWrapper(2) + ); + + do_check( + reloaded, + reference, + queries, + build_time, + stringify("initial build (", num_indices_to_add, ") points"), + true + ); + + // Make sure parameters were saved across the saving. + CATCH_REQUIRE(index.get_alpha() == reloaded.get_alpha()); + CATCH_REQUIRE(index.get_graph_max_degree() == reloaded.get_graph_max_degree()); + CATCH_REQUIRE(index.get_max_candidates() == reloaded.get_max_candidates()); + CATCH_REQUIRE( + index.get_construction_window_size() == reloaded.get_construction_window_size() + ); + CATCH_REQUIRE(index.get_prune_to() == reloaded.get_prune_to()); + CATCH_REQUIRE(index.get_full_search_history() == reloaded.get_full_search_history()); + CATCH_REQUIRE(index.size() == reloaded.size()); + // ID's preserved across runs. + index.on_ids([&](size_t e) { CATCH_REQUIRE(reloaded.has_id(e)); }); + + CATCH_REQUIRE(global_captured_logs.empty()); +} + +CATCH_TEST_CASE( + "Concurrent MutableVamanaIndex Per-Index Logging Test", "[concurrent][logging]" +) { + // Vector to store captured log messages + std::vector captured_logs; + std::vector global_captured_logs; + + // Create a callback sink to capture log messages + auto callback_sink = std::make_shared( + [&captured_logs](const spdlog::details::log_msg& msg) { + captured_logs.emplace_back(msg.payload.data(), msg.payload.size()); + } + ); + callback_sink->set_level(spdlog::level::trace); // Capture all log levels + + // Create a logger with the callback sink + auto test_logger = std::make_shared("test_logger", callback_sink); + test_logger->set_level(spdlog::level::trace); + + auto global_callback_sink = std::make_shared( + [&global_captured_logs](const spdlog::details::log_msg& msg) { + global_captured_logs.emplace_back(msg.payload.data(), msg.payload.size()); + } + ); + global_callback_sink->set_level(spdlog::level::trace); + + auto global_sink_guard = ScopedGlobalSink{svs::logging::get(), global_callback_sink}; + + // Setup index + auto data = svs::data::SimpleData::load(test_dataset::data_svs_file()); + std::vector initial_indices(data.size()); + std::iota(initial_indices.begin(), initial_indices.end(), 0); + svs::index::vamana::VamanaBuildParameters buildParams(1.2, 64, 10, 20, 10, true); + auto threadpool = svs::threads::DefaultThreadPool(1); + auto index = cc::MutableVamanaIndex( + buildParams, + std::move(data), + initial_indices, + svs::DistanceL2(), + std::move(threadpool), + test_logger + ); + + // Verify the internal log messages + CATCH_REQUIRE(global_captured_logs.empty()); + CATCH_REQUIRE(captured_logs[0].find("Vamana Build Parameters:") != std::string::npos); + CATCH_REQUIRE(captured_logs[1].find("Number of syncs:") != std::string::npos); + CATCH_REQUIRE(captured_logs[2].find("Batch Size:") != std::string::npos); +} + +CATCH_TEST_CASE( + "Concurrent MutableVamanaIndex Default Logger Test", "[concurrent][logging]" +) { + // Setup index with default logger + auto data = svs::data::SimpleData::load(test_dataset::data_svs_file()); + std::vector initial_indices(data.size()); + std::iota(initial_indices.begin(), initial_indices.end(), 0); + svs::index::vamana::VamanaBuildParameters buildParams(1.2, 64, 10, 20, 10, true); + auto threadpool = svs::threads::DefaultThreadPool(1); + auto index = cc::MutableVamanaIndex( + buildParams, + std::move(data), + initial_indices, + svs::DistanceL2(), + std::move(threadpool) + ); + + // Verify that the default logger is used + auto default_logger = svs::logging::get(); + CATCH_REQUIRE(index.get_logger() == default_logger); +} + +CATCH_TEST_CASE( + "Concurrent Vamana Index Default Parameters", "[concurrent][parameter][vamana]" +) { + using Catch::Approx; + std::filesystem::path data_path = test_dataset::data_svs_file(); + + CATCH_SECTION("L2 Distance Defaults") { + auto expected_result = test_dataset::vamana::expected_build_results( + svs::L2, svsbenchmark::Uncompressed(svs::DataType::float32) + ); + auto build_params = expected_result.build_parameters_.value(); + auto data_loader = svs::data::SimpleData::load(data_path); + + // Get IDs for all points in the dataset + std::vector indices(data_loader.size()); + std::iota(indices.begin(), indices.end(), 0); + + // Build dynamic index with L2 distance + auto index = cc::MutableVamanaIndex( + build_params, std::move(data_loader), indices, svs::distance::DistanceL2(), 2 + ); + + CATCH_REQUIRE(index.get_alpha() == Approx(svs::VAMANA_ALPHA_MINIMIZE_DEFAULT)); + } + + CATCH_SECTION("MIP Distance Defaults") { + auto expected_result = test_dataset::vamana::expected_build_results( + svs::MIP, svsbenchmark::Uncompressed(svs::DataType::float32) + ); + auto build_params = expected_result.build_parameters_.value(); + auto data_loader = svs::data::SimpleData::load(data_path); + + // Get IDs for all points in the dataset + std::vector indices(data_loader.size()); + std::iota(indices.begin(), indices.end(), 0); + + // Build dynamic index with MIP distance + auto index = cc::MutableVamanaIndex( + build_params, std::move(data_loader), indices, svs::distance::DistanceIP(), 2 + ); + + CATCH_REQUIRE(index.get_alpha() == Approx(svs::VAMANA_ALPHA_MAXIMIZE_DEFAULT)); + } + + CATCH_SECTION("Invalid Alpha for L2") { + auto expected_result = test_dataset::vamana::expected_build_results( + svs::L2, svsbenchmark::Uncompressed(svs::DataType::float32) + ); + auto build_params = expected_result.build_parameters_.value(); + build_params.alpha = 0.8f; + auto data_loader = svs::data::SimpleData::load(data_path); + + // Get IDs for all points in the dataset + std::vector indices(data_loader.size()); + std::iota(indices.begin(), indices.end(), 0); + + CATCH_REQUIRE_THROWS_WITH( + cc::MutableVamanaIndex( + build_params, + std::move(data_loader), + indices, + svs::distance::DistanceL2(), + 2 + ), + "For L2 distance, alpha must be >= 1.0" + ); + } + + CATCH_SECTION("Invalid Alpha for MIP") { + auto expected_result = test_dataset::vamana::expected_build_results( + svs::MIP, svsbenchmark::Uncompressed(svs::DataType::float32) + ); + auto build_params = expected_result.build_parameters_.value(); + build_params.alpha = 1.2f; + auto data_loader = svs::data::SimpleData::load(data_path); + + // Get IDs for all points in the dataset + std::vector indices(data_loader.size()); + std::iota(indices.begin(), indices.end(), 0); + + CATCH_REQUIRE_THROWS_WITH( + cc::MutableVamanaIndex( + build_params, + std::move(data_loader), + indices, + svs::distance::DistanceIP(), + 2 + ), + "For MIP/Cosine distance, alpha must be <= 1.0" + ); + } + + CATCH_SECTION("Invalid prune_to > graph_max_degree") { + auto expected_result = test_dataset::vamana::expected_build_results( + svs::L2, svsbenchmark::Uncompressed(svs::DataType::float32) + ); + auto build_params = expected_result.build_parameters_.value(); + build_params.prune_to = build_params.graph_max_degree + 10; + auto data_loader = svs::data::SimpleData::load(data_path); + + // Get IDs for all points in the dataset + std::vector indices(data_loader.size()); + std::iota(indices.begin(), indices.end(), 0); + + CATCH_REQUIRE_THROWS_WITH( + cc::MutableVamanaIndex( + build_params, + std::move(data_loader), + indices, + svs::distance::DistanceL2(), + 2 + ), + "prune_to must be <= graph_max_degree" + ); + } + + CATCH_SECTION("L2 Distance Empty Params") { + svs::index::vamana::VamanaBuildParameters params; + std::vector data(32); + for (size_t i = 0; i < data.size(); i++) { + data[i] = static_cast(i + 1); + } + auto data_view = svs::data::SimpleDataView(data.data(), 8, 4); + std::vector indices = {0, 1, 2, 3, 4, 5, 6, 7}; + auto index = cc::MutableVamanaIndex( + params, std::move(data_view), indices, svs::distance::DistanceL2(), 1 + ); + CATCH_REQUIRE(index.get_alpha() == Approx(svs::VAMANA_ALPHA_MINIMIZE_DEFAULT)); + CATCH_REQUIRE(index.get_graph_max_degree() == svs::VAMANA_GRAPH_MAX_DEGREE_DEFAULT); + CATCH_REQUIRE(index.get_prune_to() == svs::VAMANA_GRAPH_MAX_DEGREE_DEFAULT - 4); + CATCH_REQUIRE( + index.get_construction_window_size() == svs::VAMANA_WINDOW_SIZE_DEFAULT + ); + CATCH_REQUIRE(index.get_max_candidates() == 3 * svs::VAMANA_WINDOW_SIZE_DEFAULT); + CATCH_REQUIRE( + index.get_full_search_history() == svs::VAMANA_USE_FULL_SEARCH_HISTORY_DEFAULT + ); + } +} diff --git a/tests/svs/concurrent/graph.cpp b/tests/svs/concurrent/graph.cpp new file mode 100644 index 000000000..30cc1eb45 --- /dev/null +++ b/tests/svs/concurrent/graph.cpp @@ -0,0 +1,129 @@ +/* + * Copyright 2023 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// svs +#include "svs/concurrent/graph.h" +#include "svs/concurrent/graph_concepts.h" + +// test utils +#include "tests/utils/utils.h" + +// catch2 +#include "catch2/catch_test_macros.hpp" + +// stdlib +#include +#include +#include +#include + +// The concurrent index carries its own graph types (per-node sequence locks, +// per-node spin locks, atomic adjacency-list access, and the optional reverse-edge +// index). They shadow the same names in `svs::graphs`, which is left untouched. +namespace graphs = svs::index::vamana::concurrent::graphs; + +CATCH_TEST_CASE("Concurrent Simple Graph", "[concurrent][graphs][simple]") { + using Idx = uint32_t; + size_t n_nodes = 10; + const size_t max_degree = 5; + + auto graph = graphs::SimpleGraph(n_nodes, max_degree); + CATCH_REQUIRE(graph.n_nodes() == n_nodes); + CATCH_REQUIRE(graph.max_degree() == max_degree); + + // Ensure the constructor sets initializes the adjacency list sizes to zero. + for (Idx i = 0; i < graph.n_nodes(); ++i) { + CATCH_REQUIRE(graph.get_node_degree(i) == 0); + } + + // Test edge adding + auto check_increments = [n_nodes](auto&& span, Idx start) { + for (Idx i = 0; i < span.size(); ++i) { + CATCH_REQUIRE(span[i] == (start + i) % n_nodes); + } + }; + + for (size_t i = 0; i <= max_degree; ++i) { + // Verify that the current state of the adjacency lists is consistent. + for (Idx j = 0; j < n_nodes; ++j) { + CATCH_REQUIRE(graph.get_node_degree(j) == i); + auto l = graph.get_node(j); + CATCH_REQUIRE(l.size() == i); + check_increments(l, j + 1); + } + + for (Idx j = 0; j < n_nodes; ++j) { + bool should_be_added = (i < max_degree); + auto dst = (j + i + 1) % n_nodes; + CATCH_REQUIRE(!graph.has_edge(j, dst)); + graph.add_edge(j, dst); + + // Make sure that the edge is added (or not) depending on whether the graph + // is full or not. + if (should_be_added) { + CATCH_REQUIRE(graph.has_edge(j, dst)); + } else { + CATCH_REQUIRE(!graph.has_edge(j, dst)); + } + + // Filter out redundant assignments. + graph.add_edge(j, dst); + } + } + + // The last round should have added more than the acceptable number of neighbors. + // Make sure we didn't overwrite anything. + for (Idx j = 0; j < n_nodes; ++j) { + CATCH_REQUIRE(graph.get_node_degree(j) == max_degree); + auto l = graph.get_node(j); + CATCH_REQUIRE(l.size() == max_degree); + check_increments(l, j + 1); + } + + graph.reset(); + for (Idx j = 0; j < n_nodes; ++j) { + CATCH_REQUIRE(graph.get_node_degree(j) == 0); + } + + // Replace Node. + // Purposely use a replacement that is too large to verify the truncating logic. + { + std::vector replacement{5, 4, 3, 2, 1, 6, 7, 8, 9}; + graph.replace_node(0, std::span{replacement.data(), replacement.size()}); + CATCH_REQUIRE(graph.get_node_degree(0) == max_degree); + auto s = graph.get_node(0); + CATCH_REQUIRE(s.size() == max_degree); + std::array expected{5, 4, 3, 2, 1}; + CATCH_REQUIRE(expected.size() == max_degree); + CATCH_REQUIRE(std::equal(s.begin(), s.end(), expected.begin())); + // Make sure the next elements weren't touched. + CATCH_REQUIRE(graph.get_node_degree(1) == 0); + } + + // Now, use fewer than the max degree to make sure that works as well. + { + const size_t replacement_size = 3; + std::vector replacement{10, 7, 6}; + CATCH_REQUIRE(replacement.size() == replacement_size); + + Idx last = n_nodes - 1; + graph.replace_node(last, std::span{replacement.data(), replacement.size()}); + CATCH_REQUIRE(graph.get_node_degree(last) == replacement.size()); + auto s = graph.get_node(last); + CATCH_REQUIRE(s.size() == replacement.size()); + CATCH_REQUIRE(std::equal(s.begin(), s.end(), replacement.begin())); + } +} diff --git a/tests/svs/concurrent/iterator.cpp b/tests/svs/concurrent/iterator.cpp new file mode 100644 index 000000000..093eb3258 --- /dev/null +++ b/tests/svs/concurrent/iterator.cpp @@ -0,0 +1,422 @@ +/* + * Copyright 2024 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// header under test +#include "svs/concurrent/iterator.h" +#include "svs/concurrent/multi.h" + +// svstest +#include "tests/utils/test_dataset.h" +#include "tests/utils/vamana_reference.h" + +// catch2 +#include "catch2/catch_test_macros.hpp" + +// stl +#include +#include +#include + +// The concurrent index is a separate type from `svs::index::vamana::MutableVamanaIndex`; +// the pre-existing static and dynamic indexes are untouched by these tests. +namespace cc = svs::index::vamana::concurrent; +namespace ccg = svs::index::vamana::concurrent::graphs; + +namespace { + +// A static countdown to throwing an exception. +// This enables testing the restart functionality of the iterator. +size_t EXCEPTION_COUNTDOWN = 0; + +// A wrapper around the L2 distance that throws when the countdown reaches zero. +struct ThrowingL2 { + using compare = std::less<>; + static float compute(std::span left, std::span right) { + // If the exception countdown is active (non-zero), decrement the countdown. + // If this action caused the countdown to hit zero, throw an exception. + if (EXCEPTION_COUNTDOWN != 0) { + --EXCEPTION_COUNTDOWN; + if (EXCEPTION_COUNTDOWN == 0) { + throw ANNEXCEPTION("Exception countdown triggered!"); + } + } + auto real = svs::DistanceL2{}; + return svs::distance::compute(real, left, right); + } +}; + +} // namespace + +template <> struct svs::index::vamana::concurrent::PruneStrategy { + using type = ProgressivePruneStrategy; +}; + +namespace { + +const size_t QUERIES_TO_CHECK = 10; + +using ConcurrentGraph = ccg::SimpleBlockedGraph; +using ConcurrentData = cc::SegmentedBlockedData; + +// `test_dataset::graph()` / `test_dataset::data_f32()` return the pre-existing graph and +// dataset types. The concurrent index is built on the grow-stable variants, so load the +// same files into those instead. +auto concurrent_graph_loader() { + return svs::lib::Lazy([]() { + return ConcurrentGraph::load(test_dataset::graph_file()); + }); +} + +auto concurrent_data_loader() { + return svs::lib::Lazy([]() { + return ConcurrentData::load(test_dataset::data_svs_file()); + }); +} + +// The concurrent analogue of `test_dataset::vamana::load_dynamic_test_index`. +template auto load_concurrent_dynamic_test_index(const Distance& d) { + return cc::auto_dynamic_assemble( + test_dataset::vamana_config_file(), + concurrent_graph_loader(), + concurrent_data_loader(), + d, + 1, + true // debug_load_from_static + ); +} + +// Common test routines for the static and dynamic indexes. +template >> +void check( + Index& index, + svs::data::ConstSimpleDataView queries, + svs::data::ConstSimpleDataView groundtruth, + IDChecker& checker +) { + const size_t num_neighbors = 100; + // Through an exception during search every `throw_exception_every` batches. + const size_t throw_exception_every = 3; + const auto batchsizes = std::vector{{10, 20, 25, 50, 100}}; + + CATCH_REQUIRE(index.size() > num_neighbors); + auto p = svs::index::vamana::VamanaSearchParameters{ + {num_neighbors, num_neighbors}, false, 0, 0}; + + auto scratch = index.scratchspace(p); + + auto id_to_distance = std::unordered_map(); + auto id_buffer = std::vector(); + + CATCH_REQUIRE(checker(id_to_distance)); + + auto from_iterator = std::unordered_set(); + for (size_t query_index = 0; query_index < QUERIES_TO_CHECK; ++query_index) { + auto query = queries.get_datum(query_index); + + // Perform a single, full-precision search to obtain reference results. + index.search(query, scratch); + const auto& buffer = scratch.buffer; + + id_to_distance.clear(); + id_buffer.clear(); + for (const auto& neighbor : buffer) { + size_t id = [&]() -> size_t { + if constexpr (Index::needs_id_translation) { + return index.translate_internal_id(neighbor.id()); + } else { + return neighbor.id(); + } + }(); + id_to_distance.insert({id, neighbor.distance()}); + id_buffer.push_back(id); + } + + // Ensure we have reasonable recall between. + CATCH_REQUIRE( + svs::lib::count_intersect(id_buffer, groundtruth.get_datum(query_index)) >= + 0.9 * num_neighbors + ); + + // Begin performing batch searches. + for (auto batchsize : batchsizes) { + CATCH_REQUIRE(num_neighbors % batchsize == 0); + size_t num_batches = num_neighbors / batchsize; + + auto iterator = index.make_batch_iterator(query); + CATCH_REQUIRE(iterator.size() == 0); + iterator.next(batchsize); + + from_iterator.clear(); + size_t similar_count = 0; + + // IDs returned from the most recent batch. + // Keep track of this because we want to ensure that if an exception is thrown + // during search, the state of the iterator is unchanged. + auto ids_returned_this_batch = std::vector(); + for (size_t batch = 0; batch < num_batches; ++batch) { + // Make sure the batch number is the same. + CATCH_REQUIRE(iterator.batch_number() == batch + 1); + ids_returned_this_batch.clear(); + for (auto i : iterator) { + auto id = i.id(); + // Make sure that this ID has not been returned yet. + CATCH_REQUIRE(!from_iterator.contains(id)); + auto itr = id_to_distance.find(id); + if (itr != id_to_distance.end()) { + // Make sure the returned distances match. + CATCH_REQUIRE(itr->second == i.distance()); + ++similar_count; + } + + // Insert the ID into the `from_iterator` container to detect for + // duplicates from future calls. + from_iterator.insert(id); + ids_returned_this_batch.push_back(id); + } + + // The number of IDs returned should equal the number of IDs reported + // by the iterator. + CATCH_REQUIRE(ids_returned_this_batch.size() == iterator.size()); + CATCH_REQUIRE(ids_returned_this_batch.size() == batchsize); + + // Now we've extracted the neighbors, decide if we are going to try again + // but throw an exception. + // + // If so, we want to ensure that the buffer is left in a sane state. + // Furthermore, on the next iteration, we want to make sure we can resume + // search without incident. + if (batch % throw_exception_every == 0) { + EXCEPTION_COUNTDOWN = 50; + CATCH_REQUIRE_THROWS_AS(iterator.next(batchsize), svs::ANNException); + // The batch reported by the iterator must be unchanged. + CATCH_REQUIRE(iterator.batch_number() == batch + 1); + // The contents of the iterator should be unchanged. + CATCH_REQUIRE(iterator.size() == ids_returned_this_batch.size()); + CATCH_REQUIRE(std::equal( + iterator.begin(), + iterator.end(), + ids_returned_this_batch.begin(), + [](svs::NeighborLike auto left, size_t right) { + return left.id() == right; + } + )); + } + + iterator.next(batchsize); + } + + // Make sure the expected number of neighbors has been obtained. + CATCH_REQUIRE(from_iterator.size() == num_neighbors); + + // Ensure that the results returned by the iterator are "substantively similar" + // to those returned from the full search. + CATCH_REQUIRE(similar_count >= 0.98 * num_neighbors); + } + + // Invoke the checker on the IDs returned from the iterator. + CATCH_REQUIRE(checker(from_iterator)); + } +} + +template +void check( + Index& index, + svs::data::ConstSimpleDataView queries, + svs::data::ConstSimpleDataView groundtruth +) { + auto checker = svs::lib::Returns>(); + check(index, queries, groundtruth, checker); +} + +struct DynamicChecker { + DynamicChecker(const std::unordered_set& valid_ids) + : valid_ids_{valid_ids} {} + + // Check whether `id` is valid or not. + bool check(size_t id) { + seen_.insert(id); + return valid_ids_.contains(id); + } + + template bool operator()(const std::unordered_map& ids) { + for (const auto& itr : ids) { + if (!check(itr.first)) { + return false; + } + } + return true; + } + + template bool operator()(const std::unordered_set& ids) { + for (auto itr : ids) { + if (!check(itr)) { + return false; + } + } + return true; + } + + void clear() { seen_.clear(); } + + // Valid IDs + const std::unordered_set& valid_ids_; + std::unordered_set seen_; +}; + +} // namespace + +CATCH_TEST_CASE("Concurrent Vamana Iterator", "[concurrent][index][vamana][iterator]") { + // This tests the general behavior of the iterator for correctness. + // It is not concerned with whether the returned neighbors are accurate. + // + // That responsibility is delegated to the integration tests. + auto queries = test_dataset::queries(); + auto gt = test_dataset::groundtruth_euclidean(); + // No "Static Index" section: this stack adds a concurrent *dynamic* index only, and + // leaves `svs::index::vamana::VamanaIndex` untouched (it keeps its own test). + + // For the dynamic index, iterated search should honor the internal deleted state of + // IDs. + CATCH_SECTION("Dynamic Index") { + auto index = load_concurrent_dynamic_test_index(ThrowingL2()); + auto original = test_dataset::data_f32(); + + // Increase the number of threads to help a little with run time. + index.set_threadpool(svs::threads::DefaultThreadPool(2)); + auto itr = svs::threads::UnitRange{0, index.size()}; + auto valid_ids = std::unordered_set{itr.begin(), itr.end()}; + auto checker = DynamicChecker{valid_ids}; + check(index, queries.cview(), gt.cview(), checker); + + // Delete the best candidate for each of the test queries. + auto ids_to_delete = std::vector(); + for (size_t i = 0; i < QUERIES_TO_CHECK; ++i) { + auto nearest_neighbor = gt.get_datum(i).front(); + auto itr = + std::find(ids_to_delete.begin(), ids_to_delete.end(), nearest_neighbor); + if (itr == ids_to_delete.end()) { + ids_to_delete.push_back(nearest_neighbor); + CATCH_REQUIRE(valid_ids.erase(nearest_neighbor) == 1); + CATCH_REQUIRE(checker.seen_.contains(nearest_neighbor)); + } + } + + fmt::print("Deleting\n"); + index.delete_entries(ids_to_delete); + checker.clear(); + check(index, queries.cview(), gt.cview(), checker); + + for (auto id : ids_to_delete) { + CATCH_REQUIRE(!checker.seen_.contains(id)); + } + + // Compact and consolidate. + index.consolidate(); + index.compact(); + + fmt::print("Compacting\n"); + checker.clear(); + check(index, queries.cview(), gt.cview(), checker); + for (auto id : ids_to_delete) { + CATCH_REQUIRE(!checker.seen_.contains(id)); + } + + // Add back the points we deleted and try again. + fmt::print("Adding\n"); + auto slots = index.add_points( + svs::data::make_const_view(original, ids_to_delete), ids_to_delete + ); + + checker.clear(); + for (auto id : ids_to_delete) { + auto [_, inserted] = valid_ids.insert(id); + CATCH_REQUIRE(inserted); + } + + check(index, queries.cview(), gt.cview(), checker); + for (auto id : ids_to_delete) { + CATCH_REQUIRE(checker.seen_.contains(id)); + } + } + + // Multi-vector batch iterator should also pass non-multi-vector tests + CATCH_SECTION("Multi batch iterator") { + auto index = cc::auto_multi_dynamic_assemble( + test_dataset::vamana_config_file(), + concurrent_graph_loader(), + concurrent_data_loader(), + ThrowingL2(), + 1, + cc::MultiMutableVamanaLoad::FROM_STATIC + ); + auto original = test_dataset::data_f32(); + index.set_threadpool(svs::threads::DefaultThreadPool(2)); + auto itr = svs::threads::UnitRange{0, index.size()}; + auto valid_ids = std::unordered_set{itr.begin(), itr.end()}; + auto checker = DynamicChecker{valid_ids}; + check(index, queries.cview(), gt.cview(), checker); + + // Delete the best candidate for each of the test queries. + auto ids_to_delete = std::vector(); + for (size_t i = 0; i < QUERIES_TO_CHECK; ++i) { + auto nearest_neighbor = gt.get_datum(i).front(); + auto itr = + std::find(ids_to_delete.begin(), ids_to_delete.end(), nearest_neighbor); + if (itr == ids_to_delete.end()) { + ids_to_delete.push_back(nearest_neighbor); + CATCH_REQUIRE(valid_ids.erase(nearest_neighbor) == 1); + CATCH_REQUIRE(checker.seen_.contains(nearest_neighbor)); + } + } + + fmt::print("Deleting\n"); + index.delete_entries(ids_to_delete); + checker.clear(); + check(index, queries.cview(), gt.cview(), checker); + + for (auto id : ids_to_delete) { + CATCH_REQUIRE(!checker.seen_.contains(id)); + } + + // Compact and consolidate. + index.consolidate(); + index.compact(); + + fmt::print("Compacting\n"); + checker.clear(); + check(index, queries.cview(), gt.cview(), checker); + for (auto id : ids_to_delete) { + CATCH_REQUIRE(!checker.seen_.contains(id)); + } + + // Add back the points we deleted and try again. + fmt::print("Adding\n"); + auto slots = index.add_points( + svs::data::make_const_view(original, ids_to_delete), ids_to_delete + ); + + checker.clear(); + for (auto id : ids_to_delete) { + auto [_, inserted] = valid_ids.insert(id); + CATCH_REQUIRE(inserted); + } + + check(index, queries.cview(), gt.cview(), checker); + for (auto id : ids_to_delete) { + CATCH_REQUIRE(checker.seen_.contains(id)); + } + } +} diff --git a/tests/svs/concurrent/multi.cpp b/tests/svs/concurrent/multi.cpp new file mode 100644 index 000000000..509fe2cc3 --- /dev/null +++ b/tests/svs/concurrent/multi.cpp @@ -0,0 +1,475 @@ +/* + * Copyright 2025 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// header under test +#include "svs/concurrent/multi.h" + +// svstest +#include "tests/utils/test_dataset.h" +#include "tests/utils/vamana_reference.h" + +// catch2 +#include "catch2/catch_template_test_macros.hpp" +#include "catch2/catch_test_macros.hpp" + +// stl +#include +#include +#include +#include + +// The concurrent multi-label index is a separate type from +// `svs::index::vamana::MultiMutableVamanaIndex`; the pre-existing indexes are untouched. +namespace cc = svs::index::vamana::concurrent; +namespace ccg = svs::index::vamana::concurrent::graphs; + +namespace { + +// The concurrent index is built on grow-stable blocked storage. Copy a plain +// `SimpleData` into that type so the same source vectors can serve both as the +// reference index's dataset and as the concurrent index's. +template auto to_segmented(const Src& src) { + using T = typename Src::element_type; + auto dst = cc::SegmentedBlockedData(src.size(), src.dimensions()); + for (size_t i = 0; i < src.size(); ++i) { + dst.set_datum(i, src.get_datum(i)); + } + return dst; +} + +template float pick_alpha(Distance SVS_UNUSED(dist)) { + if constexpr (std::is_same_v) { + return 1.2; + } else if constexpr (std::is_same_v) { + return 0.95; + } else if constexpr (std::is_same_v) { + return 0.95; + } else { + throw ANNEXCEPTION("Unsupported distance type!"); + } +} + +} // namespace + +CATCH_TEMPLATE_TEST_CASE( + "Concurrent multi-vector dynamic vamana index", + "[long][concurrent][index][vamana][multi]", + svs::DistanceL2, + svs::DistanceIP, + svs::DistanceCosineSimilarity +) { + using Eltype = float; + using Distance = TestType; + const size_t N = 128; + const size_t max_degree = 64; + const float alpha = pick_alpha(Distance()); + const size_t num_threads = 4; + const size_t num_neighbors = 10; + + const auto data = svs::data::SimpleData::load(test_dataset::data_svs_file()); + const auto num_points = data.size(); + const auto queries = test_dataset::queries(); + const auto groundtruth = test_dataset::load_groundtruth(svs::distance_type_v); + + const svs::index::vamana::VamanaBuildParameters build_parameters{ + alpha, max_degree, 2 * max_degree, 1000, max_degree - 4, true}; + + const auto search_parameters = svs::index::vamana::VamanaSearchParameters(); + + const float epsilon = 0.05f; + std::vector ref_indices(num_points); + std::iota(ref_indices.begin(), ref_indices.end(), 0); + + auto ref_index = svs::index::vamana::MutableVamanaIndex( + build_parameters, data, ref_indices, Distance(), num_threads + ); + auto ref_results = svs::QueryResult(queries.size(), num_neighbors); + ref_index.search(ref_results.view(), queries.view(), search_parameters); + auto ref_recall = svs::k_recall_at_n(groundtruth, ref_results); + + // Original data label: + // 0 1 2 3 + // + // For each duplicate iteration, insert each vector with label increase by one + // Suppose we duplicate three times (i.e., num_duplicated = 3): + // 1 2 3 4 + // 2 3 4 5 + // 3 4 5 6 + // + // After deleting all the original labels, the remaining + // number of vectors will be : + // (num_duplicated * (num_duplicated + 1)) / 2 + // + // For the above examples, after deleting 0, 1, 2, 3 + // the remaining vectors becomes: + // 4 + // 4 5 + // 4 5 6 + // And the number of remaining vectors becomes + // (3 + 4) / 2 = 6 vectors + CATCH_SECTION("Insertion/Deletion in duplicated test datasets") { + const size_t num_duplicated = 3; + + std::vector test_indices(num_points); + std::iota(test_indices.begin(), test_indices.end(), 0); + + auto test_index = cc::MultiMutableVamanaIndex( + build_parameters, to_segmented(data), test_indices, Distance(), num_threads + ); + + for (size_t i = 0; i < num_duplicated; ++i) { + std::iota(test_indices.begin(), test_indices.end(), i + 1); + test_index.add_points(data, test_indices); + } + CATCH_REQUIRE(test_index.labelcount() == ref_index.size() + num_duplicated); + CATCH_REQUIRE(test_index.size() == ref_index.size() * (num_duplicated + 1)); + + std::iota(test_indices.begin(), test_indices.end(), 0); + test_index.delete_entries(test_indices); + CATCH_REQUIRE(test_index.labelcount() == num_duplicated); + CATCH_REQUIRE(test_index.size() == (num_duplicated * (num_duplicated + 1)) / 2); + } + CATCH_SECTION("Duplicated vectors with same labels") { + const size_t num_duplicated = 3; + + std::vector test_indices(num_points); + std::iota(test_indices.begin(), test_indices.end(), 0); + + auto test_index = cc::MultiMutableVamanaIndex( + build_parameters, to_segmented(data), test_indices, Distance(), num_threads + ); + + for (size_t i = 0; i < num_duplicated; ++i) { + test_index.add_points(data, test_indices); + } + CATCH_REQUIRE(test_index.labelcount() == test_indices.size()); + CATCH_REQUIRE(test_index.size() == test_indices.size() * (num_duplicated + 1)); + + auto test_results = svs::QueryResult(queries.size(), num_neighbors); + test_index.search(test_results.view(), queries.view(), search_parameters); + auto test_recall = svs::k_recall_at_n(groundtruth, test_results); + + CATCH_REQUIRE(test_recall > ref_recall - epsilon); + + test_index.delete_entries(test_indices); + CATCH_REQUIRE(test_index.labelcount() == 0); + CATCH_REQUIRE(test_index.size() == 0); + + test_index.add_points(data, test_indices); + test_index.consolidate(); + test_index.compact(); + for (size_t i = 0; i < num_duplicated; ++i) { + test_index.add_points(data, test_indices); + } + + auto test_results2 = svs::QueryResult(queries.size(), num_neighbors); + test_index.search(test_results2.view(), queries.view(), search_parameters); + auto test_recall2 = svs::k_recall_at_n(groundtruth, test_results2); + + CATCH_REQUIRE(test_recall2 > test_recall - epsilon); + CATCH_REQUIRE(test_recall2 < test_recall + epsilon); + } + + CATCH_SECTION("Step grouping") { + size_t start = 0; + size_t step = 4; + CATCH_REQUIRE(num_points % step == 0); + size_t num_groups = num_points / step; + + auto remapped_groundtruth = groundtruth; + CATCH_REQUIRE(remapped_groundtruth.size() == queries.size()); + + // It is okay to have duplicated neighbor ids in groundtruth + // as the recall is checked by counting intersect + for (size_t i = 0; i < queries.size(); ++i) { + auto arr = remapped_groundtruth.get_datum(i); + for (auto& each : arr) { + each /= step; + } + } + + std::vector test_indices(num_points); + for (size_t i = 0; i < num_points; i += step) { + for (size_t s = 0; s < step; ++s) { + test_indices[i + s] = start; + } + ++start; + } + + auto test_index = cc::MultiMutableVamanaIndex( + build_parameters, to_segmented(data), test_indices, Distance(), num_threads + ); + test_index.add_points(data, test_indices); + + auto test_results = svs::QueryResult(queries.size(), num_neighbors); + test_index.search(test_results.view(), queries.view(), search_parameters); + auto test_recall = svs::k_recall_at_n(remapped_groundtruth, test_results); + + CATCH_REQUIRE(test_recall > ref_recall - epsilon); + + // test get_distance + for (size_t i = 0; i < queries.size(); ++i) { + size_t k = std::rand() % num_groups; + double ref_distance = svs::INVALID_DISTANCE; + for (size_t s = 0; s < step; ++s) { + if constexpr (std::is_same_v) { + ref_distance = std::fmin( + ref_distance, + ref_index.get_distance( + ref_indices[k * step + s], queries.get_datum(i) + ) + ); + } else { + ref_distance = std::fmax( + ref_distance, + ref_index.get_distance( + ref_indices[k * step + s], queries.get_datum(i) + ) + ); + } + } + + double test_distance = + test_index.get_distance(test_indices[k * step], queries.get_datum(i)); + CATCH_REQUIRE(test_distance == ref_distance); + } + } + + CATCH_SECTION("Logging") { + std::vector test_indices(num_points); + std::iota(test_indices.begin(), test_indices.end(), 0); + + auto test_index = cc::MultiMutableVamanaIndex( + build_parameters, to_segmented(data), test_indices, Distance(), num_threads + ); + + CATCH_REQUIRE(ref_index.get_logger() == test_index.get_logger()); + } + + CATCH_SECTION("Save/Load") { + svs_test::prepare_temp_directory(); + auto dir = svs_test::temp_directory(); + auto config_dir = dir / "config"; + auto graph_dir = dir / "graph"; + auto data_dir = dir / "data"; + std::vector test_indices(num_points); + // Fill the test indices with labels in the range of num_labels + // to ensure that there are labels mapped to more than 1 vector. + const size_t per_label = 2; + const auto num_labels = num_points / per_label; + for (auto& i : test_indices) { + i = std::rand() % num_labels; + } + auto test_index = cc::MultiMutableVamanaIndex( + build_parameters, to_segmented(data), test_indices, Distance(), num_threads + ); + auto test_results = svs::QueryResult(queries.size(), num_neighbors); + test_index.search(test_results.view(), queries.view(), search_parameters); + auto test_recall = svs::k_recall_at_n(groundtruth, test_results); + + test_index.save(config_dir, graph_dir, data_dir); + + // `svs::GraphLoader` / `svs::VectorDataLoader` produce the pre-existing graph + // and dataset types. Reload into the concurrent stack's own types instead. + auto test_index_2 = cc::auto_multi_dynamic_assemble( + config_dir, + svs::lib::Lazy([&]() { + return ccg::SimpleBlockedGraph::load(graph_dir); + }), + svs::lib::Lazy([&]() { + return cc::SegmentedBlockedData::load(data_dir); + }), + Distance(), + svs::threads::CppAsyncThreadPool(2) + ); + auto test_results_2 = svs::QueryResult(queries.size(), num_neighbors); + test_index_2.search(test_results_2.view(), queries.view(), search_parameters); + auto test_recall_2 = svs::k_recall_at_n(groundtruth, test_results_2); + + // Check that the results are the same + CATCH_REQUIRE(test_results.n_neighbors() == test_results_2.n_neighbors()); + for (size_t i = 0; i < test_results.n_queries(); ++i) { + for (size_t j = 0; j < test_results.n_neighbors(); ++j) { + CATCH_REQUIRE( + test_results.indices().at(i, j) == test_results_2.indices().at(i, j) + ); + } + } + + CATCH_REQUIRE(test_index.size() == test_index_2.size()); + CATCH_REQUIRE(test_index.dimensions() == test_index_2.dimensions()); + // Index Properties + CATCH_REQUIRE(test_index.get_alpha() == test_index_2.get_alpha()); + CATCH_REQUIRE( + test_index.get_construction_window_size() == + test_index_2.get_construction_window_size() + ); + CATCH_REQUIRE(test_index.get_max_candidates() == test_index_2.get_max_candidates()); + CATCH_REQUIRE(test_index.max_degree() == test_index_2.max_degree()); + CATCH_REQUIRE(test_index.get_prune_to() == test_index_2.get_prune_to()); + CATCH_REQUIRE( + test_index.get_full_search_history() == test_index_2.get_full_search_history() + ); + CATCH_REQUIRE(test_index.view_data() == test_index_2.view_data()); + + CATCH_REQUIRE(test_recall_2 > test_recall - epsilon); + } +} + +CATCH_TEST_CASE( + "Concurrent MultiMutableVamana Index Save and Load", + "[concurrent][index][vamana][multi][saveload]" +) { + using Eltype = float; + using Distance = svs::DistanceL2; + const size_t N = 128; + const size_t num_threads = 4; + const size_t num_neighbors = 10; + const size_t max_degree = 64; + + const auto data = svs::data::SimpleData::load(test_dataset::data_svs_file()); + const auto num_points = data.size(); + const auto queries = test_dataset::queries(); + const auto groundtruth = test_dataset::load_groundtruth(svs::distance_type_v); + + const svs::index::vamana::VamanaBuildParameters build_parameters{ + 1.2, max_degree, 10, 20, 10, true}; + + const auto search_parameters = svs::index::vamana::VamanaSearchParameters(); + + const float epsilon = 0.05f; + + std::vector test_indices(num_points); + const size_t per_label = 2; + const auto num_labels = num_points / per_label; + for (auto& i : test_indices) { + i = std::rand() % num_labels; + } + + auto index = cc::MultiMutableVamanaIndex( + build_parameters, to_segmented(data), test_indices, Distance(), num_threads + ); + auto results = svs::QueryResult(queries.size(), num_neighbors); + index.search(results.view(), queries.view(), search_parameters); + + CATCH_SECTION("Load MultiMutableVamana Index being serialized natively to stream") { + std::stringstream stream; + index.save(stream); + { + auto deserializer = svs::lib::detail::Deserializer::build(stream); + CATCH_REQUIRE(deserializer.is_native()); + + using Data_t = cc::SegmentedBlockedData; + using GraphType = ccg::SimpleBlockedGraph; + + auto loaded = cc::auto_multi_dynamic_assemble( + stream, + [&]() -> GraphType { return GraphType::load(stream); }, + [&]() -> Data_t { return svs::lib::load_from_stream(stream); }, + Distance(), + num_threads + ); + + CATCH_REQUIRE(loaded.size() == index.size()); + CATCH_REQUIRE(loaded.dimensions() == index.dimensions()); + CATCH_REQUIRE(loaded.get_alpha() == index.get_alpha()); + CATCH_REQUIRE( + loaded.get_construction_window_size() == + index.get_construction_window_size() + ); + CATCH_REQUIRE(loaded.get_max_candidates() == index.get_max_candidates()); + CATCH_REQUIRE(loaded.max_degree() == index.max_degree()); + CATCH_REQUIRE(loaded.get_prune_to() == index.get_prune_to()); + CATCH_REQUIRE( + loaded.get_full_search_history() == index.get_full_search_history() + ); + CATCH_REQUIRE(loaded.view_data() == index.view_data()); + + auto loaded_results = svs::QueryResult(queries.size(), num_neighbors); + loaded.search(loaded_results.view(), queries.view(), search_parameters); + for (size_t i = 0; i < results.n_queries(); ++i) { + for (size_t j = 0; j < results.n_neighbors(); ++j) { + CATCH_REQUIRE( + results.indices().at(i, j) == loaded_results.indices().at(i, j) + ); + } + } + + auto loaded_recall = svs::k_recall_at_n(groundtruth, loaded_results); + auto test_recall = svs::k_recall_at_n(groundtruth, results); + CATCH_REQUIRE(loaded_recall > test_recall - epsilon); + } + } + + CATCH_SECTION("Load MultiMutableVamana Index being serialized with intermediate files" + ) { + std::stringstream stream; + svs::lib::UniqueTempDirectory tempdir{"svs_multivamana_save"}; + const auto config_dir = tempdir.get() / "config"; + const auto graph_dir = tempdir.get() / "graph"; + const auto data_dir = tempdir.get() / "data"; + std::filesystem::create_directories(config_dir); + std::filesystem::create_directories(graph_dir); + std::filesystem::create_directories(data_dir); + index.save(config_dir, graph_dir, data_dir); + svs::lib::DirectoryArchiver::pack(tempdir, stream); + { + using Data_t = cc::SegmentedBlockedData; + using GraphType = ccg::SimpleBlockedGraph; + + auto deserializer = svs::lib::detail::Deserializer::build(stream); + CATCH_REQUIRE(!deserializer.is_native()); + svs::lib::DirectoryArchiver::unpack(stream, tempdir, deserializer.magic()); + + auto loaded = cc::auto_multi_dynamic_assemble( + config_dir, + svs::lib::Lazy([&]() { return GraphType::load(graph_dir); }), + svs::lib::Lazy([&]() { return Data_t::load(data_dir); }), + Distance(), + num_threads + ); + + CATCH_REQUIRE(loaded.size() == index.size()); + CATCH_REQUIRE(loaded.dimensions() == index.dimensions()); + CATCH_REQUIRE(loaded.get_alpha() == index.get_alpha()); + CATCH_REQUIRE( + loaded.get_construction_window_size() == + index.get_construction_window_size() + ); + CATCH_REQUIRE(loaded.get_max_candidates() == index.get_max_candidates()); + CATCH_REQUIRE(loaded.max_degree() == index.max_degree()); + CATCH_REQUIRE(loaded.get_prune_to() == index.get_prune_to()); + CATCH_REQUIRE( + loaded.get_full_search_history() == index.get_full_search_history() + ); + CATCH_REQUIRE(loaded.view_data() == index.view_data()); + + auto loaded_results = svs::QueryResult(queries.size(), num_neighbors); + loaded.search(loaded_results.view(), queries.view(), search_parameters); + for (size_t i = 0; i < results.n_queries(); ++i) { + for (size_t j = 0; j < results.n_neighbors(); ++j) { + CATCH_REQUIRE( + results.indices().at(i, j) == loaded_results.indices().at(i, j) + ); + } + } + + auto loaded_recall = svs::k_recall_at_n(groundtruth, loaded_results); + auto test_recall = svs::k_recall_at_n(groundtruth, results); + CATCH_REQUIRE(loaded_recall > test_recall - epsilon); + } + } +} diff --git a/tests/svs/concurrent/prune.cpp b/tests/svs/concurrent/prune.cpp new file mode 100644 index 000000000..86f7ee198 --- /dev/null +++ b/tests/svs/concurrent/prune.cpp @@ -0,0 +1,59 @@ +/* + * Copyright 2023 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// header under test +#include "svs/concurrent/prune.h" + +// catch2 +#include "catch2/catch_test_macros.hpp" + +CATCH_TEST_CASE("Concurrent Pruning", "[concurrent][index][vamana]") { + namespace v = svs::index::vamana::concurrent; + // Protect against changes to the default strategies getting merged. + static_assert(std::is_same_v< + v::prune_strategy_t, + v::ProgressivePruneStrategy>); + static_assert(std::is_same_v< + v::prune_strategy_t, + v::IterativePruneStrategy>); + static_assert(std::is_same_v< + v::prune_strategy_t, + v::IterativePruneStrategy>); + + CATCH_SECTION("Iterative Strategy") { + CATCH_SECTION("Prune State") { + // The concurrent prune is a two-round heuristic rather than the alpha + // escalation loop: a neighbor that survives round one only conditionally is + // marked `Candidate` and reconsidered in round two, while `Pruned` is final. + // `reenable` therefore promotes `Candidate` -- not `Pruned` -- back to + // `Available`. The pre-existing `svs::index::vamana::PruneState` keeps the + // old three-state semantics and its own test. + CATCH_REQUIRE( + v::reenable(v::PruneState::Available) == v::PruneState::Available + ); + CATCH_REQUIRE(v::reenable(v::PruneState::Added) == v::PruneState::Added); + CATCH_REQUIRE(v::reenable(v::PruneState::Pruned) == v::PruneState::Pruned); + CATCH_REQUIRE( + v::reenable(v::PruneState::Candidate) == v::PruneState::Available + ); + + CATCH_REQUIRE(v::excluded(v::PruneState::Available) == false); + CATCH_REQUIRE(v::excluded(v::PruneState::Added) == true); + CATCH_REQUIRE(v::excluded(v::PruneState::Pruned) == true); + CATCH_REQUIRE(v::excluded(v::PruneState::Candidate) == true); + } + } +} diff --git a/tests/svs/concurrent/translation.cpp b/tests/svs/concurrent/translation.cpp new file mode 100644 index 000000000..d693d482c --- /dev/null +++ b/tests/svs/concurrent/translation.cpp @@ -0,0 +1,288 @@ +/* + * Copyright 2023 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// header under test. +#include "svs/concurrent/translation.h" + +// misc utilities +#include "svs/lib/narrow.h" + +// Test utils +#include "tests/utils/generators.h" + +// Catch2 +#include "catch2/catch_test_macros.hpp" + +// stl +#include + +// The concurrent index carries its own ID translator (transactional +// `replace_stale_and_insert` plus the non-throwing `_or` lookups). It is a +// distinct type from `svs::IDTranslator`, which is left untouched. +using IDTranslator = svs::index::vamana::concurrent::IDTranslator; + +namespace { +template +void check_contents( + const Begin& begin, const End& end, const External& external, const Internal& internal +) { + int64_t dist = std::distance(begin, end); + CATCH_REQUIRE(dist == svs::lib::narrow(external.size())); + CATCH_REQUIRE(dist == svs::lib::narrow(internal.size())); + + size_t count = 0; + for (auto i = begin; i != end; ++i, ++count) { + CATCH_REQUIRE(i->first == external[count]); + CATCH_REQUIRE(i->second == internal[count]); + } +} + +template +void check_contents( + const IDTranslator& translator, const External& external, const Internal& internal +) { + check_contents(translator.begin(), translator.end(), external, internal); +} + +template +void check_translation( + const IDTranslator& translator, const External& external, const Internal& internal +) { + CATCH_REQUIRE(translator.size() == external.size()); + CATCH_REQUIRE(translator.size() == internal.size()); + + for (size_t i = 0, imax = translator.size(); i < imax; ++i) { + CATCH_REQUIRE(translator.get_internal(external[i]) == internal[i]); + CATCH_REQUIRE(translator.get_external(internal[i]) == external[i]); + } +} + +void check( + const IDTranslator& translator, + std::vector& external_ids, + std::vector& internal_ids +) { + check_translation(translator, external_ids, internal_ids); + check_contents(translator, external_ids, internal_ids); +} + +} // namespace + +CATCH_TEST_CASE("Concurrent Translation Table", "[concurrent][translation]") { + CATCH_SECTION("Identity Constructor") { + auto translator = IDTranslator(IDTranslator::Identity(10)); + CATCH_REQUIRE(translator.size() == 10); + for (size_t i = 0; i < 10; ++i) { + CATCH_REQUIRE(translator.has_external(i)); + CATCH_REQUIRE(translator.has_internal(i)); + CATCH_REQUIRE(translator.get_external(i) == i); + CATCH_REQUIRE(translator.get_internal(i) == i); + } + } + + CATCH_SECTION("Basic Tests") { + auto translator = IDTranslator(); + CATCH_REQUIRE(translator.size() == 0); + + auto external_ids = std::vector{0, 2, 4, 6, 8}; + auto internal_ids = std::vector{0, 10, 20, 30, 40}; + + CATCH_REQUIRE(external_ids.size() == internal_ids.size()); + translator.insert(external_ids, internal_ids); + + check(translator, external_ids, internal_ids); + + /// + /// Adding more IDs + /// + CATCH_SECTION("Adding more entries") { + auto extra_external_ids = std::vector{10, 12, 14}; + auto extra_internal_ids = std::vector{50, 60, 70}; + translator.insert(extra_external_ids, extra_internal_ids); + external_ids.insert( + external_ids.end(), extra_external_ids.begin(), extra_external_ids.end() + ); + internal_ids.insert( + internal_ids.end(), extra_internal_ids.begin(), extra_internal_ids.end() + ); + + check(translator, external_ids, internal_ids); + } + + /// + /// Error Handling + /// + CATCH_SECTION("Mismatched number of entries") { + // Mismatched number of internal and external IDs. + auto external_mismatch_ids = std::vector{10, 12, 14}; + auto internal_mismatch_ids = std::vector{50, 60}; + CATCH_REQUIRE_THROWS_AS( + translator.insert(external_mismatch_ids, internal_mismatch_ids), + svs::ANNException + ); + + // State of the translator should be unchanged. + check(translator, external_ids, internal_ids); + } + + CATCH_SECTION("Repeat External IDs") { + auto external_mismatch_ids = std::vector{10, 12, 8}; + auto internal_mismatch_ids = std::vector{50, 60, 70}; + + CATCH_REQUIRE_THROWS_AS( + translator.insert(external_mismatch_ids, internal_mismatch_ids), + svs::ANNException + ); + + // State of the translator should be unchanged. + check(translator, external_ids, internal_ids); + } + + CATCH_SECTION("Non-unique External IDs") { + auto external_mismatch_ids = std::vector{10, 12, 10}; + auto internal_mismatch_ids = std::vector{50, 60, 70}; + + CATCH_REQUIRE_THROWS_AS( + translator.insert(external_mismatch_ids, internal_mismatch_ids), + svs::ANNException + ); + + // State of the translator should be unchanged. + check(translator, external_ids, internal_ids); + } + + CATCH_SECTION("replace_stale_and_insert is transactional on a live duplicate") { + // Batch: a fresh external id followed by a live (non-stale) duplicate. + // The whole batch must be validated before any mutation, so the throw + // leaves the translator byte-for-byte unchanged — in particular the + // fresh id `100` must NOT have been committed. + auto add_external = std::vector{100, 8}; // 8 already exists + auto add_internal = std::vector{50, 60}; + auto never_stale = [](uint32_t) { return false; }; + + CATCH_REQUIRE_THROWS_AS( + translator.replace_stale_and_insert( + add_external, add_internal, never_stale + ), + svs::ANNException + ); + + CATCH_REQUIRE_FALSE(translator.has_external(100)); // not committed + CATCH_REQUIRE_FALSE(translator.has_internal(50)); + check(translator, external_ids, internal_ids); // unchanged + } + + CATCH_SECTION("replace_stale_and_insert replaces a stale mapping") { + // Mark internal 20 (external 4) stale; re-add external 4 at a new + // internal id plus a genuinely fresh id. Both should commit. + auto add_external = std::vector{4, 100}; + auto add_internal = std::vector{50, 60}; + auto stale_if_20 = [](uint32_t i) { return i == 20; }; + + CATCH_REQUIRE_NOTHROW( + translator.replace_stale_and_insert(add_external, add_internal, stale_if_20) + ); + + CATCH_REQUIRE(translator.get_internal(4) == 50); + CATCH_REQUIRE(translator.get_internal(100) == 60); + CATCH_REQUIRE(translator.get_external(50) == 4); + CATCH_REQUIRE(translator.get_external(60) == 100); + // Old internal 20 no longer maps back to any external. + CATCH_REQUIRE_FALSE(translator.has_internal(20)); + } + + CATCH_SECTION("Repeat Internal IDs") { + auto external_mismatch_ids = std::vector{10, 12, 14}; + auto internal_mismatch_ids = std::vector{50, 10, 70}; + + CATCH_REQUIRE_THROWS_AS( + translator.insert(external_mismatch_ids, internal_mismatch_ids), + svs::ANNException + ); + + // State of the translator should be unchanged. + check(translator, external_ids, internal_ids); + } + + CATCH_SECTION("Repeat Internal IDs") { + auto external_mismatch_ids = std::vector{10, 12, 14}; + auto internal_mismatch_ids = std::vector{50, 60, 60}; + + CATCH_REQUIRE_THROWS_AS( + translator.insert(external_mismatch_ids, internal_mismatch_ids), + svs::ANNException + ); + + // State of the translator should be unchanged. + check(translator, external_ids, internal_ids); + } + + /// + /// Deletion + /// + + CATCH_SECTION("Delete External") { + auto external_id_delete = std::vector{2, 6}; + translator.delete_external(external_id_delete); + + external_ids = {0, 4, 8}; + internal_ids = {0, 20, 40}; + check(translator, external_ids, internal_ids); + } + + CATCH_SECTION("Delete External Error") { + // The external id `10` doesn't exist. + // We should get an error without modifying the underlying container. + auto external_id_delete = std::vector{2, 10}; + CATCH_REQUIRE_THROWS_AS( + translator.delete_external(external_id_delete), svs::ANNException + ); + check(translator, external_ids, internal_ids); + } + + CATCH_SECTION("Delete Internal") { + auto internal_id_delete = std::vector{0, 20, 40}; + translator.delete_internal(internal_id_delete); + + external_ids = {2, 6}; + internal_ids = {10, 30}; + check(translator, external_ids, internal_ids); + } + + CATCH_SECTION("Delete Internal Error") { + auto internal_id_delete = std::vector{0, 20, 2}; + CATCH_REQUIRE_THROWS_AS( + translator.delete_internal(internal_id_delete), svs::ANNException + ); + check(translator, external_ids, internal_ids); + } + + /// + /// Saving and loading + /// + + CATCH_SECTION("Saving and Reloading") { + svs_test::prepare_temp_directory(); + auto tempdir = svs_test::temp_directory(); + + svs::lib::save_to_disk(translator, tempdir); + auto reloaded = svs::lib::load_from_disk(tempdir); + + check(translator, external_ids, internal_ids); + check(reloaded, external_ids, internal_ids); + } + } +} diff --git a/tests/svs/lib/segmented_vector.cpp b/tests/svs/lib/segmented_vector.cpp new file mode 100644 index 000000000..6ffb1c4d5 --- /dev/null +++ b/tests/svs/lib/segmented_vector.cpp @@ -0,0 +1,237 @@ +/* + * Copyright 2026 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// header under test +#include "svs/lib/segmented_vector.h" + +// catch2 +#include "catch2/catch_test_macros.hpp" + +// stl +#include +#include +#include + +namespace { + +// Even modest element counts exercise many directory buckets, since bucket k holds +// only 1<; + +} // namespace + +CATCH_TEST_CASE("SegmentedVector basic semantics", "[core][segmented_vector]") { + CATCH_SECTION("Default construction is empty") { + SmallVec v; + CATCH_REQUIRE(v.size() == 0); + } + + CATCH_SECTION("Sized construction and read/write") { + SmallVec v(10); + CATCH_REQUIRE(v.size() == 10); + for (size_t i = 0; i < 10; ++i) { + v[i] = static_cast(i * 2); + } + for (size_t i = 0; i < 10; ++i) { + CATCH_REQUIRE(v[i] == static_cast(i * 2)); + } + } + + CATCH_SECTION("Fill construction") { + SmallVec v(7, 42); + CATCH_REQUIRE(v.size() == 7); + for (size_t i = 0; i < 7; ++i) { + CATCH_REQUIRE(v[i] == 42); + } + } + + CATCH_SECTION("resize grows and preserves existing elements") { + SmallVec v(5); + for (size_t i = 0; i < 5; ++i) { + v[i] = static_cast(i); + } + v.resize(100); + CATCH_REQUIRE(v.size() == 100); + for (size_t i = 0; i < 5; ++i) { + CATCH_REQUIRE(v[i] == static_cast(i)); + } + } + + CATCH_SECTION("resize with fill") { + SmallVec v(3, 1); + v.resize(20, 9); + CATCH_REQUIRE(v.size() == 20); + for (size_t i = 0; i < 3; ++i) { + CATCH_REQUIRE(v[i] == 1); + } + for (size_t i = 3; i < 20; ++i) { + CATCH_REQUIRE(v[i] == 9); + } + } + + CATCH_SECTION("shrink via resize destroys dropped elements; regrow default-constructs" + ) { + SmallVec v(50); + for (size_t i = 0; i < 50; ++i) { + v[i] = static_cast(i); + } + v.resize(10); + CATCH_REQUIRE(v.size() == 10); + for (size_t i = 0; i < 10; ++i) { + CATCH_REQUIRE(v[i] == static_cast(i)); + } + // Regrow: std::vector-like semantics — the dropped range was destroyed on + // shrink, so the regrown tail is default-constructed (0), not the old values. + v.resize(50); + CATCH_REQUIRE(v.size() == 50); + for (size_t i = 0; i < 10; ++i) { + CATCH_REQUIRE(v[i] == static_cast(i)); + } + for (size_t i = 10; i < 50; ++i) { + CATCH_REQUIRE(v[i] == 0); + } + } + + CATCH_SECTION("shrink_to frees segments and lowers size") { + SmallVec v(100); + for (size_t i = 0; i < 100; ++i) { + v[i] = static_cast(i); + } + v.shrink_to(7); + CATCH_REQUIRE(v.size() == 7); + for (size_t i = 0; i < 7; ++i) { + CATCH_REQUIRE(v[i] == static_cast(i)); + } + // Regrow after a real shrink: new elements default-constructed (0). + v.resize(20); + CATCH_REQUIRE(v.size() == 20); + for (size_t i = 0; i < 7; ++i) { + CATCH_REQUIRE(v[i] == static_cast(i)); + } + } +} + +CATCH_TEST_CASE("SegmentedVector copy and move", "[core][segmented_vector]") { + SmallVec v(30); + for (size_t i = 0; i < 30; ++i) { + v[i] = static_cast(i + 100); + } + + CATCH_SECTION("copy construction is a deep copy") { + SmallVec c(v); + CATCH_REQUIRE(c.size() == 30); + for (size_t i = 0; i < 30; ++i) { + CATCH_REQUIRE(c[i] == static_cast(i + 100)); + } + c[0] = -1; + CATCH_REQUIRE(v[0] == 100); // original unaffected + } + + CATCH_SECTION("move construction transfers contents") { + SmallVec m(std::move(v)); + CATCH_REQUIRE(m.size() == 30); + for (size_t i = 0; i < 30; ++i) { + CATCH_REQUIRE(m[i] == static_cast(i + 100)); + } + } +} + +// Address stability: a reference to v[i] obtained before a grow must remain valid +// (point to the same storage) after the grow. This is the core invariant Option C +// relies on. +CATCH_TEST_CASE( + "SegmentedVector grow preserves element addresses", "[core][segmented_vector]" +) { + SmallVec v(8); + for (size_t i = 0; i < 8; ++i) { + v[i] = static_cast(i); + } + std::vector addrs; + for (size_t i = 0; i < 8; ++i) { + addrs.push_back(&v[i]); + } + // Force many segment + directory-bucket allocations. + v.resize(10000); + for (size_t i = 0; i < 8; ++i) { + CATCH_REQUIRE(&v[i] == addrs[i]); // address unchanged + CATCH_REQUIRE(*addrs[i] == static_cast(i)); // value intact + } +} + +// Concurrent grow vs. read: one writer repeatedly grows the vector and fills the new +// elements, publishing a separate `published` counter ONLY after the fill — exactly the +// "publish after construct" contract the dynamic index uses (num_valid_ bumped after a +// slot is fully Valid). Readers read every index < published and assert v[i] == i with +// no tolerance. This both (a) catches use-after-free / torn outer-pointer reads on grow +// and (b) is a hard value invariant. Under TSan it must report no data races. +CATCH_TEST_CASE("SegmentedVector concurrent grow and read", "[core][segmented_vector]") { + using Vec = svs::lib::SegmentedVector; + constexpr size_t kInitial = 100; + constexpr size_t kFinal = 200000; + constexpr size_t kReaders = 8; + + Vec v(kInitial); + for (size_t i = 0; i < kInitial; ++i) { + v[i] = static_cast(i); + } + + std::atomic start{false}; + std::atomic writer_done{false}; + std::atomic published{kInitial}; + std::atomic failure{false}; + + auto reader = [&]() { + while (!start.load(std::memory_order_acquire)) {} + while (!writer_done.load(std::memory_order_acquire)) { + size_t n = published.load(std::memory_order_acquire); + size_t step = n / 256 + 1; + for (size_t i = 0; i < n; i += step) { + if (v[i] != static_cast(i)) { + failure.store(true, std::memory_order_relaxed); + } + } + } + }; + + std::vector readers; + for (size_t r = 0; r < kReaders; ++r) { + readers.emplace_back(reader); + } + + std::thread writer([&]() { + while (!start.load(std::memory_order_acquire)) {} + size_t cur = kInitial; + while (cur < kFinal) { + size_t next = std::min(cur + 137, kFinal); + v.resize(next); // allocate + publish segments (grow) + for (size_t i = cur; i < next; ++i) { // fill new elements + v[i] = static_cast(i); + } + published.store(next, std::memory_order_release); // publish after fill + cur = next; + } + writer_done.store(true, std::memory_order_release); + }); + + start.store(true, std::memory_order_release); + writer.join(); + for (auto& t : readers) { + t.join(); + } + + CATCH_REQUIRE(v.size() == kFinal); + CATCH_REQUIRE_FALSE(failure.load(std::memory_order_relaxed)); +}