Optimize clustering, batch search, and index allocation

docs/docs/index/ivf.md

@@ -83,6 +83,10 @@ vectors, and compute one-bit codes and factors. Quantized mode also computes
 float32 vectors instead. In raw mode, the index owns a copy of the input, so
 the original data can be released after construction.
 
+Construction reuses a rotation buffer of at most 32 vectors per worker. This
+temporary storage uses at most `workers * 32 * padded_dim * sizeof(float)` bytes,
+independently of cluster sizes.
+
 After construction, you can directly save the index file to disk:
 ```c++
 ivf.save(outoput_index_file);
@@ -157,6 +161,25 @@ squared L2 or `1 - dot(query, vector)` in the original coordinates; the search
 API is unchanged. Cluster selection and filtering remain approximate in both
 modes. Search returns the top `k` results after scanning the selected clusters.
 
+For multiple queries, `search_batch` distributes queries across workers.
+Input queries are contiguous row-major float32
+vectors with shape `(num_queries, dimension())`; output buffers have shape
+`(num_queries, k)`.
+
+```cpp
+ivf.search_batch(queries, num_queries, k, nprobe, ids, distances,
+                 std::nullopt, 4);  // Automatic precision, up to four workers
+```
+
+The last two arguments are `std::optional<bool> use_hacc` (default
+`std::nullopt`) and `size_t num_threads` (default `1`, using OpenMP for parallel
+batches). Pass `true` or `false`
+for an explicit precision override. Distances may be `nullptr`. An empty batch
+does no work. Input and output buffers must not overlap. Simultaneous searches
+need separate output buffers, and must not overlap index updates.
+Python's existing `index.search(queries, k, nprobe, num_threads=4)` uses this
+batch path with the same output and precision rules.
+
 ## Updating an Index
 
 Starting with 0.5.1, C++ `IVF` and Python `IvfIndex` support `add()` and `remove()`

docs/docs/index/qg.md

@@ -29,7 +29,7 @@ index.build(data, ef_construction=200, num_threads=32, init="pipnn")
 index.save("qg_example.index")
 
 loaded = SymqgIndex.load("qg_example.index")
-ids, distances = loaded.search(queries, k=10, ef=100, num_threads=1)
+ids, distances = loaded.search_batch(queries, k=10, ef=100, num_threads=1)
 ```
 
 `init` defaults to `"pipnn"`; use `"random"` for random initialization.
