From 4333e67f73ba577fa7494b4b75ad6f1103f10645 Mon Sep 17 00:00:00 2001 From: jay23606 <4602056+jay23606@users.noreply.github.com> Date: Thu, 12 Nov 2026 23:16:39 -0500 Subject: [PATCH] Speed up batched IVF flat routing --- .../rabitqlib/clustering/detail/kmeans.hpp | 108 ++++++++++++++++++ include/rabitqlib/index/ivf/ivf.hpp | 10 ++ src/index/ivf.cpp | 65 ++++++++++- tests/unit/rabitqlib/index/ivf_test.cpp | 2 +- 4 files changed, 181 insertions(+), 4 deletions(-) diff --git a/include/rabitqlib/clustering/detail/kmeans.hpp b/include/rabitqlib/clustering/detail/kmeans.hpp index 809dc689..43728182 100644 --- a/include/rabitqlib/clustering/detail/kmeans.hpp +++ b/include/rabitqlib/clustering/detail/kmeans.hpp @@ -222,6 +222,114 @@ inline void exact_assign( failure.rethrow(); } +// Assign each row to its exact nearest `nprobe` centroids. Matrix products +// screen centroid tiles for a block of rows; distances near the current top-n +// boundary are recomputed in double precision before updating the result. +// Output IDs and distances are row-major and sorted by distance, then ID. +inline void exact_assign_topn( + const float* x, + const float* centroids, + size_t n, + size_t d, + size_t k, + size_t nprobe, + bool spherical, + uint32_t num_workers, + PID* assignments, + float* distances +) { + if (n == 0 || nprobe == 0) { + return; + } + constexpr size_t kPointBlock = 32; + constexpr size_t kCentroidBlock = 256; + const double error_scale = 8.0 * static_cast(d) * + static_cast(std::numeric_limits::epsilon()); + const double error_floor = 8.0 * static_cast(d) * + std::numeric_limits::min() + + 8.0 * std::numeric_limits::epsilon(); + const size_t blocks = 1 + (n - 1) / kPointBlock; + index_detail::parallel_for( + 0, blocks, num_workers, [&](size_t block, size_t /*thread*/) { + const size_t first = block * kPointBlock; + const size_t rows = std::min(kPointBlock, n - first); + std::vector dots(rows * std::min(k, kCentroidBlock)); + std::vector point_norms(rows); + std::vector>> best(rows); + const auto exact_distance = [d, spherical](const float* vector, + const float* centroid) { + return spherical ? static_cast(dot_product_dis(vector, centroid, d)) + : static_cast(euclidean_sqr(vector, centroid, d)); + }; + for (size_t row = 0; row < rows; ++row) { + point_norms[row] = l2norm_sqr_f64(x + (first + row) * d, d); + best[row].reserve(nprobe); + } + std::vector norms(k); + for (size_t c = 0; c < k; ++c) { + norms[c] = l2norm_sqr_f64(centroids + c * d, d); + } + for (size_t first_cluster = 0; first_cluster < k; + first_cluster += kCentroidBlock) { + const size_t cols = std::min(kCentroidBlock, k - first_cluster); + simd::matrix_product_transposed( + x + first * d, centroids + first_cluster * d, dots.data(), rows, d, cols + ); + for (size_t row = 0; row < rows; ++row) { + const float* vector = x + (first + row) * d; + auto& candidates = best[row]; + for (size_t col = 0; col < cols; ++col) { + const size_t cluster = first_cluster + col; + const double norm_sum = point_norms[row] + norms[cluster]; + const double dot = dots[row * cols + col]; + const double estimate = spherical ? 1.0 - dot : norm_sum - 2.0 * dot; + const double error = error_scale * norm_sum + error_floor; + const double worst = candidates.size() < nprobe + ? std::numeric_limits::infinity() + : candidates.back().distance; + if (estimate > worst + error) { + continue; + } + AnnCandidate candidate( + static_cast(cluster), + exact_distance(vector, centroids + cluster * d) + ); + const auto position = std::lower_bound( + candidates.begin(), candidates.end(), candidate, + [](const auto& lhs, const auto& rhs) { + return lhs.distance < rhs.distance || + (lhs.distance == rhs.distance && lhs.id < rhs.id); + } + ); + if (position == candidates.end() && candidates.size() == nprobe) { + continue; + } + candidates.insert(position, candidate); + if (candidates.size() > nprobe) { + candidates.pop_back(); + } + } + } + } + for (size_t row = 0; row < rows; ++row) { + auto& candidates = best[row]; + std::sort(candidates.begin(), candidates.end(), [](const auto& lhs, + const auto& rhs) { + return lhs.distance < rhs.distance || + (lhs.distance == rhs.distance && lhs.id < rhs.id); + }); + for (size_t i = 0; i < nprobe; ++i) { + assignments[(first + row) * nprobe + i] = candidates[i].id; + distances[(first + row) * nprobe + i] = static_cast( + spherical ? candidates[i].distance + : std::sqrt(candidates[i].distance) + ); + } + } + } + ); +} + // Returns count sorted, distinct row indices in [0, n). The seed sequence keeps // this stream independent of the centroid initialization stream. [[nodiscard]] inline std::vector sample_training_points( diff --git a/include/rabitqlib/index/ivf/ivf.hpp b/include/rabitqlib/index/ivf/ivf.hpp index 745ed983..90122690 100644 --- a/include/rabitqlib/index/ivf/ivf.hpp +++ b/include/rabitqlib/index/ivf/ivf.hpp @@ -175,6 +175,16 @@ class IVF { const float* query ) const; + void search_with_candidates( + const float* query, + const float* rotated_query, + size_t k, + const std::vector>& centroid_candidates, + PID* result_ids, + float* distances, + bool use_hacc + ) const; + inline void scan_one_batch( const char* batch_data, const char* rerank_data, diff --git a/src/index/ivf.cpp b/src/index/ivf.cpp index abc91769..44475041 100644 --- a/src/index/ivf.cpp +++ b/src/index/ivf.cpp @@ -1177,11 +1177,27 @@ void IVF::search( std::vector> centroid_candidates(nprobe); initializer_->centroids_distances(rotated_query.data(), nprobe, centroid_candidates); + search_with_candidates( + query, rotated_query.data(), k, centroid_candidates, result_ids, distances, use_hacc + ); +} + +void IVF::search_with_candidates( + const float* query, + const float* rotated_query, + size_t k, + const std::vector>& centroid_candidates, + PID* result_ids, + float* distances, + bool use_hacc +) const { + const size_t nprobe = centroid_candidates.size(); + // The pool never needs more slots than points; the rest of k is padded. buffer::SearchBuffer result_pool(std::min(k, max_elements_)); SplitBatchQuery query_wrapper( - rotated_query.data(), padded_dim_, ex_bits_, metric_type_, use_hacc + rotated_query, padded_dim_, ex_bits_, metric_type_, use_hacc ); for (size_t i = 0; i < nprobe; ++i) { @@ -1193,10 +1209,10 @@ void IVF::search( query_wrapper.set_g_add(dist); } else { const float residual_norm = std::sqrt(euclidean_sqr( - rotated_query.data(), initializer_->centroid(cluster_id), padded_dim_ + rotated_query, initializer_->centroid(cluster_id), padded_dim_ )); auto query_centroid_ip = dot_product( - rotated_query.data(), initializer_->centroid(cluster_id), padded_dim_ + rotated_query, initializer_->centroid(cluster_id), padded_dim_ ); query_wrapper.set_g_add(residual_norm, query_centroid_ip); } @@ -1238,6 +1254,49 @@ void IVF::search_batch( num_queries == 1 || num_threads == 1 ? 1 : std::min(resolve_num_threads(num_threads), num_queries); + if (initializer_type_ == InitializerType::Flat && num_queries > 1) { + const size_t max_block_bytes = size_t{128} << 20; + const size_t row_bytes = sizeof(float) * padded_dim_ + + sizeof(PID) * nprobe + sizeof(float) * nprobe; + const size_t block_rows = std::max( + 1, std::min(num_queries, max_block_bytes / std::max(row_bytes, 1)) + ); + const auto* centroids = initializer_->centroid(0); + for (size_t first = 0; first < num_queries; first += block_rows) { + const size_t rows = std::min(block_rows, num_queries - first); + std::vector rotated(rows * padded_dim_); + index_detail::parallel_for( + 0, rows, num_threads, [&](size_t row, size_t /*thread*/) { + rotator_->rotate( + queries + (first + row) * dim_, rotated.data() + row * padded_dim_ + ); + } + ); + std::vector routes(rows * nprobe); + std::vector route_distances(rows * nprobe); + clustering::detail::exact_assign_topn( + rotated.data(), centroids, rows, padded_dim_, num_clusters_, nprobe, + metric_type_ == METRIC_IP, static_cast(num_workers), routes.data(), + route_distances.data() + ); + index_detail::parallel_for( + 0, rows, num_threads, [&](size_t row, size_t /*thread*/) { + std::vector> candidates(nprobe); + for (size_t i = 0; i < nprobe; ++i) { + candidates[i] = AnnCandidate( + routes[row * nprobe + i], route_distances[row * nprobe + i] + ); + } + search_with_candidates( + queries + (first + row) * dim_, rotated.data() + row * padded_dim_, k, + candidates, result_ids + (first + row) * k, + distances == nullptr ? nullptr : distances + (first + row) * k, hacc + ); + } + ); + } + return; + } if (num_workers == 1) { for (size_t i = 0; i < num_queries; ++i) { search( diff --git a/tests/unit/rabitqlib/index/ivf_test.cpp b/tests/unit/rabitqlib/index/ivf_test.cpp index cfce20c4..74bf0162 100644 --- a/tests/unit/rabitqlib/index/ivf_test.cpp +++ b/tests/unit/rabitqlib/index/ivf_test.cpp @@ -78,7 +78,7 @@ TEST(IvfSearchTest, BatchSearchMatchesIndividualQueries) { index.construct(data.data(), centroids.data(), labels.data(), false, 1); std::array expected{}, actual{}, ids_only{}; std::array expected_distances{}, actual_distances{}; - for (size_t nprobe : {1U, 99U}) { + for (size_t nprobe : {1U, 2U, 3U, 99U}) { for (std::optional hacc : {std::optional{}, std::optional(false),