Speed up batched IVF Flat routing

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<double>(d) *
+                               static_cast<double>(std::numeric_limits<float>::epsilon());
+    const double error_floor = 8.0 * static_cast<double>(d) *
+                                   std::numeric_limits<float>::min() +
+                               8.0 * std::numeric_limits<double>::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<float> dots(rows * std::min(k, kCentroidBlock));
+            std::vector<double> point_norms(rows);
+            std::vector<std::vector<AnnCandidate<double>>> best(rows);
+            const auto exact_distance = [d, spherical](const float* vector,
+                                                       const float* centroid) {
+                return spherical ? static_cast<double>(dot_product_dis(vector, centroid, d))
+                                 : static_cast<double>(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<double> 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<double>::infinity()
+                                                 : candidates.back().distance;
+                        if (estimate > worst + error) {
+                            continue;
+                        }
+                        AnnCandidate<double> candidate(
+                            static_cast<PID>(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<float>(
+                        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<size_t> sample_training_points(

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<AnnCandidate<float>>& centroid_candidates,
+        PID* result_ids,
+        float* distances,
+        bool use_hacc
+    ) const;
+
     inline void scan_one_batch(
         const char* batch_data,
         const char* rerank_data,

src/index/ivf.cpp

@@ -1177,11 +1177,27 @@ void IVF::search(
     std::vector<AnnCandidate<float>> 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<AnnCandidate<float>>& 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<float> 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<float>(
-                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<size_t>(
+            1, std::min(num_queries, max_block_bytes / std::max<size_t>(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<float> 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<PID> routes(rows * nprobe);
+            std::vector<float> route_distances(rows * nprobe);
+            clustering::detail::exact_assign_topn(
+                rotated.data(), centroids, rows, padded_dim_, num_clusters_, nprobe,
+                metric_type_ == METRIC_IP, static_cast<uint32_t>(num_workers), routes.data(),
+                route_distances.data()
+            );
+            index_detail::parallel_for(
+                0, rows, num_threads, [&](size_t row, size_t /*thread*/) {
+                    std::vector<AnnCandidate<float>> candidates(nprobe);
+                    for (size_t i = 0; i < nprobe; ++i) {
+                        candidates[i] = AnnCandidate<float>(
+                            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(

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<PID, kNumQueries * kTopK> expected{}, actual{}, ids_only{};
             std::array<float, kNumQueries * kTopK> expected_distances{}, actual_distances{};
-            for (size_t nprobe : {1U, 99U}) {
+            for (size_t nprobe : {1U, 2U, 3U, 99U}) {
                 for (std::optional<bool> hacc :
                      {std::optional<bool>{},
                       std::optional<bool>(false),