@@ -118,5 +118,31 @@ std::vector<float> distances(10);
 qg.search(query.data(), 10, ids.data(), distances.data());
 ```
 
+For contiguous row-major batches, `search_batch()` reuses scratch storage within
+each worker and preserves the results of independent `search()` calls:
+
+```cpp
+// queries contains num_queries * qg.dimension() floats in the original dimension.
+std::vector<rabitqlib::PID> batch_ids(num_queries * 10);
+std::vector<float> batch_distances(num_queries * 10);
+qg.search_batch(
+    queries.data(), num_queries, 10, batch_ids.data(), batch_distances.data(), 4
+);
+```
+
+The C++ signature is `search_batch(queries, num_queries, k, ids, distances,
+num_threads=1)`. Set `ef >= k` with `set_ef()` before searching. Both output
+buffers are required and contain `num_queries * k` elements. Empty C++ batches
+do no work and may use null buffers. Inputs and outputs must not overlap.
+Concurrent searches need separate outputs; do not modify the index or call
+`set_ef()` while searches are running.
+
+Python provides `search_batch(queries, k, ef, num_threads=1)`; the existing
+`search()` accepts the same two-dimensional array and remains an equivalent
+entry point. Both return `(ids, distances)` arrays of shape `(num_queries, k)`,
+including empty batches. Both APIs default to one worker; `num_threads=0`
+selects the available hardware thread count, capped by the number of queries.
+Queries use the original dimension, including when the internal rotation pads it.
+
 See `sample/cpp/symqg_indexing.cpp`, `sample/cpp/symqg_querying.cpp`, and their
 Python counterparts for complete examples.

include/rabitqlib/clustering/detail/kmeans.hpp

@@ -21,6 +21,7 @@
 
 #include "rabitqlib/defines.hpp"
 #include "rabitqlib/simd/matrix_dispatch.hpp"
+#include "rabitqlib/utils/cpu_features.hpp"
 #include "rabitqlib/utils/rotator.hpp"
 #include "rabitqlib/utils/space.hpp"
 #include "rabitqlib/utils/tools.hpp"
@@ -55,6 +56,8 @@ namespace detail {
 
 // Inputs have already passed clustering validation. GEMM only screens candidates;
 // final distances and tie breaking retain the scalar double-precision arithmetic.
+// Optional previous assignments contain n valid centroid IDs and may alias the
+// output. Recompute their distances against the current, updated centroids.
 inline void exact_assign(
     const float* x,
     const float* centroids,
@@ -64,27 +67,41 @@ inline void exact_assign(
     bool spherical,
     uint32_t num_threads,
     PID* assignments,
-    float* distances
+    float* distances,
+    const PID* previous_assignments = nullptr
 ) {
-    constexpr size_t kPointBlock = 64;
+    constexpr size_t kMaxPointBlock = 128;
+    // Larger row blocks amortize packing on the AVX-512 Eigen backend. Retain
+    // smaller blocks on other backends and require at least two blocks per
+    // worker to avoid doubling the work of a worker handling a small tail.
+    const size_t point_block =
+        cpu::has_avx512_core() && k >= 256 && n / num_threads >= 2 * kMaxPointBlock
+            ? kMaxPointBlock
+            : 64;
     constexpr size_t kCentroidBlock = 256;
-    const size_t blocks = 1 + (n - 1) / kPointBlock;
+    const size_t blocks = 1 + (n - 1) / point_block;
     const auto threads = static_cast<uint32_t>(std::min<size_t>(num_threads, blocks));
     std::vector<double> centroid_norms(k);
     std::vector<std::vector<float>> products(
-        threads, std::vector<float>(kPointBlock * std::min(k, kCentroidBlock))
+        threads, std::vector<float>(point_block * std::min(k, kCentroidBlock))
     );
     std::vector<std::exception_ptr> errors(threads);
 
-#pragma omp parallel for num_threads(threads) schedule(static)
-    for (std::ptrdiff_t cluster = 0; cluster < static_cast<std::ptrdiff_t>(k); ++cluster) {
-        double norm = 0.0;
-        for (size_t dim = 0; dim < d; ++dim) {
-            const double value = centroids[static_cast<size_t>(cluster) * d + dim];
-            norm += value * value;
+    const auto exact_distance = [d, spherical](const float* vector, const float* centroid) {
+        double distance = 0.0;
+        if (spherical) {
+            for (size_t dim = 0; dim < d; ++dim) {
+                distance -= static_cast<double>(vector[dim]) * centroid[dim];
+            }
+            distance += 1.0;
+        } else {
+            for (size_t dim = 0; dim < d; ++dim) {
+                const double difference = static_cast<double>(vector[dim]) - centroid[dim];
+                distance += difference * difference;
+            }
         }
-        centroid_norms[static_cast<size_t>(cluster)] = norm;
-    }
+        return distance;
+    };
 
     // For d <= 65536, this exceeds the float dot-product error bound
     // gamma_d * sum(abs(x_i*c_i)), using 2*abs(x_i*c_i) <= x_i^2+c_i^2.
@@ -101,17 +118,29 @@ inline void exact_assign(
         omp_set_num_threads(1);
         const auto thread = static_cast<size_t>(omp_get_thread_num());
         float* dots = products[thread].data();
-        std::array<double, kPointBlock> point_norms{};
-        std::array<double, kPointBlock> best_distances{};
+        std::array<double, kMaxPointBlock> point_norms{};
+        std::array<double, kMaxPointBlock> best_distances{};
+        // Reuse this worker team for norms and assignment. The implicit barrier
+        // publishes all centroid norms before any matrix tile is screened.
+#pragma omp for schedule(static)
+        for (std::ptrdiff_t cluster = 0; cluster < static_cast<std::ptrdiff_t>(k);
+             ++cluster) {
+            double norm = 0.0;
+            for (size_t dim = 0; dim < d; ++dim) {
+                const double value = centroids[static_cast<size_t>(cluster) * d + dim];
+                norm += value * value;
+            }
+            centroid_norms[static_cast<size_t>(cluster)] = norm;
+        }
 #pragma omp for schedule(static)
         for (std::ptrdiff_t block = 0; block < static_cast<std::ptrdiff_t>(blocks);
              ++block) {
             if (errors[thread]) {
                 continue;
             }
             try {
-                const size_t first_point = static_cast<size_t>(block) * kPointBlock;
-                const size_t rows = std::min(kPointBlock, n - first_point);
+                const size_t first_point = static_cast<size_t>(block) * point_block;
+                const size_t rows = std::min(point_block, n - first_point);
                 for (size_t row = 0; row < rows; ++row) {
                     const float* vector = x + (first_point + row) * d;
                     double norm = 0.0;
@@ -121,6 +150,14 @@ inline void exact_assign(
                     }
                     point_norms[row] = norm;
                     best_distances[row] = std::numeric_limits<double>::infinity();
+                    if (previous_assignments != nullptr) {
+                        const size_t point = first_point + row;
+                        const PID previous = previous_assignments[point];
+                        assert(previous < k);
+                        best_distances[row] =
+                            exact_distance(vector, centroids + previous * d);
+                        assignments[point] = previous;
+                    }
                 }
                 for (size_t first_cluster = 0; first_cluster < k;
                      first_cluster += kCentroidBlock) {
@@ -148,22 +185,12 @@ inline void exact_assign(
                                 continue;
                             }
                             const float* centroid = centroids + cluster * d;
-                            double distance = 0.0;
-                            if (spherical) {
-                                for (size_t dim = 0; dim < d; ++dim) {
-                                    distance -=
-                                        static_cast<double>(vector[dim]) * centroid[dim];
-                                }
-                                distance += 1.0;
-                            } else {
-                                for (size_t dim = 0; dim < d; ++dim) {
-                                    const double difference =
-                                        static_cast<double>(vector[dim]) - centroid[dim];
-                                    distance += difference * difference;
-                                }
-                            }
-                            // Visit in ascending ID order, retaining the first exact tie.
-                            if (distance < best_distances[row]) {
+                            const double distance = exact_distance(vector, centroid);
+                            // The hint may have a higher ID than an equally good
+                            // centroid. Preserve the original lowest-ID tie rule.
+                            if (distance < best_distances[row] ||
+                                (distance == best_distances[row] &&
+                                 cluster < assignments[point])) {
                                 assignments[point] = static_cast<PID>(cluster);
                                 best_distances[row] = distance;
                             }
@@ -289,10 +316,10 @@ class LloydKMeans : public Parameters {
         double previous_obj = std::numeric_limits<double>::infinity();
 
         for (size_t iteration = 0; iteration < niter; ++iteration) {
-            previous_centroids = centroids;
-            assigner.assign(
-                previous_centroids.data(), assignments.data(), distances.data()
-            );
+            assigner.assign(centroids.data(), assignments.data(), distances.data());
+            // Assignment has finished borrowing the current centroids. Keep them
+            // for the shift calculation and overwrite the other buffer with means.
+            previous_centroids.swap(centroids);
 
             const double obj = std::accumulate(distances.begin(), distances.end(), 0.0);
             std::fill(sums.begin(), sums.end(), 0.0);
@@ -332,7 +359,8 @@ class LloydKMeans : public Parameters {
                 spherical,
                 threads,
                 assignments.data(),
-                distances.data()
+                distances.data(),
+                assignments.data()
             );
         } else {
             assigner.assign(centroids.data(), assignments.data(), distances.data(), true);

include/rabitqlib/index/hnsw/hnsw.hpp

@@ -919,8 +919,7 @@ inline void HierarchicalNSW::construct(
 
     num_cluster_ = cluster_num;
     const size_t centroids_bytes = num_cluster_ * padded_dim_ * sizeof(float);
-    centroids_memory_ =
-        static_cast<char*>(memory::aligned_allocate_bytes(64, centroids_bytes));
+    centroids_memory_ = memory::huge_page_allocate<char>(centroids_bytes);
     if (centroids_memory_ == nullptr) {
         throw std::runtime_error("Not enough memory: HNSW failed to allocate centroids");
     }