Accelerate quantization rescale search with SIMD interval pruning

CMakeLists.txt

@@ -50,6 +50,7 @@ set(RABITQ_COMMON_SOURCES
 )
 
 set(RABITQ_AVX2_SOURCES
+    src/simd/quantization_avx2.cpp
     src/simd/pack_excode_avx2.cpp
     src/simd/space_excode_avx2.cpp
     src/simd/space_avx2.cpp
@@ -59,6 +60,7 @@ set(RABITQ_AVX2_SOURCES
 )
 
 set(RABITQ_AVX512_SOURCES
+    src/simd/quantization_avx512.cpp
     src/simd/pack_excode_avx512.cpp
     src/simd/space_excode_avx512.cpp
     src/simd/space_avx512.cpp

include/rabitqlib/quantization/rabitq_impl.hpp

@@ -9,15 +9,16 @@
 #include <cstddef>
 #include <cstdint>
 #include <functional>
-#include <queue>
 #include <stdexcept>
+#include <type_traits>
 #include <utility>
 #include <vector>
 
 #include "rabitqlib/defines.hpp"
 #include "rabitqlib/fastscan/fastscan.hpp"
 #include "rabitqlib/quantization/data_layout.hpp"
 #include "rabitqlib/quantization/pack_excode.hpp"
+#include "rabitqlib/simd/quantization_dispatch.hpp"
 #include "rabitqlib/utils/space.hpp"
 
 namespace rabitqlib::quant::rabitq_impl {
@@ -321,46 +322,90 @@ inline double best_rescale_factor(const T* o_abs, size_t dim, size_t ex_bits) {
     const double t_end = static_cast<double>(max_code + 10) / max_o;
     const double t_start = t_end * kTightStart[ex_bits];
 
+    // Both paths maximize cosine similarity between the input and quantized
+    // directions over the configured scale interval. Here a_i = o_abs[i] is a
+    // rotated, normalized residual magnitude, and q_i = code_i + 0.5.
+    // With ||a|| = 1, cosine(a, q) = dot(a, q / ||q||) = N / sqrt(S), where
+    // N = sum(a_i * q_i) and S = sum(q_i * q_i).
+    // SIMD evaluates whole vectors at selected scales and prunes intervals;
+    // its search and bounds are in src/simd/rescale_search.hpp.
+    if constexpr (std::is_same_v<T, float>) {
+        const double t = simd::best_rescale_factor(o_abs, dim, max_code, t_start, t_end);
+        if (t >= 0)
+            return t;
+    }
+
+    // Scalar event sweep, also used if SIMD is unavailable or cannot certify a
+    // winner within its work limit. Each heap entry holds a coordinate's next
+    // code-changing scale (code_i + 1) / a_i. Visiting only these events skips
+    // ranges with constant codes and updates N and S in O(1) per coordinate.
+    // Heap traversal is sequential and branch-dependent, so it offers little
+    // SIMD parallelism; the SIMD path instead prunes ranges of candidate codes.
     using Event = std::pair<double, size_t>;
-    std::priority_queue<Event, std::vector<Event>, std::greater<Event>> next_t;
+    std::vector<Event> next_t;
+    next_t.reserve(dim);
     std::vector<int> cur_o_bar(dim);
     double sqr_denominator = static_cast<double>(dim) * 0.25;
     double numerator = 0;
 
-    auto enqueue_next = [&](size_t i) {
-        const double magnitude = static_cast<double>(o_abs[i]);
-        // Never increment a saturated coordinate; skip zero coordinates.
-        if (magnitude > 0 && cur_o_bar[i] < max_code) {
-            const double next = static_cast<double>(cur_o_bar[i] + 1) / magnitude;
-            if (next < t_end)
-                next_t.emplace(next, i);
-        }
-    };
-
+    // init quantization codes for each coordinate by t_start
     for (size_t i = 0; i < dim; ++i) {
         const double magnitude = static_cast<double>(o_abs[i]);
         const int cur = quantized_level_at_scale(magnitude, t_start, max_code);
         cur_o_bar[i] = cur;
         sqr_denominator += (cur * cur) + cur;
         numerator += (cur + 0.5) * magnitude;
-        enqueue_next(i);
+        // Never increment a saturated coordinate; skip zero coordinates.
+        if (magnitude > 0 && cur < max_code) {
+            const double next = static_cast<double>(cur + 1) / magnitude;
+            if (next < t_end)
+                next_t.emplace_back(next, i);
+        }
     }
+    std::make_heap(next_t.begin(), next_t.end(), std::greater<Event>{});
 
     // The initial state may already be the best state inside the interval.
     double max_ip = numerator / std::sqrt(sqr_denominator);
     double best_t = t_start;
 
     while (!next_t.empty()) {
-        const double cur_t = next_t.top().first;
+        const double cur_t = next_t.front().first;
         // All coordinates crossing the same threshold change together.
         do {
-            const size_t i = next_t.top().second;
-            next_t.pop();
+            const size_t i = next_t.front().second;
             ++cur_o_bar[i];
+            // With the new code k, q changes from k-0.5 to k+0.5: delta S = 2k.
             sqr_denominator += 2.0 * cur_o_bar[i];
-            numerator += static_cast<double>(o_abs[i]);
-            enqueue_next(i);
-        } while (!next_t.empty() && next_t.top().first == cur_t);
+            const double magnitude = static_cast<double>(o_abs[i]);
+            numerator += magnitude;
+
+            Event next{t_end, i};
+            if (cur_o_bar[i] < max_code)
+                next.first = static_cast<double>(cur_o_bar[i] + 1) / magnitude;
+            if (next.first >= t_end) {
+                // No further event for this coordinate: fill the root's hole
+                // with the last heap entry, then restore the heap below.
+                next = next_t.back();
+                next_t.pop_back();
+            }
+            if (!next_t.empty()) {
+                // Replace the consumed root event with `next`. Its threshold
+                // cannot precede the old minimum, so one sift-down suffices;
+                // priority_queue pop() followed by push() would repair twice.
+                size_t parent = 0;
+                size_t child = 1;
+                while (child < next_t.size()) {
+                    if (child + 1 < next_t.size() && next_t[child + 1] < next_t[child])
+                        ++child;
+                    if (!(next_t[child] < next))
+                        break;
+                    next_t[parent] = next_t[child];
+                    parent = child;
+                    child = 2 * parent + 1;
+                }
+                next_t[parent] = next;
+            }
+        } while (!next_t.empty() && next_t.front().first == cur_t);
 
         const double cur_ip = numerator / std::sqrt(sqr_denominator);
         if (cur_ip > max_ip) {

include/rabitqlib/simd/quantization_dispatch.hpp

@@ -0,0 +1,20 @@
+#pragma once
+
+#include <cstddef>
+
+namespace rabitqlib::simd {
+
+// Search normalized, nonnegative float magnitudes with max_code in [0, 255].
+// A negative result requests the event sweep when the bounded search cannot
+// certify a winner within its work limit.
+double best_rescale_factor(
+    const float* magnitudes, size_t dim, int max_code, double start, double end
+);
+double best_rescale_factor_avx2(
+    const float* magnitudes, size_t dim, int max_code, double start, double end
+);
+double best_rescale_factor_avx512(
+    const float* magnitudes, size_t dim, int max_code, double start, double end
+);
+
+}  // namespace rabitqlib::simd

src/simd/dispatch.cpp

@@ -7,13 +7,31 @@
 
 #include "rabitqlib/simd/fastscan_dispatch.hpp"
 #include "rabitqlib/simd/pack_excode_dispatch.hpp"
+#include "rabitqlib/simd/quantization_dispatch.hpp"
 #include "rabitqlib/simd/rotator_dispatch.hpp"
 #include "rabitqlib/simd/space_dispatch.hpp"
 #include "rabitqlib/simd/warmup_dispatch.hpp"
 #include "rabitqlib/utils/cpu_features.hpp"
+#include "rescale_search.hpp"
 
 namespace rabitqlib::simd {
 
+namespace detail {
+
+RescaleScratch& get_thread_local_rescale_scratch(size_t dim) {
+    // Search evaluation is synchronous and does not re-enter the quantizer.
+    // Retain the largest buffers on each worker; smaller vectors overwrite only
+    // their active prefix without shrinking or zero-initializing the storage.
+    thread_local RescaleScratch scratch;
+    if (scratch.magnitudes.size() < dim)
+        scratch.magnitudes.resize(dim);
+    if (scratch.reciprocals.size() < dim)
+        scratch.reciprocals.resize(dim);
+    return scratch;
+}
+
+}  // namespace detail
+
 [[noreturn]] static void missing_feature(const char* feature_name) {
     throw std::runtime_error(
         std::string(feature_name) + " requires AVX2/FMA or AVX512 support"
@@ -29,6 +47,27 @@ static float ip_fxu0(
     return 0.0F;
 }
 
+static double request_scalar_rescale_search(const float*, size_t, int, double, double) {
+    return -1;
+}
+
+using BestRescaleFactorFn = double (*)(const float*, size_t, int, double, double);
+const BestRescaleFactorFn kBestRescaleFactorFn = [] {
+    if (cpu::has_avx512_core()) {
+        return best_rescale_factor_avx512;
+    } else if (cpu::has_avx2()) {
+        return best_rescale_factor_avx2;
+    } else {
+        return request_scalar_rescale_search;
+    }
+}();
+
+double best_rescale_factor(
+    const float* magnitudes, size_t dim, int max_code, double start, double end
+) {
+    return kBestRescaleFactorFn(magnitudes, dim, max_code, start, end);
+}
+
 static float missing_excode_ip(const float*, const uint8_t*, size_t) {
     missing_feature("excode ip functions");
 }

src/simd/quantization_avx2.cpp

@@ -0,0 +1,167 @@
+#include <immintrin.h>
+
+#include <algorithm>
+#include <cstddef>
+#include <limits>
+
+#include "rabitqlib/simd/quantization_dispatch.hpp"
+#include "rescale_search.hpp"
+namespace rabitqlib::simd::detail {
+
+// Evaluate one candidate scale, processing four coordinates per iteration.
+// For a_i = magnitudes[i], c_i = min(floor(scale*a_i), max_code), q_i = c_i+0.5,
+// return N = sum(a_i*q_i), S = sum(q_i*q_i), and the constant-code interval
+// [first, next). The shared search uses N^2/S to compare and bound squared cosine.
+// No code vector is stored here; the caller emits codes once the search finishes.
+//
+// Magnitudes are float32 values promoted exactly to double. For levels <= 255,
+// unequal comparable level/magnitude ratios differ by at least 2^-33 relatively.
+// Reciprocal multiplication has much less error, so its extreme indices remain
+// valid. Recompute those two endpoints by division to retain exact event scales.
+static RescaleSearchState evaluate_scale_state_avx2(
+    const double* magnitudes,
+    const double* reciprocals,
+    size_t dim,
+    int max_code,
+    double t_start,
+    double t_end,
+    double scale
+) {
+    __m256d z = _mm256_setzero_pd(), one = _mm256_set1_pd(1), half = _mm256_set1_pd(.5),
+            vt = _mm256_set1_pd(scale), vcap = _mm256_set1_pd(max_code);
+    __m256d guard = _mm256_set1_pd(
+                2 * std::numeric_limits<double>::epsilon() * (max_code + 1.0)
+            ),
+            vn = z, vd = z, vfirst = z,
+            vnext = _mm256_set1_pd(std::numeric_limits<double>::infinity());
+    __m256d positions = _mm256_setr_pd(0, 1, 2, 3), first_idx = _mm256_set1_pd(-1),
+            next_idx = first_idx;
+    size_t i = 0;
+    for (; i + 4 <= dim; i += 4) {
+        __m256d v = _mm256_loadu_pd(magnitudes + i), vi = _mm256_loadu_pd(reciprocals + i),
+                product = _mm256_mul_pd(v, vt),
+                c = _mm256_min_pd(_mm256_floor_pd(product), vcap);
+        // Floor is sufficient away from integer boundaries. The guard covers
+        // multiplication rounding near those boundaries; saturation removes
+        // the upper boundary. Zero magnitudes always keep code zero.
+        __m256d interior = _mm256_and_pd(
+            _mm256_cmp_pd(product, _mm256_add_pd(c, guard), _CMP_GT_OQ),
+            _mm256_or_pd(
+                _mm256_cmp_pd(c, vcap, _CMP_EQ_OQ),
+                _mm256_cmp_pd(
+                    product, _mm256_sub_pd(_mm256_add_pd(c, one), guard), _CMP_LT_OQ
+                )
+            )
+        );
+        unsigned boundary =
+            (~static_cast<unsigned>(_mm256_movemask_pd(interior))) &
+            static_cast<unsigned>(_mm256_movemask_pd(_mm256_cmp_pd(v, z, _CMP_GT_OQ))) & 15;
+        // Repair only flagged lanes using the heap's division-based thresholds.
+        // This ensures an event at k/a_i belongs to code k in both algorithms.
+        if (boundary) {
+            alignas(32) double corrected[4];
+            _mm256_store_pd(corrected, c);
+            for (unsigned mask = boundary; mask; mask &= mask - 1) {
+                unsigned k = __builtin_ctz(mask);
+                corrected[k] = quantized_code_at_scale(magnitudes[i + k], scale, max_code);
+            }
+            c = _mm256_load_pd(corrected);
+        }
+        // vn/vd hold lane-local contributions to N/S. Accumulate in double;
+        // the shared search encloses rounding in N, while S remains exact.
+        __m256d q = _mm256_add_pd(c, half);
+        vn = _mm256_add_pd(vn, _mm256_mul_pd(v, q));
+        vd = _mm256_add_pd(vd, _mm256_mul_pd(q, q));
+        // A coordinate enters its current code at c_i/a_i (only if c_i>0)
+        // and leaves at (c_i+1)/a_i (only if a_i>0 and not saturated).
+        // Inactive coordinates contribute neutral values: zero to max(first)
+        // and infinity to min(next). Reciprocals avoid per-coordinate division.
+        __m256d af = _mm256_cmp_pd(c, z, _CMP_GT_OQ),
+                an = _mm256_and_pd(
+                    _mm256_cmp_pd(c, vcap, _CMP_LT_OQ), _mm256_cmp_pd(v, z, _CMP_GT_OQ)
+                );
+        __m256d f = _mm256_blendv_pd(z, _mm256_mul_pd(c, vi), af),
+                n = _mm256_blendv_pd(
+                    _mm256_set1_pd(std::numeric_limits<double>::infinity()),
+                    _mm256_mul_pd(_mm256_add_pd(c, one), vi),
+                    an
+                );
+        __m256d fm = _mm256_and_pd(af, _mm256_cmp_pd(f, vfirst, _CMP_GT_OQ)),
+                nm = _mm256_and_pd(an, _mm256_cmp_pd(n, vnext, _CMP_LT_OQ));
+        // Retain original coordinate indices along with each lane's extrema,
+        // so the final boundary can be recomputed by division below.
+        first_idx = _mm256_blendv_pd(first_idx, positions, fm);
+        next_idx = _mm256_blendv_pd(next_idx, positions, nm);
+        vfirst = _mm256_max_pd(vfirst, f);
+        vnext = _mm256_min_pd(vnext, n);
+        positions = _mm256_add_pd(positions, _mm256_set1_pd(4));
+    }
+    // Reduce the four partial sums and select the global boundary coordinates.
+    // Index -1 means this lane never had an active boundary candidate.
+    alignas(32) double ns[4], ds[4], fs[4], ts[4];
+    _mm256_store_pd(ns, vn);
+    _mm256_store_pd(ds, vd);
+    _mm256_store_pd(fs, vfirst);
+    _mm256_store_pd(ts, vnext);
+    RescaleSearchState s{
+        ns[0] + ns[1] + ns[2] + ns[3], ds[0] + ds[1] + ds[2] + ds[3], t_start, t_end};
+    alignas(32) double first_positions[4], next_positions[4];
+    _mm256_store_pd(first_positions, first_idx);
+    _mm256_store_pd(next_positions, next_idx);
+    size_t fi[4], ni[4];
+    for (size_t k = 0; k < 4; ++k) {
+        fi[k] = first_positions[k] < 0 ? std::numeric_limits<size_t>::max()
+                                       : static_cast<size_t>(first_positions[k]);
+        ni[k] = next_positions[k] < 0 ? std::numeric_limits<size_t>::max()
+                                      : static_cast<size_t>(next_positions[k]);
+    }
+    size_t fk = std::max_element(fs, fs + 4) - fs, nk = std::min_element(ts, ts + 4) - ts;
+    // Reciprocal products selected the coordinates; division now recovers the
+    // exact heap event scales, clipped to the configured search interval.
+    if (fi[fk] != std::numeric_limits<size_t>::max())
+        s.first = std::max(
+            s.first,
+            static_cast<double>(quantized_code_at_scale(magnitudes[fi[fk]], scale, max_code)
+            ) / magnitudes[fi[fk]]
+        );
+    if (ni[nk] != std::numeric_limits<size_t>::max())
+        s.next = std::min(
+            s.next,
+            (quantized_code_at_scale(magnitudes[ni[nk]], scale, max_code) + 1.0) /
+                magnitudes[ni[nk]]
+        );
+    // Handle the remaining zero to three coordinates with the same definitions.
+    for (; i < dim; ++i) {
+        int c = quantized_code_at_scale(magnitudes[i], scale, max_code);
+        double q = c + .5;
+        s.numerator += magnitudes[i] * q;
+        s.squared_norm += q * q;
+        if (c > 0)
+            s.first = std::max(s.first, static_cast<double>(c) / magnitudes[i]);
+        if (c < max_code && magnitudes[i] > 0)
+            s.next = std::min(s.next, (c + 1.0) / magnitudes[i]);
+    }
+    return s;
+}
+
+}  // namespace rabitqlib::simd::detail
+
+namespace rabitqlib::simd {
+double best_rescale_factor_avx2(
+    const float* magnitudes, size_t dim, int max_code, double start, double end
+) {
+    // A negative result asks the caller to rerun the scalar heap search.
+    double best_t = 0;
+    return detail::try_find_best_scale_by_interval_search(
+               magnitudes,
+               dim,
+               max_code,
+               start,
+               end,
+               detail::evaluate_scale_state_avx2,
+               best_t
+           )
+               ? best_t
+               : -1;
+}
+}  // namespace rabitqlib::simd