Fix parallel index execution and zero-range quantization

include/rabitqlib/index/detail/parallel_for.hpp

@@ -13,7 +13,13 @@
 
 namespace rabitqlib::index_detail {
 template <class Function>
-inline void parallel_for(size_t start, size_t end, size_t num_threads, Function fn) {
+inline void parallel_for(
+    size_t start,
+    size_t end,
+    size_t num_threads,
+    Function fn,
+    [[maybe_unused]] bool reuse_workers = false
+) {
     if (start >= end) {
         return;
     }
@@ -25,9 +31,12 @@ inline void parallel_for(size_t start, size_t end, size_t num_threads, Function
         }
     } else {
 #if defined(_OPENMP) && _OPENMP >= 201307
-        if (omp_get_proc_bind() != omp_proc_bind_false) {
-            // Native threads inherit only the caller's place. Let OpenMP honor
-            // the configured placement of the whole team when binding is enabled.
+        if ((reuse_workers && omp_get_level() == 0) ||
+            omp_get_proc_bind() != omp_proc_bind_false) {
+            // Independent query batches can reuse workers; graph insertion keeps
+            // its existing scheduling. Binding requires OpenMP placement because
+            // native threads inherit only the caller's place. Keep unbound nested
+            // calls on the existing native-thread path.
             std::vector<std::exception_ptr> errors(num_threads);
 #pragma omp parallel for num_threads(num_threads) schedule(dynamic)
             for (size_t id = start; id < end; ++id) {

include/rabitqlib/index/ivf/initializer.hpp

@@ -7,7 +7,6 @@
 #include <fstream>
 #include <ios>
 #include <memory>
-#include <mutex>
 #include <queue>
 #include <stdexcept>
 #include <string>
@@ -192,7 +191,6 @@ class HNSWInitializer : public Initializer {
     hnswlib::SpaceInterface<float>* space_;
     MetricType metric_type_;
     std::vector<hnswlib::tableint> label_to_internal_;
-    mutable std::mutex search_mutex_;
 
     void rebuild_label_lookup() {
         if (alg_hnsw_->cur_element_count != num_cluster_) {
@@ -253,10 +251,8 @@ class HNSWInitializer : public Initializer {
     void centroids_distances(
         const float* query, size_t nprobe, std::vector<AnnCandidate<float>>& candidates
     ) const override {
-        std::lock_guard<std::mutex> lock(search_mutex_);
-        alg_hnsw_->setEf(std::max(size_t{768}, 2 * nprobe));
         std::priority_queue<std::pair<float, hnswlib::labeltype>> result =
-            alg_hnsw_->searchKnn(query, nprobe);
+            alg_hnsw_->searchKnnWithEf(query, nprobe, std::max(size_t{768}, 2 * nprobe));
         if (result.size() != nprobe) {
             throw std::runtime_error("HNSW centroid graph returned too few candidates");
         }

include/rabitqlib/third/hnswlib/hnswalg.h

@@ -1158,6 +1158,7 @@ class HierarchicalNSW : public AlgorithmInterface<dist_t> {
 
     tableint addPoint(const void *data_point, labeltype label, int level) {
         tableint cur_c = 0;
+        int curlevel = 0;
         {
             // Checking if the element with the same label already exists
             // if so, updating it *instead* of creating a new element.
@@ -1187,21 +1188,22 @@ class HierarchicalNSW : public AlgorithmInterface<dist_t> {
             cur_c = cur_element_count;
             cur_element_count++;
             label_lookup_[label] = cur_c;
+            // RaBitQ: the shared generator must advance under the registration lock.
+            curlevel = getRandomLevel(mult_);
         }
 
         std::unique_lock <std::mutex> lock_el(link_list_locks_[cur_c]);
-        int curlevel = getRandomLevel(mult_);
         if (level > 0)
             curlevel = level;
 
         element_levels_[cur_c] = curlevel;
 
         std::unique_lock <std::mutex> templock(global);
         int maxlevelcopy = maxlevel_;
-        if (curlevel <= maxlevelcopy)
-            templock.unlock();
         tableint currObj = enterpoint_node_;
         tableint enterpoint_copy = enterpoint_node_;
+        if (curlevel <= maxlevelcopy)
+            templock.unlock();
 
         memset(data_level0_memory_ + cur_c * size_data_per_element_ + offsetLevel0_, 0, size_data_per_element_);
 
@@ -1275,6 +1277,13 @@ class HierarchicalNSW : public AlgorithmInterface<dist_t> {
 
     std::priority_queue<std::pair<dist_t, labeltype >>
     searchKnn(const void *query_data, size_t k, BaseFilterFunctor* isIdAllowed = nullptr) const {
+        return searchKnnWithEf(query_data, k, ef_, isIdAllowed);
+    }
+
+
+    // RaBitQ extension: concurrent searches can choose ef without changing shared state.
+    std::priority_queue<std::pair<dist_t, labeltype >>
+    searchKnnWithEf(const void *query_data, size_t k, size_t ef, BaseFilterFunctor* isIdAllowed = nullptr) const {
         std::priority_queue<std::pair<dist_t, labeltype >> result;
         if (cur_element_count == 0) return result;
 
@@ -1312,10 +1321,10 @@ class HierarchicalNSW : public AlgorithmInterface<dist_t> {
         bool bare_bone_search = !num_deleted_ && !isIdAllowed;
         if (bare_bone_search) {
             top_candidates = searchBaseLayerST<true>(
-                    currObj, query_data, std::max(ef_, k), isIdAllowed);
+                    currObj, query_data, std::max(ef, k), isIdAllowed);
         } else {
             top_candidates = searchBaseLayerST<false>(
-                    currObj, query_data, std::max(ef_, k), isIdAllowed);
+                    currObj, query_data, std::max(ef, k), isIdAllowed);
         }
 
         while (top_candidates.size() > k) {

include/rabitqlib/utils/space.hpp

@@ -37,6 +37,10 @@ void scalar_quantize_normal(
     float lo,
     float delta
 ) {
+    if (delta == 0.0F) {
+        std::fill_n(result, dim, T{0});
+        return;
+    }
     float one_over_delta = 1.0F / delta;
 
     ConstRowMajorArrayMap<float> v0(vec0, 1, static_cast<long>(dim));
@@ -188,6 +192,7 @@ inline void data_range(const T* __restrict__ vec0, size_t dim, T& lo, T& hi) {
     hi = v0.maxCoeff();
 }
 
+// A constant range (delta == 0) is represented by zero codes.
 template <typename T, typename TD>
 void scalar_quantize(
     T* __restrict__ result, const TD* __restrict__ vec0, size_t dim, TD lo, TD delta

src/index/hnsw.cpp

@@ -951,13 +951,17 @@ std::vector<std::vector<std::pair<float, PID>>> HierarchicalNSW::search(
     }
     std::vector<std::vector<std::pair<float, PID>>> results(num_queries);
     rabitqlib::index_detail::parallel_for(
-        0, num_queries, num_threads, [&](size_t idx, size_t /*threadId*/) {
+        0,
+        num_queries,
+        num_threads,
+        [&](size_t idx, size_t /*threadId*/) {
             std::vector<float> rotated_query(padded_dim_);
             this->rotator_->rotate(queries + (idx * dim_), rotated_query.data());
             results[idx] = detail::search_knn_sorted(
                 *this, rotated_query.data(), k, std::max(ef_search, k)
             );
-        }
+        },
+        /*reuse_workers=*/true
     );
     return results;
 }

src/index/pipnn.cpp

@@ -7,6 +7,7 @@
 #include <cstddef>
 #include <cstdint>
 #include <iterator>
+#include <limits>
 #include <mutex>
 #include <numeric>
 #include <random>
@@ -28,6 +29,7 @@ namespace rabitqlib::symqg::detail {
 
 namespace pipnn_impl {
 constexpr size_t kTileSize = 2048;
+constexpr size_t kMinTileSize = 64;
 constexpr size_t kHashBits = 12;
 constexpr size_t kLocalDegree = 2;
 constexpr float kAlpha = 1.1F;
@@ -180,6 +182,33 @@ void pairwise(Workspace& work, size_t size, size_t dim, MetricType metric) {
     );
 }
 
+void select_local_neighbors(const Bucket& ids, Workspace& work) {
+    static_assert(kLocalDegree == 2);
+    work.candidates.assign(
+        ids.size() * kLocalDegree, {std::numeric_limits<float>::infinity(), kPidMax}
+    );
+    const auto keep_nearest = [&](size_t row, float distance, PID target) {
+        auto* best = work.candidates.data() + row * kLocalDegree;
+        const std::pair<float, PID> candidate(distance, target);
+        if (candidate < best[1]) {
+            best[1] = candidate;
+            if (best[1] < best[0]) {
+                std::swap(best[0], best[1]);
+            }
+        }
+    };
+    // Both metrics are symmetric. Read each stored pair once, contiguously,
+    // instead of gathering the upper half through a strided column walk.
+    for (size_t i = 1; i < ids.size(); ++i) {
+        const float* distances = work.distances + i * ids.size();
+        for (size_t j = 0; j < i; ++j) {
+            const float distance = distances[j];
+            keep_nearest(i, distance, ids[j]);
+            keep_nearest(j, distance, ids[i]);
+        }
+    }
+}
+
 std::vector<Bucket> partition(
     const float* data,
     size_t dim,
@@ -189,6 +218,15 @@ std::vector<Bucket> partition(
     ScratchPool& scratch
 ) {
     const auto& ids = job.ids;
+    // Small root partitions need enough tiles to occupy the requested team.
+    // Deeper levels already distribute independent buckets across workers.
+    const size_t tile_size =
+        job.depth == 0
+            ? std::min(
+                  kTileSize,
+                  round_up_to_multiple(div_round_up(ids.size(), threads), kMinTileSize)
+              )
+            : kTileSize;
     const size_t leader_count = std::min(
         ids.size(), std::clamp<size_t>(job.depth == 0 ? 512 : ids.size() / 200, 3, 512)
     );
@@ -210,7 +248,10 @@ std::vector<Bucket> partition(
     simd::row_norms(leader_data.data(), leader_norms.data(), leader_count, dim);
     std::vector<PID> assignments(ids.size() * fanout);
     const auto assign_tile = [&](size_t begin, Workspace& work) {
-        const size_t count = std::min(kTileSize, ids.size() - begin);
+        // Both callers run inside a PiPNN OpenMP task, including one-thread teams.
+        // Keep Eigen within that worker without changing the caller's thread budget.
+        omp_set_num_threads(1);
+        const size_t count = std::min(tile_size, ids.size() - begin);
         RowMajorMatrixMap<float> points(
             work.points, static_cast<Eigen::Index>(count), static_cast<Eigen::Index>(dim)
         );
@@ -251,13 +292,13 @@ std::vector<Bucket> partition(
     if (job.depth == 0) {
         {
             index_detail::ParallelException failure;
-#pragma omp parallel num_threads(threads)
+#pragma omp parallel num_threads(std::min(threads, div_round_up(ids.size(), tile_size)))
             {
                 auto work = scratch.worker(static_cast<size_t>(omp_get_thread_num()));
 #pragma omp for schedule(static)
                 for (std::ptrdiff_t offset = 0;
                      offset < static_cast<std::ptrdiff_t>(ids.size());
-                     offset += static_cast<std::ptrdiff_t>(kTileSize)) {
+                     offset += static_cast<std::ptrdiff_t>(tile_size)) {
                     if (failure.failed()) {
                         continue;
                     }
@@ -273,7 +314,7 @@ std::vector<Bucket> partition(
         // Already inside the bucket-level team. Eigen must see that team rather
         // than a serialized nested region when deciding its own parallelism.
         auto work = scratch.worker(static_cast<size_t>(omp_get_thread_num()));
-        for (size_t begin = 0; begin < ids.size(); begin += kTileSize) {
+        for (size_t begin = 0; begin < ids.size(); begin += tile_size) {
             assign_tile(begin, work);
         }
     }
@@ -367,7 +408,7 @@ std::vector<Bucket> cluster(
         } else {
             {
                 index_detail::ParallelException failure;
-#pragma omp parallel for num_threads(threads) schedule(dynamic)
+#pragma omp parallel for num_threads(std::min(threads, active.size())) schedule(dynamic)
                 for (std::ptrdiff_t i = 0; i < static_cast<std::ptrdiff_t>(active.size());
                      ++i) {
                     if (failure.failed()) {
@@ -424,17 +465,22 @@ InitialGraph build_initial_graph(
         projections.data()[i] = normal(random);
     }
     RowMajorMatrix<float> sketches(count, kHashBits);
+    const size_t sketch_tile = std::min(
+        kTileSize, round_up_to_multiple(div_round_up(count, threads), kMinTileSize)
+    );
     {
         index_detail::ParallelException failure;
-#pragma omp parallel for num_threads(threads) schedule(static)
+#pragma omp parallel for num_threads(std::min(threads, div_round_up(count, sketch_tile))) \
+    schedule(static)
         for (std::ptrdiff_t offset = 0; offset < static_cast<std::ptrdiff_t>(count);
-             offset += static_cast<std::ptrdiff_t>(kTileSize)) {
+             offset += static_cast<std::ptrdiff_t>(sketch_tile)) {
             if (failure.failed()) {
                 continue;
             }
             try {
                 const size_t begin = static_cast<size_t>(offset);
-                const size_t rows = std::min(kTileSize, count - begin);
+                const size_t rows = std::min(sketch_tile, count - begin);
+                omp_set_num_threads(1);
                 simd::matrix_product(
                     data + begin * dim,
                     projections.data(),
@@ -450,7 +496,7 @@ InitialGraph build_initial_graph(
     CandidateTable table(count, degree * 5 / 2);
     {
         index_detail::ParallelException failure;
-#pragma omp parallel num_threads(threads)
+#pragma omp parallel num_threads(std::min(threads, leaves.size()))
         {
             auto work = scratch.worker(static_cast<size_t>(omp_get_thread_num()));
 #pragma omp for schedule(dynamic)
@@ -468,25 +514,14 @@ InitialGraph build_initial_graph(
                     }
                     gather(data, dim, ids, work);
                     pairwise(work, ids.size(), dim, metric);
+                    select_local_neighbors(ids, work);
+                    const size_t keep = std::min(kLocalDegree, ids.size() - 1);
+                    // Retain the original point order and (distance, ID) tie order
+                    // when publishing forward and reverse proposals.
                     for (size_t i = 0; i < ids.size(); ++i) {
-                        work.candidates.clear();
-                        for (size_t j = 0; j < ids.size(); ++j) {
-                            if (i != j) {
-                                work.candidates.emplace_back(
-                                    work.distances
-                                        [std::max(i, j) * ids.size() + std::min(i, j)],
-                                    ids[j]
-                                );
-                            }
-                        }
-                        const size_t keep = std::min(kLocalDegree, work.candidates.size());
-                        std::partial_sort(
-                            work.candidates.begin(),
-                            work.candidates.begin() + static_cast<ptrdiff_t>(keep),
-                            work.candidates.end()
-                        );
                         for (size_t j = 0; j < keep; ++j) {
-                            const auto [distance, target] = work.candidates[j];
+                            const auto [distance, target] =
+                                work.candidates[i * kLocalDegree + j];
                             table.merge(ids[i], target, distance, sketches);
                             table.merge(target, ids[i], distance, sketches);
                         }

src/index/pipnn_internal.hpp

@@ -34,6 +34,9 @@ struct ScratchPool {
 
 void gather(const float*, size_t, const Bucket&, Workspace&);
 void pairwise(Workspace&, size_t, size_t, MetricType);
+// After pairwise(), write two sorted (distance, ID) candidates per local row.
+// Excludes self; a missing second neighbor retains kPidMax.
+void select_local_neighbors(const Bucket&, Workspace&);
 std::vector<Bucket> cluster(const float*, size_t, size_t, MetricType, size_t, ScratchPool&);
 
 }  // namespace rabitqlib::symqg::detail::pipnn_impl

src/simd/portable_generic.cpp

@@ -135,6 +135,10 @@ float excode_ip(const float* query, const uint8_t* compact, size_t dim) {
 
 template <typename T>
 void quantize(T* result, const float* data, size_t dim, float lo, float delta) {
+    if (delta == 0.0F) {
+        std::fill_n(result, dim, T{0});
+        return;
+    }
     const float reciprocal = 1.0F / delta;
     for (size_t i = 0; i < dim; ++i) {
         result[i] = static_cast<T>(std::round((data[i] - lo) * reciprocal));

src/simd/query_neon.cpp

@@ -30,6 +30,10 @@ inline uint32x4_t quantize4(const float* data, float32x4_t lo, float32x4_t recip
 void scalar_quantize_uint8_neon(
     uint8_t* result, const float* data, size_t dim, float lo, float delta
 ) {
+    if (delta == 0.0F) {
+        std::fill_n(result, dim, uint8_t{0});
+        return;
+    }
     const float reciprocal = 1.0F / delta;
     const float32x4_t low = vdupq_n_f32(lo), scale = vdupq_n_f32(reciprocal);
     size_t i = 0;
@@ -48,6 +52,10 @@ void scalar_quantize_uint8_neon(
 void scalar_quantize_uint16_neon(
     uint16_t* result, const float* data, size_t dim, float lo, float delta
 ) {
+    if (delta == 0.0F) {
+        std::fill_n(result, dim, uint16_t{0});
+        return;
+    }
     const float reciprocal = 1.0F / delta;
     const float32x4_t low = vdupq_n_f32(lo), scale = vdupq_n_f32(reciprocal);
     size_t i = 0;