refactor: share OpenMP, core, binding and HNSW construction helpers

include/rabitqlib/clustering/detail/kmeans.hpp

@@ -9,7 +9,6 @@
 #include <cstddef>
 #include <cstdint>
 #include <cstring>
-#include <exception>
 #include <functional>
 #include <iostream>
 #include <limits>
@@ -22,6 +21,7 @@
 
 #include "rabitqlib/core/detail/simd/matrix_dispatch.hpp"
 #include "rabitqlib/core/rotator.hpp"
+#include "rabitqlib/index/detail/parallel_exception.hpp"
 #include "rabitqlib/index/types.hpp"
 #include "rabitqlib/utils/cpu_features.hpp"
 #include "rabitqlib/utils/space.hpp"
@@ -86,7 +86,7 @@ inline void exact_assign(
     std::vector<std::vector<float>> products(
         num_workers, std::vector<float>(point_block * std::min(k, kCentroidBlock))
     );
-    std::vector<std::exception_ptr> errors(num_workers);
+    index_detail::ParallelException failure;
 
     const auto exact_distance = [d, spherical](const float* vector, const float* centroid) {
         double distance = 0.0;
@@ -136,7 +136,7 @@ inline void exact_assign(
 #pragma omp for schedule(static)
         for (std::ptrdiff_t block = 0; block < static_cast<std::ptrdiff_t>(blocks);
              ++block) {
-            if (errors[thread]) {
+            if (failure.failed()) {
                 continue;
             }
             try {
@@ -201,14 +201,10 @@ inline void exact_assign(
                 for (size_t row = 0; row < rows; ++row) {
                     distances[first_point + row] = static_cast<float>(best_distances[row]);
                 }
-            } catch (...) { errors[thread] = std::current_exception(); }
-        }
-    }
-    for (const auto& error : errors) {
-        if (error) {
-            std::rethrow_exception(error);
+            } catch (...) { failure.capture(); }
         }
     }
+    failure.rethrow();
 }
 
 template <typename Parameters, typename Assignment>

include/rabitqlib/clustering/qgkmeans.hpp

@@ -5,7 +5,6 @@
 #include <algorithm>
 #include <cstddef>
 #include <cstdint>
-#include <exception>
 #include <limits>
 #include <memory>
 #include <numeric>
@@ -17,6 +16,7 @@
 #include "rabitqlib/core/rotator.hpp"
 #include "rabitqlib/fastscan/pack_codes.hpp"
 #include "rabitqlib/fastscan/query.hpp"
+#include "rabitqlib/index/detail/parallel_exception.hpp"
 #include "rabitqlib/index/detail/visited_set.hpp"
 #include "rabitqlib/index/symqg/detail/search_scratch.hpp"
 #include "rabitqlib/index/symqg/qg.hpp"
@@ -256,15 +256,15 @@ class QGAssigner {
                     static_cast<uint32_t>(point);
             }
         }
-        std::vector<std::exception_ptr> errors(num_workers_);
+        index_detail::ParallelException failure;
 #pragma omp parallel for num_threads(num_workers_) schedule(dynamic, 64)
         for (std::ptrdiff_t loop_index = 0; loop_index < static_cast<std::ptrdiff_t>(n);
              ++loop_index) {
             const size_t i = query_order.empty()
                                  ? static_cast<size_t>(loop_index)
                                  : query_order[static_cast<size_t>(loop_index)];
             const int thread_id = omp_get_thread_num();
-            if (errors[static_cast<size_t>(thread_id)]) {
+            if (failure.failed()) {
                 continue;
             }
             try {
@@ -300,16 +300,10 @@ class QGAssigner {
                     distances[i] =
                         distance_func(x + i * dim_, centroids + labels[i] * dim_, dim_);
                 }
-            } catch (...) {
-                errors[static_cast<size_t>(thread_id)] = std::current_exception();
-            }
+            } catch (...) { failure.capture(); }
         }
 
-        for (const auto& error : errors) {
-            if (error) {
-                std::rethrow_exception(error);
-            }
-        }
+        failure.rethrow();
 
         previous_labels_.assign(labels, labels + n);
         query_sums.complete();

include/rabitqlib/clustering/rabitqkmeans.hpp

@@ -6,7 +6,6 @@
 #include <array>
 #include <cstddef>
 #include <cstdint>
-#include <exception>
 #include <limits>
 #include <stdexcept>
 #include <vector>
@@ -19,6 +18,7 @@
 #include "rabitqlib/fastscan/estimator.hpp"
 #include "rabitqlib/fastscan/pack_codes.hpp"
 #include "rabitqlib/fastscan/query.hpp"
+#include "rabitqlib/index/detail/parallel_exception.hpp"
 #include "rabitqlib/index/types.hpp"
 #include "rabitqlib/utils/space.hpp"
 
@@ -126,12 +126,12 @@ class RaBitQAssigner {
         std::vector<std::vector<float>> rotated_blocks(
             num_workers_, std::vector<float>(fastscan::kBatchSize * padded_dim_)
         );
-        std::vector<std::exception_ptr> errors(num_workers_);
+        index_detail::ParallelException failure;
 #pragma omp parallel for num_threads(num_workers_) schedule(static)
         for (std::ptrdiff_t batch = 0; batch < static_cast<std::ptrdiff_t>(num_batches_);
              ++batch) {
             const auto thread = static_cast<size_t>(omp_get_thread_num());
-            if (errors[thread]) {
+            if (failure.failed()) {
                 continue;
             }
             try {
@@ -153,13 +153,9 @@ class RaBitQAssigner {
                         block * QGBatchDataMap<float>::data_bytes(padded_dim_),
                     metric_type_
                 );
-            } catch (...) { errors[thread] = std::current_exception(); }
-        }
-        for (const auto& error : errors) {
-            if (error) {
-                std::rethrow_exception(error);
-            }
+            } catch (...) { failure.capture(); }
         }
+        failure.rethrow();
     }
 
     void assign(const float* centroids, PID* labels, float* distances, bool = false) {
@@ -175,12 +171,11 @@ class RaBitQAssigner {
             query.set_g_add(distance_func(centroid, center_.data(), dim_));
         }
 
-        std::vector<std::exception_ptr> errors(num_workers_);
+        index_detail::ParallelException failure;
 #pragma omp parallel for num_threads(num_workers_) schedule(static)
         for (std::ptrdiff_t batch = 0; batch < static_cast<std::ptrdiff_t>(num_batches_);
              ++batch) {
-            const auto thread = static_cast<size_t>(omp_get_thread_num());
-            if (errors[thread]) {
+            if (failure.failed()) {
                 continue;
             }
             try {
@@ -218,13 +213,9 @@ class RaBitQAssigner {
                         data_ + id * dim_, centroids + best_labels[point] * dim_, dim_
                     );
                 }
-            } catch (...) { errors[thread] = std::current_exception(); }
-        }
-        for (const auto& error : errors) {
-            if (error) {
-                std::rethrow_exception(error);
-            }
+            } catch (...) { failure.capture(); }
         }
+        failure.rethrow();
     }
 };
 

