diff --git a/AGENTS.md b/AGENTS.md index 9d2ed6e3..b1b37985 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -135,6 +135,12 @@ Required: - Follow `.clang-format`, `.clang-tidy`, `.editorconfig`, and surrounding naming conventions: `PascalCase` types, `snake_case` functions/variables/members, and `kPascalCase` constants. +- Across indexes, use the [index naming glossary](docs/development.md#index-variable-names) + for shared concepts. Distinguish capacity, stored/live counts, internal IDs/external labels, + original/padded dimensions, and storage sizes/offsets. Storage offsets default to bytes; + keep `*_offset_` names without a redundant `_bytes` suffix. Naming-only work must preserve + public/protected names, Python keywords, types, defaults, and serialized layouts unless + a compatibility change is explicitly approved. - Public API misuse must produce descriptive exceptions rather than `exit()`, `abort()`, or unconditional diagnostics. Assertions are for internal invariants. Keep library output opt-in. - Make ownership and non-owning pointer lifetimes explicit. Wrap legacy owning factory pointers diff --git a/docs/development.md b/docs/development.md index 12ec35a3..7412d483 100644 --- a/docs/development.md +++ b/docs/development.md @@ -71,6 +71,44 @@ To format only lines changed in the staged files, use: The complete-file formatter remains useful before the initial formatting pass or after changing `.clang-format`; CI always checks complete files. +## Index variable names + +Use the same name for the same role across IVF, HNSW and SymphonyQG. Choose the +role and unit first; similar-looking parameters can have different contracts. +These conventions apply to new internal variables and compatible renaming work. + +| Concept | Name and meaning | +| --- | --- | +| Dimensions | `dim` is the original dimension at an index boundary; `padded_dim` is the rotated, padded dimension. A lower-level helper's `dim` describes its own input shape. | +| Counts | `num_points` for an input batch or a builder's point set; `num_queries` and `num_clusters` for those input counts. A partial encoding batch uses `batch_size`. | +| Index state | `max_elements_` for configured logical capacity; `num_stored_points_` for stored records, including removals. The public `num_points()` counts live records. QG has no removal and derives both counts from capacity only after construction completes. | +| Point identity | `internal_id` or `node_id` for graph slots; `label` for a caller-visible HNSW label; `cluster_id` for centroid routing. Never assume an external label equals an internal slot. | +| Search outputs | `k` for requested result count, `result_ids` for ID-only output arrays, `distances` for their scores. `results` may hold complete distance/ID records. | +| Search state | `candidate_set` for indexing traversal candidates and `search_pool` for querying traversal candidates, `result_pool` for retained results (HNSW retains `nearest_candidates` for query results and `top_candidates` for construction/pruning candidates), `query_wrapper` for query preprocessing objects, `visited` for visited state. | +| Search windows | `ef_search` for a per-call graph window, `default_ef_search_` for a stored default, and `ef_construction` for graph-building search breadth. IVF's `nprobe` counts clusters and remains distinct. | +| Distances | `query_centroid_ip` for a dot product, `residual_norm` for a residual norm, `estimated_distances` for estimated score arrays, `lower_bounds` for pruning bounds, and `dist_k` for the current result cutoff. Do not conflate squared L2 with L2 or an IP distance with a dot product. | +| Storage | Name byte sizes after the stored object: HNSW's `level0_record_bytes_` is one complete level-0 record; `upper_layer_links_bytes_` and `level0_links_bytes_` are one node's adjacency-list sizes at the respective levels, including their headers. QG uses `record_bytes_` for one complete record. Storage offsets use `*_offset_`: bytes are implicit, so do not add a `_bytes` suffix. PiPNN scratch uses `points_buffer_size`, `distances_buffer_size` and `worker_buffer_size`, with a declaration comment specifying float-element counts. | +| Threads and encoding | `num_threads` for a requested limit, `num_workers` for the resolved worker count requested from the runtime (OpenMP may use fewer threads), and `fast_quantization` for the fixed-scale encoding option. | + +Keep short loop indices and conventional mathematical symbols where their role +is local and unambiguous. HNSW's `M` remains its algorithm-specific configuration; +its upper and base layers use `upper_layer_max_degree_` and +`base_layer_max_degree_`. QG's `max_degree_` is its fixed row degree, a multiple +of 32. These settings are not interchangeable. + +Public names and Python keywords are compatibility contracts. For example, +`Cluster::num()`, the protected `Initializer::num_cluster_`, QG's +`degree_bound()`/`set_ef()`, Python's `ef`, and `QGKMeansParameters::ef_build` +retain their existing spellings. C++ parameter names may be clarified without +changing types, order or defaults. Keep declarations, definitions and comments +in agreement; do not rename methods or literals by textual substitution. + +Storage-mode parameters also retain their contracts: IVF/HNSW use `nbits()` +with separate internal `ex_bits_`; IVF additionally accepts 32 for raw-vector +reranking. QG uses `quantization_bits()` with 0 for raw storage and 4/8 for packed +vectors. Naming work does not change these values, allocation sizes, numeric +operations, serialized field order, or ownership. + ## Static analysis clang-tidy performs semantic checks and is kept separate from clang-format. diff --git a/docs/docs/index/hnsw.md b/docs/docs/index/hnsw.md index c7d9b43e..6538ca0a 100644 --- a/docs/docs/index/hnsw.md +++ b/docs/docs/index/hnsw.md @@ -88,40 +88,44 @@ input pointer and search candidate helpers are private implementation details; see the [C++ migration guide](../modular_migration.md#public-and-implementation-interfaces). ```cpp -HierarchicalNSW::construct(size_t cluster_num, - const float* centroids, - size_t data_num, - const float* data, - PID* cluster_ids, - size_t num_threads = 0, - bool faster = false); +void HierarchicalNSW::construct( + size_t num_clusters, + const float* centroids, + size_t num_points, + const float* data, + PID* cluster_ids, + size_t num_threads = 0, + bool fast_quantization = false +); ``` - **data**: Vectors in the original input coordinates. -- **data_num**: The number of data vectors. -- **centroids**: Centroids computed by [RaBitQKMeans](../clustering.md#rabitqkmeans); the examples use `cluster_num = 16`. -- **cluster_ids**: Array of length `data_num`; every entry must be in the range `[0, cluster_num)`. +- **num_points**: The number of data vectors. +- **centroids**: Centroids computed by [RaBitQKMeans](../clustering.md#rabitqkmeans); the examples use `num_clusters = 16`. +- **cluster_ids**: Array of length `num_points`; every entry must be in the range `[0, num_clusters)`. - **num_threads**: Number of threads to use (default: 0, which auto-selects). -- **faster**: Use calibrated scaling instead of per-vector scale search; see [encoding paths](../rabitq/quantizer.md#implementations). +- **fast_quantization**: Use calibrated scaling instead of per-vector scale search; see [encoding paths](../rabitq/quantizer.md#implementations). ## Querying ```cpp -std::vector>> HierarchicalNSW::search(const float* queries, - size_t query_num, - size_t TOPK, - size_t efSearch, - size_t thread_num); +std::vector>> HierarchicalNSW::search( + const float* queries, + size_t num_queries, + size_t k, + size_t ef_search, + size_t num_threads +); ``` - **queries**: Query vectors in the original input coordinates. -- **query_num**: The number of query vectors. -- **TOPK**: The number of nearest neighbors to search. -- **efSearch**: The size of the candidate set for searching HNSW-RaBitQ base layer. -- **thread_num**: Number of threads to use. Each query is processed by one thread. +- **num_queries**: The number of query vectors. +- **k**: The number of nearest neighbors to search. +- **ef_search**: The size of the candidate set for searching HNSW-RaBitQ base layer. +- **num_threads**: Number of threads to use. Each query is processed by one thread. The search window belongs to each call, so concurrent C++ searches can use -different `efSearch` values and separate outputs while the index is unchanged. +different `ef_search` values and separate outputs while the index is unchanged. Python follows the shared [concurrency contract](usage.md#threading-and-file-paths). ## Updating an index @@ -134,21 +138,21 @@ the index file format; files with removals have ```c++ std::vector HierarchicalNSW::add( const float* data, - size_t n, + size_t num_points, const PID* cluster_ids = nullptr, - bool faster = false, + bool fast_quantization = false, size_t num_threads = 1 ); void HierarchicalNSW::resize(size_t new_max_elements); ``` -- **data**, **n**: `n` new vectors in the original coordinates. Vector `i` receives +- **data**, **num_points**: `num_points` new vectors in the original coordinates. Vector `i` receives the label `num_stored_points() + i`, so labels stay dense, and `add` returns them. - **cluster_ids**: The cluster of each new vector, in `[0, num_clusters)`. When it is `nullptr`, each vector goes to its nearest centroid, chosen the same way a query is routed. -- **faster**, **num_threads**: Same meaning as in `construct`, but `add` defaults +- **fast_quantization**, **num_threads**: Same meaning as in `construct`, but `add` defaults to one thread; pass `0` to auto-select. A point keeps the label `num_stored_points() + i` whatever order the threads finish in, because the slots are reserved as one block before the inserts start. @@ -203,10 +207,10 @@ of a rebuild. ### Removing points ```c++ -size_t HierarchicalNSW::remove(const PID* labels, size_t n); +size_t HierarchicalNSW::remove(const PID* labels, size_t num_points); ``` -- **labels**, **n**: Labels of the points to remove, as returned by `add` and +- **labels**, **num_points**: `num_points` labels of the points to remove, as returned by `add` and `search`. Every label must be in the index; nothing is removed if one is not. `remove` returns how many points were newly removed, so repeating a label or a call is safe. @@ -227,15 +231,15 @@ exclusive access as other updates. #### Choosing `ef` after removals -Removed points still take slots in the `efSearch` candidate set, so a search with -`efSearch` close to `k` finds fewer live points as removals accumulate. C++ search -uses `max(efSearch, k)`. Python uses `max(k, 10)` when `ef` is omitted or zero; +Removed points still take slots in the `ef_search` candidate set, so a search with +`ef_search` close to `k` finds fewer live points as removals accumulate. C++ search +uses `max(ef_search, k)`. Python uses `max(k, 10)` when `ef` is omitted or zero; an explicit positive `ef` is also raised to at least `k`. To keep about the same number of live candidates as before the removals, scale the effective window by the share of live points: ```text -max(efSearch, k) * num_stored_points() / (num_stored_points() - removed points) +max(ef_search, k) * num_stored_points() / (num_stored_points() - removed points) ``` Callers that track how many points they removed can apply this directly. @@ -298,9 +302,9 @@ In the upper layers of HNSW-RaBitQ, we compute the 1-bit estimated distance (usi In the base layer, we apply an adaptive re-ranking strategy: -- **candidate_set**: A distance-sorted exploration window with capacity - `max(efSearch, TOPK)`, containing both expanded and unexpanded candidates. -- **boundedKNN**: Current best TOPK candidates. +- **search_pool**: A distance-sorted exploration window with capacity + `max(ef_search, k)`, containing both expanded and unexpanded candidates. +- **nearest_candidates**: Current best k candidates. Repeat while the exploration window has an unexpanded candidate: @@ -309,21 +313,21 @@ Repeat while the exploration window has an unexpanded candidate: 2. Visit all unvisited neighbors of `e`. 3. For each neighbor: - Compute the 1-bit lower-bound distance (along with 1-bit estimated distance) using `BinData`. - - For a nonremoved neighbor, if `boundedKNN` has fewer than TOPK elements + - For a nonremoved neighbor, if `nearest_candidates` has fewer than k elements or the neighbor's lower bound is smaller than the worst result's estimated distance: 1. Refine the distance estimate using `ExData` when extra bits are present. - 2. Update `boundedKNN`. - - Insert the neighbor into `candidate_set` only if the window has space or + 2. Update `nearest_candidates`. + - Insert the neighbor into `search_pool` only if the window has space or its (possibly refined) estimated distance is no greater than the worst distance in the window. Inserting into a full window displaces its last entry. The search stops when no unexpanded candidates remain in the window, even though -the window still contains entries. `efSearch` controls this exploration budget -separately from the TOPK result count; increasing it allows a wider search at +the window still contains entries. `ef_search` controls this exploration budget +separately from the k result count; increasing it allows a wider search at additional cost, rather than requiring traversal of the entire graph. -A point hidden by `remove` never enters `boundedKNN`, so search never returns it, -but it can still enter `candidate_set` under the same window admission rule, +A point hidden by `remove` never enters `nearest_candidates`, so search never returns it, +but it can still enter `search_pool` under the same window admission rule, allowing search to walk through it. A removed entry point keeps the estimate already computed for it (the full-bits one when nbits > 1). A removed neighbor is not refined and uses its 1-bit estimate for window admission. diff --git a/docs/docs/index/ivf.md b/docs/docs/index/ivf.md index 64ad3c51..804009dd 100644 --- a/docs/docs/index/ivf.md +++ b/docs/docs/index/ivf.md @@ -113,9 +113,9 @@ using index_type = rabitqlib::ivf::IVF; size_t num_points = data.rows(); size_t dim = data.cols(); -size_t k = centroids.rows(); +size_t num_clusters = centroids.rows(); -index_type ivf(num_points, dim, k, total_bits); +index_type ivf(num_points, dim, num_clusters, total_bits); ``` ### Centroid routing @@ -131,7 +131,7 @@ print(index.initializer) # resolved routing type ``` ```c++ -index_type ivf(num_points, dim, k, total_bits, rabitqlib::METRIC_L2, +index_type ivf(num_points, dim, num_clusters, total_bits, rabitqlib::METRIC_L2, rabitqlib::RotatorType::FhtKacRotator, rabitqlib::ivf::InitializerType::FlatRaBitQ); auto routing = ivf.initializer_type(); @@ -174,10 +174,10 @@ Finally, call the construct API: ```c++ void IVF::construct( - const float* data, - const float* centroids, - const PID* cluster_ids, - bool faster = false, + const float* data, + const float* centroids, + const PID* cluster_ids, + bool fast_quantization = false, size_t num_threads = std::numeric_limits::max() ); ``` @@ -185,7 +185,7 @@ void IVF::construct( - **data**: Vectors in the original input coordinates. - **centroids**: One centroid per configured cluster, in the original input coordinates. - **cluster_ids**: One ID per input vector, in `[0, num_clusters())`. -- **faster**: Use a calibrated scaling factor instead of per-vector scale search; defaults to `false`. See [encoding paths](../rabitq/quantizer.md#implementations). +- **fast_quantization**: Use a calibrated scaling factor instead of per-vector scale search; defaults to `false`. See [encoding paths](../rabitq/quantizer.md#implementations). - **num_threads**: Maximum number of OpenMP threads used to quantize clusters. For example: @@ -207,19 +207,19 @@ Search a constructed or loaded index with: ```c++ void IVF::search( - const float* __restrict__ query, - size_t k, - size_t nprobe, - PID* __restrict__ results, - float* dists = nullptr + const float* query, + size_t k, + size_t nprobe, + PID* result_ids, + float* distances = nullptr ) const; ``` - **query**: Query vector. - **k**: Top-k. - **nprobe**: The number of closest clusters to search. -- **results**: Result buffer, size of k. -- **dists**: Optional distance buffer, size of `k`. +- **result_ids**: Output ID buffer, size of `k`. +- **distances**: Optional distance buffer, size of `k`. FastScan precision is selected automatically: HACC for 4–9 quantized bits, and standard FastScan for 1–3 bits or raw vectors (`32`). HACC reduces lookup-table @@ -271,25 +271,25 @@ or loaded index without the original dataset. Updates preserve the file format. ```c++ void IVF::add( const float* data, - size_t n, + size_t num_points, const PID* cluster_ids = nullptr, - bool faster = false, + bool fast_quantization = false, size_t num_threads = std::numeric_limits::max() ); -size_t IVF::remove(const PID* ids_to_remove, size_t n); +size_t IVF::remove(const PID* ids_to_remove, size_t num_points); void IVF::reserve(PID cluster_id, size_t num_points); ``` -- **data**, **n**: `n` new vectors, quantized like `construct` does. Vector `i` +- **data**, **num_points**: `num_points` new vectors, quantized like `construct` does. Vector `i` receives the PID `num_stored_points() + i`, using the count before the call. Existing IDs stay unchanged, and removed IDs are never reused. -- **cluster_ids**: The cluster of each new vector, in `[0, cluster_num)`. When +- **cluster_ids**: The cluster of each new vector, in `[0, num_clusters)`. When it is `nullptr`, vectors use the same centroid routing as queries. Routing uses the initializer selected at construction or restored from the index file. -- **faster**, **num_threads**: Same as in `construct`. -- **ids_to_remove**: PIDs to remove. Every PID must be below `num_stored_points()`; nothing is +- **fast_quantization**, **num_threads**: Same as in `construct`. +- **ids_to_remove**: `num_points` PIDs to remove. Every PID must be below `num_stored_points()`; nothing is removed if one is not. `remove` returns how many were newly removed and can be repeated safely. - **cluster_id**, **num_points**: `reserve` preallocates room for cluster `cluster_id` diff --git a/docs/docs/index/qg.md b/docs/docs/index/qg.md index 50aeea61..c31ccae6 100644 --- a/docs/docs/index/qg.md +++ b/docs/docs/index/qg.md @@ -94,21 +94,22 @@ class template argument deduction: ```cpp QuantizedGraph( - size_t num, size_t dim, size_t max_deg, + size_t max_elements, size_t dim, size_t max_degree, MetricType metric_type = METRIC_L2, RotatorType rotator_type = RotatorType::FhtKacRotator, size_t quantization_bits = 0, uint32_t seed = std::random_device{}() ); QGBuilder( - QuantizedGraph& index, uint32_t ef_build, const float* data, + QuantizedGraph& index, uint32_t ef_construction, const float* data, size_t num_threads = std::numeric_limits::max(), QGInitialization init = QGInitialization::PiPNN, uint32_t seed = std::random_device{}(), bool cache_vectors = false ); ``` -`data` contains `num * dim` floats; `max_deg` has Python's `max_degree` -constraints. C++ defaults to all available threads. `build()` runs one refinement +`data` contains `max_elements * dim` floats. This is the configured point count; +QG has no spare capacity for later appends. `max_degree` has the same constraints +as Python's parameter. C++ defaults to all available threads. `build()` runs one refinement for PiPNN or three passes for `QGInitialization::Random`; `build(n)` requests `n >= 2` passes. The older overload with a numeric fifth argument uses that value as the random-initialization seed. @@ -125,7 +126,7 @@ Input may be released after reset returns; rebuild before querying or saving. ### Vector cache `cache_vectors=true` retains transformed vectors and query factors -(`4 * num * padded_dim` bytes plus factors). [QGKMeans](../clustering.md#qgkmeans) enables +(`4 * max_elements * padded_dim` bytes plus factors). [QGKMeans](../clustering.md#qgkmeans) enables this for its centroid graph; the cache is not persisted. ### Random seeds @@ -162,8 +163,8 @@ qg.search_batch( ); ``` -The C++ signature is `search_batch(queries, num_queries, k, ids, distances, -num_threads=1)`. Set `ef >= k` with `set_ef()` before searching. Both output +The C++ signature is `search_batch(queries, num_queries, k, result_ids, distances, +num_threads=1)`. Set `ef_search >= k` with `set_ef()` before searching. Both output buffers are required and contain `num_queries * k` elements. Empty C++ batches do no work and may use null buffers. Inputs and outputs must not overlap. Concurrent searches need separate outputs; do not modify the index or call @@ -179,7 +180,7 @@ Python searches use per-call `ef` values. For access rules, see the [concurrency contract](usage.md#threading-and-file-paths). For C++, use -`search_batch_with_ef(queries, num_queries, k, ids, distances, ef, num_threads=1)` +`search_batch_with_ef(queries, num_queries, k, result_ids, distances, ef_search, num_threads=1)` to run concurrent calls with independent windows: ```cpp diff --git a/include/rabitqlib/index/detail/buffer.hpp b/include/rabitqlib/index/detail/buffer.hpp index 5b3b551a..c2f64cf5 100644 --- a/include/rabitqlib/index/detail/buffer.hpp +++ b/include/rabitqlib/index/detail/buffer.hpp @@ -102,16 +102,16 @@ class SearchBuffer { [[nodiscard]] size_t size() const { return size_; } - void copy_results(PID* knn) const { + void copy_results(PID* result_ids) const { for (size_t i = 0; i < size_; ++i) { - knn[i] = data_[i].id; + result_ids[i] = data_[i].id; } } - void copy_results(PID* knn, T* dist) const { + void copy_results(PID* result_ids, T* distances) const { for (size_t i = 0; i < size_; ++i) { - knn[i] = data_[i].id; - dist[i] = data_[i].distance; + result_ids[i] = data_[i].id; + distances[i] = data_[i].distance; } } diff --git a/include/rabitqlib/index/detail/exact_nn.hpp b/include/rabitqlib/index/detail/exact_nn.hpp index f2f15c12..00362a3b 100644 --- a/include/rabitqlib/index/detail/exact_nn.hpp +++ b/include/rabitqlib/index/detail/exact_nn.hpp @@ -20,7 +20,7 @@ inline PID exact_nn( T (*dist_func)(const T*, const T*, size_t) ) { #if RABITQ_USE_OPENMP - const auto thread_count = static_cast(std::max( + const auto num_workers = static_cast(std::max( 1, std::min( {resolve_num_threads(num_threads), @@ -30,12 +30,12 @@ inline PID exact_nn( )); #else (void)num_threads; - constexpr int thread_count = 1; + constexpr int num_workers = 1; #endif - std::vector> best_entries(thread_count); + std::vector> best_entries(num_workers); #if RABITQ_USE_OPENMP -#pragma omp parallel for schedule(dynamic) num_threads(thread_count) +#pragma omp parallel for schedule(dynamic) num_threads(num_workers) #endif for (std::ptrdiff_t index = 0; index < static_cast(num_points); ++index) { diff --git a/include/rabitqlib/index/detail/hnsw_dispatch.hpp b/include/rabitqlib/index/detail/hnsw_dispatch.hpp index 3b2f9d53..dabc82af 100644 --- a/include/rabitqlib/index/detail/hnsw_dispatch.hpp +++ b/include/rabitqlib/index/detail/hnsw_dispatch.hpp @@ -12,27 +12,27 @@ class HierarchicalNSW; namespace detail { // Compatibility adapter; individual backends return ordered results only. std::priority_queue> search_knn( - HierarchicalNSW&, const float*, size_t, size_t ef = 0 + HierarchicalNSW&, const float*, size_t, size_t ef_search = 0 ); // Ordered results used by batch search and backend parity tests. std::vector> search_knn_sorted( - HierarchicalNSW&, const float*, size_t, size_t ef = 0 + HierarchicalNSW&, const float*, size_t, size_t ef_search = 0 ); std::vector> search_knn_sorted_generic( - HierarchicalNSW&, const float*, size_t, size_t ef = 0 + HierarchicalNSW&, const float*, size_t, size_t ef_search = 0 ); std::vector> search_knn_sorted_neon( - HierarchicalNSW&, const float*, size_t, size_t ef = 0 + HierarchicalNSW&, const float*, size_t, size_t ef_search = 0 ); std::vector> search_knn_sorted_avx2( - HierarchicalNSW&, const float*, size_t, size_t ef = 0 + HierarchicalNSW&, const float*, size_t, size_t ef_search = 0 ); std::vector> search_knn_sorted_avx512_core( - HierarchicalNSW&, const float*, size_t, size_t ef = 0 + HierarchicalNSW&, const float*, size_t, size_t ef_search = 0 ); std::vector> search_knn_sorted_avx512_popcnt( - HierarchicalNSW&, const float*, size_t, size_t ef = 0 + HierarchicalNSW&, const float*, size_t, size_t ef_search = 0 ); } // namespace detail diff --git a/include/rabitqlib/index/hnsw/hnsw.hpp b/include/rabitqlib/index/hnsw/hnsw.hpp index 22e97c27..3b7a3de4 100644 --- a/include/rabitqlib/index/hnsw/hnsw.hpp +++ b/include/rabitqlib/index/hnsw/hnsw.hpp @@ -68,9 +68,9 @@ class HierarchicalNSW { // Logical capacity, nonremoved points, and stored records (including removals). [[nodiscard]] size_t max_elements() const { return max_elements_; } [[nodiscard]] size_t num_points() const { - return cur_element_count_.load() - num_removed_; + return num_stored_points_.load() - num_removed_; } - [[nodiscard]] size_t num_stored_points() const { return cur_element_count_; } + [[nodiscard]] size_t num_stored_points() const { return num_stored_points_; } [[nodiscard]] PID cluster_id_of(PID internal_id) const { return get_clusterid_by_internalid(internal_id); } @@ -81,8 +81,8 @@ class HierarchicalNSW { // Centroids and data use the original dimension; cluster_ids has num_points entries. // Input buffers must remain valid until this synchronous call returns or throws; // the index retains no pointers to them afterward. - // faster selects fixed-scale extra-bit quantization instead of exact scale search, - // trading quantization accuracy for speed. One-bit storage is unaffected. + // fast_quantization selects fixed-scale extra-bit quantization instead of exact scale + // search, trading quantization accuracy for speed. One-bit storage is unaffected. void construct( size_t num_clusters, const float* centroids, @@ -90,7 +90,7 @@ class HierarchicalNSW { const float* data, PID* cluster_ids, size_t num_threads = 0, - bool faster = false + bool fast_quantization = false ); /** @@ -104,20 +104,20 @@ class HierarchicalNSW { * during linking may leave a partially linked batch counted in the index, * which remains safe to search, save, load, and destroy. * - * @param data New points, `n * dim` floats in the original coordinates - * @param n Number of new points + * @param data New points, `num_points * dim` floats in the original coordinates + * @param num_points Number of new points * @param cluster_ids Cluster of each point in `[0, num_clusters)`, or nullptr to * route each one to its nearest centroid - * @param faster Use fixed-scale extra-bit quantization for speed at a possible - * loss of accuracy; has no effect on one-bit storage + * @param fast_quantization Use fixed-scale extra-bit quantization for speed at a + * possible loss of accuracy; has no effect on one-bit storage * @param num_threads Threads used to route and insert * @return The label given to each new point */ std::vector add( const float* data, - size_t n, + size_t num_points, const PID* cluster_ids = nullptr, - bool faster = false, + bool fast_quantization = false, size_t num_threads = 1 ); @@ -146,10 +146,10 @@ class HierarchicalNSW { * before anything changes. Not safe to call while another thread searches or adds. * * @param labels Labels of the points to remove, as returned by `add` or `search` - * @param n Number of labels + * @param num_points Number of labels * @return Number of points newly removed */ - size_t remove(const PID* labels, size_t n); + size_t remove(const PID* labels, size_t num_points); // Queries use the original dimension. Each result is (distance, external label), // where distance estimates squared L2 or 1 - dot according to the index metric. @@ -165,19 +165,9 @@ class HierarchicalNSW { // Borrows construct() input only while its workers run; cleared on success or failure. const float* construction_data_{nullptr}; - struct ResultRecord { - float est_dist; - float low_dist; - ResultRecord(float est_dist, float low_dist) - : est_dist(est_dist), low_dist(low_dist) {} - bool operator<(const ResultRecord& other) const { - return this->est_dist < other.est_dist; - } - }; - // Bounded priority queue implemented as a sorted vector. struct Candidate { - HierarchicalNSW::ResultRecord record; + float est_dist; PID id; }; @@ -192,7 +182,7 @@ class HierarchicalNSW { queue_.end(), cand, [](const Candidate& a, const Candidate& b) { - return a.record.est_dist < b.record.est_dist; + return a.est_dist < b.est_dist; } ); queue_.insert(it, cand); @@ -211,7 +201,7 @@ class HierarchicalNSW { private: size_t capacity_; - // Sorted in ascending order by record.est_dist so that the worst is at the back. + // Sorted in ascending order by est_dist so that the worst is at the back. std::vector queue_; }; @@ -237,15 +227,15 @@ class HierarchicalNSW { // keep num_clusters_ at or below 2^31, so a real cluster ID never has it. static constexpr PID kRemovedMask = 0x80000000U; size_t max_elements_{0}; - mutable std::atomic cur_element_count_{0}; + mutable std::atomic num_stored_points_{0}; size_t num_removed_{0}; // Recovered from stored cluster-ID markers when loading. - size_t size_data_per_element_{0}; - size_t size_links_per_element_{0}; + size_t level0_record_bytes_{0}; // One complete level-0 record. + size_t upper_layer_links_bytes_{0}; // One node's adjacency list at one upper level. size_t M_{0}; - size_t maxM_{0}; - size_t maxM0_{0}; + size_t upper_layer_max_degree_{0}; + size_t base_layer_max_degree_{0}; size_t ef_construction_{0}; - size_t ef_{0}; + size_t default_ef_search_{0}; MetricType metric_type_{METRIC_L2}; double mult_{0.0}; @@ -259,10 +249,11 @@ class HierarchicalNSW { PID entry_point_id_{0}; - size_t size_links_level0_{0}; + size_t level0_links_bytes_{ + 0}; // One node's level-0 adjacency list, including its header. size_t bin_data_offset_{0}, ex_data_offset_{0}, label_offset_{0}; - size_t size_bin_data_{0}, size_ex_data_{0}; + size_t bin_data_bytes_{0}, ex_data_bytes_{0}; size_t ex_bits_{0}; struct AlignedDeleter { @@ -304,8 +295,6 @@ class HierarchicalNSW { float (*raw_dist_func_)(const float* __restrict__, const float* __restrict__, size_t){ nullptr}; - void set_ef(size_t ef) { ef_ = ef; } - std::mutex& get_label_op_mutex(PID label) const { size_t lock_id = label & (kMaxLabelOperationLock - 1); return label_op_locks_[lock_id]; @@ -315,48 +304,53 @@ class HierarchicalNSW { PID return_label; memcpy( &return_label, - (data_level0_memory_.get() + (internal_id * size_data_per_element_) + + (data_level0_memory_.get() + (internal_id * level0_record_bytes_) + label_offset_), sizeof(PID) ); return return_label; } - void set_external_label(PID internal_id, PID label) const { + void set_external_label(PID internal_id, PID label) { memcpy( - (data_level0_memory_.get() + (internal_id * size_data_per_element_) + + (data_level0_memory_.get() + (internal_id * level0_record_bytes_) + label_offset_), &label, sizeof(PID) ); } - PID* get_external_label_pt(PID internal_id) const { + PID* get_external_label_pt(PID internal_id) { return reinterpret_cast( - data_level0_memory_.get() + (internal_id * size_data_per_element_) + - label_offset_ + data_level0_memory_.get() + (internal_id * level0_record_bytes_) + label_offset_ ); } - char* get_bindata_by_internalid(PID internal_id) const { - return reinterpret_cast( - data_level0_memory_.get() + (internal_id * size_data_per_element_) + - bin_data_offset_ - ); + char* get_bindata_by_internalid(PID internal_id) { + return data_level0_memory_.get() + (internal_id * level0_record_bytes_) + + bin_data_offset_; } - char* get_exdata_by_internalid(PID internal_id) const { - return reinterpret_cast( - data_level0_memory_.get() + (internal_id * size_data_per_element_) + - ex_data_offset_ - ); + const char* get_bindata_by_internalid(PID internal_id) const { + return data_level0_memory_.get() + (internal_id * level0_record_bytes_) + + bin_data_offset_; + } + + char* get_exdata_by_internalid(PID internal_id) { + return data_level0_memory_.get() + (internal_id * level0_record_bytes_) + + ex_data_offset_; + } + + const char* get_exdata_by_internalid(PID internal_id) const { + return data_level0_memory_.get() + (internal_id * level0_record_bytes_) + + ex_data_offset_; } // The stored cluster ID with the removal mark, as save and load see it. PID get_raw_clusterid(PID internal_id) const { - return *(reinterpret_cast( - data_level0_memory_.get() + (internal_id * size_data_per_element_) + - size_links_level0_ + return *(reinterpret_cast( + data_level0_memory_.get() + (internal_id * level0_record_bytes_) + + level0_links_bytes_ )); } @@ -368,28 +362,36 @@ class HierarchicalNSW { return (get_raw_clusterid(internal_id) & kRemovedMask) != 0; } - char* get_clusterid_pt(PID internal_id) const { + char* get_clusterid_pt(PID internal_id) { return reinterpret_cast( - data_level0_memory_.get() + (internal_id * size_data_per_element_) + - size_links_level0_ + data_level0_memory_.get() + (internal_id * level0_record_bytes_) + + level0_links_bytes_ ); } int get_random_level(double reverse_size); - size_t get_max_elements() const { return max_elements_; } + PID* get_linklist(PID internal_id, int level) { + return reinterpret_cast( + link_lists_[internal_id].get() + ((level - 1) * upper_layer_links_bytes_) + ); + } - size_t get_current_element_count() const { return cur_element_count_; } + const PID* get_linklist(PID internal_id, int level) const { + return reinterpret_cast( + link_lists_[internal_id].get() + ((level - 1) * upper_layer_links_bytes_) + ); + } - PID* get_linklist(PID internal_id, int level) const { + PID* get_linklist0(PID internal_id) { return reinterpret_cast( - link_lists_[internal_id].get() + ((level - 1) * size_links_per_element_) + data_level0_memory_.get() + (internal_id * level0_record_bytes_) ); } - PID* get_linklist0(PID internal_id) const { - return reinterpret_cast( - data_level0_memory_.get() + (internal_id * size_data_per_element_) + const PID* get_linklist0(PID internal_id) const { + return reinterpret_cast( + data_level0_memory_.get() + (internal_id * level0_record_bytes_) ); } @@ -404,33 +406,39 @@ class HierarchicalNSW { // ANN Search template void get_bin_est_direct( - std::vector&, SplitSingleQuery&, PID, HierarchicalNSW::EstimateRecord& - ); + const std::vector&, + const SplitSingleQuery&, + PID, + HierarchicalNSW::EstimateRecord& + ) const; template void get_full_est_direct( - std::vector&, SplitSingleQuery&, PID, HierarchicalNSW::EstimateRecord& + const std::vector&, + const SplitSingleQuery&, + PID, + HierarchicalNSW::EstimateRecord& ) const; template std::vector> search_knn_sorted_direct( const float*, size_t, size_t - ); + ) const; template void search_base_layer_direct( - PID ep_id, - size_t ef, + PID entry_point_id, + size_t ef_search, size_t k, - SplitSingleQuery& query_wrapper, - std::vector& q_to_centroids, + const SplitSingleQuery& query_wrapper, + const std::vector& q_to_centroids, const float* query, BoundedKNN& nearest_candidates - ); + ) const; // Construction // Currently only support index construction with non-quantized vectors - float get_data_dist(PID obj1, PID obj2) { + float get_data_dist(PID obj1, PID obj2) const { PID label1 = get_external_label(obj1); PID label2 = get_external_label(obj2); return raw_dist_func_( diff --git a/include/rabitqlib/index/ivf/cluster.hpp b/include/rabitqlib/index/ivf/cluster.hpp index e343ec76..f1268a4b 100644 --- a/include/rabitqlib/index/ivf/cluster.hpp +++ b/include/rabitqlib/index/ivf/cluster.hpp @@ -11,14 +11,16 @@ class Cluster { private: friend class IVF; - size_t num_; // Num of vectors in this cluster + size_t num_stored_points_; // Num of vectors in this cluster char* batch_data_ = nullptr; // RaBitQ code and factors char* ex_data_ = nullptr; // Extra-bit codes/factors or original float32 vectors PID* ids_ = nullptr; // PID of vectors // The owning index refreshes only views whose storage or count changed. - void reset(size_t num, char* batch_data, char* ex_data, PID* ids) noexcept { - num_ = num; + void reset( + size_t num_stored_points, char* batch_data, char* ex_data, PID* ids + ) noexcept { + num_stored_points_ = num_stored_points; batch_data_ = batch_data; ex_data_ = ex_data; ids_ = ids; @@ -38,10 +40,13 @@ class Cluster { [[nodiscard]] PID* ids() const { return this->ids_; } - [[nodiscard]] size_t num() const { return num_; } + [[nodiscard]] size_t num() const { return num_stored_points_; } }; -inline Cluster::Cluster(size_t num, char* batch_data, char* ex_data, PID* ids) - : num_(num), batch_data_(batch_data), ex_data_(ex_data), ids_(ids) {} +inline Cluster::Cluster(size_t num_stored_points, char* batch_data, char* ex_data, PID* ids) + : num_stored_points_(num_stored_points) + , batch_data_(batch_data) + , ex_data_(ex_data) + , ids_(ids) {} } // namespace rabitqlib::ivf diff --git a/include/rabitqlib/index/ivf/initializer.hpp b/include/rabitqlib/index/ivf/initializer.hpp index 3da82fb2..0f04cff5 100644 --- a/include/rabitqlib/index/ivf/initializer.hpp +++ b/include/rabitqlib/index/ivf/initializer.hpp @@ -47,7 +47,8 @@ class Initializer { size_t num_cluster_; public: - explicit Initializer(size_t d, size_t k) : dim_(d), num_cluster_(k) {} + explicit Initializer(size_t dim, size_t num_clusters) + : dim_(dim), num_cluster_(num_clusters) {} virtual ~Initializer() = 0; [[nodiscard]] virtual const float* centroid(PID) const = 0; virtual void add_vectors(const float*, size_t) = 0; @@ -66,9 +67,11 @@ class FlatInitializer : public Initializer { public: explicit FlatInitializer( - size_t d, size_t k, MetricType metric_type = MetricType::METRIC_L2 + size_t dim, size_t num_clusters, MetricType metric_type = MetricType::METRIC_L2 ) - : Initializer(d, k), centroids_(num_cluster_ * dim_), metric_type_(metric_type) { + : Initializer(dim, num_clusters) + , centroids_(num_cluster_ * dim_) + , metric_type_(metric_type) { validate_metric_type(metric_type_); } @@ -78,29 +81,31 @@ class FlatInitializer : public Initializer { return ¢roids_[id * dim_]; } - void add_vectors(const float* cent, size_t) override { - std::memcpy(centroids_.data(), cent, sizeof(float) * num_cluster_ * dim_); + void add_vectors(const float* centroids, size_t) override { + std::memcpy(centroids_.data(), centroids, sizeof(float) * num_cluster_ * dim_); } void centroids_distances( const float* query, size_t nprobe, std::vector>& candidates ) const override { - std::vector> centroid_dist(this->num_cluster_); + std::vector> centroid_candidates(this->num_cluster_); for (PID i = 0; i < num_cluster_; ++i) { - centroid_dist[i].id = i; - centroid_dist[i].distance = + centroid_candidates[i].id = i; + centroid_candidates[i].distance = metric_type_ == METRIC_IP ? dot_product_dis(query, centroid(i), dim_) : std::sqrt(euclidean_sqr(query, centroid(i), dim_)); } std::partial_sort( - centroid_dist.begin(), - centroid_dist.begin() + static_cast(nprobe), - centroid_dist.end() + centroid_candidates.begin(), + centroid_candidates.begin() + static_cast(nprobe), + centroid_candidates.end() ); std::memcpy( - candidates.data(), centroid_dist.data(), sizeof(AnnCandidate) * nprobe + candidates.data(), + centroid_candidates.data(), + sizeof(AnnCandidate) * nprobe ); } @@ -130,7 +135,9 @@ class FlatRaBitQInitializer : public Initializer { MetricType metric_type_; public: - explicit FlatRaBitQInitializer(size_t d, size_t k, MetricType metric = METRIC_L2); + explicit FlatRaBitQInitializer( + size_t dim, size_t num_clusters, MetricType metric = METRIC_L2 + ); [[nodiscard]] const float* centroid(PID id) const override; void add_vectors(const float* centroids, size_t num_threads) override; void centroids_distances( @@ -214,11 +221,11 @@ class HNSWInitializer : public Initializer { public: explicit HNSWInitializer( - size_t d, size_t k, MetricType metric_type = MetricType::METRIC_L2 + size_t dim, size_t num_clusters, MetricType metric_type = MetricType::METRIC_L2 ) - : Initializer(d, k) - , l2_space_(d) - , ip_space_(d) + : Initializer(dim, num_clusters) + , l2_space_(dim) + , ip_space_(dim) , space_( metric_type == METRIC_IP ? static_cast*>(&ip_space_) @@ -231,12 +238,12 @@ class HNSWInitializer : public Initializer { ); } - void add_vectors(const float* cent, size_t num_threads) override { + void add_vectors(const float* centroids, size_t num_threads) override { size_t start = 0; size_t rows = num_cluster_; rabitqlib::index_detail::parallel_for( start, rows, num_threads, [&](size_t row, size_t /*thread_id*/) { - alg_hnsw_->addPoint(cent + (row * dim_), row); + alg_hnsw_->addPoint(centroids + (row * dim_), row); } ); rebuild_label_lookup(); diff --git a/include/rabitqlib/index/ivf/ivf.hpp b/include/rabitqlib/index/ivf/ivf.hpp index ffe86d23..a1404539 100644 --- a/include/rabitqlib/index/ivf/ivf.hpp +++ b/include/rabitqlib/index/ivf/ivf.hpp @@ -118,8 +118,8 @@ class IVF { const quant::RabitqConfig& ) const; - // True if `n` points cannot be addressed by the storage size computations - [[nodiscard]] bool exceeds_storage(size_t n) const; + // True if `num_points` points cannot be addressed by the storage size computations + [[nodiscard]] bool exceeds_storage(size_t num_points) const; [[nodiscard]] static size_t ids_bytes(size_t cluster_size) { return sizeof(PID) * cluster_size; @@ -132,12 +132,12 @@ class IVF { [[nodiscard]] inline size_t ex_data_bytes(size_t cluster_size) const; - // Size every buffer of a cluster's storage for `n` points - void resize_storage(ClusterStorage& cs, size_t n) const; + // Size every buffer of a cluster's storage for `num_points` points + void resize_storage(ClusterStorage& cs, size_t num_points) const; - void reserve_storage(ClusterStorage& cs, size_t n) const; + void reserve_storage(ClusterStorage& cs, size_t num_points) const; - [[nodiscard]] bool has_capacity(const ClusterStorage& cs, size_t n) const; + [[nodiscard]] bool has_capacity(const ClusterStorage& cs, size_t num_points) const; void allocate_memory(const std::vector&); @@ -159,8 +159,8 @@ class IVF { const char* batch_data, const char* ex_data, const PID* ids, - const SplitBatchQuery& q_obj, - buffer::SearchBuffer& knns, + const SplitBatchQuery& query_wrapper, + buffer::SearchBuffer& result_pool, size_t num_points, bool, const float* @@ -197,13 +197,14 @@ class IVF { [[nodiscard]] RotatorType rotator_type() const { return rotator_type_; } [[nodiscard]] InitializerType initializer_type() const { return initializer_type_; } - // faster selects fixed-scale extra-bit quantization instead of exact scale search, - // trading quantization accuracy for speed. One-bit and raw storage are unaffected. + // fast_quantization selects fixed-scale extra-bit quantization instead of exact scale + // search, trading quantization accuracy for speed. One-bit and raw storage are + // unaffected. void construct( const float* data, const float* centroids, const PID* cluster_ids, - bool faster = false, + bool fast_quantization = false, size_t num_threads = std::numeric_limits::max() ); @@ -217,24 +218,24 @@ class IVF { * * Only clusters that receive points are touched. Each one is rebuilt next to its old * storage, so a call needs extra memory for, and takes time proportional to, those - * clusters plus `n`; prefer a few large calls to many small ones. A cluster with - * enough capacity from `reserve` grows in place instead. `construct` expects + * clusters plus `num_points`; prefer a few large calls to many small ones. A cluster + * with enough capacity from `reserve` grows in place instead. `construct` expects * `max_elements()` rows afterwards. * Not safe to call while another thread searches, adds or removes. * - * @param data New points (n * dim) - * @param n Number of new points + * @param data New points (num_points * dim) + * @param num_points Number of new points * @param cluster_ids Cluster of each new point, or nullptr to route each point to its * nearest centroid the same way a query is routed - * @param faster Use fixed-scale extra-bit quantization for speed at a possible - * loss of accuracy; has no effect on one-bit or raw storage + * @param fast_quantization Use fixed-scale extra-bit quantization for speed at a + * possible loss of accuracy; has no effect on one-bit or raw storage * @param num_threads Threads used for routing and quantization */ void add( const float* data, - size_t n, + size_t num_points, const PID* cluster_ids = nullptr, - bool faster = false, + bool fast_quantization = false, size_t num_threads = std::numeric_limits::max() ); @@ -262,10 +263,10 @@ class IVF { * restored. Not safe to call while another thread searches or adds. * * @param ids_to_remove PIDs to remove - * @param n Number of PIDs + * @param num_points Number of PIDs * @return Number of points newly removed */ - size_t remove(const PID* ids_to_remove, size_t n); + size_t remove(const PID* ids_to_remove, size_t num_points); void save(const char* filename) const; @@ -277,26 +278,26 @@ class IVF { const float* query, size_t k, size_t nprobe, - PID* results, + PID* result_ids, float* distances = nullptr ) const; void search( - const float* query, size_t k, size_t nprobe, PID* results, bool use_hacc + const float* query, size_t k, size_t nprobe, PID* result_ids, bool use_hacc ) const; void search( const float* query, size_t k, size_t nprobe, - PID* results, + PID* result_ids, float* distances, bool use_hacc ) const; /** * Search contiguous row-major queries using at most num_threads workers. - * queries contains num_queries * dimension() floats. results and optional + * queries contains num_queries * dimension() floats. result_ids and optional * distances each hold num_queries * k elements. Empty batches do no work. * use_hacc defaults to the same automatic policy as search(). * Input and output buffers must not overlap. Concurrent calls are safe if @@ -307,7 +308,7 @@ class IVF { size_t num_queries, size_t k, size_t nprobe, - PID* results, + PID* result_ids, float* distances = nullptr, std::optional use_hacc = std::nullopt, size_t num_threads = 1 diff --git a/include/rabitqlib/index/symqg/detail/pipnn.hpp b/include/rabitqlib/index/symqg/detail/pipnn.hpp index 859b9dec..2d7eeff5 100644 --- a/include/rabitqlib/index/symqg/detail/pipnn.hpp +++ b/include/rabitqlib/index/symqg/detail/pipnn.hpp @@ -20,7 +20,7 @@ InitialGraph build_initial_graph( size_t count, size_t dim, size_t degree, - MetricType metric = METRIC_L2, + MetricType metric_type = METRIC_L2, size_t num_threads = std::numeric_limits::max() ); diff --git a/include/rabitqlib/index/symqg/detail/search_scratch.hpp b/include/rabitqlib/index/symqg/detail/search_scratch.hpp index a6fd2785..87456143 100644 --- a/include/rabitqlib/index/symqg/detail/search_scratch.hpp +++ b/include/rabitqlib/index/symqg/detail/search_scratch.hpp @@ -30,11 +30,11 @@ struct SearchScratch { buffer::SearchBuffer search_pool; buffer::SearchBuffer result_pool; - SearchScratch(size_t padded_dim, size_t degree, size_t ef, size_t k) + SearchScratch(size_t padded_dim, size_t degree, size_t ef_search, size_t k) : rotated_query(padded_dim) , estimated_distances(degree) , lookup_table(padded_dim * 4) - , search_pool(ef) + , search_pool(ef_search) , result_pool(k) {} SearchScratchView view(VisitedSet& visited) { diff --git a/include/rabitqlib/index/symqg/qg.hpp b/include/rabitqlib/index/symqg/qg.hpp index fe0eeae4..7e6a9b01 100644 --- a/include/rabitqlib/index/symqg/qg.hpp +++ b/include/rabitqlib/index/symqg/qg.hpp @@ -57,9 +57,9 @@ struct SearchScratchView; void search_with_workspace( QuantizedGraph& graph, const float* query, - uint32_t knn, - uint32_t* results, - float* dists, + uint32_t k, + uint32_t* result_ids, + float* distances, SearchScratchView scratch, PID hint = kPidMax, const float* k1xsumq = nullptr @@ -82,12 +82,12 @@ class QuantizedGraph { friend struct QGConstructionTestAccess; private: - size_t num_points_ = 0; - size_t degree_bound_ = 0; + size_t max_elements_ = 0; + size_t max_degree_ = 0; size_t dim_ = 0; size_t padded_dim_ = 0; float (*raw_dist_func_)(const float*, const float*, size_t) = nullptr; - PID entry_point_ = 0; + PID entry_point_id_ = 0; MetricType metric_type_ = MetricType::METRIC_L2; RotatorType rotator_type_ = RotatorType::FhtKacRotator; size_t quantization_bits_ = 0; // 0: raw vectors, 4/8: packed RaBitQ vectors @@ -106,12 +106,12 @@ class QuantizedGraph { std::unique_ptr> rotator_; // Position of row data (raw vector or packed qg-quant vector), neighbor - // quantization data, and neighbor IDs. Since every degree equals degree_bound_ + // quantization data, and neighbor IDs. Since every degree equals max_degree_ // (a multiple of 32), the degree does not need to be stored per vertex. size_t batch_data_offset_ = 0; size_t neighbor_offset_ = 0; - size_t row_offset_ = 0; // length of entire row - size_t ef_ = 0; + size_t record_bytes_ = 0; // length of entire row + size_t default_ef_search_ = 0; bool ready_ = false; uint32_t seed_ = 42; @@ -140,11 +140,11 @@ class QuantizedGraph { } [[nodiscard]] float* get_vector(PID data_id) { - return data_.data() + ((row_offset_ / sizeof(float)) * data_id); + return data_.data() + ((record_bytes_ / sizeof(float)) * data_id); } [[nodiscard]] const float* get_vector(PID data_id) const { - return data_.data() + ((row_offset_ / sizeof(float)) * data_id); + return data_.data() + ((record_bytes_ / sizeof(float)) * data_id); } [[nodiscard]] char* get_quantized_vector(PID data_id) { return get_row_data(data_id); } @@ -221,9 +221,9 @@ class QuantizedGraph { public: explicit QuantizedGraph( - size_t num, + size_t max_elements, size_t dim, - size_t max_deg, + size_t max_degree, MetricType metric_type = METRIC_L2, RotatorType rotator_type = RotatorType::FhtKacRotator, size_t quantization_bits = 0, @@ -270,15 +270,15 @@ class QuantizedGraph { void search( const float* __restrict__ query, - uint32_t knn, - uint32_t* __restrict__ results, - float* __restrict__ dists + uint32_t k, + uint32_t* __restrict__ result_ids, + float* __restrict__ distances ); /** * Search contiguous row-major queries, using the window set by set_ef(). - * queries contains num_queries * dimension() floats; results and dists - * each hold num_queries * knn elements. Empty batches do no work. + * queries contains num_queries * dimension() floats; result_ids and distances + * each hold num_queries * k elements. Empty batches do no work. * Scratch is reused within each worker. num_threads defaults to one; * zero selects the available logical CPU count, capped by the batch. * Input and output buffers must not overlap. Concurrent searches require @@ -287,21 +287,21 @@ class QuantizedGraph { void search_batch( const float* queries, size_t num_queries, - uint32_t knn, - uint32_t* results, - float* dists, + uint32_t k, + uint32_t* result_ids, + float* distances, size_t num_threads = 1 ); - // Like search_batch, but ef belongs to this call and never reads or changes + // Like search_batch, but ef_search belongs to this call and never reads or changes // the default set by set_ef(). Concurrent calls may use different windows. void search_batch_with_ef( const float* queries, size_t num_queries, - uint32_t knn, - uint32_t* results, - float* dists, - size_t ef, + uint32_t k, + uint32_t* result_ids, + float* distances, + size_t ef_search, size_t num_threads = 1 ); @@ -311,7 +311,7 @@ class QuantizedGraph { * Reuses scratch and can use a previous result as a warm-start hint. * BatchQuery may allocate LUT storage on its first use or when its shape changes; * this is not the stable core v1 zero-allocation interface. Search and result - * buffers must be sized for ef and knn respectively. rotated_query_scratch and + * buffers must be sized for ef_search and k respectively. rotated_query_scratch and * lut_float_scratch need D and 4 * D floats, where D is the graph's stored padded * dimension; this advanced caller must already know D. dimension() is the original * input dimension. est_dist_scratch needs degree_bound() floats. The visited set must @@ -324,9 +324,9 @@ class QuantizedGraph { */ void search_with_scratch( const float* __restrict__ query, - uint32_t knn, - uint32_t* __restrict__ results, - float* __restrict__ dists, + uint32_t k, + uint32_t* __restrict__ result_ids, + float* __restrict__ distances, float* __restrict__ rotated_query_scratch, float* __restrict__ est_dist_scratch, float* __restrict__ lut_float_scratch, diff --git a/include/rabitqlib/index/symqg/qg_builder.hpp b/include/rabitqlib/index/symqg/qg_builder.hpp index 5fd7c75e..0b0f9885 100644 --- a/include/rabitqlib/index/symqg/qg_builder.hpp +++ b/include/rabitqlib/index/symqg/qg_builder.hpp @@ -52,11 +52,11 @@ class QGBuilder { }; QuantizedGraph& qg_; - size_t ef_build_; // size of search pool for indexing - size_t num_threads_; - size_t num_nodes_; + size_t ef_construction_; // size of search pool for indexing + size_t num_workers_; + size_t num_points_; size_t dim_; - size_t degree_bound_; // degree bound for qg, multiple of 32 + size_t max_degree_; // degree bound for qg, multiple of 32 bool seeded_ = true; static constexpr size_t kMaxCandidatePoolSize = 750; static constexpr size_t kMaxPrunedSize = 300; @@ -74,7 +74,7 @@ class QGBuilder { std::vector rotated_neighbors; std::vector neighbor_ids; - BuildWorker(size_t num_nodes, size_t visited_capacity); + BuildWorker(size_t num_points, size_t visited_capacity); }; std::vector workers_; @@ -103,7 +103,7 @@ class QGBuilder { QGBuilder( QuantizedGraph& index, - uint32_t ef_build, + uint32_t ef_construction, size_t num_threads, uint32_t seed = std::random_device{}(), bool cache_vectors = false @@ -112,7 +112,7 @@ class QGBuilder { public: explicit QGBuilder( QuantizedGraph& index, - uint32_t ef_build, + uint32_t ef_construction, const float* data, size_t num_threads = std::numeric_limits::max(), QGInitialization init = QGInitialization::PiPNN, @@ -123,7 +123,7 @@ class QGBuilder { // Preserve the seeded random-initialization overload used before PiPNN. QGBuilder( QuantizedGraph& index, - uint32_t ef_build, + uint32_t ef_construction, const float* data, size_t num_threads, uint32_t seed, diff --git a/python_bindings/hnsw_bindings.cpp b/python_bindings/hnsw_bindings.cpp index 97f44a65..318b97ae 100644 --- a/python_bindings/hnsw_bindings.cpp +++ b/python_bindings/hnsw_bindings.cpp @@ -203,23 +203,27 @@ class HnswIndex { const auto shape = std::vector{ static_cast(query_array.shape(0)), static_cast(k)}; - auto ids = py::array_t(shape); - auto dists = py::array_t(shape); - std::fill(ids.mutable_data(), ids.mutable_data() + ids.size(), rabitqlib::kPidMax); + auto result_ids = py::array_t(shape); + auto distances = py::array_t(shape); std::fill( - dists.mutable_data(), - dists.mutable_data() + dists.size(), + result_ids.mutable_data(), + result_ids.mutable_data() + result_ids.size(), + rabitqlib::kPidMax + ); + std::fill( + distances.mutable_data(), + distances.mutable_data() + distances.size(), std::numeric_limits::infinity() ); - auto ids_buf = ids.mutable_unchecked<2>(); - auto dists_buf = dists.mutable_unchecked<2>(); + auto result_ids_view = result_ids.mutable_unchecked<2>(); + auto distances_view = distances.mutable_unchecked<2>(); const auto* queries_data = query_array.data(); - const auto nq = static_cast(query_array.shape(0)); + const auto num_queries = static_cast(query_array.shape(0)); std::vector>> results; { py::gil_scoped_release release; - results = index_->search(queries_data, nq, k, ef, num_threads); + results = index_->search(queries_data, num_queries, k, ef, num_threads); } for (py::ssize_t i = 0; i < static_cast(results.size()); ++i) { @@ -227,13 +231,13 @@ class HnswIndex { k, results[static_cast(i)].size() )); ++j) { - ids_buf(i, j) = + result_ids_view(i, j) = results[static_cast(i)][static_cast(j)].second; - dists_buf(i, j) = + distances_view(i, j) = results[static_cast(i)][static_cast(j)].first; } } - return py::make_tuple(ids, dists); + return py::make_tuple(result_ids, distances); } void save(py::handle path) const { diff --git a/python_bindings/ivf_bindings.cpp b/python_bindings/ivf_bindings.cpp index a7d83faa..68304632 100644 --- a/python_bindings/ivf_bindings.cpp +++ b/python_bindings/ivf_bindings.cpp @@ -232,16 +232,18 @@ class IvfIndex { throw std::invalid_argument("nprobe must be positive"); } - const size_t nq = static_cast(query_array.shape(0)); + const size_t num_queries = static_cast(query_array.shape(0)); const auto shape = std::vector{ - static_cast(nq), static_cast(k)}; - auto ids = py::array_t(shape); - auto dists = py::array_t(shape); - auto* ids_data = ids.mutable_data(); - auto* dists_data = dists.mutable_data(); - std::fill(ids_data, ids_data + ids.size(), rabitqlib::kPidMax); + static_cast(num_queries), static_cast(k)}; + auto result_ids = py::array_t(shape); + auto distances = py::array_t(shape); + auto* ids_data = result_ids.mutable_data(); + auto* distances_data = distances.mutable_data(); + std::fill(ids_data, ids_data + result_ids.size(), rabitqlib::kPidMax); std::fill( - dists_data, dists_data + dists.size(), std::numeric_limits::infinity() + distances_data, + distances_data + distances.size(), + std::numeric_limits::infinity() ); const auto* queries_data = query_array.data(); @@ -249,17 +251,17 @@ class IvfIndex { py::gil_scoped_release release; index_->search_batch( queries_data, - nq, + num_queries, k, nprobe, ids_data, - dists_data, + distances_data, high_accuracy, num_threads ); } - return py::make_tuple(ids, dists); + return py::make_tuple(result_ids, distances); } void save(py::handle path) const { diff --git a/python_bindings/symqg_bindings.cpp b/python_bindings/symqg_bindings.cpp index be226068..94bcfdc8 100644 --- a/python_bindings/symqg_bindings.cpp +++ b/python_bindings/symqg_bindings.cpp @@ -36,7 +36,7 @@ class SymqgIndex { ) : dim_(dim) , max_degree_(max_degree) - , metric_(metric_from_string(metric)) + , metric_type_(metric_from_string(metric)) , quantization_bits_(quantization_bits) { if (max_degree == 0 || max_degree % rabitqlib::fastscan::kBatchSize != 0) { throw std::invalid_argument("max_degree must be a positive multiple of 32"); @@ -68,7 +68,7 @@ class SymqgIndex { num_points, dim_, max_degree_, - metric_, + metric_type_, rabitqlib::RotatorType::FhtKacRotator, quantization_bits_ ); @@ -108,28 +108,28 @@ class SymqgIndex { } ensure_dimension(query_array, dim_, "query dimension does not match index dim"); - const size_t nq = static_cast(query_array.shape(0)); + const size_t num_queries = static_cast(query_array.shape(0)); const auto shape = std::vector{ - static_cast(nq), static_cast(k)}; - auto ids = py::array_t(shape); - auto dists = py::array_t(shape); - auto* ids_data = ids.mutable_data(); - auto* dists_data = dists.mutable_data(); + static_cast(num_queries), static_cast(k)}; + auto result_ids = py::array_t(shape); + auto distances = py::array_t(shape); + auto* ids_data = result_ids.mutable_data(); + auto* distances_data = distances.mutable_data(); const auto* queries_data = query_array.data(); { py::gil_scoped_release release; index_->search_batch_with_ef( queries_data, - nq, + num_queries, static_cast(k), ids_data, - dists_data, + distances_data, ef, num_threads ); } - return py::make_tuple(ids, dists); + return py::make_tuple(result_ids, distances); } void save(py::handle path) const { @@ -152,7 +152,7 @@ class SymqgIndex { } wrapper.dim_ = wrapper.index_->dimension(); wrapper.max_degree_ = wrapper.index_->degree_bound(); - wrapper.metric_ = wrapper.index_->metric_type(); + wrapper.metric_type_ = wrapper.index_->metric_type(); wrapper.quantization_bits_ = wrapper.index_->quantization_bits(); wrapper.built_ = true; return wrapper; @@ -161,7 +161,7 @@ class SymqgIndex { [[nodiscard]] std::string repr() const { const auto operation = access_.read(); return "SymqgIndex(dim=" + std::to_string(dim_) + ", metric='" + - metric_to_string(metric_) + + metric_to_string(metric_type_) + "', quantization_bits=" + std::to_string(quantization_bits_) + ", num_points=" + std::to_string(index_ ? index_->num_points() : 0) + ", max_degree=" + std::to_string(max_degree_) + @@ -194,7 +194,7 @@ class SymqgIndex { } [[nodiscard]] std::string metric() const { const auto operation = access_.read(); - return metric_to_string(metric_); + return metric_to_string(metric_type_); } [[nodiscard]] size_t quantization_bits() const { const auto operation = access_.read(); @@ -206,7 +206,7 @@ class SymqgIndex { size_t dim_ = 0; size_t max_degree_ = 0; - rabitqlib::MetricType metric_ = rabitqlib::METRIC_L2; + rabitqlib::MetricType metric_type_ = rabitqlib::METRIC_L2; size_t quantization_bits_ = 0; bool built_ = false; mutable IndexAccess access_; diff --git a/src/index/hnsw.cpp b/src/index/hnsw.cpp index 7c66e881..aef1adcc 100644 --- a/src/index/hnsw.cpp +++ b/src/index/hnsw.cpp @@ -16,6 +16,7 @@ #include #include #include +#include #include #include #include @@ -108,25 +109,25 @@ HierarchicalNSW::HierarchicalNSW( M_ = 10000; } - maxM_ = M_; - maxM0_ = M_ * 2; + upper_layer_max_degree_ = M_; + base_layer_max_degree_ = M_ * 2; ef_construction_ = std::max(ef_construction, M_); - ef_ = 10; + default_ef_search_ = 10; - size_bin_data_ = BinDataMap::data_bytes(padded_dim_); - size_ex_data_ = ExDataMap::data_bytes(padded_dim_, ex_bits_); - size_links_level0_ = maxM0_ * sizeof(PID) + sizeof(PID); + bin_data_bytes_ = BinDataMap::data_bytes(padded_dim_); + ex_data_bytes_ = ExDataMap::data_bytes(padded_dim_, ex_bits_); + level0_links_bytes_ = base_layer_max_degree_ * sizeof(PID) + sizeof(PID); label_offset_ = - size_links_level0_ + sizeof(PID); // (# of edges + edges) + (cluster_id) + level0_links_bytes_ + sizeof(PID); // (# of edges + edges) + (cluster_id) bin_data_offset_ = label_offset_ + sizeof(PID); // (# of edges + edges) + (cluster_id) + (external label) ex_data_offset_ = - bin_data_offset_ + size_bin_data_; // (# of edges + edges) + (cluster_id) - // + (external label) + (BinData) - size_data_per_element_ = ex_data_offset_ + size_ex_data_; + bin_data_offset_ + bin_data_bytes_; // (# of edges + edges) + (cluster_id) + // + (external label) + (BinData) + level0_record_bytes_ = ex_data_offset_ + ex_data_bytes_; data_level0_memory_.reset( - memory::huge_page_allocate(max_elements_ * size_data_per_element_) + memory::huge_page_allocate(max_elements_ * level0_record_bytes_) ); if (data_level0_memory_ == nullptr) { throw std::runtime_error("Not enough memory"); @@ -134,7 +135,7 @@ HierarchicalNSW::HierarchicalNSW( level_generator_.seed(random_seed); - cur_element_count_ = 0; + num_stored_points_ = 0; visited_list_pool_ = std::make_unique(1, max_elements_); @@ -143,7 +144,7 @@ HierarchicalNSW::HierarchicalNSW( max_level_ = -1; link_lists_.reserve(max_elements_); - size_links_per_element_ = maxM_ * sizeof(PID) + sizeof(PID); + upper_layer_links_bytes_ = upper_layer_max_degree_ * sizeof(PID) + sizeof(PID); mult_ = 1 / log(1.0 * static_cast(M_)); this->query_config_ = @@ -157,7 +158,7 @@ void HierarchicalNSW::save(const char* filename) const { } output.write(reinterpret_cast(&max_elements_), sizeof(size_t)); - output.write(reinterpret_cast(&cur_element_count_), sizeof(size_t)); + output.write(reinterpret_cast(&num_stored_points_), sizeof(size_t)); output.write(reinterpret_cast(&dim_), sizeof(size_t)); output.write(reinterpret_cast(&padded_dim_), sizeof(size_t)); @@ -165,21 +166,21 @@ void HierarchicalNSW::save(const char* filename) const { output.write(reinterpret_cast(&ex_bits_), sizeof(size_t)); output.write(reinterpret_cast(&metric_type_), sizeof(metric_type_)); - output.write(reinterpret_cast(&size_bin_data_), sizeof(size_t)); - output.write(reinterpret_cast(&size_ex_data_), sizeof(size_t)); - output.write(reinterpret_cast(&size_links_level0_), sizeof(size_t)); + output.write(reinterpret_cast(&bin_data_bytes_), sizeof(size_t)); + output.write(reinterpret_cast(&ex_data_bytes_), sizeof(size_t)); + output.write(reinterpret_cast(&level0_links_bytes_), sizeof(size_t)); output.write(reinterpret_cast(&bin_data_offset_), sizeof(size_t)); output.write(reinterpret_cast(&ex_data_offset_), sizeof(size_t)); output.write(reinterpret_cast(&label_offset_), sizeof(PID)); - output.write(reinterpret_cast(&size_data_per_element_), sizeof(size_t)); - output.write(reinterpret_cast(&size_links_per_element_), sizeof(size_t)); + output.write(reinterpret_cast(&level0_record_bytes_), sizeof(size_t)); + output.write(reinterpret_cast(&upper_layer_links_bytes_), sizeof(size_t)); output.write(reinterpret_cast(&max_level_), sizeof(int)); output.write(reinterpret_cast(&entry_point_id_), sizeof(PID)); output.write(reinterpret_cast(&M_), sizeof(size_t)); - output.write(reinterpret_cast(&maxM_), sizeof(size_t)); - output.write(reinterpret_cast(&maxM0_), sizeof(size_t)); + output.write(reinterpret_cast(&upper_layer_max_degree_), sizeof(size_t)); + output.write(reinterpret_cast(&base_layer_max_degree_), sizeof(size_t)); output.write(reinterpret_cast(&mult_), sizeof(double)); output.write(reinterpret_cast(&ef_construction_), sizeof(size_t)); @@ -190,12 +191,12 @@ void HierarchicalNSW::save(const char* filename) const { output.write( reinterpret_cast(data_level0_memory_.get()), - static_cast(cur_element_count_ * size_data_per_element_) + static_cast(num_stored_points_ * level0_record_bytes_) ); - for (size_t i = 0; i < cur_element_count_; i++) { + for (size_t i = 0; i < num_stored_points_; i++) { unsigned int link_list_size = - element_levels_[i] > 0 ? size_links_per_element_ * element_levels_[i] : 0; + element_levels_[i] > 0 ? upper_layer_links_bytes_ * element_levels_[i] : 0; output.write(reinterpret_cast(&link_list_size), sizeof(unsigned int)); if (link_list_size != 0) { output.write(link_lists_[i].get(), link_list_size); @@ -245,20 +246,20 @@ void HierarchicalNSW::load(const char* filename) { read_header(loaded.num_clusters_); read_header(loaded.ex_bits_); read_header(loaded.metric_type_); - read_header(loaded.size_bin_data_); - read_header(loaded.size_ex_data_); - read_header(loaded.size_links_level0_); + read_header(loaded.bin_data_bytes_); + read_header(loaded.ex_data_bytes_); + read_header(loaded.level0_links_bytes_); read_header(loaded.bin_data_offset_); read_header(loaded.ex_data_offset_); read_header(label_offset); loaded.label_offset_ = label_offset; - read_header(loaded.size_data_per_element_); - read_header(loaded.size_links_per_element_); + read_header(loaded.level0_record_bytes_); + read_header(loaded.upper_layer_links_bytes_); read_header(loaded.max_level_); read_header(loaded.entry_point_id_); read_header(loaded.M_); - read_header(loaded.maxM_); - read_header(loaded.maxM0_); + read_header(loaded.upper_layer_max_degree_); + read_header(loaded.base_layer_max_degree_); read_header(loaded.mult_); read_header(loaded.ef_construction_); @@ -271,8 +272,9 @@ void HierarchicalNSW::load(const char* filename) { loaded.padded_dim_ < loaded.dim_ || loaded.padded_dim_ % 32 != 0 || loaded.padded_dim_ > rotator_impl::FhtKacRotator::kMaxDim || loaded.ex_bits_ > 8 || (loaded.metric_type_ != METRIC_L2 && loaded.metric_type_ != METRIC_IP) || - loaded.M_ < 2 || loaded.M_ > 10000 || loaded.maxM_ != loaded.M_ || - loaded.maxM0_ != 2 * loaded.M_ || loaded.ef_construction_ < loaded.M_ || + loaded.M_ < 2 || loaded.M_ > 10000 || loaded.upper_layer_max_degree_ != loaded.M_ || + loaded.base_layer_max_degree_ != 2 * loaded.M_ || + loaded.ef_construction_ < loaded.M_ || (element_count == 0 && (loaded.max_level_ != -1 || loaded.entry_point_id_ != kPidMax)) || (element_count != 0 && @@ -280,21 +282,22 @@ void HierarchicalNSW::load(const char* filename) { invalid_file(); } - const size_t expected_level0 = (loaded.maxM0_ + 1) * sizeof(PID); + const size_t expected_level0 = (loaded.base_layer_max_degree_ + 1) * sizeof(PID); const size_t expected_bin = BinDataMap::data_bytes(loaded.padded_dim_); const size_t expected_ex = ExDataMap::data_bytes(loaded.padded_dim_, loaded.ex_bits_); const size_t expected_label = expected_level0 + sizeof(PID); const size_t expected_bin_offset = expected_label + sizeof(PID); const size_t expected_ex_offset = expected_bin_offset + expected_bin; - const size_t expected_stride = expected_ex_offset + expected_ex; - if (loaded.size_bin_data_ != expected_bin || loaded.size_ex_data_ != expected_ex || - loaded.size_links_level0_ != expected_level0 || + const size_t expected_level0_record_bytes = expected_ex_offset + expected_ex; + if (loaded.bin_data_bytes_ != expected_bin || loaded.ex_data_bytes_ != expected_ex || + loaded.level0_links_bytes_ != expected_level0 || loaded.label_offset_ != expected_label || loaded.bin_data_offset_ != expected_bin_offset || loaded.ex_data_offset_ != expected_ex_offset || - loaded.size_data_per_element_ != expected_stride || - loaded.size_links_per_element_ != (loaded.maxM_ + 1) * sizeof(PID) || + loaded.level0_record_bytes_ != expected_level0_record_bytes || + loaded.upper_layer_links_bytes_ != + (loaded.upper_layer_max_degree_ + 1) * sizeof(PID) || !std::isfinite(loaded.mult_) || loaded.mult_ <= 0.0) { invalid_file(); } @@ -303,8 +306,8 @@ void HierarchicalNSW::load(const char* filename) { checked_product(loaded.num_clusters_, loaded.padded_dim_), sizeof(float) ); const size_t data_capacity_bytes = - checked_product(loaded.max_elements_, loaded.size_data_per_element_); - const size_t data_bytes = checked_product(element_count, loaded.size_data_per_element_); + checked_product(loaded.max_elements_, loaded.level0_record_bytes_); + const size_t data_bytes = checked_product(element_count, loaded.level0_record_bytes_); // Validate pointer-table capacity before allocating it. (void)checked_product(loaded.max_elements_, sizeof(LinkStorage)); const size_t link_headers_bytes = checked_product(element_count, sizeof(unsigned int)); @@ -360,15 +363,15 @@ void HierarchicalNSW::load(const char* filename) { loaded.link_lists_.resize(element_count); for (size_t i = 0; i < element_count; ++i) { const char* row = - loaded.data_level0_memory_.get() + i * loaded.size_data_per_element_; + loaded.data_level0_memory_.get() + i * loaded.level0_record_bytes_; PID cluster_id = 0; PID label = 0; unsigned short level0_degree = 0; std::memcpy(&level0_degree, row, sizeof(level0_degree)); - std::memcpy(&cluster_id, row + loaded.size_links_level0_, sizeof(cluster_id)); + std::memcpy(&cluster_id, row + loaded.level0_links_bytes_, sizeof(cluster_id)); std::memcpy(&label, row + loaded.label_offset_, sizeof(label)); if ((cluster_id & ~kRemovedMask) >= loaded.num_clusters_ || - level0_degree > loaded.maxM0_) { + level0_degree > loaded.base_layer_max_degree_) { invalid_file(); } loaded.num_removed_ += (cluster_id & kRemovedMask) != 0; @@ -386,8 +389,8 @@ void HierarchicalNSW::load(const char* filename) { unsigned int link_list_size = 0; read_payload(reinterpret_cast(&link_list_size), sizeof(link_list_size)); if (link_list_size != 0) { - if (link_list_size % loaded.size_links_per_element_ != 0 || - link_list_size / loaded.size_links_per_element_ > + if (link_list_size % loaded.upper_layer_links_bytes_ != 0 || + link_list_size / loaded.upper_layer_links_bytes_ > static_cast(loaded.max_level_) || link_list_size > remaining) { invalid_file(); @@ -397,17 +400,17 @@ void HierarchicalNSW::load(const char* filename) { throw std::bad_alloc(); } loaded.element_levels_[i] = - static_cast(link_list_size / loaded.size_links_per_element_); + static_cast(link_list_size / loaded.upper_layer_links_bytes_); } - loaded.cur_element_count_ = i + 1; + loaded.num_stored_points_ = i + 1; if (link_list_size != 0) { read_payload(loaded.link_lists_[i].get(), link_list_size); for (int level = 0; level < loaded.element_levels_[i]; ++level) { const char* link = - loaded.link_lists_[i].get() + level * loaded.size_links_per_element_; + loaded.link_lists_[i].get() + level * loaded.upper_layer_links_bytes_; unsigned short degree = 0; std::memcpy(°ree, link, sizeof(degree)); - if (degree > loaded.maxM_) { + if (degree > loaded.upper_layer_max_degree_) { invalid_file(); } for (size_t edge = 0; edge < degree; ++edge) { @@ -426,7 +429,7 @@ void HierarchicalNSW::load(const char* filename) { for (int level = 1; level <= loaded.element_levels_[i]; ++level) { const char* link = loaded.link_lists_[i].get() + - (static_cast(level) - 1) * loaded.size_links_per_element_; + (static_cast(level) - 1) * loaded.upper_layer_links_bytes_; unsigned short degree = 0; std::memcpy(°ree, link, sizeof(degree)); for (size_t edge = 0; edge < degree; ++edge) { @@ -461,7 +464,7 @@ void HierarchicalNSW::load(const char* filename) { loaded.raw_dist_func_ = (loaded.metric_type_ == METRIC_IP) ? dot_product_dis : euclidean_sqr; loaded.ip_func_ = select_excode_ipfunc(loaded.ex_bits_); - loaded.ef_ = 10; + loaded.default_ef_search_ = 10; loaded.query_config_ = quant::faster_config(loaded.padded_dim_, SplitSingleQuery::kNumBits); @@ -469,29 +472,29 @@ void HierarchicalNSW::load(const char* filename) { using std::swap; swap(construction_data_, loaded.construction_data_); swap(max_elements_, loaded.max_elements_); - const size_t old_count = cur_element_count_.load(); - cur_element_count_ = loaded.cur_element_count_.load(); - loaded.cur_element_count_ = old_count; + const size_t old_count = num_stored_points_.load(); + num_stored_points_ = loaded.num_stored_points_.load(); + loaded.num_stored_points_ = old_count; swap(num_removed_, loaded.num_removed_); - swap(size_data_per_element_, loaded.size_data_per_element_); - swap(size_links_per_element_, loaded.size_links_per_element_); + swap(level0_record_bytes_, loaded.level0_record_bytes_); + swap(upper_layer_links_bytes_, loaded.upper_layer_links_bytes_); swap(M_, loaded.M_); - swap(maxM_, loaded.maxM_); - swap(maxM0_, loaded.maxM0_); + swap(upper_layer_max_degree_, loaded.upper_layer_max_degree_); + swap(base_layer_max_degree_, loaded.base_layer_max_degree_); swap(ef_construction_, loaded.ef_construction_); - swap(ef_, loaded.ef_); + swap(default_ef_search_, loaded.default_ef_search_); swap(metric_type_, loaded.metric_type_); swap(mult_, loaded.mult_); swap(max_level_, loaded.max_level_); label_op_locks_.swap(loaded.label_op_locks_); link_list_locks_.swap(loaded.link_list_locks_); swap(entry_point_id_, loaded.entry_point_id_); - swap(size_links_level0_, loaded.size_links_level0_); + swap(level0_links_bytes_, loaded.level0_links_bytes_); swap(bin_data_offset_, loaded.bin_data_offset_); swap(ex_data_offset_, loaded.ex_data_offset_); swap(label_offset_, loaded.label_offset_); - swap(size_bin_data_, loaded.size_bin_data_); - swap(size_ex_data_, loaded.size_ex_data_); + swap(bin_data_bytes_, loaded.bin_data_bytes_); + swap(ex_data_bytes_, loaded.ex_data_bytes_); swap(ex_bits_, loaded.ex_bits_); swap(data_level0_memory_, loaded.data_level0_memory_); swap(link_lists_, loaded.link_lists_); @@ -515,9 +518,9 @@ void HierarchicalNSW::construct( const float* data, PID* cluster_ids, size_t num_threads, - bool faster + bool fast_quantization ) { - if (num_clusters_ != 0 || cur_element_count_ != 0) { + if (num_clusters_ != 0 || num_stored_points_ != 0) { throw std::logic_error("HNSW index is already constructed or loaded"); } if (num_clusters == 0 || num_clusters > buffer::kSearchBufferMaxPointCount) { @@ -557,7 +560,7 @@ void HierarchicalNSW::construct( } quant::RabitqConfig config; - if (faster) { + if (fast_quantization) { config = quant::faster_config(padded_dim_, ex_bits_ + 1); } @@ -570,9 +573,9 @@ void HierarchicalNSW::construct( } catch (...) { // parallel_for joins every worker before propagating an exception. // Discard the partial graph, retaining capacity and rotation for a retry. - std::fill_n(element_levels_.begin(), cur_element_count_.load(), 0); + std::fill_n(element_levels_.begin(), num_stored_points_.load(), 0); link_lists_.clear(); - cur_element_count_ = 0; + num_stored_points_ = 0; num_removed_ = 0; label_lookup_.clear(); entry_point_id_ = kPidMax; @@ -588,12 +591,12 @@ void HierarchicalNSW::construct( void HierarchicalNSW::add_point( PID label, PID cluster_id, const quant::RabitqConfig& config ) { - std::unique_lock lock_label(get_label_op_mutex(label)); + std::unique_lock label_lock(get_label_op_mutex(label)); int node_level = 0; - PID cur_c = 0; + PID new_node_id = 0; { - std::unique_lock lock_table(label_lookup_lock_); + std::unique_lock label_lookup_lock(label_lookup_lock_); if (label_lookup_.find(label) != label_lookup_.end()) { throw std::runtime_error( "Currently not support replacement of existing elements, only support " @@ -601,37 +604,37 @@ void HierarchicalNSW::add_point( ); } - if (cur_element_count_ >= max_elements_) { + if (num_stored_points_ >= max_elements_) { throw std::runtime_error("The number of elements exceeds the specified limit"); } - cur_c = cur_element_count_; + new_node_id = num_stored_points_; link_lists_.emplace_back(); - cur_element_count_++; - label_lookup_[label] = cur_c; + num_stored_points_++; + label_lookup_[label] = new_node_id; node_level = get_random_level(mult_); } - std::unique_lock lock_el(link_list_locks_[cur_c]); + std::unique_lock node_lock(link_list_locks_[new_node_id]); - element_levels_[cur_c] = node_level; - std::unique_lock templock(global_); + element_levels_[new_node_id] = node_level; + std::unique_lock entry_point_lock(global_); int previous_max_level = max_level_; PID current_id = entry_point_id_; if (node_level <= previous_max_level) { - templock.unlock(); + entry_point_lock.unlock(); } // initialize the current memory. memset( - data_level0_memory_.get() + (cur_c * size_data_per_element_), + data_level0_memory_.get() + (new_node_id * level0_record_bytes_), 0, - size_data_per_element_ + level0_record_bytes_ ); // Initialisation of label and cluster id - memcpy(get_external_label_pt(cur_c), &label, sizeof(PID)); - memcpy(get_clusterid_pt(cur_c), &cluster_id, sizeof(PID)); + memcpy(get_external_label_pt(new_node_id), &label, sizeof(PID)); + memcpy(get_clusterid_pt(new_node_id), &cluster_id, sizeof(PID)); // Quantize raw data and initialize quantized data std::vector rotated_data(padded_dim_); @@ -641,46 +644,53 @@ void HierarchicalNSW::add_point( reinterpret_cast(centroids_memory_.get()) + (cluster_id * padded_dim_), padded_dim_, ex_bits_, - get_bindata_by_internalid(cur_c), - get_exdata_by_internalid(cur_c), + get_bindata_by_internalid(new_node_id), + get_exdata_by_internalid(new_node_id), metric_type_, config ); // If the current vertex is at level >0, it needs some space to store the extra edges. if (node_level > 0) { - link_lists_[cur_c].reset( - static_cast(std::malloc((size_links_per_element_ * node_level) + 1)) + link_lists_[new_node_id].reset( + static_cast(std::malloc((upper_layer_links_bytes_ * node_level) + 1)) ); - if (link_lists_[cur_c] == nullptr) { + if (link_lists_[new_node_id] == nullptr) { throw std::runtime_error( "Not enough memory: add_point failed to allocate linklist" ); } - memset(link_lists_[cur_c].get(), 0, (size_links_per_element_ * node_level) + 1); + memset( + link_lists_[new_node_id].get(), 0, (upper_layer_links_bytes_ * node_level) + 1 + ); } if (static_cast(current_id) != -1) { if (node_level < previous_max_level) { - float curdist = get_data_dist(current_id, cur_c); + float current_distance = get_data_dist(current_id, new_node_id); for (int level = previous_max_level; level > node_level; level--) { bool changed = true; while (changed) { changed = false; - unsigned int* data; std::unique_lock lock(link_list_locks_[current_id]); - data = get_linklist(current_id, level); - int size = get_list_count(data); + const PID* data = get_linklist(current_id, level); + const size_t size = get_list_count(data); - auto* datal = static_cast(data + 1); - for (int i = 0; i < size; i++) { + const PID* datal = data + 1; + for (size_t i = 0; i < size; i++) { PID cand = datal[i]; if (cand > max_elements_) { - throw std::runtime_error("cand error"); + throw std::runtime_error( + "HNSW construct upper-layer traversal: invalid candidate " + "ID " + + std::to_string(cand) + " at level " + + std::to_string(level) + + " (capacity=" + std::to_string(max_elements_) + ")" + ); } - float d = get_data_dist(cand, cur_c); - if (d < curdist) { - curdist = d; + float d = get_data_dist(cand, new_node_id); + if (d < current_distance) { + current_distance = d; current_id = cand; changed = true; } @@ -691,8 +701,8 @@ void HierarchicalNSW::add_point( for (int level = std::min(node_level, previous_max_level); level >= 0; level--) { maxheap> top_candidates = - search_base_layer(current_id, cur_c, level); - current_id = mutually_connect_new_element(cur_c, top_candidates, level); + search_base_layer(current_id, new_node_id, level); + current_id = mutually_connect_new_element(new_node_id, top_candidates, level); } } else { // Do nothing for the first element @@ -702,23 +712,23 @@ void HierarchicalNSW::add_point( // Releasing lock for the maximum level if (node_level > previous_max_level) { - entry_point_id_ = cur_c; + entry_point_id_ = new_node_id; max_level_ = node_level; } } maxheap> HierarchicalNSW::search_base_layer( - PID ep_id, PID cur_c, int layer + PID entry_point_id, PID new_node_id, int layer ) { - auto vl = visited_list_pool_->get_free_vislist(); + auto visited = visited_list_pool_->get_free_vislist(); maxheap> top_candidates; minheap> candidate_set; - float lower_bound = get_data_dist(ep_id, cur_c); - top_candidates.emplace(lower_bound, ep_id); - candidate_set.emplace(lower_bound, ep_id); - vl->set(ep_id); + float lower_bound = get_data_dist(entry_point_id, new_node_id); + top_candidates.emplace(lower_bound, entry_point_id); + candidate_set.emplace(lower_bound, entry_point_id); + visited->set(entry_point_id); while (!candidate_set.empty()) { std::pair curr_el_pair = candidate_set.top(); @@ -727,18 +737,14 @@ maxheap> HierarchicalNSW::search_base_layer( } candidate_set.pop(); - PID cur_node_num = curr_el_pair.second; + PID current_id = curr_el_pair.second; - std::unique_lock lock(link_list_locks_[cur_node_num]); + std::unique_lock lock(link_list_locks_[current_id]); - int* data; - if (layer == 0) { - data = reinterpret_cast(get_linklist0(cur_node_num)); - } else { - data = reinterpret_cast(get_linklist(cur_node_num, layer)); - } - size_t size = get_list_count(reinterpret_cast(data)); - auto* datal = reinterpret_cast(data + 1); + const PID* data = + (layer == 0) ? get_linklist0(current_id) : get_linklist(current_id, layer); + const size_t size = get_list_count(data); + const PID* datal = data + 1; rabitqlib::memory::mem_prefetch_l1( reinterpret_cast( @@ -756,10 +762,10 @@ maxheap> HierarchicalNSW::search_base_layer( for (size_t j = 0; j < size; j++) { PID candidate_id = *(datal + j); - if (vl->get(candidate_id)) { + if (visited->get(candidate_id)) { continue; } - vl->set(candidate_id); + visited->set(candidate_id); if (j < size - 1) { rabitqlib::memory::mem_prefetch_l1( @@ -770,10 +776,11 @@ maxheap> HierarchicalNSW::search_base_layer( ); } - float dist1 = get_data_dist(candidate_id, cur_c); - if (top_candidates.size() < ef_construction_ || lower_bound > dist1) { - candidate_set.emplace(dist1, candidate_id); - top_candidates.emplace(dist1, candidate_id); + float candidate_distance = get_data_dist(candidate_id, new_node_id); + if (top_candidates.size() < ef_construction_ || + lower_bound > candidate_distance) { + candidate_set.emplace(candidate_distance, candidate_id); + top_candidates.emplace(candidate_distance, candidate_id); if (top_candidates.size() > ef_construction_) { top_candidates.pop(); } @@ -783,7 +790,7 @@ maxheap> HierarchicalNSW::search_base_layer( } } } - visited_list_pool_->release_vis_list(std::move(vl)); + visited_list_pool_->release_vis_list(std::move(visited)); return top_candidates; } @@ -808,13 +815,13 @@ void prune_neighbors( break; } std::pair current_pair = queue_closest.top(); - float dist_to_query = current_pair.first; + float query_distance = current_pair.first; queue_closest.pop(); bool good = true; for (std::pair second_pair : return_list) { - float curdist = distance(second_pair.second, current_pair.second); - if (curdist < dist_to_query) { + float current_distance = distance(second_pair.second, current_pair.second); + if (current_distance < query_distance) { good = false; break; } @@ -832,13 +839,13 @@ void prune_neighbors( template PID HierarchicalNSW::mutually_connect( - PID cur_c, + PID new_node_id, maxheap>& top_candidates, int level, Distance distance, ReverseEdgeDistance reverse_edge_distance ) { - size_t max_m = level > 0 ? maxM_ : maxM0_; + size_t max_m = level > 0 ? upper_layer_max_degree_ : base_layer_max_degree_; detail::prune_neighbors(top_candidates, M_, distance); if (top_candidates.size() > M_) { throw std::runtime_error( @@ -858,7 +865,8 @@ PID HierarchicalNSW::mutually_connect( PID next_closest_entry_point = selected_neighbors.back(); { - PID* ll_cur = (level == 0) ? get_linklist0(cur_c) : get_linklist(cur_c, level); + PID* ll_cur = + (level == 0) ? get_linklist0(new_node_id) : get_linklist(new_node_id, level); if (*ll_cur > 0) { throw std::runtime_error( "The newly inserted element should have blank link list" @@ -888,7 +896,7 @@ PID HierarchicalNSW::mutually_connect( if (sz_link_list_other > max_m) { throw std::runtime_error("Bad value of sz_link_list_other"); } - if (selected_neighbor == cur_c) { + if (selected_neighbor == new_node_id) { throw std::runtime_error("Trying to connect an element to itself"); } if (level > element_levels_[selected_neighbor]) { @@ -899,7 +907,7 @@ PID HierarchicalNSW::mutually_connect( bool is_cur_c_present = false; for (size_t j = 0; j < sz_link_list_other; j++) { - if (data[j] == cur_c) { + if (data[j] == new_node_id) { is_cur_c_present = true; break; } @@ -907,12 +915,12 @@ PID HierarchicalNSW::mutually_connect( if (!is_cur_c_present) { if (sz_link_list_other < max_m) { - data[sz_link_list_other] = cur_c; + data[sz_link_list_other] = new_node_id; set_list_count(ll_other, sz_link_list_other + 1); } else { - float d_max = reverse_edge_distance(selected_neighbor); + float new_neighbor_distance = reverse_edge_distance(selected_neighbor); maxheap> candidates; - candidates.emplace(d_max, cur_c); + candidates.emplace(new_neighbor_distance, new_node_id); for (size_t j = 0; j < sz_link_list_other; j++) { candidates.emplace(distance(data[j], selected_neighbor), data[j]); } @@ -934,15 +942,15 @@ PID HierarchicalNSW::mutually_connect( } PID HierarchicalNSW::mutually_connect_new_element( - PID cur_c, maxheap>& top_candidates, int level + PID new_node_id, maxheap>& top_candidates, int level ) { return mutually_connect( - cur_c, + new_node_id, top_candidates, level, [this](PID target, PID query) { return get_data_dist(target, query); }, // Preserve the original operand order for a newly proposed reverse edge. - [this, cur_c](PID neighbor) { return get_data_dist(neighbor, cur_c); } + [this, new_node_id](PID neighbor) { return get_data_dist(neighbor, new_node_id); } ); } @@ -969,12 +977,12 @@ std::vector>> HierarchicalNSW::search( return results; } -size_t HierarchicalNSW::remove(const PID* labels, size_t n) { +size_t HierarchicalNSW::remove(const PID* labels, size_t num_points) { if (data_level0_memory_ == nullptr || centroids_memory_ == nullptr || rotator_ == nullptr || num_clusters_ == 0) { throw std::logic_error("HNSW index must be constructed or loaded before remove"); } - if (n == 0) { + if (num_points == 0) { return 0; } if (labels == nullptr) { @@ -983,10 +991,10 @@ size_t HierarchicalNSW::remove(const PID* labels, size_t n) { // Resolve every label before marking any point, so an unknown label removes // nothing. - std::vector internal_ids(n); + std::vector internal_ids(num_points); { - std::unique_lock lock_table(label_lookup_lock_); - for (size_t i = 0; i < n; ++i) { + std::unique_lock label_lookup_lock(label_lookup_lock_); + for (size_t i = 0; i < num_points; ++i) { const auto found = label_lookup_.find(labels[i]); if (found == label_lookup_.end()) { throw std::invalid_argument("HNSW remove label is not in the index"); @@ -1102,17 +1110,17 @@ float HierarchicalNSW::get_quant_dist( } maxheap> HierarchicalNSW::search_base_layer_quant( - PID ep_id, PID cur_c, int layer, detail::QuantQuery& query + PID entry_point_id, PID new_node_id, int layer, detail::QuantQuery& query ) { - auto vl = visited_list_pool_->get_free_vislist(); + auto visited = visited_list_pool_->get_free_vislist(); maxheap> top_candidates; minheap> candidate_set; - float lower_bound = get_quant_dist(ep_id, cur_c, query); - top_candidates.emplace(lower_bound, ep_id); - candidate_set.emplace(lower_bound, ep_id); - vl->set(ep_id); + float lower_bound = get_quant_dist(entry_point_id, new_node_id, query); + top_candidates.emplace(lower_bound, entry_point_id); + candidate_set.emplace(lower_bound, entry_point_id); + visited->set(entry_point_id); while (!candidate_set.empty()) { std::pair curr_el_pair = candidate_set.top(); @@ -1121,28 +1129,33 @@ maxheap> HierarchicalNSW::search_base_layer_quant( } candidate_set.pop(); - PID cur_node_num = curr_el_pair.second; - std::unique_lock lock(link_list_locks_[cur_node_num]); + PID current_id = curr_el_pair.second; + std::unique_lock lock(link_list_locks_[current_id]); - PID* data = - (layer == 0) ? get_linklist0(cur_node_num) : get_linklist(cur_node_num, layer); - size_t size = get_list_count(data); - auto* datal = data + 1; + const PID* data = + (layer == 0) ? get_linklist0(current_id) : get_linklist(current_id, layer); + const size_t size = get_list_count(data); + const PID* datal = data + 1; for (size_t j = 0; j < size; j++) { PID candidate_id = datal[j]; - if (candidate_id >= cur_element_count_) { - throw std::runtime_error("cand error"); + if (candidate_id >= num_stored_points_) { + throw std::runtime_error( + "HNSW add layer search: invalid candidate ID " + + std::to_string(candidate_id) + " at level " + std::to_string(layer) + + " (stored_points=" + std::to_string(num_stored_points_.load()) + ")" + ); } - if (vl->get(candidate_id)) { + if (visited->get(candidate_id)) { continue; } - vl->set(candidate_id); + visited->set(candidate_id); - float dist1 = get_quant_dist(candidate_id, cur_c, query); - if (top_candidates.size() < ef_construction_ || lower_bound > dist1) { - candidate_set.emplace(dist1, candidate_id); - top_candidates.emplace(dist1, candidate_id); + float candidate_distance = get_quant_dist(candidate_id, new_node_id, query); + if (top_candidates.size() < ef_construction_ || + lower_bound > candidate_distance) { + candidate_set.emplace(candidate_distance, candidate_id); + top_candidates.emplace(candidate_distance, candidate_id); if (top_candidates.size() > ef_construction_) { top_candidates.pop(); } @@ -1152,26 +1165,26 @@ maxheap> HierarchicalNSW::search_base_layer_quant( } } } - visited_list_pool_->release_vis_list(std::move(vl)); + visited_list_pool_->release_vis_list(std::move(visited)); return top_candidates; } PID HierarchicalNSW::mutually_connect_quant( - PID cur_c, + PID new_node_id, maxheap>& top_candidates, int level, detail::QuantQuery& query ) { return mutually_connect( - cur_c, + new_node_id, top_candidates, level, [this, &query](PID target, PID source) { return get_quant_dist(target, source, query); }, // Quantized scoring is directional: reconstruct the existing neighbor. - [this, cur_c, &query](PID neighbor) { - return get_quant_dist(cur_c, neighbor, query); + [this, new_node_id, &query](PID neighbor) { + return get_quant_dist(new_node_id, neighbor, query); } ); } @@ -1179,41 +1192,47 @@ PID HierarchicalNSW::mutually_connect_quant( // The slot, the level and the link list are all decided by the caller, so this // can run on many threads at once and the labels do not depend on their order. void HierarchicalNSW::add_point_quant( - PID cur_c, int node_level, LinkStorage& link_list, detail::QuantQuery& query + PID new_node_id, int node_level, LinkStorage& link_list, detail::QuantQuery& query ) { - std::unique_lock lock_el(link_list_locks_[cur_c]); + std::unique_lock node_lock(link_list_locks_[new_node_id]); - std::unique_lock templock(global_); + std::unique_lock entry_point_lock(global_); int previous_max_level = max_level_; PID current_id = entry_point_id_; if (node_level <= previous_max_level) { - templock.unlock(); + entry_point_lock.unlock(); } const PID previous_entry_point = current_id; - link_lists_[cur_c] = std::move(link_list); - element_levels_[cur_c] = node_level; + link_lists_[new_node_id] = std::move(link_list); + element_levels_[new_node_id] = node_level; try { if (static_cast(current_id) != -1) { if (node_level < previous_max_level) { - float curdist = get_quant_dist(current_id, cur_c, query); + float current_distance = get_quant_dist(current_id, new_node_id, query); for (int level = previous_max_level; level > node_level; level--) { bool changed = true; while (changed) { changed = false; std::unique_lock lock(link_list_locks_[current_id]); - PID* data = get_linklist(current_id, level); - int size = get_list_count(data); - auto* datal = data + 1; - for (int i = 0; i < size; i++) { + const PID* data = get_linklist(current_id, level); + const size_t size = get_list_count(data); + const PID* datal = data + 1; + for (size_t i = 0; i < size; i++) { PID cand = datal[i]; - if (cand >= cur_element_count_) { - throw std::runtime_error("cand error"); + if (cand >= num_stored_points_) { + throw std::runtime_error( + "HNSW add upper-layer traversal: invalid candidate " + "ID " + + std::to_string(cand) + " at level " + + std::to_string(level) + " (stored_points=" + + std::to_string(num_stored_points_.load()) + ")" + ); } - float d = get_quant_dist(cand, cur_c, query); - if (d < curdist) { - curdist = d; + float d = get_quant_dist(cand, new_node_id, query); + if (d < current_distance) { + current_distance = d; current_id = cand; changed = true; } @@ -1225,22 +1244,23 @@ void HierarchicalNSW::add_point_quant( for (int level = std::min(node_level, previous_max_level); level >= 0; level--) { maxheap> top_candidates = - search_base_layer_quant(current_id, cur_c, level, query); - current_id = mutually_connect_quant(cur_c, top_candidates, level, query); + search_base_layer_quant(current_id, new_node_id, level, query); + current_id = + mutually_connect_quant(new_node_id, top_candidates, level, query); } } else { - entry_point_id_ = cur_c; + entry_point_id_ = new_node_id; max_level_ = node_level; } if (node_level > previous_max_level) { - entry_point_id_ = cur_c; + entry_point_id_ = new_node_id; max_level_ = node_level; } } catch (...) { // Upper layers may already expose this node. Keep a route to the old // graph if base-layer linking was interrupted before choosing neighbors. - PID* base_links = get_linklist0(cur_c); + PID* base_links = get_linklist0(new_node_id); if (previous_entry_point != kPidMax && get_list_count(base_links) == 0) { base_links[1] = previous_entry_point; set_list_count(base_links, 1); @@ -1248,9 +1268,9 @@ void HierarchicalNSW::add_point_quant( // An interrupted promotion has no incoming links above the previous // maximum level. Discard those levels so save/load sees a valid graph. if (node_level > previous_max_level) { - element_levels_[cur_c] = std::max(0, previous_max_level); - if (element_levels_[cur_c] == 0) { - link_lists_[cur_c].reset(); + element_levels_[new_node_id] = std::max(0, previous_max_level); + if (element_levels_[new_node_id] == 0) { + link_lists_[new_node_id].reset(); } } throw; @@ -1260,10 +1280,10 @@ void HierarchicalNSW::add_point_quant( // A capacity change reallocates the level-0 block, invalidating pointers into it. // This must not run while another thread searches; equal capacity is a no-op. void HierarchicalNSW::resize(size_t new_max_elements) { - if (size_data_per_element_ == 0) { + if (level0_record_bytes_ == 0) { throw std::logic_error("HNSW index must be initialized before resize"); } - if (new_max_elements < cur_element_count_) { + if (new_max_elements < num_stored_points_) { throw std::invalid_argument("HNSW resize: capacity below the element count"); } if (new_max_elements > buffer::kSearchBufferMaxPointCount) { @@ -1272,21 +1292,21 @@ void HierarchicalNSW::resize(size_t new_max_elements) { if (new_max_elements == max_elements_) { return; } - if (new_max_elements > std::numeric_limits::max() / size_data_per_element_) { + if (new_max_elements > std::numeric_limits::max() / level0_record_bytes_) { throw std::invalid_argument("HNSW resize: capacity exceeds addressable storage"); } // Everything new is built before anything old is released, so a failure // leaves the index as it was. realloc is not an option: the level-0 block // comes from aligned_alloc, which realloc does not keep aligned. - const size_t level0_bytes = new_max_elements * size_data_per_element_; + const size_t level0_bytes = new_max_elements * level0_record_bytes_; AlignedStorage level0(memory::huge_page_allocate(level0_bytes)); if (level0 == nullptr) { throw std::bad_alloc(); } std::memset(level0.get(), 0, level0_bytes); std::memcpy( - level0.get(), data_level0_memory_.get(), cur_element_count_ * size_data_per_element_ + level0.get(), data_level0_memory_.get(), num_stored_points_ * level0_record_bytes_ ); std::vector links; @@ -1311,30 +1331,34 @@ void HierarchicalNSW::resize(size_t new_max_elements) { } std::vector HierarchicalNSW::add( - const float* data, size_t n, const PID* cluster_ids, bool faster, size_t num_threads + const float* data, + size_t num_points, + const PID* cluster_ids, + bool fast_quantization, + size_t num_threads ) { if (data_level0_memory_ == nullptr || centroids_memory_ == nullptr || rotator_ == nullptr || num_clusters_ == 0) { throw std::logic_error("HNSW index must be constructed or loaded before add"); } - if (n == 0) { + if (num_points == 0) { return {}; } if (data == nullptr) { throw std::invalid_argument("HNSW add inputs must not be null"); } - if (n > max_elements_ - cur_element_count_) { + if (num_points > max_elements_ - num_stored_points_) { throw std::invalid_argument("HNSW add: not enough capacity"); } - if (n > buffer::kSearchBufferMaxPointCount - cur_element_count_) { + if (num_points > buffer::kSearchBufferMaxPointCount - num_stored_points_) { throw std::invalid_argument("HNSW add: point count exceeds the supported ID range"); } - std::vector labels(n); + std::vector labels(num_points); // Either the caller names the cluster of every point, or each one goes to its // nearest centroid, ordered the way a query is routed. - std::vector assigned(n, 0); + std::vector assigned(num_points, 0); if (cluster_ids != nullptr) { - for (size_t i = 0; i < n; ++i) { + for (size_t i = 0; i < num_points; ++i) { if (cluster_ids[i] >= num_clusters_) { throw std::invalid_argument("HNSW cluster ID is out of range"); } @@ -1342,43 +1366,45 @@ std::vector HierarchicalNSW::add( } } else { const auto* centroids = reinterpret_cast(centroids_memory_.get()); - rabitqlib::index_detail::parallel_for(0, n, num_threads, [&](size_t i, size_t) { - std::vector rotated(padded_dim_); - rotator_->rotate(data + (i * dim_), rotated.data()); - float best = std::numeric_limits::max(); - for (size_t c = 0; c < num_clusters_; ++c) { - const float dist = raw_dist_func_( - rotated.data(), centroids + (c * padded_dim_), padded_dim_ - ); - if (dist < best) { - best = dist; - assigned[i] = static_cast(c); + rabitqlib::index_detail::parallel_for( + 0, num_points, num_threads, [&](size_t i, size_t) { + std::vector rotated(padded_dim_); + rotator_->rotate(data + (i * dim_), rotated.data()); + float best = std::numeric_limits::max(); + for (size_t c = 0; c < num_clusters_; ++c) { + const float dist = raw_dist_func_( + rotated.data(), centroids + (c * padded_dim_), padded_dim_ + ); + if (dist < best) { + best = dist; + assigned[i] = static_cast(c); + } } } - }); + ); } quant::RabitqConfig config; - if (faster) { + if (fast_quantization) { config = quant::faster_config(padded_dim_, ex_bits_ + 1); } // Drawn here because the generator is shared: the inserts below run in // parallel and must not touch it. - std::vector levels(n); - for (size_t i = 0; i < n; ++i) { + std::vector levels(num_points); + for (size_t i = 0; i < num_points; ++i) { levels[i] = get_random_level(mult_); } // Every link list is allocated before any of them is installed, so running // out of memory leaves the index untouched. std::vector link_lists; - link_lists.reserve(n); - for (size_t i = 0; i < n; ++i) { + link_lists.reserve(num_points); + for (size_t i = 0; i < num_points; ++i) { char* raw = nullptr; if (levels[i] > 0) { raw = static_cast( - std::calloc(1, (size_links_per_element_ * levels[i]) + 1) + std::calloc(1, (upper_layer_links_bytes_ * levels[i]) + 1) ); if (raw == nullptr) { throw std::runtime_error( @@ -1389,44 +1415,46 @@ std::vector HierarchicalNSW::add( link_lists.emplace_back(raw); } - const size_t workers = std::min(resolve_num_threads(num_threads), n); - std::vector queries(workers); + const size_t num_workers = std::min(resolve_num_threads(num_threads), num_points); + std::vector queries(num_workers); // Prepare every code before publishing the slots. If quantization or worker // creation fails, the graph and point count are unchanged. { - std::unique_lock lock_table(label_lookup_lock_); - if (cur_element_count_ + n > max_elements_) { + std::unique_lock label_lookup_lock(label_lookup_lock_); + if (num_stored_points_ + num_points > max_elements_) { throw std::runtime_error("The number of elements exceeds the specified limit"); } - const auto first = static_cast(cur_element_count_.load()); - rabitqlib::index_detail::parallel_for(0, n, num_threads, [&](size_t i, size_t) { - const auto id = static_cast(first + i); - labels[i] = id; - std::memset( - data_level0_memory_.get() + (id * size_data_per_element_), - 0, - size_data_per_element_ - ); - std::memcpy(get_external_label_pt(id), &id, sizeof(PID)); - std::memcpy(get_clusterid_pt(id), &assigned[i], sizeof(PID)); - std::vector rotated_data(padded_dim_); - rotator_->rotate(data + (i * dim_), rotated_data.data()); - quant::quantize_split_single( - rotated_data.data(), - reinterpret_cast(centroids_memory_.get()) + - (assigned[i] * padded_dim_), - padded_dim_, - ex_bits_, - get_bindata_by_internalid(id), - get_exdata_by_internalid(id), - metric_type_, - config - ); - }); + const auto first = static_cast(num_stored_points_.load()); + rabitqlib::index_detail::parallel_for( + 0, num_points, num_threads, [&](size_t i, size_t) { + const auto id = static_cast(first + i); + labels[i] = id; + std::memset( + data_level0_memory_.get() + (id * level0_record_bytes_), + 0, + level0_record_bytes_ + ); + std::memcpy(get_external_label_pt(id), &id, sizeof(PID)); + std::memcpy(get_clusterid_pt(id), &assigned[i], sizeof(PID)); + std::vector rotated_data(padded_dim_); + rotator_->rotate(data + (i * dim_), rotated_data.data()); + quant::quantize_split_single( + rotated_data.data(), + reinterpret_cast(centroids_memory_.get()) + + (assigned[i] * padded_dim_), + padded_dim_, + ex_bits_, + get_bindata_by_internalid(id), + get_exdata_by_internalid(id), + metric_type_, + config + ); + } + ); size_t registered = 0; try { - for (; registered < n; ++registered) { + for (; registered < num_points; ++registered) { label_lookup_[labels[registered]] = labels[registered]; } } catch (...) { @@ -1437,16 +1465,18 @@ std::vector HierarchicalNSW::add( } // Unstarted inserts are valid level-0 nodes with owned codes. Their // upper-layer lists remain under RAII until an insert takes ownership. - link_lists_.resize(cur_element_count_ + n); - for (size_t i = 0; i < n; ++i) { + link_lists_.resize(num_stored_points_ + num_points); + for (size_t i = 0; i < num_points; ++i) { element_levels_[labels[i]] = 0; } - cur_element_count_ += n; + num_stored_points_ += num_points; } - rabitqlib::index_detail::parallel_for(0, n, workers, [&](size_t i, size_t worker) { - add_point_quant(labels[i], levels[i], link_lists[i], queries[worker]); - }); + rabitqlib::index_detail::parallel_for( + 0, num_points, num_workers, [&](size_t i, size_t worker) { + add_point_quant(labels[i], levels[i], link_lists[i], queries[worker]); + } + ); return labels; } diff --git a/src/index/hnsw_dispatch.cpp b/src/index/hnsw_dispatch.cpp index f24f8e5a..9f85a6ae 100644 --- a/src/index/hnsw_dispatch.cpp +++ b/src/index/hnsw_dispatch.cpp @@ -16,7 +16,7 @@ using SearchKnnFn = // The core variant uses AVX2 warmup; the popcount variant has its own stricter tier. } // namespace std::vector> search_knn_sorted( - HierarchicalNSW& index, const float* query, size_t topk, size_t ef + HierarchicalNSW& index, const float* query, size_t k, size_t ef_search ) { #if defined(__x86_64__) || defined(_M_X64) static const SearchKnnFn kSearchKnnFn = @@ -33,12 +33,12 @@ std::vector> search_knn_sorted( ); #endif - return kSearchKnnFn(index, query, topk, ef); + return kSearchKnnFn(index, query, k, ef_search); } std::priority_queue> search_knn( - HierarchicalNSW& index, const float* query, size_t topk, size_t ef + HierarchicalNSW& index, const float* query, size_t k, size_t ef_search ) { - auto results = search_knn_sorted(index, query, topk, ef); + auto results = search_knn_sorted(index, query, k, ef_search); return {results.begin(), results.end()}; } } // namespace rabitqlib::hnsw::detail diff --git a/src/index/hnsw_search_avx2.cpp b/src/index/hnsw_search_avx2.cpp index 360cfd99..820982af 100644 --- a/src/index/hnsw_search_avx2.cpp +++ b/src/index/hnsw_search_avx2.cpp @@ -30,10 +30,10 @@ struct HnswAvx2Kernel { }; std::vector> search_knn_sorted_avx2( - HierarchicalNSW& index, const float* rotated_query, size_t topk, size_t ef + HierarchicalNSW& index, const float* rotated_query, size_t k, size_t ef_search ) { return index.search_knn_sorted_direct( - rotated_query, topk, ef == 0 ? index.ef_ : ef + rotated_query, k, ef_search == 0 ? index.default_ef_search_ : ef_search ); } diff --git a/src/index/hnsw_search_avx512_core.cpp b/src/index/hnsw_search_avx512_core.cpp index 5a80e902..c6c80002 100644 --- a/src/index/hnsw_search_avx512_core.cpp +++ b/src/index/hnsw_search_avx512_core.cpp @@ -31,10 +31,10 @@ struct HnswAvx512CoreKernel { }; std::vector> search_knn_sorted_avx512_core( - HierarchicalNSW& index, const float* rotated_query, size_t topk, size_t ef + HierarchicalNSW& index, const float* rotated_query, size_t k, size_t ef_search ) { return index.search_knn_sorted_direct( - rotated_query, topk, ef == 0 ? index.ef_ : ef + rotated_query, k, ef_search == 0 ? index.default_ef_search_ : ef_search ); } diff --git a/src/index/hnsw_search_avx512_popcnt.cpp b/src/index/hnsw_search_avx512_popcnt.cpp index db4c6ddd..0602e087 100644 --- a/src/index/hnsw_search_avx512_popcnt.cpp +++ b/src/index/hnsw_search_avx512_popcnt.cpp @@ -30,10 +30,10 @@ struct HnswAvx512PopcntKernel { }; std::vector> search_knn_sorted_avx512_popcnt( - HierarchicalNSW& index, const float* rotated_query, size_t topk, size_t ef + HierarchicalNSW& index, const float* rotated_query, size_t k, size_t ef_search ) { return index.search_knn_sorted_direct( - rotated_query, topk, ef == 0 ? index.ef_ : ef + rotated_query, k, ef_search == 0 ? index.default_ef_search_ : ef_search ); } diff --git a/src/index/hnsw_search_generic.cpp b/src/index/hnsw_search_generic.cpp index fa7abc84..77b7fba1 100644 --- a/src/index/hnsw_search_generic.cpp +++ b/src/index/hnsw_search_generic.cpp @@ -33,10 +33,10 @@ struct HnswGenericKernel { }; std::vector> search_knn_sorted_generic( - HierarchicalNSW& index, const float* rotated_query, size_t topk, size_t ef + HierarchicalNSW& index, const float* rotated_query, size_t k, size_t ef_search ) { return index.search_knn_sorted_direct( - rotated_query, topk, ef == 0 ? index.ef_ : ef + rotated_query, k, ef_search == 0 ? index.default_ef_search_ : ef_search ); } diff --git a/src/index/hnsw_search_impl.hpp b/src/index/hnsw_search_impl.hpp index 17f1e5b9..48f85f34 100644 --- a/src/index/hnsw_search_impl.hpp +++ b/src/index/hnsw_search_impl.hpp @@ -4,6 +4,7 @@ #include #include #include +#include #include #include @@ -21,78 +22,82 @@ namespace rabitqlib::hnsw { template inline void HierarchicalNSW::get_bin_est_direct( - std::vector& q_to_centroids, - SplitSingleQuery& query_wrapper, + const std::vector& q_to_centroids, + const SplitSingleQuery& query_wrapper, PID candidate_id, - HierarchicalNSW::EstimateRecord& res -) { + HierarchicalNSW::EstimateRecord& estimate +) const { if (metric_type_ == METRIC_IP) { - float norm = q_to_centroids[get_clusterid_by_internalid(candidate_id)]; - float error = + const float query_centroid_ip = + q_to_centroids[get_clusterid_by_internalid(candidate_id)]; + const float residual_norm = q_to_centroids[get_clusterid_by_internalid(candidate_id) + num_clusters_]; - res = estimate_split_single_direct( + estimate = estimate_split_single_direct( get_bindata_by_internalid(candidate_id), query_wrapper, padded_dim_, - -norm, - error + -query_centroid_ip, + residual_norm ); } else { // L2 distance - float norm = q_to_centroids[get_clusterid_by_internalid(candidate_id)]; - res = estimate_split_single_direct( + const float residual_norm = + q_to_centroids[get_clusterid_by_internalid(candidate_id)]; + estimate = estimate_split_single_direct( get_bindata_by_internalid(candidate_id), query_wrapper, padded_dim_, - norm * norm, - norm + residual_norm * residual_norm, + residual_norm ); } } template inline void HierarchicalNSW::get_full_est_direct( - std::vector& q_to_centroids, - SplitSingleQuery& query_wrapper, + const std::vector& q_to_centroids, + const SplitSingleQuery& query_wrapper, PID candidate_id, - HierarchicalNSW::EstimateRecord& res + HierarchicalNSW::EstimateRecord& estimate ) const { if (metric_type_ == METRIC_IP) { - float norm = q_to_centroids[get_clusterid_by_internalid(candidate_id)]; - float error = + const float query_centroid_ip = + q_to_centroids[get_clusterid_by_internalid(candidate_id)]; + const float residual_norm = q_to_centroids[get_clusterid_by_internalid(candidate_id) + num_clusters_]; - res = estimate_split_single_full_direct( + estimate = estimate_split_single_full_direct( get_bindata_by_internalid(candidate_id), get_exdata_by_internalid(candidate_id), ip_func_, query_wrapper, padded_dim_, ex_bits_, - -norm, - error + -query_centroid_ip, + residual_norm ); } else { // L2 distance - float norm = q_to_centroids[get_clusterid_by_internalid(candidate_id)]; - res = estimate_split_single_full_direct( + const float residual_norm = + q_to_centroids[get_clusterid_by_internalid(candidate_id)]; + estimate = estimate_split_single_full_direct( get_bindata_by_internalid(candidate_id), get_exdata_by_internalid(candidate_id), ip_func_, query_wrapper, padded_dim_, ex_bits_, - norm * norm, - norm + residual_norm * residual_norm, + residual_norm ); } } template inline std::vector> HierarchicalNSW::search_knn_sorted_direct( - const float* rotated_query, size_t k, size_t ef -) { + const float* rotated_query, size_t k, size_t ef_search +) const { std::vector> result; - if (cur_element_count_ == 0) { + if (num_stored_points_ == 0) { return result; } @@ -107,53 +112,57 @@ inline std::vector> HierarchicalNSW::search_knn_sorted_dir for (size_t i = 0; i < num_clusters_; i++) { q_to_centroids[i] = std::sqrt(raw_dist_func_( rotated_query, - reinterpret_cast(centroids_memory_.get()) + (i * padded_dim_), + reinterpret_cast(centroids_memory_.get()) + (i * padded_dim_), padded_dim_ )); } } else if (metric_type_ == METRIC_IP) { q_to_centroids.resize(2 * num_clusters_); - // first half as g_add, second half as g_error + // First half stores dot(query, centroid); second half stores ||query - centroid||. for (size_t i = 0; i < num_clusters_; i++) { q_to_centroids[i] = dot_product( rotated_query, - reinterpret_cast(centroids_memory_.get()) + (i * padded_dim_), + reinterpret_cast(centroids_memory_.get()) + (i * padded_dim_), padded_dim_ ); q_to_centroids[i + num_clusters_] = std::sqrt(euclidean_sqr( rotated_query, - reinterpret_cast(centroids_memory_.get()) + (i * padded_dim_), + reinterpret_cast(centroids_memory_.get()) + (i * padded_dim_), padded_dim_ )); } } PID current_id = entry_point_id_; - EstimateRecord curest; + EstimateRecord current_estimate; - get_bin_est_direct(q_to_centroids, query_wrapper, current_id, curest); + get_bin_est_direct(q_to_centroids, query_wrapper, current_id, current_estimate); for (int level = max_level_; level > 0; level--) { bool changed = true; while (changed) { changed = false; - unsigned int* data; - - data = static_cast(get_linklist(current_id, level)); - int size = get_list_count(data); + const PID* data = get_linklist(current_id, level); + const size_t size = get_list_count(data); - PID* datal = static_cast(data + 1); - for (int i = 0; i < size; i++) { + const PID* datal = data + 1; + for (size_t i = 0; i < size; i++) { PID cand = datal[i]; if (cand > max_elements_) { - throw std::runtime_error("cand error"); + throw std::runtime_error( + "HNSW search upper-layer traversal: invalid candidate ID " + + std::to_string(cand) + " at level " + std::to_string(level) + + " (capacity=" + std::to_string(max_elements_) + ")" + ); } - EstimateRecord candest; - get_bin_est_direct(q_to_centroids, query_wrapper, cand, candest); + EstimateRecord candidate_estimate; + get_bin_est_direct( + q_to_centroids, query_wrapper, cand, candidate_estimate + ); - if (candest.distance < curest.distance) { - curest = candest; + if (candidate_estimate.distance < current_estimate.distance) { + current_estimate = candidate_estimate; current_id = cand; changed = true; } @@ -164,7 +173,7 @@ inline std::vector> HierarchicalNSW::search_knn_sorted_dir BoundedKNN nearest_candidates(k); search_base_layer_direct( current_id, - std::max(ef, k), + std::max(ef_search, k), k, query_wrapper, q_to_centroids, @@ -173,7 +182,7 @@ inline std::vector> HierarchicalNSW::search_knn_sorted_dir ); result.reserve(nearest_candidates.size()); for (const auto& candidate : nearest_candidates.candidates()) { - result.emplace_back(candidate.record.est_dist, get_external_label(candidate.id)); + result.emplace_back(candidate.est_dist, get_external_label(candidate.id)); } // Candidate selection uses distance only. Preserve that selection, then match // the former heap's (distance, external label) ordering within each tie group. @@ -190,51 +199,50 @@ inline std::vector> HierarchicalNSW::search_knn_sorted_dir template inline void HierarchicalNSW::search_base_layer_direct( - PID ep_id, - size_t ef, + PID entry_point_id, + size_t ef_search, size_t k, - SplitSingleQuery& query_wrapper, - std::vector& q_to_centroids, + const SplitSingleQuery& query_wrapper, + const std::vector& q_to_centroids, [[maybe_unused]] const float* query, BoundedKNN& nearest_candidates -) { - auto vl = visited_list_pool_->get_free_vislist(); +) const { + auto visited = visited_list_pool_->get_free_vislist(); // Use our bounded priority queue instead of the maxheap. - buffer::SearchBuffer candidate_set(ef); + buffer::SearchBuffer search_pool(ef_search); - float distk = 1e10; + float dist_k = 1e10; EstimateRecord start_estimate_record; if (ex_bits_ > 0) { get_full_est_direct( - q_to_centroids, query_wrapper, ep_id, start_estimate_record + q_to_centroids, query_wrapper, entry_point_id, start_estimate_record ); } else { get_bin_est_direct( - q_to_centroids, query_wrapper, ep_id, start_estimate_record + q_to_centroids, query_wrapper, entry_point_id, start_estimate_record ); } float est_dist = start_estimate_record.distance; - float low_dist = start_estimate_record.lower_bound; // Insert initial candidate. A removed point is still explored, never returned. - if (!is_removed(ep_id)) { - nearest_candidates.insert({ResultRecord(est_dist, low_dist), ep_id}); - distk = est_dist; + if (!is_removed(entry_point_id)) { + nearest_candidates.insert({est_dist, entry_point_id}); + dist_k = est_dist; } - candidate_set.insert(ep_id, est_dist); + search_pool.insert(entry_point_id, est_dist); - vl->set(ep_id); + visited->set(entry_point_id); const size_t prefetch_size = (((padded_dim_ / 8) + 63) / 64) + 1; const size_t prefetch_lookahead = 4; // Number of neighbors to prefetch in advance. - while (candidate_set.has_next()) { + while (search_pool.has_next()) { // Step 1 - get the next node to explore. - PID current_node_id = candidate_set.pop(); - int* data = (int*)get_linklist0(current_node_id); - size_t size = get_list_count((PID*)data); + PID current_node_id = search_pool.pop(); + const PID* data = get_linklist0(current_node_id); + const size_t size = get_list_count(data); for (size_t p = 0; p < prefetch_lookahead; ++p) { rabitqlib::memory::mem_prefetch_l1( @@ -243,7 +251,7 @@ inline void HierarchicalNSW::search_base_layer_direct( } // Iterate over neighbors. (List starts at index 1.) for (size_t j = 1; j <= size; j++) { - int candidate_id = *(data + j); + const PID candidate_id = *(data + j); if (j + prefetch_lookahead <= size) { rabitqlib::memory::mem_prefetch_l1( @@ -252,46 +260,44 @@ inline void HierarchicalNSW::search_base_layer_direct( ); } - if (vl->get(candidate_id)) { + if (visited->get(candidate_id)) { continue; } - vl->set(candidate_id); + visited->set(candidate_id); - EstimateRecord candest; + EstimateRecord candidate_estimate; get_bin_est_direct( - q_to_centroids, query_wrapper, candidate_id, candest + q_to_centroids, query_wrapper, candidate_id, candidate_estimate ); - // A removed neighbor keeps its binary estimate and stays in candidate_set. + // A removed neighbor keeps its binary estimate and stays in search_pool. bool update_results = - !is_removed(static_cast(candidate_id)) && - (nearest_candidates.size() < k || candest.lower_bound < distk); + !is_removed(candidate_id) && + (nearest_candidates.size() < k || candidate_estimate.lower_bound < dist_k); if (update_results) { // Compute the full estimate if promising. if (ex_bits_ > 0) { get_full_est_direct( - q_to_centroids, query_wrapper, candidate_id, candest + q_to_centroids, query_wrapper, candidate_id, candidate_estimate ); } - Candidate cand{ - ResultRecord(candest.distance, candest.lower_bound), - static_cast(candidate_id)}; + Candidate cand{candidate_estimate.distance, candidate_id}; nearest_candidates.insert(cand); - distk = nearest_candidates.worst().record.est_dist; + dist_k = nearest_candidates.worst().est_dist; } - if (!candidate_set.is_full(candest.distance)) { - candidate_set.insert(candidate_id, candest.distance); + if (!search_pool.is_full(candidate_estimate.distance)) { + search_pool.insert(candidate_id, candidate_estimate.distance); } rabitqlib::memory::mem_prefetch_l2( - (char*)get_linklist0(candidate_set.next_id()), 2 + reinterpret_cast(get_linklist0(search_pool.next_id())), 2 ); } } - visited_list_pool_->release_vis_list(std::move(vl)); + visited_list_pool_->release_vis_list(std::move(visited)); } } // namespace rabitqlib::hnsw diff --git a/src/index/hnsw_search_neon.cpp b/src/index/hnsw_search_neon.cpp index c0c8f782..a07e9506 100644 --- a/src/index/hnsw_search_neon.cpp +++ b/src/index/hnsw_search_neon.cpp @@ -34,10 +34,10 @@ struct HnswNeonKernel { }; std::vector> search_knn_sorted_neon( - HierarchicalNSW& index, const float* rotated_query, size_t topk, size_t ef + HierarchicalNSW& index, const float* rotated_query, size_t k, size_t ef_search ) { return index.search_knn_sorted_direct( - rotated_query, topk, ef == 0 ? index.ef_ : ef + rotated_query, k, ef_search == 0 ? index.default_ef_search_ : ef_search ); } diff --git a/src/index/ivf.cpp b/src/index/ivf.cpp index fa1a3714..7b920533 100644 --- a/src/index/ivf.cpp +++ b/src/index/ivf.cpp @@ -45,12 +45,14 @@ namespace rabitqlib::ivf { IVF::IVF() noexcept = default; IVF::~IVF() = default; -bool IVF::exceeds_storage(size_t n) const { +bool IVF::exceeds_storage(size_t num_points) const { const size_t max_size = std::numeric_limits::max(); const size_t batch_bytes = BatchDataMap::data_bytes(padded_dim_); const size_t rerank_bytes = rerank_vector_bytes(); - return n > max_size / sizeof(PID) || n > max_size / sizeof(float) / padded_dim_ || - n > max_size / batch_bytes || (rerank_bytes != 0 && n > max_size / rerank_bytes); + return num_points > max_size / sizeof(PID) || + num_points > max_size / sizeof(float) / padded_dim_ || + num_points > max_size / batch_bytes || + (rerank_bytes != 0 && num_points > max_size / rerank_bytes); } size_t IVF::batch_data_bytes(size_t cluster_size) const { @@ -69,39 +71,40 @@ inline size_t IVF::ex_data_bytes(size_t cluster_size) const { return rerank_vector_bytes() * cluster_size; } -void IVF::resize_storage(ClusterStorage& cs, size_t n) const { - cs.batch_storage.resize(batch_data_bytes(n)); +void IVF::resize_storage(ClusterStorage& cs, size_t num_points) const { + cs.batch_storage.resize(batch_data_bytes(num_points)); if (rerank_vector_bytes() > 0) { if (raw_reranking_) { - cs.raw_storage.resize(n * dim_); + cs.raw_storage.resize(num_points * dim_); } else { - cs.ex_storage.resize(ex_data_bytes(n)); + cs.ex_storage.resize(ex_data_bytes(num_points)); } } - cs.id_storage.resize(n); + cs.id_storage.resize(num_points); } -void IVF::reserve_storage(ClusterStorage& cs, size_t n) const { - cs.batch_storage.reserve(batch_data_bytes(n)); +void IVF::reserve_storage(ClusterStorage& cs, size_t num_points) const { + cs.batch_storage.reserve(batch_data_bytes(num_points)); if (rerank_vector_bytes() > 0) { if (raw_reranking_) { - cs.raw_storage.reserve(n * dim_); + cs.raw_storage.reserve(num_points * dim_); } else { - cs.ex_storage.reserve(ex_data_bytes(n)); + cs.ex_storage.reserve(ex_data_bytes(num_points)); } } - cs.id_storage.reserve(n); + cs.id_storage.reserve(num_points); } -bool IVF::has_capacity(const ClusterStorage& cs, size_t n) const { - if (cs.batch_storage.capacity() < batch_data_bytes(n) || cs.id_storage.capacity() < n) { +bool IVF::has_capacity(const ClusterStorage& cs, size_t num_points) const { + if (cs.batch_storage.capacity() < batch_data_bytes(num_points) || + cs.id_storage.capacity() < num_points) { return false; } if (rerank_vector_bytes() == 0) { return true; } - return raw_reranking_ ? cs.raw_storage.capacity() >= n * dim_ - : cs.ex_storage.capacity() >= ex_data_bytes(n); + return raw_reranking_ ? cs.raw_storage.capacity() >= num_points * dim_ + : cs.ex_storage.capacity() >= ex_data_bytes(num_points); } void IVF::free_memory() { @@ -185,7 +188,7 @@ void IVF::construct( const float* data, const float* centroids, const PID* cluster_ids, - bool faster, + bool fast_quantization, size_t num_threads ) { if (max_elements_ == 0 || dim_ == 0 || num_clusters_ == 0 || rotator_ == nullptr) { @@ -199,12 +202,12 @@ void IVF::construct( std::vector counts(num_clusters_, 0); std::vector> id_lists(num_clusters_); for (size_t i = 0; i < max_elements_; ++i) { - PID cid = cluster_ids[i]; - if (cid >= num_clusters_) { + PID cluster_id = cluster_ids[i]; + if (cluster_id >= num_clusters_) { throw std::invalid_argument("Cluster ID is out of range"); } - id_lists[cid].push_back(static_cast(i)); - counts[cid] += 1; + id_lists[cluster_id].push_back(static_cast(i)); + counts[cluster_id] += 1; } allocate_memory(counts); @@ -216,7 +219,7 @@ void IVF::construct( std::vector rotated_centroids(num_clusters_ * padded_dim_); quant::RabitqConfig config; - if (faster) { + if (fast_quantization) { config = quant::faster_config(padded_dim_, ex_bits_ + 1); } @@ -319,8 +322,8 @@ void IVF::quantize_cluster( char* batch_data = cp.batch_data(); char* ex_data = cp.ex_data(); for (size_t i = 0; i < num_points; i += fastscan::kBatchSize) { - size_t n = std::min(fastscan::kBatchSize, num_points - i); - for (size_t j = 0; j < n; ++j) { + size_t batch_size = std::min(fastscan::kBatchSize, num_points - i); + for (size_t j = 0; j < batch_size; ++j) { const float* vector = data + (point_ids[i + j] * dim_); rotator_->rotate(vector, rotated_data + (j * padded_dim_)); if (raw_reranking_) { @@ -335,7 +338,7 @@ void IVF::quantize_cluster( quant::quantize_split_batch( rotated_data, rotated_centroid, - n, + batch_size, padded_dim_, ex_bits_, batch_data, @@ -346,36 +349,40 @@ void IVF::quantize_cluster( batch_data += BatchDataMap::data_bytes(padded_dim_); if (ex_bits_ > 0) { - ex_data += ExDataMap::data_bytes(padded_dim_, ex_bits_) * n; + ex_data += ExDataMap::data_bytes(padded_dim_, ex_bits_) * batch_size; } } } void IVF::add( - const float* data, size_t n, const PID* cluster_ids, bool faster, size_t num_threads + const float* data, + size_t num_points, + const PID* cluster_ids, + bool fast_quantization, + size_t num_threads ) { if (!ready_ || initializer_ == nullptr || rotator_ == nullptr || clusters_.size() != num_clusters_ || cluster_storage_.size() != num_clusters_) { throw std::logic_error("IVF index must be constructed or loaded before add"); } - if (n == 0) { + if (num_points == 0) { return; } if (data == nullptr) { throw std::invalid_argument("IVF add data must not be null"); } - if (n > buffer::kSearchBufferMaxPointCount - max_elements_) { + if (num_points > buffer::kSearchBufferMaxPointCount - max_elements_) { throw std::invalid_argument("IVF point count exceeds the supported ID range"); } - const size_t new_num_points = max_elements_ + n; + const size_t new_num_points = max_elements_ + num_points; if (exceeds_storage(new_num_points)) { throw std::invalid_argument("IVF configuration exceeds addressable storage"); } /* Assign every new point to a cluster */ - std::vector assigned(n); + std::vector assigned(num_points); if (cluster_ids != nullptr) { - for (size_t i = 0; i < n; ++i) { + for (size_t i = 0; i < num_points; ++i) { if (cluster_ids[i] >= num_clusters_) { throw std::invalid_argument("Cluster ID is out of range"); } @@ -383,7 +390,7 @@ void IVF::add( } } else { rabitqlib::index_detail::parallel_for( - 0, n, num_threads, [&](size_t row, size_t /*thread_id*/) { + 0, num_points, num_threads, [&](size_t row, size_t /*thread_id*/) { std::vector rotated(padded_dim_); std::vector> nearest(1); rotator_->rotate(data + (row * dim_), rotated.data()); @@ -405,7 +412,7 @@ void IVF::add( // Group only incoming rows, preserving their order within each cluster. // A small append must not allocate or scan num_clusters_ entries. std::unordered_map positions; - for (size_t i = 0; i < n; ++i) { + for (size_t i = 0; i < num_points; ++i) { const auto [position, inserted] = positions.try_emplace(assigned[i], additions.size()); if (inserted) { @@ -417,10 +424,10 @@ void IVF::add( } } - const size_t workers = std::min(resolve_num_threads(num_threads), additions.size()); - static_cast(workers); // Consumed by the OpenMP clause below. + const size_t num_workers = std::min(resolve_num_threads(num_threads), additions.size()); + static_cast(num_workers); // Consumed by the OpenMP clause below. quant::RabitqConfig config; - if (faster) { + if (fast_quantization) { config = quant::faster_config(padded_dim_, ex_bits_ + 1); } @@ -450,7 +457,7 @@ void IVF::add( // An exception must not leave an OpenMP region, so keep the first one and rethrow it. index_detail::ParallelException failure; -#pragma omp parallel for schedule(dynamic) num_threads(workers) +#pragma omp parallel for schedule(dynamic) num_threads(num_workers) for (std::ptrdiff_t index = 0; index < static_cast(additions.size()); ++index) { auto& addition = additions[static_cast(index)]; @@ -654,12 +661,12 @@ void IVF::reserve(PID cluster_id, size_t num_points) { init_clusters(cluster_sizes); } -size_t IVF::remove(const PID* ids_to_remove, size_t n) { +size_t IVF::remove(const PID* ids_to_remove, size_t num_points) { if (!ready_ || initializer_ == nullptr || rotator_ == nullptr || clusters_.size() != num_clusters_ || cluster_storage_.size() != num_clusters_) { throw std::logic_error("IVF index must be constructed or loaded before remove"); } - if (n == 0) { + if (num_points == 0) { return 0; } if (ids_to_remove == nullptr) { @@ -668,7 +675,7 @@ size_t IVF::remove(const PID* ids_to_remove, size_t n) { // Validate every ID before changing anything std::vector marked(max_elements_, false); - for (size_t i = 0; i < n; ++i) { + for (size_t i = 0; i < num_points; ++i) { if (ids_to_remove[i] >= max_elements_) { throw std::invalid_argument("IVF remove ID is out of range"); } @@ -737,8 +744,8 @@ void IVF::save(const char* filename) const { /* Save number of vectors of each cluster */ std::vector cluster_sizes; cluster_sizes.reserve(num_clusters_); - for (const auto& cur_cluster : clusters_) { - cluster_sizes.push_back(cur_cluster.num()); + for (const auto& cluster : clusters_) { + cluster_sizes.push_back(cluster.num()); } output.write( reinterpret_cast(cluster_sizes.data()), @@ -970,75 +977,77 @@ void IVF::load(const char* filename) { } void IVF::search( - const float* query, size_t k, size_t nprobe, PID* results, float* distances + const float* query, size_t k, size_t nprobe, PID* result_ids, float* distances ) const { - search(query, k, nprobe, results, distances, !raw_reranking_ && ex_bits_ > 2); + search(query, k, nprobe, result_ids, distances, !raw_reranking_ && ex_bits_ > 2); } void IVF::search( const float* __restrict__ query, size_t k, size_t nprobe, - PID* __restrict__ results, + PID* __restrict__ result_ids, bool use_hacc ) const { - this->search(query, k, nprobe, results, nullptr, use_hacc); + this->search(query, k, nprobe, result_ids, nullptr, use_hacc); } void IVF::search( const float* __restrict__ query, size_t k, size_t nprobe, - PID* __restrict__ results, + PID* __restrict__ result_ids, float* __restrict__ distances, bool use_hacc ) const { - validate_search(query, k, nprobe, results); + validate_search(query, k, nprobe, result_ids); nprobe = std::min(nprobe, num_clusters_); // corner case std::vector rotated_query(padded_dim_); this->rotator_->rotate(query, rotated_query.data()); // use initer to get closest nprobe centroids - std::vector> centroid_dist(nprobe); - this->initializer_->centroids_distances(rotated_query.data(), nprobe, centroid_dist); + std::vector> centroid_candidates(nprobe); + this->initializer_->centroids_distances( + rotated_query.data(), nprobe, centroid_candidates + ); - buffer::SearchBuffer knns(k); + buffer::SearchBuffer result_pool(k); - SplitBatchQuery q_obj( + SplitBatchQuery query_wrapper( rotated_query.data(), padded_dim_, ex_bits_, metric_type_, use_hacc ); for (size_t i = 0; i < nprobe; ++i) { - PID cid = centroid_dist[i].id; - float dist = centroid_dist[i].distance; - const Cluster& cur_cluster = clusters_[cid]; + PID cluster_id = centroid_candidates[i].id; + float dist = centroid_candidates[i].distance; + const Cluster& cluster = clusters_[cluster_id]; if (metric_type_ == METRIC_L2) { - q_obj.set_g_add(dist); + query_wrapper.set_g_add(dist); } else if (metric_type_ == METRIC_IP) { const float residual_norm = std::sqrt(euclidean_sqr( - rotated_query.data(), initializer_->centroid(cid), padded_dim_ + rotated_query.data(), initializer_->centroid(cluster_id), padded_dim_ )); - auto g_add_ip = dot_product( - rotated_query.data(), initializer_->centroid(cid), padded_dim_ + auto query_centroid_ip = dot_product( + rotated_query.data(), initializer_->centroid(cluster_id), padded_dim_ ); - q_obj.set_g_add(residual_norm, g_add_ip); + query_wrapper.set_g_add(residual_norm, query_centroid_ip); } else { // unsupported throw std::invalid_argument("Quantization only supports L2 and IP metrics"); } - search_cluster(cur_cluster, q_obj, knns, use_hacc, query); + search_cluster(cluster, query_wrapper, result_pool, use_hacc, query); } - const size_t found = knns.size(); + const size_t found = result_pool.size(); if (distances != nullptr) { - knns.copy_results(results, distances); + result_pool.copy_results(result_ids, distances); std::fill(distances + found, distances + k, std::numeric_limits::infinity()); } else { - knns.copy_results(results); + result_pool.copy_results(result_ids); } - std::fill(results + found, results + k, kPidMax); + std::fill(result_ids + found, result_ids + k, kPidMax); } void IVF::search_batch( @@ -1046,7 +1055,7 @@ void IVF::search_batch( size_t num_queries, size_t k, size_t nprobe, - PID* results, + PID* result_ids, float* distances, std::optional use_hacc, size_t num_threads @@ -1054,23 +1063,24 @@ void IVF::search_batch( if (num_queries == 0) { return; } - validate_search(queries, k, nprobe, results); + validate_search(queries, k, nprobe, result_ids); const size_t max_values = std::numeric_limits::max() / sizeof(float); if (num_queries > max_values / dim_ || num_queries > max_values / k) { throw std::length_error("IVF query batch is too large"); } nprobe = std::min(nprobe, num_clusters_); const bool hacc = use_hacc.value_or(!raw_reranking_ && ex_bits_ > 2); - const size_t workers = num_queries == 1 || num_threads == 1 - ? 1 - : std::min(resolve_num_threads(num_threads), num_queries); - if (workers == 1) { + const size_t num_workers = + num_queries == 1 || num_threads == 1 + ? 1 + : std::min(resolve_num_threads(num_threads), num_queries); + if (num_workers == 1) { for (size_t i = 0; i < num_queries; ++i) { search( queries + i * dim_, k, nprobe, - results + i * k, + result_ids + i * k, distances == nullptr ? nullptr : distances + i * k, hacc ); @@ -1078,7 +1088,7 @@ void IVF::search_batch( return; } index_detail::ParallelException failure; -#pragma omp parallel for num_threads(workers) schedule(dynamic) +#pragma omp parallel for num_threads(num_workers) schedule(dynamic) for (std::ptrdiff_t query_index = 0; query_index < static_cast(num_queries); ++query_index) { @@ -1091,7 +1101,7 @@ void IVF::search_batch( queries + i * dim_, k, nprobe, - results + i * k, + result_ids + i * k, distances == nullptr ? nullptr : distances + i * k, hacc ); @@ -1101,7 +1111,7 @@ void IVF::search_batch( } inline void IVF::validate_search( - const float* query, size_t k, size_t nprobe, const PID* results + const float* query, size_t k, size_t nprobe, const PID* result_ids ) const { if (!ready_ || initializer_ == nullptr || rotator_ == nullptr || clusters_.size() != num_clusters_ || cluster_storage_.size() != num_clusters_) { @@ -1110,7 +1120,7 @@ inline void IVF::validate_search( if (query == nullptr) { throw std::invalid_argument("IVF search query must not be null"); } - if (results == nullptr) { + if (result_ids == nullptr) { throw std::invalid_argument("IVF search results must not be null"); } if (k == 0 || k > max_elements_ || k > buffer::kSearchBufferMaxPointCount) { @@ -1122,23 +1132,30 @@ inline void IVF::validate_search( } inline void IVF::search_cluster( - const Cluster& cur_cluster, - const SplitBatchQuery& q_obj, - buffer::SearchBuffer& knns, + const Cluster& cluster, + const SplitBatchQuery& query_wrapper, + buffer::SearchBuffer& result_pool, bool use_hacc, const float* query ) const { - size_t iter = cur_cluster.num() / fastscan::kBatchSize; - size_t remain = cur_cluster.num() - (iter * fastscan::kBatchSize); + size_t iter = cluster.num() / fastscan::kBatchSize; + size_t remain = cluster.num() - (iter * fastscan::kBatchSize); - const char* batch_data = cur_cluster.batch_data(); - const char* ex_data = cur_cluster.ex_data(); - const PID* ids = cur_cluster.ids(); + const char* batch_data = cluster.batch_data(); + const char* ex_data = cluster.ex_data(); + const PID* ids = cluster.ids(); /* Compute distances block by block */ for (size_t i = 0; i < iter; ++i) { scan_one_batch( - batch_data, ex_data, ids, q_obj, knns, fastscan::kBatchSize, use_hacc, query + batch_data, + ex_data, + ids, + query_wrapper, + result_pool, + fastscan::kBatchSize, + use_hacc, + query ); batch_data += BatchDataMap::data_bytes(padded_dim_); @@ -1150,7 +1167,9 @@ inline void IVF::search_cluster( if (remain > 0) { // scan the last block - scan_one_batch(batch_data, ex_data, ids, q_obj, knns, remain, use_hacc, query); + scan_one_batch( + batch_data, ex_data, ids, query_wrapper, result_pool, remain, use_hacc, query + ); } } @@ -1158,52 +1177,53 @@ inline void IVF::scan_one_batch( const char* batch_data, const char* ex_data, const PID* ids, - const SplitBatchQuery& q_obj, - buffer::SearchBuffer& knns, + const SplitBatchQuery& query_wrapper, + buffer::SearchBuffer& result_pool, size_t num_points, bool use_hacc, const float* query ) const { - std::array est_distance; - std::array low_distance; + std::array estimated_distances; + std::array lower_bounds; std::array ip_x0_qr; // inner product of the 1st bit split_batch_estdist( batch_data, - q_obj, + query_wrapper, padded_dim_, - est_distance.data(), - low_distance.data(), + estimated_distances.data(), + lower_bounds.data(), ip_x0_qr.data(), use_hacc ); - float distk = knns.top_dist(); + float dist_k = result_pool.top_dist(); // Without reranking data, return the one-bit estimates directly. if (ex_bits_ == 0 && !raw_reranking_) { - detail::insert_candidates(knns, ids, est_distance.data(), num_points); + detail::insert_candidates(result_pool, ids, estimated_distances.data(), num_points); return; } // Refine candidates passing the probabilistic one-bit filter. Quantized storage // reuses ip_x0_qr with the extra bits; raw storage scores the original vector. for (size_t i = 0; i < num_points; ++i) { - float lower_dist = low_distance[i]; - if (lower_dist < distk) { + float lower_bound = lower_bounds[i]; + if (lower_bound < dist_k) { PID id = ids[i]; - float ex_dist; + float refined_distance; if (raw_reranking_) { const auto* vector = reinterpret_cast(ex_data); - ex_dist = metric_type_ == METRIC_L2 ? euclidean_sqr(query, vector, dim_) - : dot_product_dis(query, vector, dim_); + refined_distance = metric_type_ == METRIC_L2 + ? euclidean_sqr(query, vector, dim_) + : dot_product_dis(query, vector, dim_); } else { - ex_dist = split_distance_boosting( - ex_data, ip_func_, q_obj, padded_dim_, ex_bits_, ip_x0_qr[i] + refined_distance = split_distance_boosting( + ex_data, ip_func_, query_wrapper, padded_dim_, ex_bits_, ip_x0_qr[i] ); } - knns.insert(id, ex_dist); - distk = knns.top_dist(); + result_pool.insert(id, refined_distance); + dist_k = result_pool.top_dist(); } ex_data += rerank_vector_bytes(); } diff --git a/src/index/ivf_initializer.cpp b/src/index/ivf_initializer.cpp index 95356fc3..9aa5e34b 100644 --- a/src/index/ivf_initializer.cpp +++ b/src/index/ivf_initializer.cpp @@ -23,21 +23,24 @@ #include "rabitqlib/utils/tools.hpp" namespace rabitqlib::ivf { -FlatRaBitQInitializer::FlatRaBitQInitializer(size_t d, size_t k, MetricType metric) - : Initializer(d, k), metric_type_(metric) { +FlatRaBitQInitializer::FlatRaBitQInitializer( + size_t dim, size_t num_clusters, MetricType metric +) + : Initializer(dim, num_clusters), metric_type_(metric) { validate_metric_type(metric); const size_t max_size = std::numeric_limits::max(); - if (d == 0 || d % 32 != 0 || d > max_size / 32 || k == 0 || - k > buffer::kSearchBufferMaxPointCount || k > max_size / sizeof(float) / d) { + if (dim == 0 || dim % 32 != 0 || dim > max_size / 32 || num_clusters == 0 || + num_clusters > buffer::kSearchBufferMaxPointCount || + num_clusters > max_size / sizeof(float) / dim) { throw std::invalid_argument("Invalid Flat RaBitQ initializer dimensions"); } - const size_t batches = div_round_up(k, fastscan::kBatchSize); - const size_t batch_bytes = BatchDataMap::data_bytes(d); + const size_t batches = div_round_up(num_clusters, fastscan::kBatchSize); + const size_t batch_bytes = BatchDataMap::data_bytes(dim); if (batches > max_size / batch_bytes) { throw std::invalid_argument("Flat RaBitQ initializer storage overflow"); } - centroids_.resize(k * d); - global_centroid_.resize(d); + centroids_.resize(num_clusters * dim); + global_centroid_.resize(dim); batch_data_.resize(batches * batch_bytes); } @@ -78,24 +81,24 @@ void FlatRaBitQInitializer::add_vectors(const float* centroids, size_t num_threa void FlatRaBitQInitializer::centroids_distances( const float* query, size_t nprobe, std::vector>& candidates ) const { - const size_t topk = std::min(nprobe, num_cluster_); - candidates.resize(topk); - if (topk == 0) { + const size_t nprobed = std::min(nprobe, num_cluster_); + candidates.resize(nprobed); + if (nprobed == 0) { return; } const auto exact_distance = [&](PID id) { return metric_type_ == METRIC_L2 ? euclidean_sqr(query, centroid(id), dim_) : dot_product_dis(query, centroid(id), dim_); }; - if (topk == num_cluster_) { - for (size_t i = 0; i < topk; ++i) { + if (nprobed == num_cluster_) { + for (size_t i = 0; i < nprobed; ++i) { candidates[i] = AnnCandidate(static_cast(i), exact_distance(i)); } std::sort(candidates.begin(), candidates.end()); } else { - SplitBatchQuery q_obj(query, dim_, 0, metric_type_, false); + SplitBatchQuery query_wrapper(query, dim_, 0, metric_type_, false); const float norm = std::sqrt(euclidean_sqr(query, global_centroid_.data(), dim_)); - q_obj.set_g_add( + query_wrapper.set_g_add( norm, metric_type_ == METRIC_IP ? dot_product(query, global_centroid_.data(), dim_) : 0.0F @@ -108,7 +111,7 @@ void FlatRaBitQInitializer::centroids_distances( split_batch_estdist( reinterpret_cast(batch_data_.data()) + (first / fastscan::kBatchSize) * batch_bytes, - q_obj, + query_wrapper, dim_, estimates.data(), lows.data(), @@ -126,21 +129,21 @@ void FlatRaBitQInitializer::centroids_distances( } } std::nth_element( - upper.begin(), upper.begin() + static_cast(topk), upper.end() + upper.begin(), upper.begin() + static_cast(nprobed), upper.end() ); - buffer::SearchBuffer knns(topk); + buffer::SearchBuffer result_pool(nprobed); // Seed with exact distances: an estimated upper bound need not hold for // every vector. Always fill the result before pruning other candidates. - for (size_t i = 0; i < topk; ++i) { - knns.insert(upper[i].id, exact_distance(upper[i].id)); + for (size_t i = 0; i < nprobed; ++i) { + result_pool.insert(upper[i].id, exact_distance(upper[i].id)); } - for (size_t i = topk; i < num_cluster_; ++i) { + for (size_t i = nprobed; i < num_cluster_; ++i) { const PID id = upper[i].id; - if (lower[id] <= knns.top_dist()) { - knns.insert(id, exact_distance(id)); + if (lower[id] <= result_pool.top_dist()) { + result_pool.insert(id, exact_distance(id)); } } - std::copy_n(knns.data().begin(), topk, candidates.begin()); + std::copy_n(result_pool.data().begin(), nprobed, candidates.begin()); } if (metric_type_ == METRIC_L2) { for (auto& candidate : candidates) { diff --git a/src/index/ivf_search.cpp b/src/index/ivf_search.cpp index eaed4877..1abd8d55 100644 --- a/src/index/ivf_search.cpp +++ b/src/index/ivf_search.cpp @@ -9,10 +9,13 @@ namespace rabitqlib::ivf::detail { // Keep the scalar top-k loop out of the large inlined IVF search body. This avoids // register spills in wheel builds without changing candidate order or tie handling. void insert_candidates( - buffer::SearchBuffer& knns, const PID* ids, const float* distances, size_t count + buffer::SearchBuffer& result_pool, + const PID* ids, + const float* distances, + size_t count ) { for (size_t i = 0; i < count; ++i) { - knns.insert(ids[i], distances[i]); + result_pool.insert(ids[i], distances[i]); } } } // namespace rabitqlib::ivf::detail diff --git a/src/index/pipnn.cpp b/src/index/pipnn.cpp index d9794371..0d96608f 100644 --- a/src/index/pipnn.cpp +++ b/src/index/pipnn.cpp @@ -120,14 +120,14 @@ class CandidateTable { } // Export only after every row has been pruned. - InitialGraph export_graph(size_t threads) const { + InitialGraph export_graph(size_t num_workers) const { InitialGraph result; result.offsets.resize(sizes_.size() + 1, 0); for (size_t i = 0; i < sizes_.size(); ++i) { result.offsets[i + 1] = result.offsets[i] + sizes_[i]; } result.neighbors.resize(result.offsets.back()); -#pragma omp parallel for num_threads(threads) +#pragma omp parallel for num_threads(num_workers) for (std::ptrdiff_t index = 0; index < static_cast(sizes_.size()); ++index) { const size_t i = static_cast(index); @@ -153,21 +153,27 @@ struct BucketJob { uint32_t seed = 555; }; -ScratchPool::ScratchPool(size_t count, size_t dim, size_t threads) - : point_size(std::min(count, kTileSize) * dim) - , distance_size( +ScratchPool::ScratchPool(size_t count, size_t dim, size_t num_workers) + : points_buffer_size(std::min(count, kTileSize) * dim) + , distances_buffer_size( std::max( std::min(count, kTileSize) * std::min(count, size_t{512}), std::min(count, kLeafSize) * std::min(count, kLeafSize) ) ) - , stride(point_size + distance_size + std::min(count, kTileSize)) + , worker_buffer_size( + points_buffer_size + distances_buffer_size + std::min(count, kTileSize) + ) // Preserve uninitialized allocation and worker first-touch placement. - , storage(new float[threads * stride]) {} + , storage(new float[num_workers * worker_buffer_size]) {} Workspace ScratchPool::worker(size_t id) { - float* base = storage.get() + id * stride; - return {base, base + point_size, base + point_size + distance_size, {}}; + float* base = storage.get() + id * worker_buffer_size; + return { + base, + base + points_buffer_size, + base + points_buffer_size + distances_buffer_size, + {}}; } void gather(const float* data, size_t dim, const Bucket& ids, Workspace& work) { @@ -176,9 +182,9 @@ void gather(const float* data, size_t dim, const Bucket& ids, Workspace& work) { } } -void pairwise(Workspace& work, size_t size, size_t dim, MetricType metric) { +void pairwise(Workspace& work, size_t size, size_t dim, MetricType metric_type) { simd::pairwise_distances_lower( - work.points, work.distances, work.norms, size, dim, metric == METRIC_IP + work.points, work.distances, work.norms, size, dim, metric_type == METRIC_IP ); } @@ -213,8 +219,8 @@ std::vector partition( const float* data, size_t dim, const BucketJob& job, - MetricType metric, - size_t threads, + MetricType metric_type, + size_t num_workers, ScratchPool& scratch ) { const auto& ids = job.ids; @@ -224,7 +230,7 @@ std::vector partition( job.depth == 0 ? std::min( kTileSize, - round_up_to_multiple(div_round_up(ids.size(), threads), kMinTileSize) + round_up_to_multiple(div_round_up(ids.size(), num_workers), kMinTileSize) ) : kTileSize; const size_t leader_count = std::min( @@ -274,7 +280,7 @@ std::vector partition( const float dot = distances(static_cast(i), static_cast(j)); const float distance = - metric == METRIC_L2 + metric_type == METRIC_L2 ? std::max(0.0F, norms.data()[i] + leader_norms.data()[j] - 2 * dot) : -dot; work.candidates[j] = {distance, static_cast(j)}; @@ -292,7 +298,7 @@ std::vector partition( if (job.depth == 0) { { index_detail::ParallelException failure; -#pragma omp parallel num_threads(std::min(threads, div_round_up(ids.size(), tile_size))) +#pragma omp parallel num_threads(std::min(num_workers, div_round_up(ids.size(), tile_size))) { auto work = scratch.worker(static_cast(omp_get_thread_num())); #pragma omp for schedule(static) @@ -331,8 +337,8 @@ std::vector cluster( const float* data, size_t count, size_t dim, - MetricType metric, - size_t threads, + MetricType metric_type, + size_t num_workers, ScratchPool& scratch ) { Bucket all(count); @@ -350,7 +356,7 @@ std::vector cluster( finished[i].push_back(std::move(job.ids)); return; } - auto children = partition(data, dim, job, metric, threads, scratch); + auto children = partition(data, dim, job, metric_type, num_workers, scratch); std::sort(children.begin(), children.end(), [](const auto& a, const auto& b) { return a.size() < b.size(); }); @@ -408,7 +414,7 @@ std::vector cluster( } else { { index_detail::ParallelException failure; -#pragma omp parallel for num_threads(std::min(threads, active.size())) schedule(dynamic) +#pragma omp parallel for num_threads(std::min(num_workers, active.size())) schedule(dynamic) for (std::ptrdiff_t i = 0; i < static_cast(active.size()); ++i) { if (failure.failed()) { @@ -443,10 +449,10 @@ InitialGraph build_initial_graph( size_t count, size_t dim, size_t degree, - MetricType metric, + MetricType metric_type, size_t num_threads ) { - validate_metric_type(metric); + validate_metric_type(metric_type); if (data == nullptr || dim == 0 || degree == 0 || degree >= count || count >= buffer::kSearchBufferMaxPointCount || degree % 32 != 0) { throw std::invalid_argument( @@ -455,9 +461,9 @@ InitialGraph build_initial_graph( ); } using namespace pipnn_impl; - const size_t threads = resolve_num_threads(num_threads); - ScratchPool scratch(count, dim, threads); - auto leaves = cluster(data, count, dim, metric, threads, scratch); + const size_t num_workers = resolve_num_threads(num_threads); + ScratchPool scratch(count, dim, num_workers); + auto leaves = cluster(data, count, dim, metric_type, num_workers, scratch); RowMajorMatrix projections(dim, kHashBits); std::mt19937 random(555); std::normal_distribution normal; @@ -466,12 +472,13 @@ InitialGraph build_initial_graph( } RowMajorMatrix sketches(count, kHashBits); const size_t sketch_tile = std::min( - kTileSize, round_up_to_multiple(div_round_up(count, threads), kMinTileSize) + kTileSize, round_up_to_multiple(div_round_up(count, num_workers), kMinTileSize) ); { index_detail::ParallelException failure; -#pragma omp parallel for num_threads(std::min(threads, div_round_up(count, sketch_tile))) \ - schedule(static) +#pragma omp parallel for num_threads( \ + std::min(num_workers, div_round_up(count, sketch_tile)) \ +) schedule(static) for (std::ptrdiff_t offset = 0; offset < static_cast(count); offset += static_cast(sketch_tile)) { if (failure.failed()) { @@ -496,7 +503,7 @@ InitialGraph build_initial_graph( CandidateTable table(count, degree * 5 / 2); { index_detail::ParallelException failure; -#pragma omp parallel num_threads(std::min(threads, leaves.size())) +#pragma omp parallel num_threads(std::min(num_workers, leaves.size())) { auto work = scratch.worker(static_cast(omp_get_thread_num())); #pragma omp for schedule(dynamic) @@ -513,7 +520,7 @@ InitialGraph build_initial_graph( continue; } gather(data, dim, ids, work); - pairwise(work, ids.size(), dim, metric); + pairwise(work, ids.size(), dim, metric_type); select_local_neighbors(ids, work); const size_t keep = std::min(kLocalDegree, ids.size() - 1); // Retain the original point order and (distance, ID) tie order @@ -537,7 +544,7 @@ InitialGraph build_initial_graph( // Prune in-place: retain selected IDs in the front of each reservoir row. { index_detail::ParallelException failure; -#pragma omp parallel num_threads(threads) +#pragma omp parallel num_threads(num_workers) { std::vector candidates; Bucket selected; @@ -562,7 +569,7 @@ InitialGraph build_initial_graph( for (const auto& candidate : candidates) { bool occluded = false; for (PID accepted : selected) { - const float distance = metric == METRIC_L2 + const float distance = metric_type == METRIC_L2 ? euclidean_sqr( data + accepted * dim, data + candidate.id * dim, @@ -591,6 +598,6 @@ InitialGraph build_initial_graph( } failure.rethrow(); } - return table.export_graph(threads); + return table.export_graph(num_workers); } } // namespace rabitqlib::symqg::detail diff --git a/src/index/pipnn_internal.hpp b/src/index/pipnn_internal.hpp index d4ceb3f2..99471c97 100644 --- a/src/index/pipnn_internal.hpp +++ b/src/index/pipnn_internal.hpp @@ -23,12 +23,13 @@ struct Workspace { }; struct ScratchPool { - size_t point_size; - size_t distance_size; - size_t stride; + // Buffer sizes in float elements. + size_t points_buffer_size; + size_t distances_buffer_size; + size_t worker_buffer_size; std::unique_ptr storage; - ScratchPool(size_t count, size_t dim, size_t threads); + ScratchPool(size_t count, size_t dim, size_t num_workers); Workspace worker(size_t id); }; diff --git a/src/index/qg.cpp b/src/index/qg.cpp index bc21f59b..12920132 100644 --- a/src/index/qg.cpp +++ b/src/index/qg.cpp @@ -75,19 +75,21 @@ size_t QuantizedGraph::padded_dimension(size_t dim) { return (checked_add(dim, 31) / 32) * 32; } -size_t QuantizedGraph::num_vertices() const { return this->num_points_; } +size_t QuantizedGraph::num_vertices() const { return this->max_elements_; } -size_t QuantizedGraph::max_elements() const { return num_points_; } +size_t QuantizedGraph::max_elements() const { return max_elements_; } -size_t QuantizedGraph::num_points() const { return ready_ ? num_points_ : 0; } +size_t QuantizedGraph::num_points() const { return ready_ ? max_elements_ : 0; } -size_t QuantizedGraph::num_stored_points() const { return ready_ ? num_points_ : 0; } +size_t QuantizedGraph::num_stored_points() const { + return ready_ ? max_elements_ : 0; +} size_t QuantizedGraph::dimension() const { return this->dim_; } -size_t QuantizedGraph::degree_bound() const { return this->degree_bound_; } +size_t QuantizedGraph::degree_bound() const { return this->max_degree_; } -PID QuantizedGraph::entry_point() const { return this->entry_point_; } +PID QuantizedGraph::entry_point() const { return this->entry_point_id_; } MetricType QuantizedGraph::metric_type() const { return this->metric_type_; } @@ -96,10 +98,10 @@ size_t QuantizedGraph::quantization_bits() const { return this->quantizat bool QuantizedGraph::is_quantized() const { return quantization_bits_ != 0; } void QuantizedGraph::set_ep(PID entry) { - if (entry >= num_points_) { + if (entry >= max_elements_) { throw std::invalid_argument("QuantizedGraph entry point is out of range"); } - entry_point_ = entry; + entry_point_id_ = entry; } QuantizedGraph::QuantizedGraph() = default; @@ -112,16 +114,16 @@ QuantizedGraph& QuantizedGraph::operator=(QuantizedGraph&&) default; QuantizedGraph::QuantizedGraph( - size_t num, + size_t max_elements, size_t dim, - size_t max_deg, + size_t max_degree, MetricType metric_type, RotatorType rotator_type, size_t quantization_bits, uint32_t seed ) - : num_points_(num) - , degree_bound_(max_deg) + : max_elements_(max_elements) + , max_degree_(max_degree) , dim_(dim) , padded_dim_(dim) , raw_dist_func_( @@ -145,22 +147,22 @@ void QuantizedGraph::validate_configuration() const { rotator_type_ != RotatorType::FhtKacRotator) { throw std::invalid_argument("QuantizedGraph rotator type is invalid"); } - if (degree_bound_ == 0 || degree_bound_ % fastscan::kBatchSize != 0) { + if (max_degree_ == 0 || max_degree_ % fastscan::kBatchSize != 0) { throw std::invalid_argument( "QuantizedGraph degree bound must be a positive multiple of 32" ); } - if (degree_bound_ >= num_points_) { + if (max_degree_ >= max_elements_) { throw std::invalid_argument( "QuantizedGraph degree bound must be smaller than the number of points" ); } - if (num_points_ > buffer::kSearchBufferMaxPointCount) { + if (max_elements_ > buffer::kSearchBufferMaxPointCount) { throw std::invalid_argument( "QuantizedGraph point count exceeds the search-buffer ID limit" ); } - if (entry_point_ >= num_points_) { + if (entry_point_id_ >= max_elements_) { throw std::invalid_argument("QuantizedGraph entry point is out of range"); } if (quantization_bits_ != 0 && quantization_bits_ != 4 && quantization_bits_ != 8) { @@ -170,15 +172,15 @@ void QuantizedGraph::validate_configuration() const { } } -void QuantizedGraph::copy_vectors(const float* data, size_t num_threads) { - const int thread_count = static_cast(num_threads); +void QuantizedGraph::copy_vectors(const float* data, size_t num_workers) { + const int omp_num_threads = static_cast(num_workers); if (quantization_bits_ != 0) { if (centroid_.size() != padded_dim_) { throw std::logic_error("qg-quant centroid must be set before copying vectors"); } { index_detail::ParallelException failure; -#pragma omp parallel num_threads(thread_count) +#pragma omp parallel num_threads(omp_num_threads) { std::vector rotated_data; std::vector quantized_data; @@ -188,7 +190,7 @@ void QuantizedGraph::copy_vectors(const float* data, size_t num_threads) } catch (...) { failure.capture(); } #pragma omp for schedule(static) for (std::ptrdiff_t index = 0; - index < static_cast(num_points_); + index < static_cast(max_elements_); ++index) { if (failure.failed()) { continue; @@ -228,8 +230,8 @@ void QuantizedGraph::copy_vectors(const float* data, size_t num_threads) } return; } -#pragma omp parallel for schedule(static) num_threads(thread_count) - for (std::ptrdiff_t index = 0; index < static_cast(num_points_); +#pragma omp parallel for schedule(static) num_threads(omp_num_threads) + for (std::ptrdiff_t index = 0; index < static_cast(max_elements_); ++index) { const size_t i = static_cast(index); const float* src = data + (dim_ * i); @@ -270,11 +272,11 @@ void QuantizedGraph::save(const char* filename) const { } /* Basic variants */ - output.write(reinterpret_cast(&num_points_), sizeof(size_t)); - output.write(reinterpret_cast(°ree_bound_), sizeof(size_t)); + output.write(reinterpret_cast(&max_elements_), sizeof(size_t)); + output.write(reinterpret_cast(&max_degree_), sizeof(size_t)); output.write(reinterpret_cast(&dim_), sizeof(size_t)); output.write(reinterpret_cast(&padded_dim_), sizeof(size_t)); - output.write(reinterpret_cast(&entry_point_), sizeof(PID)); + output.write(reinterpret_cast(&entry_point_id_), sizeof(PID)); output.write(reinterpret_cast(&rotator_type_), sizeof(RotatorType)); output.write(reinterpret_cast(&metric_type_), sizeof(MetricType)); if (quantization_bits_ != 0) { @@ -290,7 +292,7 @@ void QuantizedGraph::save(const char* filename) const { /* Data */ output.write( get_row_data(0), - static_cast(checked_multiply(num_points_, row_offset_)) + static_cast(checked_multiply(max_elements_, record_bytes_)) ); /* Rotator */ @@ -350,11 +352,11 @@ void QuantizedGraph::load(const char* filename) { QuantizedGraph loaded; size_t stored_padded_dim = 0; - read_value(loaded.num_points_, "point count"); - read_value(loaded.degree_bound_, "degree bound"); + read_value(loaded.max_elements_, "point count"); + read_value(loaded.max_degree_, "degree bound"); read_value(loaded.dim_, "dimension"); read_value(stored_padded_dim, "padded dimension"); - read_value(loaded.entry_point_, "entry point"); + read_value(loaded.entry_point_id_, "entry point"); read_value(loaded.rotator_type_, "rotator type"); read_value(loaded.metric_type_, "metric type"); if (magic == kFormatMagic) { @@ -377,7 +379,7 @@ void QuantizedGraph::load(const char* filename) { const size_t centroid_bytes = loaded.quantization_bits_ == 0 ? 0 : checked_multiply(loaded.padded_dim_, sizeof(float)); - const size_t data_bytes = checked_multiply(loaded.num_points_, loaded.row_offset_); + const size_t data_bytes = checked_multiply(loaded.max_elements_, loaded.record_bytes_); const size_t rotator_bytes = loaded.rotator_type_ == RotatorType::MatrixRotator ? checked_multiply( @@ -407,10 +409,10 @@ void QuantizedGraph::load(const char* filename) { read_exact(loaded.get_row_data(0), data_bytes, "graph data"); - for (PID source = 0; source < loaded.num_points_; ++source) { + for (PID source = 0; source < loaded.max_elements_; ++source) { const auto neighbors = loaded.get_neighbors(source); - for (size_t i = 0; i < loaded.degree_bound_; ++i) { - if (neighbors[i] >= loaded.num_points_) { + for (size_t i = 0; i < loaded.max_degree_; ++i) { + if (neighbors[i] >= loaded.max_elements_) { throw std::runtime_error("Invalid QuantizedGraph neighbor ID"); } } @@ -422,92 +424,94 @@ void QuantizedGraph::load(const char* filename) { } input.close(); - // ef is a runtime search setting rather than persisted index state. Preserve + // ef_search is a runtime search setting rather than persisted index state. Preserve // the target object's value, matching the previous in-place load behavior. - loaded.ef_ = ef_; + loaded.default_ef_search_ = default_ef_search_; loaded.ready_ = true; *this = std::move(loaded); } -void QuantizedGraph::set_ef(size_t cur_ef) { - if (cur_ef == 0) { +void QuantizedGraph::set_ef(size_t ef_search) { + if (ef_search == 0) { throw std::invalid_argument("QuantizedGraph ef must be positive"); } - this->ef_ = cur_ef; + this->default_ef_search_ = ef_search; } void QuantizedGraph::search( const float* __restrict__ query, uint32_t k, - uint32_t* __restrict__ results, - float* __restrict__ dists + uint32_t* __restrict__ result_ids, + float* __restrict__ distances ) { - search_impl(query, k, results, dists, ef_); + search_impl(query, k, result_ids, distances, default_ef_search_); } void QuantizedGraph::search_impl( - const float* query, uint32_t k, uint32_t* results, float* dists, size_t ef + const float* query, uint32_t k, uint32_t* result_ids, float* distances, size_t ef_search ) { - validate_search(query, k, results, dists, ef); - detail::SearchScratch scratch(padded_dim_, degree_bound_, ef, k); + validate_search(query, k, result_ids, distances, ef_search); + detail::SearchScratch scratch(padded_dim_, max_degree_, ef_search, k); thread_local VisitedSet visited; thread_local size_t visited_size = 0; - if (visited_size != num_points_) { - visited.initialize(num_points_, num_points_ / 10); - visited_size = num_points_; + if (visited_size != max_elements_) { + visited.initialize(max_elements_, max_elements_ / 10); + visited_size = max_elements_; } search_with_scratch_impl( - query, k, results, dists, scratch.view(visited), kPidMax, nullptr + query, k, result_ids, distances, scratch.view(visited), kPidMax, nullptr ); } void QuantizedGraph::search_batch( const float* queries, size_t num_queries, - uint32_t knn, - uint32_t* results, - float* dists, + uint32_t k, + uint32_t* result_ids, + float* distances, size_t num_threads ) { - search_batch_with_ef(queries, num_queries, knn, results, dists, ef_, num_threads); + search_batch_with_ef( + queries, num_queries, k, result_ids, distances, default_ef_search_, num_threads + ); } void QuantizedGraph::search_batch_with_ef( const float* queries, size_t num_queries, - uint32_t knn, - uint32_t* results, - float* dists, - size_t ef, + uint32_t k, + uint32_t* result_ids, + float* distances, + size_t ef_search, size_t num_threads ) { if (num_queries == 0) { return; } - validate_search(queries, knn, results, dists, ef); + validate_search(queries, k, result_ids, distances, ef_search); const size_t max_values = std::numeric_limits::max() / sizeof(float); - if (num_queries > max_values / dim_ || num_queries > max_values / knn) { + if (num_queries > max_values / dim_ || num_queries > max_values / k) { throw std::length_error("QuantizedGraph query batch is too large"); } if (num_queries == 1) { - search_impl(queries, knn, results, dists, ef); + search_impl(queries, k, result_ids, distances, ef_search); return; } - const auto workers = static_cast(std::min( + const auto num_workers = static_cast(std::min( {resolve_num_threads(num_threads), num_queries, static_cast(std::numeric_limits::max())} )); index_detail::ParallelException failure; -#pragma omp parallel num_threads(workers) if (workers > 1) +#pragma omp parallel num_threads(num_workers) if (num_workers > 1) { // Allocate and release scratch on its worker, retaining it across queries. std::optional scratch; VisitedSet visited; try { - scratch.emplace(padded_dim_, degree_bound_, ef, knn); - visited.initialize(num_points_, num_points_ / 10); + scratch.emplace(padded_dim_, max_degree_, ef_search, k); + visited.initialize(max_elements_, max_elements_ / 10); } catch (...) { failure.capture(); } #pragma omp for schedule(dynamic) for (std::ptrdiff_t query_index = 0; @@ -521,9 +525,9 @@ void QuantizedGraph::search_batch_with_ef( auto& local = *scratch; search_with_scratch_impl( queries + i * dim_, - knn, - results + i * knn, - dists + i * knn, + k, + result_ids + i * k, + distances + i * k, local.view(visited), kPidMax, nullptr @@ -535,18 +539,22 @@ void QuantizedGraph::search_batch_with_ef( } void QuantizedGraph::validate_search( - const float* query, uint32_t k, const uint32_t* results, const float* dists, size_t ef + const float* query, + uint32_t k, + const uint32_t* result_ids, + const float* distances, + size_t ef_search ) const { if (!ready_ || rotator_ == nullptr) { throw std::logic_error("QuantizedGraph must be built or loaded before search"); } - if (query == nullptr || results == nullptr || dists == nullptr) { + if (query == nullptr || result_ids == nullptr || distances == nullptr) { throw std::invalid_argument("QuantizedGraph search buffers must not be null"); } - if (k == 0 || k > num_points_) { + if (k == 0 || k > max_elements_) { throw std::invalid_argument("QuantizedGraph k must be between 1 and num_points"); } - if (ef < k) { + if (ef_search < k) { throw std::invalid_argument("QuantizedGraph ef must be at least k"); } } @@ -563,24 +571,24 @@ float QuantizedGraph::point_distance( // Scan a data row and store estimated neighbor distances. The caller scores the current // vertex from either its raw vector (vanilla QG) or its 4/8-bit code (qg-quant). void QuantizedGraph::scan_neighbors( - const BatchQuery& q_obj, + const BatchQuery& query_wrapper, PID data_id, - float* est_dist, + float* estimated_distances, buffer::SearchBuffer& search_pool, - VisitedSet& vis, - size_t cur_degree + VisitedSet& visited, + size_t num_neighbors ) const { const auto* batch_data = get_batch_data(data_id); const auto neighbors = get_neighbors(data_id); - for (size_t begin = 0; begin < cur_degree; begin += fastscan::kBatchSize) { + for (size_t begin = 0; begin < num_neighbors; begin += fastscan::kBatchSize) { const float threshold = search_pool.top_dist(); uint32_t candidate_mask = simd::qg_batch_estdist_mask( - batch_data, q_obj, padded_dim_, est_dist + begin, threshold + batch_data, query_wrapper, padded_dim_, estimated_distances + begin, threshold ); batch_data += QGBatchDataMap::data_bytes(padded_dim_); // Construction can leave a partial batch with stale IDs in its unused lanes. - const size_t remaining = cur_degree - begin; + const size_t remaining = num_neighbors - begin; if (remaining < fastscan::kBatchSize) { candidate_mask &= (uint32_t{1} << remaining) - 1; } @@ -589,13 +597,13 @@ void QuantizedGraph::scan_neighbors( const auto lane = static_cast(bitops::countr_zero32(candidate_mask)); candidate_mask &= candidate_mask - 1; const size_t i = begin + lane; - PID cur_neighbor = neighbors[i]; - float dist = est_dist[i]; + PID neighbor_id = neighbors[i]; + float dist = estimated_distances[i]; - if (search_pool.is_full(dist) || vis.get(cur_neighbor)) { + if (search_pool.is_full(dist) || visited.get(neighbor_id)) { continue; } - search_pool.insert(cur_neighbor, dist); // update search buffer + search_pool.insert(neighbor_id, dist); // update search buffer memory::mem_prefetch_l2(get_row_data(search_pool.next_id()), 10); } } @@ -603,7 +611,7 @@ void QuantizedGraph::scan_neighbors( void QuantizedGraph::update_results( buffer::SearchBuffer& result_pool, - VisitedSet& vis, + VisitedSet& visited, const float* query, const QuantizedQuery* quantized_query ) { @@ -617,12 +625,12 @@ void QuantizedGraph::update_results( ); for (const auto& record : data) { auto neighbors = get_neighbors(record.id); - for (uint32_t i = 0; i < this->degree_bound_; ++i) { - PID cur_neighbor = neighbors[i]; - if (!vis.get(cur_neighbor)) { - vis.set(cur_neighbor); + for (uint32_t i = 0; i < this->max_degree_; ++i) { + PID neighbor_id = neighbors[i]; + if (!visited.get(neighbor_id)) { + visited.set(neighbor_id); result_pool.insert( - cur_neighbor, point_distance(query, quantized_query, cur_neighbor) + neighbor_id, point_distance(query, quantized_query, neighbor_id) ); } } @@ -642,10 +650,10 @@ void QuantizedGraph::initialize_layout() { const size_t batch_bytes = QGBatchDataMap::data_bytes(padded_dim_); neighbor_offset_ = checked_add( batch_data_offset_, - checked_multiply(batch_bytes, degree_bound_ / fastscan::kBatchSize) + checked_multiply(batch_bytes, max_degree_ / fastscan::kBatchSize) ); - row_offset_ = - checked_add(neighbor_offset_, checked_multiply(degree_bound_, sizeof(PID))); + record_bytes_ = + checked_add(neighbor_offset_, checked_multiply(max_degree_, sizeof(PID))); } void QuantizedGraph::initialize(bool randomize_rotation) { @@ -660,8 +668,8 @@ void QuantizedGraph::initialize(bool randomize_rotation) { initialize_layout(); - const size_t data_bytes = checked_multiply(num_points_, row_offset_); - assert(row_offset_ % sizeof(float) == 0); + const size_t data_bytes = checked_multiply(max_elements_, record_bytes_); + assert(record_bytes_ % sizeof(float) == 0); data_ = RowStorage(data_bytes / sizeof(float)); if (quantization_bits_ != 0) { @@ -716,20 +724,20 @@ void QuantizedGraph::encode_neighbors( const float* rotated_data, const float* rotated_centroid ) { - size_t cur_degree = new_neighbors.size(); + size_t num_neighbors = new_neighbors.size(); - if (cur_degree == 0) { + if (num_neighbors == 0) { return; } auto neighbor_ptr = get_neighbors(cur_id); - for (size_t i = 0; i < cur_degree; ++i) { + for (size_t i = 0; i < num_neighbors; ++i) { neighbor_ptr[i] = new_neighbors[i]; } // quantize batches for current vertex auto* batch_data = get_batch_data(cur_id); - for (size_t i = 0; i < cur_degree; i += fastscan::kBatchSize) { - const size_t batch_size = std::min(cur_degree - i, fastscan::kBatchSize); + for (size_t i = 0; i < num_neighbors; i += fastscan::kBatchSize) { + const size_t batch_size = std::min(num_neighbors - i, fastscan::kBatchSize); quant::quantize_qg_batch( rotated_data + (i * padded_dim_), rotated_centroid, @@ -757,9 +765,9 @@ void QuantizedGraph::encode_neighbors( } void QuantizedGraph::search_with_scratch( const float* __restrict__ query, - uint32_t knn, - uint32_t* __restrict__ results, - float* __restrict__ dists, + uint32_t k, + uint32_t* __restrict__ result_ids, + float* __restrict__ distances, float* __restrict__ rotated_query_scratch, float* __restrict__ est_dist_scratch, float* __restrict__ lut_float_scratch, @@ -773,9 +781,9 @@ void QuantizedGraph::search_with_scratch( detail::search_with_workspace( *this, query, - knn, - results, - dists, + k, + result_ids, + distances, {rotated_query_scratch, est_dist_scratch, lut_float_scratch, @@ -791,22 +799,22 @@ void QuantizedGraph::search_with_scratch( void detail::search_with_workspace( QuantizedGraph& graph, const float* query, - uint32_t knn, - uint32_t* results, - float* dists, + uint32_t k, + uint32_t* result_ids, + float* distances, detail::SearchScratchView scratch, PID hint, const float* k1xsumq ) { - graph.validate_search(query, knn, results, dists, graph.ef_); - graph.search_with_scratch_impl(query, knn, results, dists, scratch, hint, k1xsumq); + graph.validate_search(query, k, result_ids, distances, graph.default_ef_search_); + graph.search_with_scratch_impl(query, k, result_ids, distances, scratch, hint, k1xsumq); } void QuantizedGraph::search_with_scratch_impl( const float* __restrict__ query, - uint32_t knn, - uint32_t* __restrict__ results, - float* __restrict__ dists, + uint32_t k, + uint32_t* __restrict__ result_ids, + float* __restrict__ distances, detail::SearchScratchView scratch, PID hint, const float* k1xsumq @@ -815,7 +823,7 @@ void QuantizedGraph::search_with_scratch_impl( scratch.lookup_table == nullptr) { throw std::invalid_argument("QuantizedGraph search buffers must not be null"); } - if (hint != kPidMax && hint >= num_points_) { + if (hint != kPidMax && hint >= max_elements_) { throw std::invalid_argument("SymphonyQG search hint is outside the graph"); } if (!scratch.visited.initialized()) { @@ -837,10 +845,10 @@ void QuantizedGraph::search_with_scratch_impl( scratch.result_pool.clear(); scratch.visited.clear(); - if (hint != kPidMax && hint != entry_point_) { + if (hint != kPidMax && hint != entry_point_id_) { scratch.search_pool.insert(hint, point_distance(query, quantized, hint)); } - scratch.search_pool.insert(entry_point_, std::numeric_limits::max()); + scratch.search_pool.insert(entry_point_id_, std::numeric_limits::max()); while (scratch.search_pool.has_next()) { const PID current = scratch.search_pool.pop(); @@ -857,16 +865,16 @@ void QuantizedGraph::search_with_scratch_impl( scratch.estimated_distances, scratch.search_pool, scratch.visited, - degree_bound_ + max_degree_ ); scratch.result_pool.insert(current, vertex_distance); } update_results(scratch.result_pool, scratch.visited, query, quantized); - if (scratch.result_pool.size() != knn) { + if (scratch.result_pool.size() != k) { throw std::runtime_error("QuantizedGraph search could not produce k results"); } - scratch.result_pool.copy_results(results, dists); + scratch.result_pool.copy_results(result_ids, distances); } } // namespace rabitqlib::symqg diff --git a/src/index/qg_builder.cpp b/src/index/qg_builder.cpp index bb6c13b1..9c92c673 100644 --- a/src/index/qg_builder.cpp +++ b/src/index/qg_builder.cpp @@ -34,25 +34,25 @@ QGBuilder::BuildQuery::BuildQuery( ) : graph_(graph), cache_(cache) {} -QGBuilder::BuildWorker::BuildWorker(size_t num_nodes, size_t visited_capacity) - : visited(num_nodes, visited_capacity) {} +QGBuilder::BuildWorker::BuildWorker(size_t num_points, size_t visited_capacity) + : visited(num_points, visited_capacity) {} const QGBuilder::BuildCache* QGBuilder::build_cache() const { return cache_ ? &*cache_ : nullptr; } QGBuilder::QGBuilder( QuantizedGraph& index, - uint32_t ef_build, + uint32_t ef_construction, size_t num_threads, uint32_t seed, bool cache_vectors ) : qg_{index} - , ef_build_{ef_build} - , num_threads_{resolve_num_threads(num_threads)} - , num_nodes_{qg_.num_vertices()} + , ef_construction_{ef_construction} + , num_workers_{resolve_num_threads(num_threads)} + , num_points_{qg_.num_vertices()} , dim_{qg_.dimension()} - , degree_bound_(qg_.degree_bound()) + , max_degree_(qg_.degree_bound()) , seed_(seed) { if (cache_vectors) { cache_.emplace(); @@ -61,14 +61,14 @@ QGBuilder::QGBuilder( QGBuilder::QGBuilder( QuantizedGraph& index, - uint32_t ef_build, + uint32_t ef_construction, const float* data, size_t num_threads, QGInitialization init, uint32_t seed, bool cache_vectors ) - : QGBuilder(index, ef_build, num_threads, seed, cache_vectors) { + : QGBuilder(index, ef_construction, num_threads, seed, cache_vectors) { if (init != QGInitialization::PiPNN && init != QGInitialization::Random) { throw std::invalid_argument("Unknown QG initialization"); } @@ -78,14 +78,20 @@ QGBuilder::QGBuilder( QGBuilder::QGBuilder( QuantizedGraph& index, - uint32_t ef_build, + uint32_t ef_construction, const float* data, size_t num_threads, uint32_t seed, bool cache_vectors ) : QGBuilder( - index, ef_build, data, num_threads, QGInitialization::Random, seed, cache_vectors + index, + ef_construction, + data, + num_threads, + QGInitialization::Random, + seed, + cache_vectors ) {} void QGBuilder::reset(const float* data) { @@ -94,7 +100,7 @@ void QGBuilder::reset(const float* data) { } if (seeded_) { auto initial = detail::build_initial_graph( - data, num_nodes_, dim_, degree_bound_, qg_.metric_type_, num_threads_ + data, num_points_, dim_, max_degree_, qg_.metric_type_, num_workers_ ); initialize_seed(data, initial.offsets, initial.neighbors); } else { @@ -106,21 +112,21 @@ void QGBuilder::reset(const float* data) { void QGBuilder::initialize_seed( const float* data, const std::vector& offsets, const std::vector& neighbors ) { - if (offsets.size() != num_nodes_ + 1 || offsets.front() != 0 || + if (offsets.size() != num_points_ + 1 || offsets.front() != 0 || offsets.back() != neighbors.size()) { throw std::invalid_argument("Seed graph offsets must delimit every vertex"); } - for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); + for (std::ptrdiff_t index = 0; index < static_cast(num_points_); ++index) { const size_t i = static_cast(index); if (offsets[i] > offsets[i + 1] || offsets[i + 1] > neighbors.size() || - offsets[i + 1] - offsets[i] > degree_bound_) { + offsets[i + 1] - offsets[i] > max_degree_) { throw std::invalid_argument( "Seed graph row exceeds degree bound or has invalid offsets" ); } for (size_t j = offsets[i]; j < offsets[i + 1]; ++j) { - if (neighbors[j] >= num_nodes_ || neighbors[j] == i) { + if (neighbors[j] >= num_points_ || neighbors[j] == i) { throw std::invalid_argument( "Seed graph IDs must be in range and exclude self" ); @@ -130,15 +136,15 @@ void QGBuilder::initialize_seed( initialize_storage(data); { index_detail::ParallelException failure; -#pragma omp parallel num_threads(num_threads_) +#pragma omp parallel num_threads(num_workers_) { auto& worker = workers_[static_cast(omp_get_thread_num())]; auto& ids = worker.neighbor_ids; try { - ids.reserve(degree_bound_); + ids.reserve(max_degree_); } catch (...) { failure.capture(); } #pragma omp for schedule(dynamic) - for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); + for (std::ptrdiff_t index = 0; index < static_cast(num_points_); ++index) { if (failure.failed()) { continue; @@ -175,8 +181,8 @@ bool QGBuilder::check_dup() const { std::atomic flag(false); { index_detail::ParallelException failure; -#pragma omp parallel for num_threads(num_threads_) - for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); +#pragma omp parallel for num_threads(num_workers_) + for (std::ptrdiff_t index = 0; index < static_cast(num_points_); ++index) { if (failure.failed()) { continue; @@ -199,7 +205,7 @@ bool QGBuilder::check_dup() const { float QGBuilder::avg_degree() const { size_t degrees = std::accumulate(degrees_.begin(), degrees_.end(), 0U); - return static_cast(degrees) / static_cast(num_nodes_); + return static_cast(degrees) / static_cast(num_points_); } void QGBuilder::run_passes(size_t num_iter) { qg_.ready_ = false; @@ -214,13 +220,13 @@ void QGBuilder::run_passes(size_t num_iter) { void QGBuilder::write_row(PID id, BuildWorker& worker) { if (!worker.neighbor_ids.empty()) { const auto& new_neighbors = worker.neighbor_ids; - const size_t cur_degree = new_neighbors.size(); + const size_t num_neighbors = new_neighbors.size(); const auto* cache = build_cache(); const size_t padded_dim = qg_.padded_dim_; auto& rotated_data = worker.rotated_neighbors; - rotated_data.resize((cur_degree + 1) * padded_dim); - float* rotated_centroid = rotated_data.data() + cur_degree * padded_dim; - for (size_t i = 0; i < cur_degree; ++i) { + rotated_data.resize((num_neighbors + 1) * padded_dim); + float* rotated_centroid = rotated_data.data() + num_neighbors * padded_dim; + for (size_t i = 0; i < num_neighbors; ++i) { if (cache != nullptr) { std::copy_n( cache->rotated_vectors.data() + new_neighbors[i] * padded_dim, @@ -266,9 +272,9 @@ void QGBuilder::initialize_storage(const float* data) { } // Allocate refinement scratch after PiPNN releases its dense workspace. // Clear old graph-dependent state on reset while retaining vector capacities. - new_neighbors_.resize(num_nodes_); - reverse_buffer_.resize(num_nodes_); - for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); + new_neighbors_.resize(num_points_); + reverse_buffer_.resize(num_points_); + for (std::ptrdiff_t index = 0; index < static_cast(num_points_); ++index) { const size_t i = static_cast(index); new_neighbors_[i].clear(); @@ -276,24 +282,26 @@ void QGBuilder::initialize_storage(const float* data) { } if (workers_.empty()) { workers_ = std::vector( - num_threads_, - BuildWorker(num_nodes_, std::min(ef_build_ * ef_build_, num_nodes_ / 10)) + num_workers_, + BuildWorker( + num_points_, std::min(ef_construction_ * ef_construction_, num_points_ / 10) + ) ); } - degrees_.assign(num_nodes_, degree_bound_); - std::vector centroid = compute_centroid(data, num_nodes_, dim_, num_threads_); + degrees_.assign(num_points_, max_degree_); + std::vector centroid = compute_centroid(data, num_points_, dim_, num_workers_); qg_.ready_ = false; qg_.set_quantization_centroid(centroid.data()); - qg_.copy_vectors(data, num_threads_); + qg_.copy_vectors(data, num_workers_); if (cache_) { const size_t padded_dim = qg_.padded_dim_; - cache_->rotated_vectors.resize(num_nodes_ * padded_dim); + cache_->rotated_vectors.resize(num_points_ * padded_dim); { index_detail::ParallelException failure; -#pragma omp parallel for num_threads(num_threads_) schedule(static) - for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); +#pragma omp parallel for num_threads(num_workers_) schedule(static) + for (std::ptrdiff_t index = 0; index < static_cast(num_points_); ++index) { if (failure.failed()) { continue; @@ -312,8 +320,8 @@ void QGBuilder::initialize_storage(const float* data) { } cache_->queries.clear(); if (qg_.is_quantized()) { - cache_->queries.reserve(num_nodes_); - for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); + cache_->queries.reserve(num_points_); + for (std::ptrdiff_t index = 0; index < static_cast(num_points_); ++index) { const size_t id = static_cast(index); cache_->queries.emplace_back( @@ -332,7 +340,7 @@ void QGBuilder::initialize_storage(const float* data) { qg_.centroid_.data(), qg_.centroid_.data(), qg_.padded_dim_, qg_.metric_type_ ); float best = std::numeric_limits::max(); - for (PID id = 0; id < num_nodes_; ++id) { + for (PID id = 0; id < num_points_; ++id) { const float distance = qg_.quantized_distance(query, id); if (distance < best) { best = distance; @@ -341,7 +349,7 @@ void QGBuilder::initialize_storage(const float* data) { } } else { entry_point = exact_nn( - data, centroid.data(), num_nodes_, dim_, num_threads_, qg_.raw_dist_func_ + data, centroid.data(), num_points_, dim_, num_workers_, qg_.raw_dist_func_ ); } @@ -359,13 +367,13 @@ void QGBuilder::add_pruned_edges( new_result = result; std::unordered_set nei_set; - nei_set.reserve(degree_bound_); + nei_set.reserve(max_degree_); for (const auto& nei : result) { nei_set.emplace(nei.id); } BuildQuery query(qg_, build_cache()); - while (new_result.size() < degree_bound_ && start < pruned_list.size()) { + while (new_result.size() < max_degree_ && start < pruned_list.size()) { const auto& cur = pruned_list[start]; bool occlude = false; query.reset(cur.id); @@ -407,7 +415,7 @@ void QGBuilder::heuristic_prune( } pruned_results.clear(); size_t poolsize = pool.size(); - if (poolsize <= degree_bound_) { + if (poolsize <= max_degree_) { pruned_results = pool; return; } @@ -416,7 +424,7 @@ void QGBuilder::heuristic_prune( size_t start = 0; BuildQuery query(qg_, build_cache()); - while (pruned_results.size() < degree_bound_ && start < poolsize) { + while (pruned_results.size() < max_degree_ && start < poolsize) { auto candidate_id = pool[start].id; if (pruned[start]) { ++start; @@ -456,40 +464,40 @@ void QGBuilder::heuristic_prune( void QGBuilder::search_new_neighbors() { { index_detail::ParallelException failure; -#pragma omp parallel for num_threads(num_threads_) schedule(dynamic) - for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); +#pragma omp parallel for num_threads(num_workers_) schedule(dynamic) + for (std::ptrdiff_t index = 0; index < static_cast(num_points_); ++index) { if (failure.failed()) { continue; } try { const size_t i = static_cast(index); - PID cur_id = i; + PID point_id = i; auto& worker = workers_[static_cast(omp_get_thread_num())]; CandidateList& candidates = worker.candidates; candidates.clear(); - VisitedSet& vis = worker.visited; + VisitedSet& visited = worker.visited; candidates.reserve(2 * kMaxCandidatePoolSize); - vis.clear(); - find_candidates(cur_id, candidates, vis); + visited.clear(); + find_candidates(point_id, candidates, visited); // Seeded construction keeps its initial edges only in QG. Materialize // their scores on demand, after the caller can release the input/CSR. - if (new_neighbors_[cur_id].empty() && degrees_[cur_id] != 0) { + if (new_neighbors_[point_id].empty() && degrees_[point_id] != 0) { BuildQuery query(qg_, build_cache()); - query.reset(cur_id); - const auto ids = qg_.get_neighbors(cur_id); - for (size_t j = 0; j < degrees_[cur_id]; ++j) { - if (ids[j] != cur_id && !vis.get(ids[j])) { + query.reset(point_id); + const auto ids = qg_.get_neighbors(point_id); + for (size_t j = 0; j < degrees_[point_id]; ++j) { + if (ids[j] != point_id && !visited.get(ids[j])) { candidates.emplace_back(ids[j], query.distance_to(ids[j])); } } } // Add current neighbors retained by preceding iterations. - for (auto& nei : new_neighbors_[cur_id]) { + for (auto& nei : new_neighbors_[point_id]) { auto neighbor_id = nei.id; - if (neighbor_id != cur_id && !vis.get(neighbor_id)) { + if (neighbor_id != point_id && !visited.get(neighbor_id)) { candidates.emplace_back(nei); } } @@ -503,8 +511,8 @@ void QGBuilder::search_new_neighbors() { candidates.resize(min_size); // prune and update qg - new_neighbors_[cur_id].reserve(degree_bound_); - heuristic_prune(candidates, new_neighbors_[cur_id], /*recorded=*/nullptr); + new_neighbors_[point_id].reserve(max_degree_); + heuristic_prune(candidates, new_neighbors_[point_id], /*recorded=*/nullptr); } catch (...) { failure.capture(); } } failure.rethrow(); @@ -512,7 +520,7 @@ void QGBuilder::search_new_neighbors() { } void QGBuilder::add_reverse_edges(bool final_pass) { - std::vector locks(num_nodes_); + std::vector locks(num_points_); for (auto& row : reverse_buffer_) { row.clear(); } @@ -521,8 +529,8 @@ void QGBuilder::add_reverse_edges(bool final_pass) { // destination row here races with another worker reading that row as its source. { index_detail::ParallelException failure; -#pragma omp parallel for num_threads(num_threads_) schedule(dynamic) - for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); +#pragma omp parallel for num_threads(num_workers_) schedule(dynamic) + for (std::ptrdiff_t index = 0; index < static_cast(num_points_); ++index) { if (failure.failed()) { continue; @@ -546,8 +554,7 @@ void QGBuilder::add_reverse_edges(bool final_pass) { std::lock_guard lock(locks[destination]); const size_t missing_slots = - degree_bound_ - - std::min(degree_bound_, destination_neighbors.size()); + max_degree_ - std::min(max_degree_, destination_neighbors.size()); if (reverse_buffer_[destination].size() < kMaxCandidatePoolSize + missing_slots) { reverse_buffer_[destination].emplace_back(data_id, nei.distance); @@ -560,37 +567,39 @@ void QGBuilder::add_reverse_edges(bool final_pass) { { index_detail::ParallelException failure; -#pragma omp parallel for num_threads(num_threads_) schedule(dynamic) - for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); +#pragma omp parallel for num_threads(num_workers_) schedule(dynamic) + for (std::ptrdiff_t index = 0; index < static_cast(num_points_); ++index) { if (failure.failed()) { continue; } try { const PID data_id = static_cast(index); - CandidateList& tmp_pool = reverse_buffer_[data_id]; - if (qg_.is_quantized() && !tmp_pool.empty()) { + CandidateList& candidate_pool = reverse_buffer_[data_id]; + if (qg_.is_quantized() && !candidate_pool.empty()) { // RaBitQ estimates are directional: score destination -> source afresh. BuildQuery query(qg_, build_cache()); query.reset(data_id); - for (auto& candidate : tmp_pool) { + for (auto& candidate : candidate_pool) { candidate.distance = query.distance_to(candidate.id); } } - tmp_pool.reserve(tmp_pool.size() + degree_bound_); - tmp_pool.insert( - tmp_pool.end(), + candidate_pool.reserve(candidate_pool.size() + max_degree_); + candidate_pool.insert( + candidate_pool.end(), new_neighbors_[data_id].begin(), new_neighbors_[data_id].end() ); - std::sort(tmp_pool.begin(), tmp_pool.end()); + std::sort(candidate_pool.begin(), candidate_pool.end()); auto& worker = workers_[static_cast(omp_get_thread_num())]; worker.pruned.clear(); if (final_pass) { worker.pruned.reserve(kMaxPrunedSize); } heuristic_prune( - tmp_pool, new_neighbors_[data_id], final_pass ? &worker.pruned : nullptr + candidate_pool, + new_neighbors_[data_id], + final_pass ? &worker.pruned : nullptr ); // Reverse collection has finished. Complete this row against the old // graph before reusing its worker scratch; graph writes stay deferred. @@ -605,11 +614,11 @@ void QGBuilder::add_reverse_edges(bool final_pass) { void QGBuilder::random_init() { const PID min_id = 0; - const PID max_id = num_nodes_ - 1; + const PID max_id = num_points_ - 1; { index_detail::ParallelException failure; -#pragma omp parallel for num_threads(num_threads_) schedule(dynamic) - for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); +#pragma omp parallel for num_threads(num_workers_) schedule(dynamic) + for (std::ptrdiff_t index = 0; index < static_cast(num_points_); ++index) { if (failure.failed()) { continue; @@ -620,8 +629,8 @@ void QGBuilder::random_init() { std::mt19937 rng(seed_ ^ static_cast(i)); std::uniform_int_distribution distribution(min_id, max_id); std::unordered_set neighbor_set; - neighbor_set.reserve(degree_bound_); - while (neighbor_set.size() < degree_bound_) { + neighbor_set.reserve(max_degree_); + while (neighbor_set.size() < max_degree_) { const PID rand_id = distribution(rng); if (rand_id != i) { neighbor_set.emplace(rand_id); @@ -630,7 +639,7 @@ void QGBuilder::random_init() { BuildQuery query(qg_, build_cache()); query.reset(i); - new_neighbors_[i].reserve(degree_bound_); + new_neighbors_[i].reserve(max_degree_); for (PID cur_neigh : neighbor_set) { new_neighbors_[i].emplace_back(cur_neigh, query.distance_to(cur_neigh)); } @@ -649,30 +658,30 @@ void QGBuilder::random_init() { // Complete one final-pass row using only this worker's reverse-prune candidates. void QGBuilder::complete_degree(PID i, BuildWorker& worker) { CandidateList& cur_neighbors = new_neighbors_[i]; - size_t cur_degree = cur_neighbors.size(); + size_t num_neighbors = cur_neighbors.size(); // skip vertices with enough neighbors - if (cur_degree >= degree_bound_) { + if (num_neighbors >= max_degree_) { return; } CandidateList& pruned_list = worker.pruned; CandidateList& new_result = worker.completion_result; - new_result.reserve(degree_bound_); + new_result.reserve(max_degree_); std::sort(pruned_list.begin(), pruned_list.end()); // use binary search to get refined results CandidateList& successful_result = worker.successful_result; successful_result.clear(); - successful_result.reserve(degree_bound_); + successful_result.reserve(max_degree_); float left = 0.5; float right = 1.0; size_t iter = 0; while (iter++ < kMaxBsIter) { float mid = (left + right) / 2; add_pruned_edges(cur_neighbors, pruned_list, new_result, mid); - if (new_result.size() < degree_bound_) { + if (new_result.size() < max_degree_) { left = mid; } else { right = mid; @@ -689,19 +698,19 @@ void QGBuilder::complete_degree(PID i, BuildWorker& worker) { } // if the vertex still doesn't have enough neighbors, use random vertices - if (new_result.size() < degree_bound_) { + if (new_result.size() < max_degree_) { std::mt19937 rng(seed_ ^ static_cast(i) ^ 0x80000000U); std::uniform_int_distribution distribution( - 0, static_cast(num_nodes_) - 1 + 0, static_cast(num_points_) - 1 ); std::unordered_set ids; - ids.reserve(degree_bound_); + ids.reserve(max_degree_); for (auto& neighbor : new_result) { ids.emplace(neighbor.id); } BuildQuery query(qg_, build_cache()); query.reset(i); - while (new_result.size() < degree_bound_) { + while (new_result.size() < max_degree_) { const PID rand_id = distribution(rng); if (rand_id != static_cast(i) && ids.find(rand_id) == ids.end()) { new_result.emplace_back(rand_id, query.distance_to(rand_id)); @@ -720,8 +729,8 @@ void QGBuilder::run_pass(bool final_pass) { // update qg { index_detail::ParallelException failure; -#pragma omp parallel for num_threads(num_threads_) schedule(dynamic) - for (std::ptrdiff_t index = 0; index < static_cast(num_nodes_); +#pragma omp parallel for num_threads(num_workers_) schedule(dynamic) + for (std::ptrdiff_t index = 0; index < static_cast(num_points_); ++index) { if (failure.failed()) { continue; @@ -775,27 +784,38 @@ const float* QGBuilder::BuildQuery::rotated_query() { return rotated_; } -// find candidate neighbors for cur_id, exclude the vertex itself -void QGBuilder::find_candidates(PID cur_id, CandidateList& results, VisitedSet& vis) const { +// find candidate neighbors for point_id, exclude the vertex itself +void QGBuilder::find_candidates( + PID point_id, CandidateList& results, VisitedSet& visited +) const { BuildQuery query(qg_, build_cache()); - query.reset(cur_id); - BatchQuery q_obj(query.rotated_query(), qg_.padded_dim_, qg_.metric_type_); - buffer::SearchBuffer tmp_pool(ef_build_); - tmp_pool.insert(qg_.entry_point_, std::numeric_limits::max()); - - std::vector est_dist(degree_bound_); - while (tmp_pool.has_next()) { - auto cur_candi = tmp_pool.pop(); - if (vis.get(cur_candi)) { + query.reset(point_id); + BatchQuery query_wrapper( + query.rotated_query(), qg_.padded_dim_, qg_.metric_type_ + ); + buffer::SearchBuffer candidate_set(ef_construction_); + candidate_set.insert(qg_.entry_point_id_, std::numeric_limits::max()); + + std::vector estimated_distances(max_degree_); + while (candidate_set.has_next()) { + auto candidate_id = candidate_set.pop(); + if (visited.get(candidate_id)) { continue; } - vis.set(cur_candi); - auto cur_degree = degrees_[cur_candi]; - const float vertex_distance = query.distance_to(cur_candi); - q_obj.set_g_add(vertex_distance); - qg_.scan_neighbors(q_obj, cur_candi, est_dist.data(), tmp_pool, vis, cur_degree); - if (cur_candi != cur_id) { - results.emplace_back(cur_candi, vertex_distance); + visited.set(candidate_id); + auto num_neighbors = degrees_[candidate_id]; + const float vertex_distance = query.distance_to(candidate_id); + query_wrapper.set_g_add(vertex_distance); + qg_.scan_neighbors( + query_wrapper, + candidate_id, + estimated_distances.data(), + candidate_set, + visited, + num_neighbors + ); + if (candidate_id != point_id) { + results.emplace_back(candidate_id, vertex_distance); } } } diff --git a/tests/unit/rabitqlib/index/hnsw_test.cpp b/tests/unit/rabitqlib/index/hnsw_test.cpp index 4e59a3f4..1dca9d79 100644 --- a/tests/unit/rabitqlib/index/hnsw_test.cpp +++ b/tests/unit/rabitqlib/index/hnsw_test.cpp @@ -45,11 +45,13 @@ struct HnswPruningTestAccess { static std::vector connect(HierarchicalNSW& index, PID incoming) { PID* neighbors = index.get_linklist0(0); - index.set_list_count(neighbors, index.maxM0_); - for (size_t i = 0; i < index.maxM0_; ++i) { + index.set_list_count(neighbors, index.base_layer_max_degree_); + for (size_t i = 0; i < index.base_layer_max_degree_; ++i) { neighbors[i + 1] = static_cast(i + 1); } - std::fill_n(index.get_linklist0(incoming), index.maxM0_ + 1, PID{0}); + std::fill_n( + index.get_linklist0(incoming), index.base_layer_max_degree_ + 1, PID{0} + ); detail::QuantQuery cached; maxheap> candidates; candidates.emplace(index.get_quant_dist(0, incoming, cached), 0); diff --git a/tests/unit/rabitqlib/index/qg_test.cpp b/tests/unit/rabitqlib/index/qg_test.cpp index 4e1e356e..ec254f4c 100644 --- a/tests/unit/rabitqlib/index/qg_test.cpp +++ b/tests/unit/rabitqlib/index/qg_test.cpp @@ -57,7 +57,7 @@ struct QGConstructionTestAccess { static CandidateList record_collinear_prune(QGBuilder& builder) { CandidateList pool; - for (PID id = 1; id < builder.num_nodes_; ++id) { + for (PID id = 1; id < builder.num_points_; ++id) { pool.emplace_back(id, build_distance(builder, 0, id)); } CandidateList result; @@ -71,7 +71,7 @@ struct QGConstructionTestAccess { for (auto& row : builder.new_neighbors_) { row.clear(); } - for (PID id = 1; id < builder.num_nodes_; ++id) { + for (PID id = 1; id < builder.num_points_; ++id) { builder.new_neighbors_[id].emplace_back(0, build_distance(builder, id, 0)); } } @@ -134,7 +134,7 @@ struct QGConstructionTestAccess { // Only one slot is missing; independently reproduce its seeded choice. std::mt19937 random(builder.seed_ ^ 0x80000000U); std::uniform_int_distribution distribution( - 0, static_cast(builder.num_nodes_) - 1 + 0, static_cast(builder.num_points_) - 1 ); PID selected = 0; do { @@ -172,9 +172,9 @@ struct QGConstructionTestAccess { static auto rows(const QuantizedGraph& graph) { std::vector result; - for (PID id = 0; id < graph.num_points_; ++id) { + for (PID id = 0; id < graph.max_elements_; ++id) { const char* row = graph.get_quantized_vector(id); - result.insert(result.end(), row, row + graph.row_offset_); + result.insert(result.end(), row, row + graph.record_bytes_); } return result; } @@ -187,11 +187,11 @@ struct QGConstructionTestAccess { ? ExDataMap::data_bytes(graph.padded_dim_, graph.quantization_bits_) : graph.dim_ * sizeof(float); const size_t batch_bytes = QGBatchDataMap::data_bytes(graph.padded_dim_) * - (graph.degree_bound_ / fastscan::kBatchSize); + (graph.max_degree_ / fastscan::kBatchSize); const size_t row_bytes = - vector_bytes + batch_bytes + (graph.degree_bound_ * sizeof(PID)); + vector_bytes + batch_bytes + (graph.max_degree_ * sizeof(PID)); const char* base = graph.get_row_data(0); - for (PID id = 0; id < graph.num_points_; ++id) { + for (PID id = 0; id < graph.max_elements_; ++id) { const char* row = graph.get_row_data(id); EXPECT_EQ(row, base + (id * row_bytes)); EXPECT_EQ(graph.get_batch_data(id), row + vector_bytes); @@ -207,7 +207,7 @@ struct QGConstructionTestAccess { EXPECT_EQ(stored_id, id); } if (!graph.is_quantized()) { - for (PID id = 0; id < graph.num_points_; ++id) { + for (PID id = 0; id < graph.max_elements_; ++id) { EXPECT_FLOAT_EQ( graph.get_vector(id)[graph.dim_ - 1], static_cast(id) ); @@ -229,7 +229,7 @@ struct QGConstructionTestAccess { } static auto codes(const QuantizedGraph& graph) { std::vector result; - for (PID id = 0; id < graph.num_points_; ++id) { + for (PID id = 0; id < graph.max_elements_; ++id) { const char* code = graph.is_quantized() ? graph.get_quantized_vector(id) : reinterpret_cast(graph.get_vector(id)); @@ -243,7 +243,7 @@ struct QGConstructionTestAccess { static const auto& degrees(const QGBuilder& builder) { return builder.degrees_; } static void retain_seed_scores(QGBuilder& builder) { auto& graph = builder.qg_; - for (PID i = 0; i < builder.num_nodes_; ++i) { + for (PID i = 0; i < builder.num_points_; ++i) { QGBuilder::BuildQuery query(graph); query.reset(i); for (size_t j = 0; j < builder.degrees_[i]; ++j) {