From 7bc170a56aca273d41c84359e6d0d2207783245c Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Thu, 29 Oct 2026 15:36:41 +0800 Subject: [PATCH] refactor: align numerical naming and documentation --- docs/development.md | 24 +++ docs/docs/index/hnsw.md | 7 +- docs/docs/index/ivf.md | 8 +- docs/docs/rabitq/estimator.md | 12 +- docs/docs/rabitq/quantizer.md | 44 +++--- .../rabitqlib/clustering/detail/kmeans.hpp | 44 +++--- include/rabitqlib/clustering/qgkmeans.hpp | 97 ++++++------ include/rabitqlib/clustering/rabitqkmeans.hpp | 61 ++++---- include/rabitqlib/core/data_layout.hpp | 14 +- include/rabitqlib/core/detail/rabitq_impl.hpp | 8 +- .../core/detail/simd/rotator_dispatch.hpp | 28 ++-- include/rabitqlib/core/estimator.hpp | 145 +++++++++--------- include/rabitqlib/core/rotator.hpp | 8 +- include/rabitqlib/fastscan/detail/encoder.hpp | 4 +- .../fastscan/detail/estimator_dispatch.hpp | 52 +++---- include/rabitqlib/fastscan/detail/simd.hpp | 8 +- include/rabitqlib/fastscan/estimator.hpp | 53 ++++--- include/rabitqlib/fastscan/fastscan.hpp | 2 +- .../rabitqlib/fastscan/highacc_fastscan.hpp | 2 +- python_bindings/__init__.pyi | 5 +- python_bindings/ivf_bindings.cpp | 5 +- src/fastscan/estimator_avx2.cpp | 28 ++-- src/fastscan/estimator_avx512.cpp | 20 ++- src/fastscan/estimator_generic.cpp | 28 ++-- src/fastscan/estimator_kernels.hpp | 41 ++--- src/fastscan/fastscan_avx2.cpp | 28 ++-- src/fastscan/fastscan_avx512.cpp | 22 +-- src/fastscan/fastscan_dispatch.cpp | 36 +++-- src/simd/dispatch.cpp | 6 +- src/simd/portable_generic.cpp | 34 ++-- src/simd/rotator_avx2.cpp | 6 +- src/simd/rotator_avx512.cpp | 6 +- src/simd/rotator_kernels.hpp | 30 ++-- src/simd/rotator_neon.cpp | 30 ++-- 34 files changed, 519 insertions(+), 427 deletions(-) diff --git a/docs/development.md b/docs/development.md index 7412d483..6f6c7820 100644 --- a/docs/development.md +++ b/docs/development.md @@ -109,6 +109,30 @@ reranking. QG uses `quantization_bits()` with 0 for raw storage and 4/8 for pack vectors. Naming work does not change these values, allocation sizes, numeric operations, serialized field order, or ownership. +## Numerical module variable names + +Apply the same role-based naming to Core, FastScan, clustering and SIMD. Preserve +public and protected names, Python keywords, field order and numerical operations. + +- Use `query_wrapper` for prepared query objects. Retain the existing `G_*` + private correction-field names where they identify mathematical coefficients. +- Distinguish a scalar `estimated_distance` or `lower_bound` from batch outputs + `estimated_distances` and `lower_bounds`. Keep `ip_x0_qr` where it identifies the + mathematical intermediate used by refinement. +- Use `lookup_table` for FastScan table pointers and `accumulated_values` for + integer accumulation outputs. Retain `lut` in compact local kernels where its + role is clear. +- Clustering's private state uses `dim_`, `num_points_`, `num_clusters_`, + `num_workers_` and `data_`. Keep public `d`, `k`, `num_threads` and graph parameter + names compatible. Retain the existing private `niter_` graph iteration name. +- FHT backend arguments use `rotated_data`, `padded_dim`, `trunc_dim` and + `scale_factor` consistently across declarations, dispatch and implementations. + A lower-level kernel's `dim` still denotes that kernel's own input shape. + +Keep short mathematical symbols, loop indices and SIMD register names when their +meaning is local and clear. Naming cleanup must not alter formulas, arithmetic +order, packed layouts, ISA selection or runtime behavior. + ## 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 6538ca0a..a92cd2ae 100644 --- a/docs/docs/index/hnsw.md +++ b/docs/docs/index/hnsw.md @@ -282,10 +282,13 @@ Each indexed element is stored in the following layout: [edges] [cluster ID] (its high bit marks a removed point, see Removing points) [external label] -[BinData (1-bit * dim + factors)] -[ExData (ex-bits * dim + factors)] +[BinData (1-bit * padded_dim + factors)] +[ExData (ex-bits * padded_dim + factors)] ``` +Here `padded_dim` is the stored rotated dimension, which may exceed the original +input `dim`. Loading preserves the index's saved padded dimension. + ### Search algorithm We first pre-process the query: diff --git a/docs/docs/index/ivf.md b/docs/docs/index/ivf.md index 804009dd..f43dea45 100644 --- a/docs/docs/index/ivf.md +++ b/docs/docs/index/ivf.md @@ -415,10 +415,10 @@ Extra-bit layouts retain full 64-coordinate blocks and store any final IVF-RaBitQ stores the following data: ```c++ -[batch data] // 1-bit code and factors -[ex_data] // code for remaining bits, or original float32 vectors -[ids] // PID of vectors (organized by clusters) -[cluster_lst] // List of clusters' metadata in IVF-RaBitQ +[batch data] // 1-bit code and factors +[ex_data] // code for remaining bits, or original float32 vectors +[ids] // PID of vectors (organized by clusters) +[cluster metadata] // List of clusters' metadata in IVF-RaBitQ ``` In raw mode, `ex_data` holds `4 * num_points * dim` bytes in cluster order, without padding or extra-bit factors. The one-bit filtering codes and factors diff --git a/docs/docs/rabitq/estimator.md b/docs/docs/rabitq/estimator.md index b06a80eb..048e7a5e 100644 --- a/docs/docs/rabitq/estimator.md +++ b/docs/docs/rabitq/estimator.md @@ -151,15 +151,15 @@ when $B=1$. // Compute the estimated distance // Note that G_add is dependent on the centroid. -float est_dist = F_add + G_add + F_rescale * (ip + G_kBxSumq); +float estimated_distance = F_add + G_add + F_rescale * (ip + G_kBxSumq); // Compute the error bound // Note that G_error is dependent on the centroid. float error_bound = F_error * G_error; // Compute the lower and upper bounds of the estimated distance -float lb_dist = est_dist - error_bound; -float ub_dist = est_dist + error_bound; +float lb_dist = estimated_distance - error_bound; +float ub_dist = estimated_distance + error_bound; ``` ### Incremental distance estimation @@ -187,10 +187,10 @@ that additional error. ```cpp // 1-bit dist -float est_dist = F_add + G_add + F_rescale * (ip_bin + G_k1xSumq); +float estimated_distance = F_add + G_add + F_rescale * (ip_bin + G_k1xSumq); float bound = F_error * G_error; -float ub_dist = est_dist + bound; -float lb_dist = est_dist - bound; +float ub_dist = estimated_distance + bound; +float lb_dist = estimated_distance - bound; // boost to full-bit dist float ex_scale = static_cast(1 << (bits - 1)); diff --git a/docs/docs/rabitq/quantizer.md b/docs/docs/rabitq/quantizer.md index 422b12ce..e0e31d59 100644 --- a/docs/docs/rabitq/quantizer.md +++ b/docs/docs/rabitq/quantizer.md @@ -187,7 +187,7 @@ Insert the following before `return 0` in the indexing example above. float k1xsumq = c_1 * std::accumulate(query.begin(), query.end(), 0.F); float g_add = rabitqlib::euclidean_sqr(query.data(), centroid.data(), dim); - float est_dist = rabitqlib::quant::full_est_dist( + float estimated_distance = rabitqlib::quant::full_est_dist( code.data(), query.data(), rabitqlib::excode_ipimpl::ip_fxi, @@ -200,7 +200,7 @@ Insert the following before `return 0` in the indexing example above. ); float gt_dist = rabitqlib::euclidean_sqr(query.data(), vector.data(), dim); - std::cout << "Estimated distance: " << est_dist << '\n'; + std::cout << "Estimated distance: " << estimated_distance << '\n'; std::cout << "GT distance: " << gt_dist << '\n'; ``` @@ -261,22 +261,22 @@ initialized `batch_data`, `centroid`, `dim`, and `batch_size`. std::vector query(dim, 1.0F); - rabitqlib::BatchQuery processed_query(query.data(), dim); + rabitqlib::BatchQuery query_wrapper(query.data(), dim); // G_add is the squared L2 distance between the query and the centroid. // This standalone example omits rotation, so both operands are raw vectors. // In an index that applies rotation, rotate both with the same transformation. - processed_query.set_g_add( + query_wrapper.set_g_add( rabitqlib::euclidean_sqr(query.data(), centroid.data(), dim) ); - std::vector est_distance(batch_size); // store the estimated distances + std::vector estimated_distances(batch_size); // store the estimated distances // QG scoring returns estimates only. Use the IVF-RaBitQ batch format below // when error intervals are needed. rabitqlib::qg_batch_estdist( - batch_data.data(), processed_query, dim, est_distance.data() + batch_data.data(), query_wrapper, dim, estimated_distances.data() ); ``` @@ -361,7 +361,7 @@ Insert this before `return 0` in the indexing example, reusing its data and dime dim, rabitqlib::SplitSingleQuery::kNumBits ); - rabitqlib::SplitSingleQuery processed_query( + rabitqlib::SplitSingleQuery query_wrapper( query.data(), dim, bits - 1, query_config, rabitqlib::METRIC_L2 ); @@ -372,16 +372,16 @@ Insert this before `return 0` in the indexing example, reusing its data and dime // Compute estimated distances based on binary codes float ip_x0_qr; - float est_dist; - float low_dist; + float estimated_distance; + float lower_bound; rabitqlib::split_single_estdist( bin_data.data(), - processed_query, + query_wrapper, dim, ip_x0_qr, - est_dist, - low_dist, + estimated_distance, + lower_bound, distance_sqr, distance_norm ); @@ -398,7 +398,7 @@ Insert this before `return 0` in the indexing example, reusing its data and dime bin_data.data(), ex_data.data(), ip_func, - processed_query, + query_wrapper, dim, bits - 1, est_dist_ex, @@ -498,7 +498,7 @@ Insert this before `return 0` in the indexing example, reusing its buffers and d // IVF-RaBitQ instead selects HACC for 4-9 total bits and standard FastScan for // 1-3 total bits or raw-vector storage, unless the caller overrides it. - rabitqlib::SplitBatchQuery processed_query( + rabitqlib::SplitBatchQuery query_wrapper( query.data(), dim, bits - 1, rabitqlib::METRIC_L2, true ); @@ -506,10 +506,10 @@ Insert this before `return 0` in the indexing example, reusing its buffers and d // This example omits rotation, so query and centroid are both raw vectors. float distance_norm = std::sqrt(rabitqlib::euclidean_sqr(query.data(), centroid.data(), dim)); - processed_query.set_g_add(distance_norm); + query_wrapper.set_g_add(distance_norm); - std::vector est_distance(batch_size); // store the estimated distances - std::vector low_distance(batch_size); // store the lower bound + std::vector estimated_distances(batch_size); // store the estimated distances + std::vector lower_bounds(batch_size); // store the lower bound std::vector ip_x0_qr(batch_size ); // store the intermediate result of inner product @@ -517,10 +517,10 @@ Insert this before `return 0` in the indexing example, reusing its buffers and d rabitqlib::split_batch_estdist( batch_data.data(), - processed_query, + query_wrapper, dim, - est_distance.data(), - low_distance.data(), + estimated_distances.data(), + lower_bounds.data(), ip_x0_qr.data(), true ); @@ -530,10 +530,10 @@ Insert this before `return 0` in the indexing example, reusing its buffers and d size_t i = 15; // Replace this candidate's one-bit estimate with the refined distance. - est_distance[i] = rabitqlib::split_distance_boosting( + estimated_distances[i] = rabitqlib::split_distance_boosting( ex_data.data() + (i * rabitqlib::ExDataMap::data_bytes(dim, bits - 1)), ip_func, - processed_query, + query_wrapper, dim, bits - 1, ip_x0_qr[i] diff --git a/include/rabitqlib/clustering/detail/kmeans.hpp b/include/rabitqlib/clustering/detail/kmeans.hpp index d92bf862..5cde801b 100644 --- a/include/rabitqlib/clustering/detail/kmeans.hpp +++ b/include/rabitqlib/clustering/detail/kmeans.hpp @@ -81,12 +81,12 @@ inline void exact_assign( : 64; constexpr size_t kCentroidBlock = 256; const size_t blocks = 1 + (n - 1) / point_block; - const auto threads = static_cast(std::min(num_threads, blocks)); + const auto num_workers = static_cast(std::min(num_threads, blocks)); std::vector centroid_norms(k); std::vector> products( - threads, std::vector(point_block * std::min(k, kCentroidBlock)) + num_workers, std::vector(point_block * std::min(k, kCentroidBlock)) ); - std::vector errors(threads); + std::vector errors(num_workers); const auto exact_distance = [d, spherical](const float* vector, const float* centroid) { double distance = 0.0; @@ -112,7 +112,7 @@ inline void exact_assign( const double error_floor = 8.0 * static_cast(d) * std::numeric_limits::min() + 8.0 * std::numeric_limits::epsilon(); -#pragma omp parallel num_threads(threads) +#pragma omp parallel num_threads(num_workers) { // Set only this OpenMP task's ICV, not Eigen's process-global thread limit. // This also prevents nested GEMM teams when the outer team has one worker. @@ -244,8 +244,8 @@ class LloydKMeans : public Parameters { std::string(Assignment::kName) + " input dimensions overflow size_t" ); } - const uint32_t threads = effective_num_threads(); - static_cast(threads); + const uint32_t num_workers = effective_num_threads(); + static_cast(num_workers); // Leave headroom for products of squared distances in graph pruning, // as well as quantized reconstruction and rotated residuals. const double max_coordinate = std::sqrt( @@ -254,7 +254,7 @@ class LloydKMeans : public Parameters { ); bool in_range = true; bool finite = true; -#pragma omp parallel for num_threads(threads) reduction(&& : finite, in_range) \ +#pragma omp parallel for num_threads(num_workers) reduction(&& : finite, in_range) \ schedule(static) for (std::ptrdiff_t loop_index = 0; loop_index < static_cast(n * d); ++loop_index) { @@ -289,7 +289,7 @@ class LloydKMeans : public Parameters { x = input_copy.data(); } - Assignment assigner(d, k, x, n, threads, *this); + Assignment assigner(d, k, x, n, num_workers, *this); centroids.resize(k * d); assignments.resize(n); @@ -304,7 +304,7 @@ class LloydKMeans : public Parameters { std::vector previous_centroids(centroids.size()); std::vector cluster_sizes(k); std::vector sums(centroids.size()); - const size_t update_workers = std::min(k, static_cast(threads)); + const size_t update_workers = std::min(k, static_cast(num_workers)); std::vector update_partitions(update_workers); std::vector update_worker_for_cluster(k); for (size_t worker = 0; worker < update_workers; ++worker) { @@ -360,7 +360,7 @@ class LloydKMeans : public Parameters { d, k, spherical, - threads, + num_workers, assignments.data(), distances.data(), assignments.data() @@ -447,9 +447,9 @@ class LloydKMeans : public Parameters { for (size_t point = 0; point < n; ++point) { partitions[worker_for_cluster[assignments[point]]].points.push_back(point); } - const uint32_t threads = effective_num_threads(); - static_cast(threads); -#pragma omp parallel for num_threads(threads) schedule(static) + const uint32_t num_workers = effective_num_threads(); + static_cast(num_workers); +#pragma omp parallel for num_threads(num_workers) schedule(static) for (std::ptrdiff_t loop_worker = 0; loop_worker < static_cast(partitions.size()); ++loop_worker) { @@ -485,8 +485,8 @@ class LloydKMeans : public Parameters { std::vector& cluster_sizes, const std::vector& previous_centroids ) { - const uint32_t threads = effective_num_threads(); - static_cast(threads); + const uint32_t num_workers = effective_num_threads(); + static_cast(num_workers); // Means, normalization, and displacement share one worker team. Preserve // the per-centroid arithmetic; empty rows are measured after refilling. const auto cluster_shift = [&](size_t cluster) { @@ -500,7 +500,7 @@ class LloydKMeans : public Parameters { return result; }; double shift = 0; -#pragma omp parallel for num_threads(threads) reduction(+ : shift) +#pragma omp parallel for num_threads(num_workers) reduction(+ : shift) for (std::ptrdiff_t loop_index = 0; loop_index < static_cast(k); ++loop_index) { const size_t cluster = static_cast(loop_index); @@ -548,8 +548,8 @@ class LloydKMeans : public Parameters { }; if (nsplit == 1) { // A single empty centroid needs only one eligible maximum per worker. - std::vector best_points(threads, n); -#pragma omp parallel num_threads(threads) + std::vector best_points(num_workers, n); +#pragma omp parallel num_threads(num_workers) { size_t best = n; #pragma omp for schedule(static) @@ -584,7 +584,7 @@ class LloydKMeans : public Parameters { return {1, shift}; } -#pragma omp parallel for num_threads(threads) schedule(static) +#pragma omp parallel for num_threads(num_workers) schedule(static) for (std::ptrdiff_t loop_index = 0; loop_index < static_cast(n); ++loop_index) { const size_t point = static_cast(loop_index); @@ -624,9 +624,9 @@ class LloydKMeans : public Parameters { } void normalize_centroids() { - const uint32_t threads = effective_num_threads(); - static_cast(threads); -#pragma omp parallel for num_threads(threads) + const uint32_t num_workers = effective_num_threads(); + static_cast(num_workers); +#pragma omp parallel for num_threads(num_workers) for (std::ptrdiff_t loop_index = 0; loop_index < static_cast(k); ++loop_index) { const size_t cluster = static_cast(loop_index); diff --git a/include/rabitqlib/clustering/qgkmeans.hpp b/include/rabitqlib/clustering/qgkmeans.hpp index d008d759..c25c21e8 100644 --- a/include/rabitqlib/clustering/qgkmeans.hpp +++ b/include/rabitqlib/clustering/qgkmeans.hpp @@ -140,27 +140,33 @@ class QGAssigner { MetricType metric_type = METRIC_L2, uint32_t quantization_bits = 0 ) - : d_(d) - , k_(k) - , M_(M) + : dim_(d) + , num_clusters_(k) + , max_degree_(M) , ef_construction_(ef_construction) , ef_search_(ef_search) , niter_(niter) - , num_threads_(num_threads) + , num_workers_(num_threads) , seed_(seed) { validate_configuration(); - num_threads_ = static_cast(resolve_num_threads(num_threads_)); + num_workers_ = static_cast(resolve_num_threads(num_workers_)); graph_ = std::make_unique>( - k_, d_, M_, metric_type, RotatorType::FhtKacRotator, quantization_bits, seed_ + num_clusters_, + dim_, + max_degree_, + metric_type, + RotatorType::FhtKacRotator, + quantization_bits, + seed_ ); graph_->set_ef(ef_search_); - scratch_.reserve(num_threads_); - for (uint32_t thread = 0; thread < num_threads_; ++thread) { + scratch_.reserve(num_workers_); + for (uint32_t thread = 0; thread < num_workers_; ++thread) { scratch_.emplace_back( std::make_unique( - graph_->padded_dim(), M_, ef_search_, 1 + graph_->padded_dim(), max_degree_, ef_search_, 1 ) ); } @@ -192,7 +198,8 @@ class QGAssigner { // Amortize their short phases without reducing assignment parallelism. constexpr size_t kMinVerticesPerBuildThread = 128; const size_t build_threads = std::min( - num_threads_, std::max(1, k_ / kMinVerticesPerBuildThread) + num_workers_, + std::max(1, num_clusters_ / kMinVerticesPerBuildThread) ); builder_ = std::make_unique( *graph_, @@ -226,7 +233,7 @@ class QGAssigner { ); } - if (n > static_cast(std::numeric_limits::max()) / d_) { + if (n > static_cast(std::numeric_limits::max()) / dim_) { throw std::invalid_argument("QGAssigner input dimensions overflow ptrdiff_t"); } @@ -239,7 +246,7 @@ class QGAssigner { std::vector query_order; if (previous_labels_.size() == n && n <= std::numeric_limits::max()) { query_order.resize(n); - std::vector offsets(k_ + 1, 0); + std::vector offsets(num_clusters_ + 1, 0); for (PID label : previous_labels_) { ++offsets[label + 1]; } @@ -249,8 +256,8 @@ class QGAssigner { static_cast(point); } } - std::vector errors(num_threads_); -#pragma omp parallel for num_threads(num_threads_) schedule(dynamic, 64) + std::vector errors(num_workers_); +#pragma omp parallel for num_threads(num_workers_) schedule(dynamic, 64) for (std::ptrdiff_t loop_index = 0; loop_index < static_cast(n); ++loop_index) { const size_t i = query_order.empty() @@ -266,18 +273,19 @@ class QGAssigner { thread_local VisitedSet visited; thread_local size_t visited_size = 0; thread_local size_t visited_capacity = 0; - const size_t capacity = - std::min(static_cast(ef_search_) * ef_search_, k_ / 10); - if (visited_size != k_ || visited_capacity != capacity) { - visited.initialize(k_, capacity); - visited_size = k_; + const size_t capacity = std::min( + static_cast(ef_search_) * ef_search_, num_clusters_ / 10 + ); + if (visited_size != num_clusters_ || visited_capacity != capacity) { + visited.initialize(num_clusters_, capacity); + visited_size = num_clusters_; visited_capacity = capacity; } const PID hint = previous_labels_.size() == n ? previous_labels_[i] : kPidMax; symqg::detail::search_with_workspace( *graph_, - x + i * d_, + x + i * dim_, 1, labels + i, distances + i, @@ -290,7 +298,7 @@ class QGAssigner { // Keep objectives and returned distances relative to the original // centroids; graph search scores are quantized estimates. distances[i] = - distance_func(x + i * d_, centroids + labels[i] * d_, d_); + distance_func(x + i * dim_, centroids + labels[i] * dim_, dim_); } } catch (...) { errors[static_cast(thread_id)] = std::current_exception(); @@ -308,20 +316,20 @@ class QGAssigner { } void validate_configuration() const { - if (num_threads_ > static_cast(std::numeric_limits::max())) { + if (num_workers_ > static_cast(std::numeric_limits::max())) { throw std::invalid_argument("QGAssigner num_threads must fit in an OpenMP int"); } - if (d_ < 64 || d_ > rotator_impl::FhtKacRotator::kMaxDim) { + if (dim_ < 64 || dim_ > rotator_impl::FhtKacRotator::kMaxDim) { throw std::invalid_argument( "QGKMeans currently requires dimensions in the range [64, 65536]" ); } - if (M_ == 0 || M_ % fastscan::kBatchSize != 0) { + if (max_degree_ == 0 || max_degree_ % fastscan::kBatchSize != 0) { throw std::invalid_argument( "QGKMeans graph degree must be a positive multiple of 32" ); } - if (k_ <= M_) { + if (num_clusters_ <= max_degree_) { throw std::invalid_argument( "QGKMeans requires more centroids than the graph degree" ); @@ -334,13 +342,13 @@ class QGAssigner { } } - size_t d_; - size_t k_; - uint32_t M_; + size_t dim_; + size_t num_clusters_; + uint32_t max_degree_; uint32_t ef_construction_; uint32_t ef_search_; uint32_t niter_; - uint32_t num_threads_; + uint32_t num_workers_; uint32_t seed_; std::unique_ptr> graph_; std::unique_ptr builder_; @@ -352,10 +360,10 @@ namespace detail { // Retain graph-specific state for one synchronous training call. class QGAssignment { - size_t d_; - const float* x_; - size_t n_; - uint32_t num_threads_; + size_t dim_; + const float* data_; + size_t num_points_; + uint32_t num_workers_; MetricType metric_type_; bool quantized_; QGAssigner assigner_; @@ -379,10 +387,10 @@ class QGAssignment { uint32_t num_threads, const QGKMeansParameters& cp ) - : d_(d) - , x_(x) - , n_(n) - , num_threads_(num_threads) + : dim_(d) + , data_(x) + , num_points_(n) + , num_workers_(num_threads) , metric_type_(cp.spherical ? METRIC_IP : METRIC_L2) , quantized_(cp.quantization_bits != 0) , assigner_( @@ -400,12 +408,14 @@ class QGAssignment { , query_sums_(assigner_.graph_->padded_dim(), cp) {} void assign(const float* centroids, PID* labels, float* distances, bool final = false) { - query_sums_.prepare(n_); + query_sums_.prepare(num_points_); std::vector previous_labels; if (final && quantized_) { - previous_labels.assign(labels, labels + n_); + previous_labels.assign(labels, labels + num_points_); } - assigner_.assign_with_cache(centroids, x_, n_, labels, distances, query_sums_); + assigner_.assign_with_cache( + centroids, data_, num_points_, labels, distances, query_sums_ + ); if (previous_labels.empty()) { return; } @@ -414,8 +424,9 @@ class QGAssignment { // Lloyd iterations: uphill approximate moves can improve final centroids. const auto distance_func = metric_type_ == METRIC_IP ? dot_product_dis : euclidean_sqr; -#pragma omp parallel for num_threads(num_threads_) schedule(static) - for (std::ptrdiff_t loop_index = 0; loop_index < static_cast(n_); +#pragma omp parallel for num_threads(num_workers_) schedule(static) + for (std::ptrdiff_t loop_index = 0; + loop_index < static_cast(num_points_); ++loop_index) { const size_t point = static_cast(loop_index); const PID previous = previous_labels[point]; @@ -423,7 +434,7 @@ class QGAssignment { continue; } const float distance = - distance_func(x_ + point * d_, centroids + previous * d_, d_); + distance_func(data_ + point * dim_, centroids + previous * dim_, dim_); if (distance < distances[point]) { labels[point] = previous; distances[point] = distance; diff --git a/include/rabitqlib/clustering/rabitqkmeans.hpp b/include/rabitqlib/clustering/rabitqkmeans.hpp index 5717b694..5f4d4296 100644 --- a/include/rabitqlib/clustering/rabitqkmeans.hpp +++ b/include/rabitqlib/clustering/rabitqkmeans.hpp @@ -47,12 +47,12 @@ namespace detail { // Encode immutable training points once; each moving centroid supplies one LUT. // RaBitQKMeans validates inputs and keeps the borrowed x buffer alive through training. class RaBitQAssigner { - size_t d_; - size_t k_; - const float* x_; - size_t n_; + size_t dim_; + size_t num_clusters_; + const float* data_; + size_t num_points_; size_t num_batches_; - uint32_t num_threads_; + uint32_t num_workers_; MetricType metric_type_; size_t padded_dim_; rotator_impl::FhtKacRotator rotator_; @@ -81,12 +81,12 @@ class RaBitQAssigner { uint32_t num_threads, const RaBitQKMeansParameters& cp ) - : d_(d) - , k_(k) - , x_(x) - , n_(n) + : dim_(d) + , num_clusters_(k) + , data_(x) + , num_points_(n) , num_batches_(1 + (n - 1) / fastscan::kBatchSize) - , num_threads_(static_cast(std::min(num_threads, num_batches_))) + , num_workers_(static_cast(std::min(num_threads, num_batches_))) , metric_type_(cp.spherical ? METRIC_IP : METRIC_L2) , padded_dim_(rotator_impl::padding_requirement(d, RotatorType::FhtKacRotator)) , rotator_(d, padded_dim_, cp.seed) @@ -97,36 +97,37 @@ class RaBitQAssigner { , lookup_table_(padded_dim_ * 4) { { std::vector> partial_sums( - num_threads_, std::vector(d_, 0.0) + num_workers_, std::vector(dim_, 0.0) ); -#pragma omp parallel num_threads(num_threads_) +#pragma omp parallel num_threads(num_workers_) { auto& sums = partial_sums[static_cast(omp_get_thread_num())]; #pragma omp for schedule(static) - for (std::ptrdiff_t point = 0; point < static_cast(n_); + for (std::ptrdiff_t point = 0; + point < static_cast(num_points_); ++point) { - const float* vector = x_ + static_cast(point) * d_; - for (size_t dim = 0; dim < d_; ++dim) { + const float* vector = data_ + static_cast(point) * dim_; + for (size_t dim = 0; dim < dim_; ++dim) { sums[dim] += static_cast(vector[dim]); } } } - for (size_t dim = 0; dim < d_; ++dim) { + for (size_t dim = 0; dim < dim_; ++dim) { double sum = 0.0; for (const auto& partial : partial_sums) { sum += partial[dim]; } - center_[dim] = static_cast(sum / static_cast(n_)); + center_[dim] = static_cast(sum / static_cast(num_points_)); } } std::vector rotated_center(padded_dim_); rotator_.rotate(center_.data(), rotated_center.data()); std::vector> rotated_blocks( - num_threads_, std::vector(fastscan::kBatchSize * padded_dim_) + num_workers_, std::vector(fastscan::kBatchSize * padded_dim_) ); - std::vector errors(num_threads_); -#pragma omp parallel for num_threads(num_threads_) schedule(static) + std::vector errors(num_workers_); +#pragma omp parallel for num_threads(num_workers_) schedule(static) for (std::ptrdiff_t batch = 0; batch < static_cast(num_batches_); ++batch) { const auto thread = static_cast(omp_get_thread_num()); @@ -136,11 +137,11 @@ class RaBitQAssigner { try { const auto block = static_cast(batch); const size_t first = block * fastscan::kBatchSize; - const size_t count = std::min(fastscan::kBatchSize, n_ - first); + const size_t count = std::min(fastscan::kBatchSize, num_points_ - first); auto& rotated = rotated_blocks[thread]; for (size_t point = 0; point < count; ++point) { rotator_.rotate( - x_ + (first + point) * d_, rotated.data() + point * padded_dim_ + data_ + (first + point) * dim_, rotated.data() + point * padded_dim_ ); } quant::quantize_qg_batch( @@ -164,18 +165,18 @@ class RaBitQAssigner { void assign(const float* centroids, PID* labels, float* distances, bool = false) { const auto distance_func = metric_type_ == METRIC_IP ? dot_product_dis : euclidean_sqr; - for (size_t cluster = 0; cluster < k_; ++cluster) { - const float* centroid = centroids + cluster * d_; + for (size_t cluster = 0; cluster < num_clusters_; ++cluster) { + const float* centroid = centroids + cluster * dim_; rotator_.rotate(centroid, rotated_query_.data()); auto& query = centroid_queries_[cluster]; query.reset( rotated_query_.data(), padded_dim_, lookup_table_.data(), metric_type_ ); - query.set_g_add(distance_func(centroid, center_.data(), d_)); + query.set_g_add(distance_func(centroid, center_.data(), dim_)); } - std::vector errors(num_threads_); -#pragma omp parallel for num_threads(num_threads_) schedule(static) + std::vector errors(num_workers_); +#pragma omp parallel for num_threads(num_workers_) schedule(static) for (std::ptrdiff_t batch = 0; batch < static_cast(num_batches_); ++batch) { const auto thread = static_cast(omp_get_thread_num()); @@ -185,14 +186,14 @@ class RaBitQAssigner { try { const auto block = static_cast(batch); const size_t first = block * fastscan::kBatchSize; - const size_t count = std::min(fastscan::kBatchSize, n_ - first); + const size_t count = std::min(fastscan::kBatchSize, num_points_ - first); const char* codes = batch_data_.data() + block * QGBatchDataMap::data_bytes(padded_dim_); std::array estimated_distances{}; std::array best_distances; best_distances.fill(std::numeric_limits::infinity()); std::array best_labels{}; - for (size_t cluster = 0; cluster < k_; ++cluster) { + for (size_t cluster = 0; cluster < num_clusters_; ++cluster) { qg_batch_estdist( codes, centroid_queries_[cluster], @@ -214,7 +215,7 @@ class RaBitQAssigner { // Scores choose the centroid; objectives retain original-space // distances. distances[id] = distance_func( - x_ + id * d_, centroids + best_labels[point] * d_, d_ + data_ + id * dim_, centroids + best_labels[point] * dim_, dim_ ); } } catch (...) { errors[thread] = std::current_exception(); } diff --git a/include/rabitqlib/core/data_layout.hpp b/include/rabitqlib/core/data_layout.hpp index 005a41da..8a1afa30 100644 --- a/include/rabitqlib/core/data_layout.hpp +++ b/include/rabitqlib/core/data_layout.hpp @@ -29,7 +29,8 @@ struct ExDataMap { explicit ExDataMap(char* data, size_t padded_dim, size_t ex_bits) : ex_code_(reinterpret_cast(data)) , f_add_ex_(data + detail::SingleCodeOffsets::add(padded_dim, ex_bits)) - , f_recale_ex_(data + detail::SingleCodeOffsets::rescale(padded_dim, ex_bits)) {} + , f_rescale_ex_(data + detail::SingleCodeOffsets::rescale(padded_dim, ex_bits)) { + } static size_t data_bytes(size_t padded_dim, size_t ex_bits) { if (ex_bits == 0) { @@ -49,12 +50,12 @@ struct ExDataMap { [[nodiscard]] uint8_t* ex_code() { return ex_code_; } [[nodiscard]] detail::PackedValueRef f_add_ex() { return f_add_ex_; } - [[nodiscard]] detail::PackedValueRef f_rescale_ex() { return f_recale_ex_; } + [[nodiscard]] detail::PackedValueRef f_rescale_ex() { return f_rescale_ex_; } private: uint8_t* ex_code_; detail::PackedValueRef f_add_ex_; - detail::PackedValueRef f_recale_ex_; + detail::PackedValueRef f_rescale_ex_; }; template @@ -63,18 +64,19 @@ struct ConstExDataMap { explicit ConstExDataMap(const char* data, size_t padded_dim, size_t ex_bits) : ex_code_(reinterpret_cast(data)) , f_add_ex_(data + detail::SingleCodeOffsets::add(padded_dim, ex_bits)) - , f_recale_ex_(data + detail::SingleCodeOffsets::rescale(padded_dim, ex_bits)) {} + , f_rescale_ex_(data + detail::SingleCodeOffsets::rescale(padded_dim, ex_bits)) { + } [[nodiscard]] const uint8_t* ex_code() const { return ex_code_; } [[nodiscard]] T f_add_ex() const { return detail::load_packed_value(f_add_ex_); } [[nodiscard]] T f_rescale_ex() const { - return detail::load_packed_value(f_recale_ex_); + return detail::load_packed_value(f_rescale_ex_); } private: const uint8_t* ex_code_; const char* f_add_ex_; - const char* f_recale_ex_; + const char* f_rescale_ex_; }; template diff --git a/include/rabitqlib/core/detail/rabitq_impl.hpp b/include/rabitqlib/core/detail/rabitq_impl.hpp index 1992291d..ef49bc64 100644 --- a/include/rabitqlib/core/detail/rabitq_impl.hpp +++ b/include/rabitqlib/core/detail/rabitq_impl.hpp @@ -225,7 +225,7 @@ inline void one_bit_compact_code_to_bytes( size_t padded_dim, uint8_t* compact_code, T& f_add, - T& f_recale, + T& f_rescale, T& f_error, MetricType metric_type = METRIC_L2 ) { @@ -237,7 +237,7 @@ inline void one_bit_compact_code_to_bytes( padded_dim, binary_code.data(), f_add, - f_recale, + f_rescale, f_error, metric_type ); @@ -252,7 +252,7 @@ inline void one_bit_compact_code( size_t padded_dim, TC* compact_code, T& f_add, - T& f_recale, + T& f_rescale, T& f_error, MetricType metric_type = METRIC_L2 ) { @@ -262,7 +262,7 @@ inline void one_bit_compact_code( padded_dim, reinterpret_cast(compact_code), f_add, - f_recale, + f_rescale, f_error, metric_type ); diff --git a/include/rabitqlib/core/detail/simd/rotator_dispatch.hpp b/include/rabitqlib/core/detail/simd/rotator_dispatch.hpp index f6bc0db7..02d17c08 100644 --- a/include/rabitqlib/core/detail/simd/rotator_dispatch.hpp +++ b/include/rabitqlib/core/detail/simd/rotator_dispatch.hpp @@ -11,11 +11,11 @@ void kacs_walk(float* data, size_t len); void fht_rotate( const float* data, - float* rotated_vec, + float* rotated_data, size_t dim, size_t padded_dim, size_t trunc_dim, - float fac, + float scale_factor, const uint8_t* flip ); @@ -23,11 +23,11 @@ void flip_sign_neon(const uint8_t* flip, float* data, size_t dim); void kacs_walk_neon(float* data, size_t len); void fht_rotate_neon( const float* data, - float* output, + float* rotated_data, size_t dim, - size_t padded, - size_t trunc, - float factor, + size_t padded_dim, + size_t trunc_dim, + float scale_factor, const uint8_t* flip ); @@ -39,20 +39,20 @@ void kacs_walk_avx2(float* data, size_t len); void fht_rotate_generic( const float* data, - float* output, + float* rotated_data, size_t dim, - size_t padded, - size_t trunc, - float factor, + size_t padded_dim, + size_t trunc_dim, + float scale_factor, const uint8_t* flip ); void fht_rotate_avx2( const float* data, - float* rotated_vec, + float* rotated_data, size_t dim, size_t padded_dim, size_t trunc_dim, - float fac, + float scale_factor, const uint8_t* flip ); @@ -62,11 +62,11 @@ void kacs_walk_avx512(float* data, size_t len); void fht_rotate_avx512( const float* data, - float* rotated_vec, + float* rotated_data, size_t dim, size_t padded_dim, size_t trunc_dim, - float fac, + float scale_factor, const uint8_t* flip ); diff --git a/include/rabitqlib/core/estimator.hpp b/include/rabitqlib/core/estimator.hpp index a3fbac35..fca698c9 100644 --- a/include/rabitqlib/core/estimator.hpp +++ b/include/rabitqlib/core/estimator.hpp @@ -31,18 +31,18 @@ template inline float split_distance_boosting( const char* ex_data, float (*ip_func)(const float*, const uint8_t*, size_t), - const Query& q_obj, + const Query& query_wrapper, size_t padded_dim, size_t ex_bits, float ip_x0_qr ) { ConstExDataMap cur_ex(ex_data, padded_dim, ex_bits); - float ex_dist = - cur_ex.f_add_ex() + q_obj.g_add() + - (cur_ex.f_rescale_ex() * - (static_cast(1 << ex_bits) * ip_x0_qr + - ip_func(q_obj.rotated_query(), cur_ex.ex_code(), padded_dim) + q_obj.kbxsumq())); + float ex_dist = cur_ex.f_add_ex() + query_wrapper.g_add() + + (cur_ex.f_rescale_ex() * + (static_cast(1 << ex_bits) * ip_x0_qr + + ip_func(query_wrapper.rotated_query(), cur_ex.ex_code(), padded_dim) + + query_wrapper.kbxsumq())); return ex_dist; } @@ -52,11 +52,11 @@ inline void split_single_fulldist_direct( const char* bin_data, const char* ex_data, float (*ip_func)(const float*, const uint8_t*, size_t), - const Query& q_obj, + const Query& query_wrapper, size_t padded_dim, size_t ex_bits, - float& est_dist, - float& low_dist, + float& estimated_distance, + float& lower_bound, float& ip_x0_qr, float g_add, float g_error @@ -67,29 +67,31 @@ inline void split_single_fulldist_direct( // Direct HNSW kernels retain their legacy word-typed ABI. Generic kernels use the // byte-oriented overload and do not manufacture a typed view of packed storage. ip_x0_qr = Kernel::mask_ip_x0_q( - q_obj.rotated_query(), + query_wrapper.rotated_query(), reinterpret_cast(cur_bin.bin_code()), padded_dim ); - est_dist = + estimated_distance = cur_ex.f_add_ex() + g_add + (cur_ex.f_rescale_ex() * (static_cast(1 << ex_bits) * ip_x0_qr + - ip_func(q_obj.rotated_query(), cur_ex.ex_code(), padded_dim) + q_obj.kbxsumq())); + ip_func(query_wrapper.rotated_query(), cur_ex.ex_code(), padded_dim) + + query_wrapper.kbxsumq())); - low_dist = est_dist - (cur_bin.f_error() * g_error / static_cast(1 << ex_bits)); + lower_bound = estimated_distance - + (cur_bin.f_error() * g_error / static_cast(1 << ex_bits)); } inline void split_single_fulldist( const char* bin_data, const char* ex_data, float (*ip_func)(const float*, const uint8_t*, size_t), - const SplitSingleQuery& q_obj, + const SplitSingleQuery& query_wrapper, size_t padded_dim, size_t ex_bits, - float& est_dist, - float& low_dist, + float& estimated_distance, + float& lower_bound, float& ip_x0_qr, float g_add, float g_error @@ -97,24 +99,26 @@ inline void split_single_fulldist( ConstBinDataMap cur_bin(bin_data, padded_dim); ConstExDataMap cur_ex(ex_data, padded_dim, ex_bits); - ip_x0_qr = mask_ip_x0_q(q_obj.rotated_query(), cur_bin.bin_code(), padded_dim); + ip_x0_qr = mask_ip_x0_q(query_wrapper.rotated_query(), cur_bin.bin_code(), padded_dim); - est_dist = + estimated_distance = cur_ex.f_add_ex() + g_add + (cur_ex.f_rescale_ex() * (static_cast(1 << ex_bits) * ip_x0_qr + - ip_func(q_obj.rotated_query(), cur_ex.ex_code(), padded_dim) + q_obj.kbxsumq())); + ip_func(query_wrapper.rotated_query(), cur_ex.ex_code(), padded_dim) + + query_wrapper.kbxsumq())); - low_dist = est_dist - (cur_bin.f_error() * g_error / static_cast(1 << ex_bits)); + lower_bound = estimated_distance - + (cur_bin.f_error() * g_error / static_cast(1 << ex_bits)); } inline void split_single_estdist( const char* bin_data, - const SplitSingleQuery& q_obj, + const SplitSingleQuery& query_wrapper, size_t padded_dim, float& ip_x0_qr, - float& est_dist, - float& low_dist, + float& estimated_distance, + float& lower_bound, float g_add = 0, float g_error = 0 ) { @@ -122,27 +126,27 @@ inline void split_single_estdist( ip_x0_qr = warmup_ip_x0_q_512( cur_bin.bin_code(), - q_obj.query_bin(), - q_obj.delta(), - q_obj.vl(), + query_wrapper.query_bin(), + query_wrapper.delta(), + query_wrapper.vl(), padded_dim, - q_obj.num_bits() + query_wrapper.num_bits() ); - est_dist = - cur_bin.f_add() + g_add + (cur_bin.f_rescale() * (ip_x0_qr + q_obj.k1xsumq())); + estimated_distance = cur_bin.f_add() + g_add + + (cur_bin.f_rescale() * (ip_x0_qr + query_wrapper.k1xsumq())); - low_dist = est_dist - (cur_bin.f_error() * g_error); + lower_bound = estimated_distance - (cur_bin.f_error() * g_error); }; template inline void split_single_estdist_direct( const char* bin_data, - const SplitSingleQuery& q_obj, + const SplitSingleQuery& query_wrapper, size_t padded_dim, float& ip_x0_qr, - float& est_dist, - float& low_dist, + float& estimated_distance, + float& lower_bound, float g_add = 0, float g_error = 0 ) { @@ -151,39 +155,39 @@ inline void split_single_estdist_direct( // Direct HNSW kernels retain their legacy word-typed ABI; see the note above. ip_x0_qr = Kernel::warmup_ip_x0_q_512( reinterpret_cast(cur_bin.bin_code()), - q_obj.query_bin(), - q_obj.delta(), - q_obj.vl(), + query_wrapper.query_bin(), + query_wrapper.delta(), + query_wrapper.vl(), padded_dim, - q_obj.num_bits() + query_wrapper.num_bits() ); - est_dist = - cur_bin.f_add() + g_add + (cur_bin.f_rescale() * (ip_x0_qr + q_obj.k1xsumq())); + estimated_distance = cur_bin.f_add() + g_add + + (cur_bin.f_rescale() * (ip_x0_qr + query_wrapper.k1xsumq())); - low_dist = est_dist - (cur_bin.f_error() * g_error); + lower_bound = estimated_distance - (cur_bin.f_error() * g_error); }; inline void xy_single_base_dist( const char* base_data, float (*base_ip_func)(const float*, const uint8_t*, size_t), - const SplitSingleQuery& q_obj, + const SplitSingleQuery& query_wrapper, size_t padded_dim, size_t base_bits, float& ip_base, - float& est_dist, - float& low_dist, + float& estimated_distance, + float& lower_bound, float g_add = 0, float g_error = 0 ) { ConstBaseDataMap cur_base(base_data, padded_dim, base_bits); - ip_base = base_ip_func(q_obj.rotated_query(), cur_base.base_code(), padded_dim); + ip_base = base_ip_func(query_wrapper.rotated_query(), cur_base.base_code(), padded_dim); - est_dist = - cur_base.f_add() + g_add + (cur_base.f_rescale() * (ip_base + q_obj.kbase_sumq())); + estimated_distance = cur_base.f_add() + g_add + + (cur_base.f_rescale() * (ip_base + query_wrapper.kbase_sumq())); - low_dist = est_dist - (cur_base.f_error() * g_error); + lower_bound = estimated_distance - (cur_base.f_error() * g_error); } inline void xy_single_full_dist( @@ -191,34 +195,37 @@ inline void xy_single_full_dist( const char* ex_data, float (*base_ip_func)(const float*, const uint8_t*, size_t), float (*ex_ip_func)(const float*, const uint8_t*, size_t), - const SplitSingleQuery& q_obj, + const SplitSingleQuery& query_wrapper, size_t padded_dim, size_t base_bits, size_t ex_bits, - float& est_dist, - float& low_dist, + float& estimated_distance, + float& lower_bound, float& ip_base, float g_add = 0, float g_error = 0 ) { ConstBaseDataMap cur_base(base_data, padded_dim, base_bits); - ip_base = base_ip_func(q_obj.rotated_query(), cur_base.base_code(), padded_dim); + ip_base = base_ip_func(query_wrapper.rotated_query(), cur_base.base_code(), padded_dim); if (ex_bits == 0) { - est_dist = cur_base.f_add() + g_add + - (cur_base.f_rescale() * (ip_base + q_obj.kbase_sumq())); - low_dist = est_dist - (cur_base.f_error() * g_error); + estimated_distance = + cur_base.f_add() + g_add + + (cur_base.f_rescale() * (ip_base + query_wrapper.kbase_sumq())); + lower_bound = estimated_distance - (cur_base.f_error() * g_error); return; } ConstExDataMap cur_ex(ex_data, padded_dim, ex_bits); float ip = (static_cast(1 << ex_bits) * ip_base) + - ex_ip_func(q_obj.rotated_query(), cur_ex.ex_code(), padded_dim); + ex_ip_func(query_wrapper.rotated_query(), cur_ex.ex_code(), padded_dim); - est_dist = cur_ex.f_add_ex() + g_add + (cur_ex.f_rescale_ex() * (ip + q_obj.kbxsumq())); + estimated_distance = cur_ex.f_add_ex() + g_add + + (cur_ex.f_rescale_ex() * (ip + query_wrapper.kbxsumq())); - low_dist = est_dist - (cur_base.f_error() * g_error / static_cast(1 << ex_bits)); + lower_bound = estimated_distance - + (cur_base.f_error() * g_error / static_cast(1 << ex_bits)); } // Named-result functions keep the numerical implementation and legacy output-reference @@ -228,7 +235,7 @@ inline detail::DistanceEstimate estimate_split_single_full_direct( const char* bin_data, const char* ex_data, float (*ip_func)(const float*, const uint8_t*, size_t), - const Query& q_obj, + const Query& query_wrapper, size_t padded_dim, size_t ex_bits, float g_add, @@ -239,7 +246,7 @@ inline detail::DistanceEstimate estimate_split_single_full_direct( bin_data, ex_data, ip_func, - q_obj, + query_wrapper, padded_dim, ex_bits, result.distance, @@ -255,7 +262,7 @@ inline detail::DistanceEstimate estimate_split_single_full( const char* bin_data, const char* ex_data, float (*ip_func)(const float*, const uint8_t*, size_t), - const SplitSingleQuery& q_obj, + const SplitSingleQuery& query_wrapper, size_t padded_dim, size_t ex_bits, float g_add, @@ -266,7 +273,7 @@ inline detail::DistanceEstimate estimate_split_single_full( bin_data, ex_data, ip_func, - q_obj, + query_wrapper, padded_dim, ex_bits, result.distance, @@ -280,7 +287,7 @@ inline detail::DistanceEstimate estimate_split_single_full( inline detail::DistanceEstimate estimate_split_single( const char* bin_data, - const SplitSingleQuery& q_obj, + const SplitSingleQuery& query_wrapper, size_t padded_dim, float g_add = 0, float g_error = 0 @@ -288,7 +295,7 @@ inline detail::DistanceEstimate estimate_split_single( detail::DistanceEstimate result; split_single_estdist( bin_data, - q_obj, + query_wrapper, padded_dim, result.inner_product, result.distance, @@ -302,7 +309,7 @@ inline detail::DistanceEstimate estimate_split_single( template inline detail::DistanceEstimate estimate_split_single_direct( const char* bin_data, - const SplitSingleQuery& q_obj, + const SplitSingleQuery& query_wrapper, size_t padded_dim, float g_add = 0, float g_error = 0 @@ -310,7 +317,7 @@ inline detail::DistanceEstimate estimate_split_single_direct( detail::DistanceEstimate result; split_single_estdist_direct( bin_data, - q_obj, + query_wrapper, padded_dim, result.inner_product, result.distance, @@ -324,7 +331,7 @@ inline detail::DistanceEstimate estimate_split_single_direct( inline detail::DistanceEstimate estimate_xy_base( const char* base_data, float (*base_ip_func)(const float*, const uint8_t*, size_t), - const SplitSingleQuery& q_obj, + const SplitSingleQuery& query_wrapper, size_t padded_dim, size_t base_bits, float g_add = 0, @@ -334,7 +341,7 @@ inline detail::DistanceEstimate estimate_xy_base( xy_single_base_dist( base_data, base_ip_func, - q_obj, + query_wrapper, padded_dim, base_bits, result.inner_product, @@ -351,7 +358,7 @@ inline detail::DistanceEstimate estimate_xy_full( const char* ex_data, float (*base_ip_func)(const float*, const uint8_t*, size_t), float (*ex_ip_func)(const float*, const uint8_t*, size_t), - const SplitSingleQuery& q_obj, + const SplitSingleQuery& query_wrapper, size_t padded_dim, size_t base_bits, size_t ex_bits, @@ -364,7 +371,7 @@ inline detail::DistanceEstimate estimate_xy_full( ex_data, base_ip_func, ex_ip_func, - q_obj, + query_wrapper, padded_dim, base_bits, ex_bits, diff --git a/include/rabitqlib/core/rotator.hpp b/include/rabitqlib/core/rotator.hpp index 5abaa5f2..0452d371 100644 --- a/include/rabitqlib/core/rotator.hpp +++ b/include/rabitqlib/core/rotator.hpp @@ -159,7 +159,7 @@ class FhtKacRotator : public Rotator { private: std::vector flip_; size_t trunc_dim_ = 0; - float fac_ = 0; + float scale_factor_ = 0; static constexpr size_t kByteLen = 8; @@ -185,7 +185,7 @@ class FhtKacRotator : public Rotator { flip_.resize((padded_dim / kByteLen) * 4); size_t bottom_log_dim = floor_log2(dim); trunc_dim_ = size_t{1} << bottom_log_dim; - fac_ = 1.0F / std::sqrt(static_cast(trunc_dim_)); + scale_factor_ = 1.0F / std::sqrt(static_cast(trunc_dim_)); } FhtKacRotator() = default; ~FhtKacRotator() override = default; @@ -222,7 +222,7 @@ class FhtKacRotator : public Rotator { this->dim_ = other.dim_; this->padded_dim_ = other.padded_dim_; trunc_dim_ = other.trunc_dim_; - fac_ = other.fac_; + scale_factor_ = other.scale_factor_; return *this; } @@ -230,7 +230,7 @@ class FhtKacRotator : public Rotator { void rotate(const float* data, float* rotated_vec) const override { simd::fht_rotate( - data, rotated_vec, dim_, padded_dim_, trunc_dim_, fac_, flip_.data() + data, rotated_vec, dim_, padded_dim_, trunc_dim_, scale_factor_, flip_.data() ); } }; diff --git a/include/rabitqlib/fastscan/detail/encoder.hpp b/include/rabitqlib/fastscan/detail/encoder.hpp index f91e74cd..f31768a0 100644 --- a/include/rabitqlib/fastscan/detail/encoder.hpp +++ b/include/rabitqlib/fastscan/detail/encoder.hpp @@ -25,7 +25,7 @@ inline void one_bit_batch_code( size_t padded_dim, uint8_t* packed_code, FAdd f_add, - FRescale f_recale, + FRescale f_rescale, FError f_error, MetricType metric_type = METRIC_L2 ) { @@ -38,7 +38,7 @@ inline void one_bit_batch_code( padded_dim, compact_codes.data(), f_add, - f_recale, + f_rescale, f_error, metric_type ); diff --git a/include/rabitqlib/fastscan/detail/estimator_dispatch.hpp b/include/rabitqlib/fastscan/detail/estimator_dispatch.hpp index 065b1ae5..cb34f043 100644 --- a/include/rabitqlib/fastscan/detail/estimator_dispatch.hpp +++ b/include/rabitqlib/fastscan/detail/estimator_dispatch.hpp @@ -12,67 +12,67 @@ class BatchQuery; namespace simd { void split_batch_estdist( const char* batch_data, - const SplitBatchQuery& q_obj, + const SplitBatchQuery& query_wrapper, size_t padded_dim, - float* est_distance, - float* low_distance, + float* estimated_distances, + float* lower_bounds, float* ip_x0_qr, bool use_hacc ); void split_batch_estdist_generic( const char* batch_data, - const SplitBatchQuery& q_obj, + const SplitBatchQuery& query_wrapper, size_t padded_dim, - float* est_distance, - float* low_distance, + float* estimated_distances, + float* lower_bounds, float* ip_x0_qr, bool use_hacc ); void split_batch_estdist_avx2( const char* batch_data, - const SplitBatchQuery& q_obj, + const SplitBatchQuery& query_wrapper, size_t padded_dim, - float* est_distance, - float* low_distance, + float* estimated_distances, + float* lower_bounds, float* ip_x0_qr, bool use_hacc ); void split_batch_estdist_avx512( const char* batch_data, - const SplitBatchQuery& q_obj, + const SplitBatchQuery& query_wrapper, size_t padded_dim, - float* est_distance, - float* low_distance, + float* estimated_distances, + float* lower_bounds, float* ip_x0_qr, bool use_hacc ); void qg_batch_estdist( const char* batch_data, - const BatchQuery& q_obj, + const BatchQuery& query_wrapper, size_t padded_dim, - float* est_distance + float* estimated_distances ); void qg_batch_estdist_generic( const char* batch_data, - const BatchQuery& q_obj, + const BatchQuery& query_wrapper, size_t padded_dim, - float* est_distance + float* estimated_distances ); void qg_batch_estdist_avx2( const char* batch_data, - const BatchQuery& q_obj, + const BatchQuery& query_wrapper, size_t padded_dim, - float* est_distance + float* estimated_distances ); void qg_batch_estdist_avx512( const char* batch_data, - const BatchQuery& q_obj, + const BatchQuery& query_wrapper, size_t padded_dim, - float* est_distance + float* estimated_distances ); uint32_t qg_candidate_mask_generic(const float* distances, float threshold); @@ -80,23 +80,23 @@ uint32_t qg_candidate_mask_avx2(const float* distances, float threshold); uint32_t qg_batch_estdist_mask( const char* batch_data, - const BatchQuery& q_obj, + const BatchQuery& query_wrapper, size_t padded_dim, - float* est_distance, + float* estimated_distances, float threshold ); uint32_t qg_batch_estdist_mask_generic( const char* batch_data, - const BatchQuery& q_obj, + const BatchQuery& query_wrapper, size_t padded_dim, - float* est_distance, + float* estimated_distances, float threshold ); uint32_t qg_batch_estdist_mask_avx2( const char* batch_data, - const BatchQuery& q_obj, + const BatchQuery& query_wrapper, size_t padded_dim, - float* est_distance, + float* estimated_distances, float threshold ); } // namespace simd diff --git a/include/rabitqlib/fastscan/detail/simd.hpp b/include/rabitqlib/fastscan/detail/simd.hpp index 7c61d264..eaa6ffbf 100644 --- a/include/rabitqlib/fastscan/detail/simd.hpp +++ b/include/rabitqlib/fastscan/detail/simd.hpp @@ -39,7 +39,7 @@ void accumulate_unsupported( ); void accumulate_avx2( const uint8_t* __restrict__ codes, - const uint8_t* __restrict__ lp_table, + const uint8_t* __restrict__ lookup_table, int32_t* __restrict__ result, size_t dim ); @@ -54,13 +54,13 @@ void accumulate_hacc_generic( void accumulate_hacc_avx2( const uint8_t* __restrict__ codes, const uint8_t* __restrict__ hc_lut, - int32_t* accu_res, + int32_t* accumulated_values, size_t dim ); void accumulate_avx512( const uint8_t* __restrict__ codes, - const uint8_t* __restrict__ lp_table, + const uint8_t* __restrict__ lookup_table, int32_t* __restrict__ result, size_t dim ); @@ -68,7 +68,7 @@ void transfer_lut_hacc_avx512(const uint16_t* lut, size_t dim, uint8_t* hc_lut); void accumulate_hacc_avx512( const uint8_t* __restrict__ codes, const uint8_t* __restrict__ hc_lut, - int32_t* accu_res, + int32_t* accumulated_values, size_t dim ); diff --git a/include/rabitqlib/fastscan/estimator.hpp b/include/rabitqlib/fastscan/estimator.hpp index f7ed71b5..9d96d3ce 100644 --- a/include/rabitqlib/fastscan/estimator.hpp +++ b/include/rabitqlib/fastscan/estimator.hpp @@ -18,13 +18,13 @@ namespace rabitqlib { * @brief Use FastScan to estimate batch distance * * @param batch_data codes and factors in fastscan/data_layout.hpp's BatchDataMap layout - * @param q_obj prepared query with the matching centroid correction already set + * @param query_wrapper prepared query with the matching centroid correction already set * @param padded_dim dim, must be multiple of 16 - * @param est_distance estimated distance - * @param low_distance lower endpoint of the probabilistic distance interval; + * @param estimated_distances estimated distance + * @param lower_bounds lower endpoint of the probabilistic distance interval; * not a deterministic bound for every query * @param ip_x0_qr intermediate result for re-ranking - * @param use_hacc must match the mode used to prepare q_obj + * @param use_hacc must match the mode used to prepare query_wrapper * All three output buffers need 32 floats, including for a partial batch. * Throws std::invalid_argument for an unprepared or mismatched query. * The margin is f_error * g_error; see docs/docs/rabitq/estimator.md, @@ -32,50 +32,61 @@ namespace rabitqlib { */ inline void split_batch_estdist( const char* batch_data, - const SplitBatchQuery& q_obj, + const SplitBatchQuery& query_wrapper, size_t padded_dim, - float* est_distance, - float* low_distance, + float* estimated_distances, + float* lower_bounds, float* ip_x0_qr, bool use_hacc ) { - if (padded_dim == 0 || q_obj.padded_dim() != padded_dim) { + if (padded_dim == 0 || query_wrapper.padded_dim() != padded_dim) { throw std::invalid_argument( "FastScan scoring requires a prepared query with matching dimension" ); } - if (q_obj.use_hacc() != use_hacc) { + if (query_wrapper.use_hacc() != use_hacc) { throw std::invalid_argument("FastScan scoring mode must match the prepared query"); } simd::split_batch_estdist( - batch_data, q_obj, padded_dim, est_distance, low_distance, ip_x0_qr, use_hacc + batch_data, + query_wrapper, + padded_dim, + estimated_distances, + lower_bounds, + ip_x0_qr, + use_hacc ); } /** * Estimate distances from QGBatchDataMap without producing refinement or error terms. - * est_distance needs 32 values of T, even for a partial batch. + * estimated_distances needs 32 values of T, even for a partial batch. * Throws std::invalid_argument for an unprepared or mismatched query. */ template inline void qg_batch_estdist( - const char* batch_data, const BatchQuery& q_obj, size_t padded_dim, T* est_distance + const char* batch_data, + const BatchQuery& query_wrapper, + size_t padded_dim, + T* estimated_distances ) { static_assert(std::is_same_v, "FastScan accumulation requires int32_t"); - if (padded_dim == 0 || q_obj.padded_dim() != padded_dim) { + if (padded_dim == 0 || query_wrapper.padded_dim() != padded_dim) { throw std::invalid_argument( "FastScan scoring requires a prepared query with matching dimension" ); } if constexpr (std::is_same_v) { - simd::qg_batch_estdist(batch_data, q_obj, padded_dim, est_distance); + simd::qg_batch_estdist(batch_data, query_wrapper, padded_dim, estimated_distances); return; } ConstQGBatchDataMap cur_batch(batch_data, padded_dim); - std::array accu_res{}; - fastscan::accumulate(cur_batch.bin_code(), q_obj.lut(), accu_res.data(), padded_dim); - ConstRowMajorArrayMap ip_arr(accu_res.data(), 1, fastscan::kBatchSize); + std::array accumulated_values{}; + fastscan::accumulate( + cur_batch.bin_code(), query_wrapper.lut(), accumulated_values.data(), padded_dim + ); + ConstRowMajorArrayMap ip_arr(accumulated_values.data(), 1, fastscan::kBatchSize); std::array f_add_values; std::array f_rescale_values; cur_batch.f_add().copy_to(f_add_values.data(), f_add_values.size()); @@ -85,11 +96,11 @@ inline void qg_batch_estdist( f_rescale_values.data(), 1, fastscan::kBatchSize ); - RowMajorArrayMap est_dist_arr(est_distance, 1, fastscan::kBatchSize); + RowMajorArrayMap est_dist_arr(estimated_distances, 1, fastscan::kBatchSize); - est_dist_arr = f_add_arr + q_obj.g_add() + - (f_rescale_arr * (q_obj.delta() * (ip_arr.template cast()) + - q_obj.sum_vl_lut() + q_obj.k1xsumq())); + est_dist_arr = f_add_arr + query_wrapper.g_add() + + (f_rescale_arr * (query_wrapper.delta() * (ip_arr.template cast()) + + query_wrapper.sum_vl_lut() + query_wrapper.k1xsumq())); } } // namespace rabitqlib diff --git a/include/rabitqlib/fastscan/fastscan.hpp b/include/rabitqlib/fastscan/fastscan.hpp index 423e0bf5..e63957dd 100644 --- a/include/rabitqlib/fastscan/fastscan.hpp +++ b/include/rabitqlib/fastscan/fastscan.hpp @@ -33,7 +33,7 @@ constexpr static std::array kPos = { // Requires AVX2/FMA, AVX-512, or ARM NEON; otherwise throws std::runtime_error. void accumulate( const uint8_t* __restrict__ codes, - const uint8_t* __restrict__ lp_table, + const uint8_t* __restrict__ lookup_table, int32_t* __restrict__ result, size_t dim ); diff --git a/include/rabitqlib/fastscan/highacc_fastscan.hpp b/include/rabitqlib/fastscan/highacc_fastscan.hpp index 390c1f44..af7d3c38 100644 --- a/include/rabitqlib/fastscan/highacc_fastscan.hpp +++ b/include/rabitqlib/fastscan/highacc_fastscan.hpp @@ -16,7 +16,7 @@ void transfer_lut_hacc(const uint16_t* lut, size_t dim, uint8_t* hc_lut); void accumulate_hacc( const uint8_t* __restrict__ codes, const uint8_t* __restrict__ hc_lut, - int32_t* accu_res, + int32_t* accumulated_values, size_t dim ); } // namespace rabitqlib::fastscan diff --git a/python_bindings/__init__.pyi b/python_bindings/__init__.pyi index 1725d6cb..22247a3b 100644 --- a/python_bindings/__init__.pyi +++ b/python_bindings/__init__.pyi @@ -236,8 +236,9 @@ class IvfIndex: Vectors in the original, unpadded coordinates. Converted to contiguous float32; a copy may be made. cluster_ids : array_like of int, shape (n,), optional - Existing centroid IDs. None assigns each vector to its nearest - centroid. Centroids do not move. + Existing centroid IDs. None routes each vector using the index's + configured centroid initializer; routing may be approximate. + Centroids do not move. num_threads : int, optional Worker count. Default is 1; 0 selects the available CPU count. fast_quantization : bool, optional diff --git a/python_bindings/ivf_bindings.cpp b/python_bindings/ivf_bindings.cpp index 0f96da2c..ae51719e 100644 --- a/python_bindings/ivf_bindings.cpp +++ b/python_bindings/ivf_bindings.cpp @@ -444,8 +444,9 @@ data : array_like of float, shape (n, dim) Vectors in the original, unpadded coordinates. Converted to contiguous float32; a copy may be made. cluster_ids : array_like of int, shape (n,), optional - Existing centroid IDs. None assigns each vector to its nearest - centroid. Centroids do not move. + Existing centroid IDs. None routes each vector using the index's + configured centroid initializer; routing may be approximate. + Centroids do not move. num_threads : int, optional Worker count. Default is 1; 0 selects the available CPU count. fast_quantization : bool, optional diff --git a/src/fastscan/estimator_avx2.cpp b/src/fastscan/estimator_avx2.cpp index 0be3878e..f41a28ed 100644 --- a/src/fastscan/estimator_avx2.cpp +++ b/src/fastscan/estimator_avx2.cpp @@ -24,28 +24,34 @@ uint32_t candidate_mask_impl(const float* distances, float threshold) { void split_batch_estdist_avx2( const char* batch_data, - const SplitBatchQuery& q_obj, + const SplitBatchQuery& query_wrapper, size_t padded_dim, - float* est_distance, - float* low_distance, + float* estimated_distances, + float* lower_bounds, float* ip_x0_qr, bool use_hacc ) { // Explicit backend entry points also accept query LUTs prepared by the host // backend. In particular, AVX2 correction can receive AVX-512 HACC bytes. split_batch_estdist_impl( - batch_data, q_obj, padded_dim, est_distance, low_distance, ip_x0_qr, use_hacc + batch_data, + query_wrapper, + padded_dim, + estimated_distances, + lower_bounds, + ip_x0_qr, + use_hacc ); } void qg_batch_estdist_avx2( const char* batch_data, - const BatchQuery& q_obj, + const BatchQuery& query_wrapper, size_t padded_dim, - float* est_distance + float* estimated_distances ) { qg_batch_estdist_impl( - batch_data, q_obj, padded_dim, est_distance + batch_data, query_wrapper, padded_dim, estimated_distances ); } @@ -55,15 +61,15 @@ uint32_t qg_candidate_mask_avx2(const float* distances, float threshold) { uint32_t qg_batch_estdist_mask_avx2( const char* batch_data, - const BatchQuery& q_obj, + const BatchQuery& query_wrapper, size_t padded_dim, - float* est_distance, + float* estimated_distances, float threshold ) { qg_batch_estdist_impl( - batch_data, q_obj, padded_dim, est_distance + batch_data, query_wrapper, padded_dim, estimated_distances ); - return candidate_mask_impl(est_distance, threshold); + return candidate_mask_impl(estimated_distances, threshold); } } // namespace rabitqlib::simd diff --git a/src/fastscan/estimator_avx512.cpp b/src/fastscan/estimator_avx512.cpp index 32550968..a7e8a092 100644 --- a/src/fastscan/estimator_avx512.cpp +++ b/src/fastscan/estimator_avx512.cpp @@ -8,10 +8,10 @@ namespace rabitqlib::simd { void split_batch_estdist_avx512( const char* batch_data, - const SplitBatchQuery& q_obj, + const SplitBatchQuery& query_wrapper, size_t padded_dim, - float* est_distance, - float* low_distance, + float* estimated_distances, + float* lower_bounds, float* ip_x0_qr, bool use_hacc ) { @@ -26,21 +26,27 @@ void split_batch_estdist_avx512( split_batch_estdist_impl< fastscan::simd::accumulate_avx512, fastscan::simd::accumulate_hacc_avx512>( - batch_data, q_obj, padded_dim, est_distance, low_distance, ip_x0_qr, use_hacc + batch_data, + query_wrapper, + padded_dim, + estimated_distances, + lower_bounds, + ip_x0_qr, + use_hacc ); } void qg_batch_estdist_avx512( const char* batch_data, - const BatchQuery& q_obj, + const BatchQuery& query_wrapper, size_t padded_dim, - float* est_distance + float* estimated_distances ) { if (padded_dim == 0 || padded_dim % 16 != 0) { throw std::invalid_argument("FastScan dimension must be a positive multiple of 16"); } qg_batch_estdist_impl( - batch_data, q_obj, padded_dim, est_distance + batch_data, query_wrapper, padded_dim, estimated_distances ); } diff --git a/src/fastscan/estimator_generic.cpp b/src/fastscan/estimator_generic.cpp index cf7cf82f..5601ecde 100644 --- a/src/fastscan/estimator_generic.cpp +++ b/src/fastscan/estimator_generic.cpp @@ -19,28 +19,34 @@ uint32_t candidate_mask_impl(const float* distances, float threshold) { void split_batch_estdist_generic( const char* batch_data, - const SplitBatchQuery& q_obj, + const SplitBatchQuery& query_wrapper, size_t padded_dim, - float* est_distance, - float* low_distance, + float* estimated_distances, + float* lower_bounds, float* ip_x0_qr, bool use_hacc ) { // Explicit backend entry points also accept query LUTs prepared by the host // backend. In particular, AVX2 correction can receive AVX-512 HACC bytes. split_batch_estdist_impl( - batch_data, q_obj, padded_dim, est_distance, low_distance, ip_x0_qr, use_hacc + batch_data, + query_wrapper, + padded_dim, + estimated_distances, + lower_bounds, + ip_x0_qr, + use_hacc ); } void qg_batch_estdist_generic( const char* batch_data, - const BatchQuery& q_obj, + const BatchQuery& query_wrapper, size_t padded_dim, - float* est_distance + float* estimated_distances ) { qg_batch_estdist_impl( - batch_data, q_obj, padded_dim, est_distance + batch_data, query_wrapper, padded_dim, estimated_distances ); } @@ -50,15 +56,15 @@ uint32_t qg_candidate_mask_generic(const float* distances, float threshold) { uint32_t qg_batch_estdist_mask_generic( const char* batch_data, - const BatchQuery& q_obj, + const BatchQuery& query_wrapper, size_t padded_dim, - float* est_distance, + float* estimated_distances, float threshold ) { qg_batch_estdist_impl( - batch_data, q_obj, padded_dim, est_distance + batch_data, query_wrapper, padded_dim, estimated_distances ); - return candidate_mask_impl(est_distance, threshold); + return candidate_mask_impl(estimated_distances, threshold); } } // namespace rabitqlib::simd diff --git a/src/fastscan/estimator_kernels.hpp b/src/fastscan/estimator_kernels.hpp index fb44a6c6..8d3cedae 100644 --- a/src/fastscan/estimator_kernels.hpp +++ b/src/fastscan/estimator_kernels.hpp @@ -16,10 +16,10 @@ namespace { template inline void split_batch_estdist_impl( const char* batch_data, - const SplitBatchQuery& q_obj, + const SplitBatchQuery& query_wrapper, size_t padded_dim, - float* est_distance, - float* low_distance, + float* estimated_distances, + float* lower_bounds, float* ip_x0_qr, bool use_hacc ) { @@ -27,9 +27,13 @@ inline void split_batch_estdist_impl( std::array accu_values{}; RowMajorArrayMap accu_arr(accu_values.data(), 1, fastscan::kBatchSize); if (use_hacc) { - AccumulateHacc(cur_batch.bin_code(), q_obj.lut(), accu_values.data(), padded_dim); + AccumulateHacc( + cur_batch.bin_code(), query_wrapper.lut(), accu_values.data(), padded_dim + ); } else { - Accumulate(cur_batch.bin_code(), q_obj.lut(), accu_values.data(), padded_dim); + Accumulate( + cur_batch.bin_code(), query_wrapper.lut(), accu_values.data(), padded_dim + ); } std::array f_add_values; @@ -46,29 +50,30 @@ inline void split_batch_estdist_impl( f_error_values.data(), 1, fastscan::kBatchSize ); - RowMajorArrayMap est_dist_arr(est_distance, 1, fastscan::kBatchSize); + RowMajorArrayMap est_dist_arr(estimated_distances, 1, fastscan::kBatchSize); RowMajorArrayMap ip_x0_qr_arr(ip_x0_qr, 1, fastscan::kBatchSize); - RowMajorArrayMap low_dist_arr(low_distance, 1, fastscan::kBatchSize); + RowMajorArrayMap low_dist_arr(lower_bounds, 1, fastscan::kBatchSize); - ip_x0_qr_arr = q_obj.delta() * (accu_arr.template cast()) + q_obj.sum_vl_lut(); + ip_x0_qr_arr = query_wrapper.delta() * (accu_arr.template cast()) + + query_wrapper.sum_vl_lut(); - est_dist_arr = - f_add_arr + q_obj.g_add() + f_rescale_arr * (ip_x0_qr_arr + q_obj.k1xsumq()); + est_dist_arr = f_add_arr + query_wrapper.g_add() + + f_rescale_arr * (ip_x0_qr_arr + query_wrapper.k1xsumq()); - low_dist_arr = est_dist_arr - f_error_arr * q_obj.g_error(); + low_dist_arr = est_dist_arr - f_error_arr * query_wrapper.g_error(); } template inline void qg_batch_estdist_impl( const char* batch_data, - const BatchQuery& q_obj, + const BatchQuery& query_wrapper, size_t padded_dim, - float* est_distance + float* estimated_distances ) { using T = float; ConstQGBatchDataMap cur_batch(batch_data, padded_dim); std::array accu_values{}; - Accumulate(cur_batch.bin_code(), q_obj.lut(), accu_values.data(), padded_dim); + Accumulate(cur_batch.bin_code(), query_wrapper.lut(), accu_values.data(), padded_dim); ConstRowMajorArrayMap ip_arr(accu_values.data(), 1, fastscan::kBatchSize); std::array f_add_values; @@ -80,11 +85,11 @@ inline void qg_batch_estdist_impl( f_rescale_values.data(), 1, fastscan::kBatchSize ); - RowMajorArrayMap est_dist_arr(est_distance, 1, fastscan::kBatchSize); + RowMajorArrayMap est_dist_arr(estimated_distances, 1, fastscan::kBatchSize); - est_dist_arr = f_add_arr + q_obj.g_add() + - (f_rescale_arr * (q_obj.delta() * (ip_arr.template cast()) + - q_obj.sum_vl_lut() + q_obj.k1xsumq())); + est_dist_arr = f_add_arr + query_wrapper.g_add() + + (f_rescale_arr * (query_wrapper.delta() * (ip_arr.template cast()) + + query_wrapper.sum_vl_lut() + query_wrapper.k1xsumq())); } } // namespace diff --git a/src/fastscan/fastscan_avx2.cpp b/src/fastscan/fastscan_avx2.cpp index d10db4cc..f187e0e9 100644 --- a/src/fastscan/fastscan_avx2.cpp +++ b/src/fastscan/fastscan_avx2.cpp @@ -43,7 +43,7 @@ constexpr size_t kBatchSize = 32; void accumulate_chunk_avx2( const uint8_t* __restrict__ codes, - const uint8_t* __restrict__ lp_table, + const uint8_t* __restrict__ lookup_table, int32_t* __restrict__ result, size_t dim ) { @@ -58,7 +58,7 @@ void accumulate_chunk_avx2( for (size_t i = 0; i < code_length; i += 64) { c = _mm256_loadu_si256((__m256i*)&codes[i]); - lut = _mm256_loadu_si256((__m256i*)&lp_table[i]); + lut = _mm256_loadu_si256((__m256i*)&lookup_table[i]); lo = _mm256_and_si256(c, low_mask); hi = _mm256_and_si256(_mm256_srli_epi16(c, 4), low_mask); @@ -71,7 +71,7 @@ void accumulate_chunk_avx2( accu3 = _mm256_add_epi16(accu3, _mm256_srli_epi16(res_hi, 8)); c = _mm256_loadu_si256((__m256i*)&codes[i + 32]); - lut = _mm256_loadu_si256((__m256i*)&lp_table[i + 32]); + lut = _mm256_loadu_si256((__m256i*)&lookup_table[i + 32]); lo = _mm256_and_si256(c, low_mask); hi = _mm256_and_si256(_mm256_srli_epi16(c, 4), low_mask); @@ -113,12 +113,12 @@ void accumulate_chunk_avx2( void accumulate_avx2( const uint8_t* __restrict__ codes, - const uint8_t* __restrict__ lp_table, + const uint8_t* __restrict__ lookup_table, int32_t* __restrict__ result, size_t dim ) { if (dim <= kChunkDim) { - accumulate_chunk_avx2(codes, lp_table, result, dim); + accumulate_chunk_avx2(codes, lookup_table, result, dim); return; } @@ -127,7 +127,7 @@ void accumulate_avx2( for (size_t offset = 0; offset < dim;) { const size_t chunk_dim = std::min(kChunkDim, dim - offset); accumulate_chunk_avx2( - codes + offset * 4, lp_table + offset * 4, chunk_result.data(), chunk_dim + codes + offset * 4, lookup_table + offset * 4, chunk_result.data(), chunk_dim ); for (size_t lane = 0; lane < kBatchSize; ++lane) { totals[lane] += chunk_result[lane]; @@ -176,7 +176,7 @@ namespace { void accumulate_hacc_chunk_avx2( const uint8_t* __restrict__ codes, const uint8_t* __restrict__ hc_lut, - int32_t* accu_res, + int32_t* accumulated_values, size_t dim ) { __m256i low_mask = _mm256_set1_epi8(0xf); @@ -243,10 +243,10 @@ void accumulate_hacc_chunk_avx2( add_shiftl8(dis0[0], dis0[1], res[0], res[1]); // res for vec 0 to 15 add_shiftl8(dis1[0], dis1[1], res[2], res[3]); // res for vec 16 to 31 - _mm256_storeu_si256((__m256i*)(accu_res), res[0]); - _mm256_storeu_si256((__m256i*)(accu_res + 8), res[1]); - _mm256_storeu_si256((__m256i*)(accu_res + 16), res[2]); - _mm256_storeu_si256((__m256i*)(accu_res + 24), res[3]); + _mm256_storeu_si256((__m256i*)(accumulated_values), res[0]); + _mm256_storeu_si256((__m256i*)(accumulated_values + 8), res[1]); + _mm256_storeu_si256((__m256i*)(accumulated_values + 16), res[2]); + _mm256_storeu_si256((__m256i*)(accumulated_values + 24), res[3]); } } // namespace @@ -254,11 +254,11 @@ void accumulate_hacc_chunk_avx2( void accumulate_hacc_avx2( const uint8_t* __restrict__ codes, const uint8_t* __restrict__ hc_lut, - int32_t* accu_res, + int32_t* accumulated_values, size_t dim ) { if (dim <= kChunkDim) { - accumulate_hacc_chunk_avx2(codes, hc_lut, accu_res, dim); + accumulate_hacc_chunk_avx2(codes, hc_lut, accumulated_values, dim); return; } @@ -280,7 +280,7 @@ void accumulate_hacc_avx2( } } for (size_t lane = 0; lane < kBatchSize; ++lane) { - accu_res[lane] = static_cast(totals[lane]); + accumulated_values[lane] = static_cast(totals[lane]); } } diff --git a/src/fastscan/fastscan_avx512.cpp b/src/fastscan/fastscan_avx512.cpp index bafaf454..e4e48f7f 100644 --- a/src/fastscan/fastscan_avx512.cpp +++ b/src/fastscan/fastscan_avx512.cpp @@ -37,7 +37,7 @@ constexpr size_t kBatchSize = 32; void accumulate_chunk_avx512( const uint8_t* __restrict__ codes, - const uint8_t* __restrict__ lp_table, + const uint8_t* __restrict__ lookup_table, int32_t* __restrict__ result, size_t dim ) { @@ -59,7 +59,7 @@ void accumulate_chunk_avx512( // ! multiple of 16 for (size_t i = 0; i < code_length; i += 64) { c = _mm512_loadu_si512(&codes[i]); - lut = _mm512_loadu_si512(&lp_table[i]); + lut = _mm512_loadu_si512(&lookup_table[i]); lo = _mm512_and_si512(c, lo_mask); // code of vector 0 to 15 hi = _mm512_and_si512(_mm512_srli_epi16(c, 4), lo_mask); // code of vector 16 to 31 @@ -105,12 +105,12 @@ void accumulate_chunk_avx512( void accumulate_avx512( const uint8_t* __restrict__ codes, - const uint8_t* __restrict__ lp_table, + const uint8_t* __restrict__ lookup_table, int32_t* __restrict__ result, size_t dim ) { if (dim <= kChunkDim) { - accumulate_chunk_avx512(codes, lp_table, result, dim); + accumulate_chunk_avx512(codes, lookup_table, result, dim); return; } @@ -119,7 +119,7 @@ void accumulate_avx512( for (size_t offset = 0; offset < dim;) { const size_t chunk_dim = std::min(kChunkDim, dim - offset); accumulate_chunk_avx512( - codes + offset * 4, lp_table + offset * 4, chunk_result.data(), chunk_dim + codes + offset * 4, lookup_table + offset * 4, chunk_result.data(), chunk_dim ); for (size_t lane = 0; lane < kBatchSize; ++lane) { totals[lane] += chunk_result[lane]; @@ -168,7 +168,7 @@ namespace { void accumulate_hacc_chunk_avx512( const uint8_t* __restrict__ codes, const uint8_t* __restrict__ hc_lut, - int32_t* accu_res, + int32_t* accumulated_values, size_t dim ) { __m512i low_mask = _mm512_set1_epi8(0xf); @@ -244,8 +244,8 @@ void accumulate_hacc_chunk_avx512( res[1] = _mm512_add_epi32(dis1[0], _mm512_slli_epi32(dis1[1], 8)); // res for vec 16 to 31 - _mm512_storeu_epi32(accu_res, res[0]); - _mm512_storeu_epi32(accu_res + 16, res[1]); + _mm512_storeu_epi32(accumulated_values, res[0]); + _mm512_storeu_epi32(accumulated_values + 16, res[1]); } } // namespace @@ -253,11 +253,11 @@ void accumulate_hacc_chunk_avx512( void accumulate_hacc_avx512( const uint8_t* __restrict__ codes, const uint8_t* __restrict__ hc_lut, - int32_t* accu_res, + int32_t* accumulated_values, size_t dim ) { if (dim <= kChunkDim) { - accumulate_hacc_chunk_avx512(codes, hc_lut, accu_res, dim); + accumulate_hacc_chunk_avx512(codes, hc_lut, accumulated_values, dim); return; } @@ -279,7 +279,7 @@ void accumulate_hacc_avx512( } } for (size_t lane = 0; lane < kBatchSize; ++lane) { - accu_res[lane] = static_cast(totals[lane]); + accumulated_values[lane] = static_cast(totals[lane]); } } diff --git a/src/fastscan/fastscan_dispatch.cpp b/src/fastscan/fastscan_dispatch.cpp index 71267d05..46634c9a 100644 --- a/src/fastscan/fastscan_dispatch.cpp +++ b/src/fastscan/fastscan_dispatch.cpp @@ -78,40 +78,48 @@ namespace rabitqlib::simd { void qg_batch_estdist( const char* batch_data, - const BatchQuery& q_obj, + const BatchQuery& query_wrapper, size_t padded_dim, - float* est_distance + float* estimated_distances ) { static const auto kQgBatchEstdistFn = fastscan::detail::selected_backend().qg; - kQgBatchEstdistFn(batch_data, q_obj, padded_dim, est_distance); + kQgBatchEstdistFn(batch_data, query_wrapper, padded_dim, estimated_distances); } uint32_t qg_batch_estdist_mask( const char* batch_data, - const BatchQuery& q_obj, + const BatchQuery& query_wrapper, size_t padded_dim, - float* est_distance, + float* estimated_distances, float threshold ) { static const auto kQgBatchEstdistMaskFn = fastscan::detail::selected_backend().qg_mask; - return kQgBatchEstdistMaskFn(batch_data, q_obj, padded_dim, est_distance, threshold); + return kQgBatchEstdistMaskFn( + batch_data, query_wrapper, padded_dim, estimated_distances, threshold + ); } void split_batch_estdist( const char* batch_data, - const SplitBatchQuery& q_obj, + const SplitBatchQuery& query_wrapper, size_t padded_dim, - float* est_distance, - float* low_distance, + float* estimated_distances, + float* lower_bounds, float* ip_x0_qr, bool use_hacc ) { static const auto kSplitBatchEstdistFn = fastscan::detail::selected_backend().split; kSplitBatchEstdistFn( - batch_data, q_obj, padded_dim, est_distance, low_distance, ip_x0_qr, use_hacc + batch_data, + query_wrapper, + padded_dim, + estimated_distances, + lower_bounds, + ip_x0_qr, + use_hacc ); } @@ -128,7 +136,7 @@ void pack_lut(size_t dim, const float* __restrict__ query, float* __restr void accumulate( const uint8_t* __restrict__ codes, - const uint8_t* __restrict__ lp_table, + const uint8_t* __restrict__ lookup_table, int32_t* __restrict__ result, size_t dim ) { @@ -137,7 +145,7 @@ void accumulate( if (dim == 0 || dim % 16 != 0) { throw std::invalid_argument("FastScan dimension must be a positive multiple of 16"); } - kAccumulateFn(codes, lp_table, result, dim); + kAccumulateFn(codes, lookup_table, result, dim); } void transfer_lut_hacc(const uint16_t* lut, size_t dim, uint8_t* hc_lut) { @@ -154,7 +162,7 @@ void transfer_lut_hacc(const uint16_t* lut, size_t dim, uint8_t* hc_lut) { void accumulate_hacc( const uint8_t* __restrict__ codes, const uint8_t* __restrict__ hc_lut, - int32_t* accu_res, + int32_t* accumulated_values, size_t dim ) { static const auto kAccumulateHaccFn = detail::selected_backend().accumulate_hacc; @@ -164,7 +172,7 @@ void accumulate_hacc( "high-accuracy FastScan dimension must be a positive multiple of 16" ); } - kAccumulateHaccFn(codes, hc_lut, accu_res, dim); + kAccumulateHaccFn(codes, hc_lut, accumulated_values, dim); } } // namespace rabitqlib::fastscan diff --git a/src/simd/dispatch.cpp b/src/simd/dispatch.cpp index 7049b898..66c9772e 100644 --- a/src/simd/dispatch.cpp +++ b/src/simd/dispatch.cpp @@ -179,11 +179,11 @@ double best_rescale_factor_with_workspace( void fht_rotate( const float* data, - float* rotated_vec, + float* rotated_data, size_t dim, size_t padded_dim, size_t trunc_dim, - float fac, + float scale_factor, const uint8_t* flip ) { #if defined(__aarch64__) || defined(_M_ARM64) @@ -194,7 +194,7 @@ void fht_rotate( RABITQ_RESOLVE(fht_rotate_avx512, fht_rotate_avx2, fht_rotate_generic); #endif - kFhtRotateFn(data, rotated_vec, dim, padded_dim, trunc_dim, fac, flip); + kFhtRotateFn(data, rotated_data, dim, padded_dim, trunc_dim, scale_factor, flip); } ExcodeIpTable resolve_excode_ip_table() { diff --git a/src/simd/portable_generic.cpp b/src/simd/portable_generic.cpp index 3392abe8..1c2e54c1 100644 --- a/src/simd/portable_generic.cpp +++ b/src/simd/portable_generic.cpp @@ -338,30 +338,30 @@ void kacs_walk_generic(float* data, size_t len) { } void fht_rotate_generic( const float* data, - float* output, + float* rotated_data, size_t dim, - size_t padded, - size_t trunc, - float factor, + size_t padded_dim, + size_t trunc_dim, + float scale_factor, const uint8_t* flip ) { - if (trunc < 64 || trunc > 65536 || (trunc & (trunc - 1)) != 0) { + if (trunc_dim < 64 || trunc_dim > 65536 || (trunc_dim & (trunc_dim - 1)) != 0) { throw std::invalid_argument("Unsupported dimension for FhtKacRotator"); } - std::memmove(output, data, dim * sizeof(float)); - std::fill(output + dim, output + padded, 0.0F); + std::memmove(rotated_data, data, dim * sizeof(float)); + std::fill(rotated_data + dim, rotated_data + padded_dim, 0.0F); for (size_t pass = 0; pass < 4; ++pass) { - flip_sign_generic(flip + pass * padded / 8, output, padded); - const size_t start = pass % 2 == 0 ? 0 : padded - trunc; - fht(output + start, trunc); - for (size_t i = start; i < start + trunc; ++i) - output[i] *= factor; - if (padded != trunc) - kacs_walk_generic(output, padded); + flip_sign_generic(flip + pass * padded_dim / 8, rotated_data, padded_dim); + const size_t start = pass % 2 == 0 ? 0 : padded_dim - trunc_dim; + fht(rotated_data + start, trunc_dim); + for (size_t i = start; i < start + trunc_dim; ++i) + rotated_data[i] *= scale_factor; + if (padded_dim != trunc_dim) + kacs_walk_generic(rotated_data, padded_dim); } - if (padded != trunc) { - for (size_t i = 0; i < padded; ++i) - output[i] *= 0.25F; + if (padded_dim != trunc_dim) { + for (size_t i = 0; i < padded_dim; ++i) + rotated_data[i] *= 0.25F; } } } // namespace rabitqlib::simd diff --git a/src/simd/rotator_avx2.cpp b/src/simd/rotator_avx2.cpp index 82b51838..191c05dd 100644 --- a/src/simd/rotator_avx2.cpp +++ b/src/simd/rotator_avx2.cpp @@ -55,15 +55,15 @@ void kacs_walk_avx2(float* data, size_t len) { void fht_rotate_avx2( const float* data, - float* rotated_vec, + float* rotated_data, size_t dim, size_t padded_dim, size_t trunc_dim, - float fac, + float scale_factor, const uint8_t* flip ) { fht_rotate_impl( - data, rotated_vec, dim, padded_dim, trunc_dim, fac, flip + data, rotated_data, dim, padded_dim, trunc_dim, scale_factor, flip ); } diff --git a/src/simd/rotator_avx512.cpp b/src/simd/rotator_avx512.cpp index f72ba705..82fbc9cd 100644 --- a/src/simd/rotator_avx512.cpp +++ b/src/simd/rotator_avx512.cpp @@ -110,15 +110,15 @@ void kacs_walk_avx512(float* data, size_t len) { void fht_rotate_avx512( const float* data, - float* rotated_vec, + float* rotated_data, size_t dim, size_t padded_dim, size_t trunc_dim, - float fac, + float scale_factor, const uint8_t* flip ) { fht_rotate_impl( - data, rotated_vec, dim, padded_dim, trunc_dim, fac, flip + data, rotated_data, dim, padded_dim, trunc_dim, scale_factor, flip ); } diff --git a/src/simd/rotator_kernels.hpp b/src/simd/rotator_kernels.hpp index 1128d1a4..4befb336 100644 --- a/src/simd/rotator_kernels.hpp +++ b/src/simd/rotator_kernels.hpp @@ -40,11 +40,11 @@ inline void kacs_walk_rescale(float* data, size_t dim, float factor) { template void fht_rotate_impl( const float* data, - float* rotated_vec, + float* rotated_data, size_t dim, size_t padded_dim, size_t trunc_dim, - float fac, + float scale_factor, const uint8_t* flip ) { void (*fht)(float*, float) = nullptr; @@ -69,36 +69,36 @@ void fht_rotate_impl( } #undef RABITQ_FHT_CASE - std::memcpy(rotated_vec, data, sizeof(float) * dim); - std::fill(rotated_vec + dim, rotated_vec + padded_dim, 0.0F); + std::memcpy(rotated_data, data, sizeof(float) * dim); + std::fill(rotated_data + dim, rotated_data + padded_dim, 0.0F); if (trunc_dim == padded_dim) { for (size_t round = 0; round < 4; ++round) { - flip_sign(flip + round * (padded_dim / 8), rotated_vec, padded_dim); - fht(rotated_vec, fac); + flip_sign(flip + round * (padded_dim / 8), rotated_data, padded_dim); + fht(rotated_data, scale_factor); } return; } const auto walk_and_flip = [&](const uint8_t* signs) { if constexpr (kacs_flip != nullptr) { - kacs_flip(signs, rotated_vec, padded_dim); + kacs_flip(signs, rotated_data, padded_dim); } else { - kacs_walk(rotated_vec, padded_dim); - flip_sign(signs, rotated_vec, padded_dim); + kacs_walk(rotated_data, padded_dim); + flip_sign(signs, rotated_data, padded_dim); } }; size_t start = padded_dim - trunc_dim; - flip_sign(flip, rotated_vec, padded_dim); - fht(rotated_vec, fac); + flip_sign(flip, rotated_data, padded_dim); + fht(rotated_data, scale_factor); walk_and_flip(flip + (padded_dim / 8)); - fht(rotated_vec + start, fac); + fht(rotated_data + start, scale_factor); walk_and_flip(flip + (2 * padded_dim / 8)); - fht(rotated_vec, fac); + fht(rotated_data, scale_factor); walk_and_flip(flip + (3 * padded_dim / 8)); - fht(rotated_vec + start, fac); - kacs_walk_rescale(rotated_vec, padded_dim, 0.25F); + fht(rotated_data + start, scale_factor); + kacs_walk_rescale(rotated_data, padded_dim, 0.25F); } } // namespace } // namespace rabitqlib::simd diff --git a/src/simd/rotator_neon.cpp b/src/simd/rotator_neon.cpp index 2c90b3c7..44c087ea 100644 --- a/src/simd/rotator_neon.cpp +++ b/src/simd/rotator_neon.cpp @@ -124,29 +124,29 @@ void kacs_walk_neon(float* data, size_t len) { void fht_rotate_neon( const float* data, - float* output, + float* rotated_data, size_t dim, - size_t padded, - size_t trunc, - float factor, + size_t padded_dim, + size_t trunc_dim, + float scale_factor, const uint8_t* flip ) { - if (trunc < 64 || trunc > 65536 || (trunc & (trunc - 1)) != 0) { + if (trunc_dim < 64 || trunc_dim > 65536 || (trunc_dim & (trunc_dim - 1)) != 0) { throw std::invalid_argument("Unsupported dimension for FhtKacRotator"); } - std::memmove(output, data, dim * sizeof(float)); - std::fill(output + dim, output + padded, 0.0F); + std::memmove(rotated_data, data, dim * sizeof(float)); + std::fill(rotated_data + dim, rotated_data + padded_dim, 0.0F); for (size_t pass = 0; pass < 4; ++pass) { - flip_sign_neon(flip + pass * padded / 8, output, padded); - const size_t start = pass % 2 == 0 ? 0 : padded - trunc; - fht(output + start, trunc); - rescale(output + start, trunc, factor); - if (padded != trunc) { - kacs_walk_neon(output, padded); + flip_sign_neon(flip + pass * padded_dim / 8, rotated_data, padded_dim); + const size_t start = pass % 2 == 0 ? 0 : padded_dim - trunc_dim; + fht(rotated_data + start, trunc_dim); + rescale(rotated_data + start, trunc_dim, scale_factor); + if (padded_dim != trunc_dim) { + kacs_walk_neon(rotated_data, padded_dim); } } - if (padded != trunc) { - rescale(output, padded, 0.25F); + if (padded_dim != trunc_dim) { + rescale(rotated_data, padded_dim, 0.25F); } } } // namespace rabitqlib::simd