include/rabitqlib/core/detail/rabitq_impl.hpp

@@ -19,6 +19,7 @@
 #include "rabitqlib/core/metric.hpp"
 #include "rabitqlib/core/pack_excode.hpp"
 #include "rabitqlib/core/types.hpp"
+#include "rabitqlib/utils/bitops.hpp"
 #include "rabitqlib/utils/openmp.hpp"  // IWYU pragma: keep (configuration macros)
 #include "rabitqlib/utils/space.hpp"
 
@@ -90,6 +91,29 @@ inline CodeFactorTerms<T> code_factor_terms(
     throw std::invalid_argument(invalid_metric_message);
 }
 
+// Rescale factor for a code centered at its own width, given code_factor_terms() of
+// that centered code. Normalized ex-bit codes use their own rescale expression.
+template <typename T>
+inline T direct_rescale(T l2_sqr, T residual_code_ip, MetricType metric_type) {
+    return l2_sqr == 0                ? T{0}
+           : metric_type == METRIC_L2 ? -2 * l2_sqr / residual_code_ip
+                                      : -l2_sqr / residual_code_ip;
+}
+
+/**
+ * Inverse of pack_binary_to_bytes<uint64_t>: write each coordinate's sign bit (0 or
+ * 1) to signs. Full 64-coordinate words come first, followed by a final 32-bit word.
+ */
+inline void unpack_sign_bits(const uint8_t* packed, size_t dim, uint16_t* signs) {
+    for (size_t start = 0; start < dim; start += 64) {
+        const size_t width = std::min(size_t{64}, dim - start);
+        const uint64_t word = bitops::load_word(packed + start / 8, width / 8);
+        for (size_t i = 0; i < width; ++i) {
+            signs[start + i] = static_cast<uint16_t>((word >> (width - 1 - i)) & 1U);
+        }
+    }
+}
+
 namespace one_bit {
 
 /**
@@ -147,9 +171,7 @@ inline void one_bit_code_with_factor_into(
         "Unsupported metric type in quantization"
     );
     f_add = factors.add;
-    f_rescale = l2_sqr == 0                ? T{0}
-                : metric_type == METRIC_L2 ? -2 * l2_sqr / factors.residual_code_ip
-                                           : -l2_sqr / factors.residual_code_ip;
+    f_rescale = direct_rescale(l2_sqr, factors.residual_code_ip, metric_type);
     f_error = factors.error;
 }
 
@@ -976,9 +998,7 @@ inline void code_factors(
         "Unsupported metric type in code_factors()"
     );
     f_add = factors.add;
-    f_rescale = l2_sqr == 0                ? T{0}
-                : metric_type == METRIC_L2 ? -2 * l2_sqr / factors.residual_code_ip
-                                           : -l2_sqr / factors.residual_code_ip;
+    f_rescale = direct_rescale(l2_sqr, factors.residual_code_ip, metric_type);
     f_error = factors.error;
 }
 

include/rabitqlib/index/detail/exact_nn.hpp

@@ -20,14 +20,7 @@ inline PID exact_nn(
     T (*dist_func)(const T*, const T*, size_t)
 ) {
 #if RABITQ_USE_OPENMP
-    const auto num_workers = static_cast<int>(std::max<size_t>(
-        1,
-        std::min(
-            {resolve_num_threads(num_threads),
-             std::max<size_t>(num_points, 1),
-             static_cast<size_t>(std::numeric_limits<int>::max())}
-        )
-    ));
+    const int num_workers = detail::openmp_team_size(num_threads, num_points);
 #else
     (void)num_threads;
     constexpr int num_workers = 1;

include/rabitqlib/index/detail/parallel_for.hpp

@@ -8,6 +8,7 @@
 #include <thread>
 #include <vector>
 
+#include "rabitqlib/index/detail/parallel_exception.hpp"
 #include "rabitqlib/utils/openmp.hpp"
 #include "rabitqlib/utils/tools.hpp"
 
@@ -37,21 +38,17 @@ inline void parallel_for(
             // 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);
+            ParallelException failure;
 #pragma omp parallel for num_threads(num_threads) schedule(dynamic)
             for (size_t id = start; id < end; ++id) {
-                const size_t worker = static_cast<size_t>(omp_get_thread_num());
-                if (!errors[worker]) {
-                    try {
-                        fn(id, worker);
-                    } catch (...) { errors[worker] = std::current_exception(); }
-                }
-            }
-            for (const auto& error : errors) {
-                if (error) {
-                    std::rethrow_exception(error);
+                if (failure.failed()) {
+                    continue;
                 }
+                try {
+                    fn(id, static_cast<size_t>(omp_get_thread_num()));
+                } catch (...) { failure.capture(); }
             }
+            failure.rethrow();
             return;
         }
 #endif
@@ -113,4 +110,24 @@ inline void parallel_for(
     }
 }
 
+#if RABITQ_USE_OPENMP
+// Runs fn(i, worker) for i in [0, count) on an OpenMP team with dynamic scheduling.
+// After an iteration throws, the remaining ones are skipped; the first exception is
+// rethrown once the team has joined.
+template <class Function>
+inline void omp_for_dynamic(size_t count, size_t num_workers, Function fn) {
+    ParallelException failure;
+#pragma omp parallel for num_threads(num_workers) schedule(dynamic)
+    for (std::ptrdiff_t index = 0; index < static_cast<std::ptrdiff_t>(count); ++index) {
+        if (failure.failed()) {
+            continue;
+        }
+        try {
+            fn(static_cast<size_t>(index), static_cast<size_t>(omp_get_thread_num()));
+        } catch (...) { failure.capture(); }
+    }
+    failure.rethrow();
+}
+#endif
+
 }  // namespace rabitqlib::index_detail

include/rabitqlib/index/hnsw/hnsw.hpp

@@ -438,7 +438,12 @@ class HierarchicalNSW {
 
     void add_point(PID, PID, const quant::RabitqConfig&);
 
-    maxheap<std::pair<float, PID>> search_base_layer(PID, PID, int);
+    // Insertion steps shared by raw and quantized scoring. distance(id) scores an
+    // existing node against the node being inserted; prefetch(id) may warm its data.
+    template <class Distance>
+    PID greedy_descend(PID, int from_level, int to_level, Distance);
+    template <class Distance, class Prefetch>
+    maxheap<std::pair<float, PID>> search_layer(PID, int, Distance, Prefetch);
 
     PID mutually_connect_new_element(PID, maxheap<std::pair<float, PID>>&, int);
 
@@ -451,9 +456,6 @@ class HierarchicalNSW {
     void reconstruct_rotated(PID, float*) const;
     float get_quant_dist(PID target, PID query, detail::QuantQuery&) const;
     void add_point_quant(PID, int, LinkStorage&, detail::QuantQuery&);
-    maxheap<std::pair<float, PID>> search_base_layer_quant(
-        PID, PID, int, detail::QuantQuery&
-    );
     PID mutually_connect_quant(
         PID, maxheap<std::pair<float, PID>>&, int, detail::QuantQuery&
     );

include/rabitqlib/utils/space.hpp

@@ -17,10 +17,6 @@
 #include "rabitqlib/utils/tools.hpp"
 #include "rabitqlib/utils/warmup_space.hpp"
 
-#if RABITQ_USE_OPENMP
-#include <limits>
-#endif
-
 namespace rabitqlib {
 
 inline uint64_t load_unaligned_u64(const void* data) {
@@ -191,14 +187,7 @@ inline std::vector<T> compute_centroid(
     const T* data, size_t num_points, size_t dim, size_t num_threads
 ) {
 #if RABITQ_USE_OPENMP
-    const auto thread_count = static_cast<int>(std::max<size_t>(
-        1,
-        std::min(
-            {resolve_num_threads(num_threads),
-             std::max<size_t>(num_points, 1),
-             static_cast<size_t>(std::numeric_limits<int>::max())}
-        )
-    ));
+    const int thread_count = detail::openmp_team_size(num_threads, num_points);
 #else
     (void)num_threads;
     constexpr int thread_count = 1;

include/rabitqlib/utils/tools.hpp

@@ -5,6 +5,7 @@
 #include <cstddef>
 #include <cstdint>
 #include <functional>
+#include <limits>
 #include <queue>
 #include <random>
 #include <thread>
@@ -93,6 +94,20 @@ inline size_t resolve_num_threads(size_t requested) {
     return requested == 0 ? available : std::min(requested, available);
 }
 
+namespace detail {
+// OpenMP team size for work_items items: at least one and representable as int.
+inline int openmp_team_size(size_t requested, size_t work_items) {
+    return static_cast<int>(std::max<size_t>(
+        1,
+        std::min(
+            {resolve_num_threads(requested),
+             work_items,
+             static_cast<size_t>(std::numeric_limits<int>::max())}
+        )
+    ));
+}
+}  // namespace detail
+
 template <typename T, typename TP>
 float distance_ratio(
     const T* data,