perf: faster LUT and query preparation, QGKMeans cache and visited set

include/rabitqlib/clustering/qgkmeans.hpp

@@ -74,7 +74,8 @@ class QuerySumCache {
 
    public:
     QuerySumCache(size_t padded_dim, const QGKMeansParameters& parameters) {
-        if (padded_dim >= 1024 &&
+        constexpr size_t kMinCachedDim = 512;
+        if (padded_dim >= kMinCachedDim &&
             (parameters.niter > 1 ||
              parameters.final_assignment != FinalAssignmentMode::Exact)) {
             state_ = State::Pending;

include/rabitqlib/core/detail/simd/space_dispatch.hpp

@@ -154,6 +154,11 @@ void scalar_quantize_uint8_generic(
 void scalar_quantize_uint8_avx2(
     uint8_t* result, const float* vec0, size_t dim, float lo, float delta
 );
+// Only for (value - lo) / delta >= 0, such as FastScan LUT entries at least their
+// minimum; there it matches scalar_quantize_uint8_avx2 with fewer operations.
+void scalar_quantize_nonnegative_uint8_avx2(
+    uint8_t* result, const float* vec0, size_t dim, float lo, float delta
+);
 void scalar_quantize_uint16_generic(
     uint16_t* result, const float* data, size_t dim, float lo, float delta
 );
@@ -172,6 +177,10 @@ float mask_ip_x0_q_avx512(const float* query, const uint64_t* data, size_t padde
 void scalar_quantize_uint8_avx512(
     uint8_t* result, const float* vec0, size_t dim, float lo, float delta
 );
+// Only for (value - lo) / delta >= 0; see scalar_quantize_nonnegative_uint8_avx2.
+void scalar_quantize_nonnegative_uint8_avx512(
+    uint8_t* result, const float* vec0, size_t dim, float lo, float delta
+);
 void scalar_quantize_uint16_avx512(
     uint16_t* result, const float* vec0, size_t dim, float lo, float delta
 );

include/rabitqlib/fastscan/detail/lut.hpp

@@ -9,6 +9,7 @@
 #include <type_traits>
 #include <vector>
 
+#include "rabitqlib/fastscan/detail/lut_prepare.hpp"
 #include "rabitqlib/fastscan/fastscan.hpp"
 #include "rabitqlib/fastscan/highacc_fastscan.hpp"
 #include "rabitqlib/utils/space.hpp"
@@ -95,18 +96,35 @@ class Lut {
         uint16_t* lut_u16_scratch,
         bool use_hacc
     ) {
-        fastscan::pack_lut(padded_dim, rotated_query, lut_scratch);
         T vl_lut;
         T vr_lut;
-        data_range(lut_scratch, table_length_, vl_lut, vr_lut);
+        bool finite_lut = false;
+        if constexpr (std::is_same_v<T, float>) {
+            finite_lut = fastscan::detail::pack_lut_with_range(
+                padded_dim, rotated_query, lut_scratch, vl_lut, vr_lut
+            );
+        } else {
+            fastscan::pack_lut(padded_dim, rotated_query, lut_scratch);
+            data_range(lut_scratch, table_length_, vl_lut, vr_lut);
+        }
 
         if (use_hacc) {
             delta_ = (vr_lut - vl_lut) / ((1 << kNumBitsHacc) - 1);
             quantize_entries(lut_u16_scratch, lut_scratch, vl_lut);
             fastscan::transfer_lut_hacc(lut_u16_scratch, padded_dim, lut_.data());
         } else {
             delta_ = (vr_lut - vl_lut) / ((1 << kNumBits) - 1);
-            quantize_entries(lut_.data(), lut_scratch, vl_lut);
+            // finite_lut is only set for float, the only type the uint8 kernels take.
+            if constexpr (std::is_same_v<T, float>) {
+                if (finite_lut) {
+                    fastscan::detail::quantize_lut_uint8(
+                        lut_.data(), lut_scratch, table_length_, vl_lut, delta_
+                    );
+                }
+            }
+            if (!finite_lut) {
+                quantize_entries(lut_.data(), lut_scratch, vl_lut);
+            }
         }
 
         const size_t num_tables = table_length_ / 16;

include/rabitqlib/fastscan/detail/lut_prepare.hpp

@@ -0,0 +1,19 @@
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+
+namespace rabitqlib::fastscan::detail {
+// Caller validates a positive dimension divisible by 16 and non-overlapping scratch.
+// Return whether every generated LUT entry is finite. Exceptional and zero bounds
+// retain the original range reduction, including its signed-zero behavior.
+bool pack_lut_with_range(
+    size_t dim, const float* query, float* lut, float& minimum, float& maximum
+);
+
+// Use only after pack_lut_with_range returns true, with its minimum and range delta.
+// General quantization and high-accuracy uint16 quantization remain unchanged.
+void quantize_lut_uint8(
+    uint8_t* result, const float* values, size_t count, float lo, float delta
+);
+}  // namespace rabitqlib::fastscan::detail

include/rabitqlib/fastscan/detail/simd.hpp

@@ -24,6 +24,20 @@ void pack_lut_generic(size_t dim, const float* query, float* lut);
 void pack_lut_avx2(size_t dim, const float* query, float* lut);
 void pack_lut_avx512(size_t dim, const float* query, float* lut);
 
+// Raw fused kernels; the wrapper preserves exceptional/zero-bound reduction behavior.
+void pack_lut_range_generic(
+    size_t dim, const float* query, float* lut, float& minimum, float& maximum
+);
+void pack_lut_range_avx2(
+    size_t dim, const float* query, float* lut, float& minimum, float& maximum
+);
+void pack_lut_range_avx512(
+    size_t dim, const float* query, float* lut, float& minimum, float& maximum
+);
+void pack_lut_range_neon(
+    size_t dim, const float* query, float* lut, float& minimum, float& maximum
+);
+
 // Scalar reference for backend correctness tests; never selected by runtime dispatch.
 void accumulate_generic(
     const uint8_t* __restrict__ codes,

include/rabitqlib/index/detail/visited_set_hash.hpp

@@ -77,7 +77,8 @@ class HashBasedVisitedSet {
             bit_size++;
             bit >>= 1;
         }
-        size_t bucket_size = static_cast<size_t>(0x1) << ((bit_size + 4) / 2 + 3);
+        // Limit per-visit spillover allocations while keeping the direct table sublinear.
+        size_t bucket_size = static_cast<size_t>(0x1) << ((bit_size + 4) / 2 + 4);
         initialize_table(bucket_size);
     }
 

src/core/api.cpp

@@ -4,6 +4,7 @@
 #include <cmath>
 #include <cstddef>
 #include <cstdint>
+#include <cstring>
 #include <limits>
 #include <memory_resource>
 #include <numeric>
@@ -89,10 +90,19 @@ void check_input(const float* input, size_t dim) {
     check_buffer<alignof(float)>(
         {input, multiply(dim, sizeof(float))}, multiply(dim, sizeof(float))
     );
+    static_assert(
+        sizeof(float) == sizeof(uint32_t) && std::numeric_limits<float>::is_iec559
+    );
+    // An integer reduction lets portable builds check several float32 values
+    // per instruction without relaxing floating-point compiler options.
+    uint32_t nonfinite = 0;
     for (size_t i = 0; i < dim; ++i) {
-        if (!std::isfinite(input[i])) {
-            throw std::invalid_argument("Input vectors must contain only finite values");
-        }
+        uint32_t bits;
+        std::memcpy(&bits, input + i, sizeof(bits));
+        nonfinite |= static_cast<uint32_t>((bits & 0x7f800000U) == 0x7f800000U);
+    }
+    if (nonfinite != 0) {
+        throw std::invalid_argument("Input vectors must contain only finite values");
     }
 }
 

src/fastscan/fastscan_avx2.cpp

@@ -11,7 +11,26 @@
 
 namespace rabitqlib::fastscan::simd {
 
-void pack_lut_avx2(size_t dim, const float* query, float* lut) {
+namespace {
+// Evaluated at compile time: unoptimized builds would otherwise export these library
+// functions compiled with this file's ISA flags for callers on other CPUs.
+constexpr float kInfinity = std::numeric_limits<float>::infinity();
+constexpr float kQuietNaN = std::numeric_limits<float>::quiet_NaN();
+
+template <bool ComputeRange>
+void pack_lut_impl(
+    size_t dim, const float* query, float* lut, float& minimum, float& maximum
+) {
+    auto min_values = _mm256_set1_ps(kInfinity), max_values = _mm256_set1_ps(-kInfinity);
+    unsigned nan_mask = 0;
+    const auto update_range = [&](__m256 values) {
+        min_values = _mm256_min_ps(min_values, values);
+        max_values = _mm256_max_ps(max_values, values);
+        nan_mask |= static_cast<unsigned>(
+            _mm256_movemask_ps(_mm256_cmp_ps(values, values, _CMP_UNORD_Q))
+        );
+    };
+
     // Keep the initial +0 addition observable for signed-zero inputs.
     volatile float zero = 0;
     for (size_t group = 0; group < dim / 4; ++group) {
@@ -27,11 +46,43 @@ void pack_lut_avx2(size_t dim, const float* query, float* lut) {
         hi = _mm256_blend_ps(hi, _mm256_add_ps(hi, q2), 0xCC);
         lo = _mm256_blend_ps(lo, _mm256_add_ps(lo, q3), 0xAA);
         hi = _mm256_blend_ps(hi, _mm256_add_ps(hi, q3), 0xAA);
+        if constexpr (ComputeRange) {
+            update_range(lo);
+            update_range(hi);
+        }
         _mm256_storeu_ps(lut, lo);
         _mm256_storeu_ps(lut + 8, hi);
         query += 4;
         lut += 16;
     }
+    if constexpr (ComputeRange) {
+        float min_lanes[8], max_lanes[8];
+        _mm256_storeu_ps(min_lanes, min_values);
+        _mm256_storeu_ps(max_lanes, max_values);
+        // Plain comparisons avoid std algorithms, for the reason given at kInfinity.
+        minimum = min_lanes[0];
+        maximum = max_lanes[0];
+        for (size_t lane = 1; lane < 8; ++lane) {
+            minimum = min_lanes[lane] < minimum ? min_lanes[lane] : minimum;
+            maximum = max_lanes[lane] > maximum ? max_lanes[lane] : maximum;
+        }
+        if (nan_mask)
+            minimum = kQuietNaN;
+    }
+}
+}  // namespace
+
+void pack_lut_avx2(size_t dim, const float* query, float* lut) {
+    float unused_minimum = 0, unused_maximum = 0;
+    pack_lut_impl<false>(dim, query, lut, unused_minimum, unused_maximum);
+}
+
+// Preserve the existing ordered subset additions; reduce their values in registers.
+// A NaN minimum signals that the wrapper must use the original reduction.
+void pack_lut_range_avx2(
+    size_t dim, const float* query, float* lut, float& minimum, float& maximum
+) {
+    pack_lut_impl<true>(dim, query, lut, minimum, maximum);
 }
 
 namespace {

src/fastscan/fastscan_avx512.cpp

@@ -11,7 +11,19 @@
 
 namespace rabitqlib::fastscan::simd {
 
-void pack_lut_avx512(size_t dim, const float* query, float* lut) {
+namespace {
+// Evaluated at compile time: unoptimized builds would otherwise export these library
+// functions compiled with this file's ISA flags for callers on other CPUs.
+constexpr float kInfinity = std::numeric_limits<float>::infinity();
+constexpr float kQuietNaN = std::numeric_limits<float>::quiet_NaN();
+
+template <bool ComputeRange>
+void pack_lut_impl(
+    size_t dim, const float* query, float* lut, float& minimum, float& maximum
+) {
+    auto min_values = _mm512_set1_ps(kInfinity), max_values = _mm512_set1_ps(-kInfinity);
+    unsigned nan_mask = 0;
+
     // Keep the initial +0 addition observable for signed-zero inputs.
     volatile float zero = 0;
     for (size_t group = 0; group < dim / 4; ++group) {
@@ -22,10 +34,33 @@ void pack_lut_avx512(size_t dim, const float* query, float* lut) {
         values = _mm512_mask_add_ps(values, 0xF0F0, values, _mm512_set1_ps(query[1]));
         values = _mm512_mask_add_ps(values, 0xCCCC, values, _mm512_set1_ps(query[2]));
         values = _mm512_mask_add_ps(values, 0xAAAA, values, _mm512_set1_ps(query[3]));
+        if constexpr (ComputeRange) {
+            min_values = _mm512_min_ps(min_values, values);
+            max_values = _mm512_max_ps(max_values, values);
+            nan_mask |= _mm512_cmp_ps_mask(values, values, _CMP_UNORD_Q);
+        }
         _mm512_storeu_ps(lut, values);
         query += 4;
         lut += 16;
     }
+    if constexpr (ComputeRange) {
+        minimum = nan_mask ? kQuietNaN : _mm512_reduce_min_ps(min_values);
+        maximum = _mm512_reduce_max_ps(max_values);
+    }
+}
+}  // namespace
+
+void pack_lut_avx512(size_t dim, const float* query, float* lut) {
+    float unused_minimum = 0, unused_maximum = 0;
+    pack_lut_impl<false>(dim, query, lut, unused_minimum, unused_maximum);
+}
+
+// Preserve the existing ordered subset additions; reduce their values in registers.
+// A NaN minimum signals that the wrapper must use the original reduction.
+void pack_lut_range_avx512(
+    size_t dim, const float* query, float* lut, float& minimum, float& maximum
+) {
+    pack_lut_impl<true>(dim, query, lut, minimum, maximum);
 }
 
 namespace {

src/fastscan/fastscan_backend.hpp

@@ -1,5 +1,6 @@
 #pragma once
 
+#include "rabitqlib/core/detail/simd/space_dispatch.hpp"
 #include "rabitqlib/fastscan/detail/estimator_dispatch.hpp"
 #include "rabitqlib/fastscan/detail/simd.hpp"
 
@@ -16,6 +17,8 @@ struct Backend {
     const char* name;
     HaccLutLayout hacc_layout;
     decltype(&simd::pack_lut_generic) pack_lut;
+    decltype(&simd::pack_lut_range_generic) pack_lut_range;
+    decltype(&rabitqlib::simd::scalar_quantize_uint8_generic) quantize_lut;
     decltype(&simd::transfer_lut_hacc_generic) transfer_hacc;
     decltype(&simd::accumulate_generic) accumulate;
     decltype(&simd::accumulate_hacc_generic) accumulate_hacc;

src/fastscan/fastscan_dispatch.cpp

@@ -1,13 +1,16 @@
+#include <cmath>
 #include <cstddef>
 #include <cstdint>
 #include <stdexcept>
 
 #include "../simd/dispatch_resolver.hpp"
 #include "fastscan_backend.hpp"
 #include "rabitqlib/fastscan/detail/estimator_dispatch.hpp"
+#include "rabitqlib/fastscan/detail/lut_prepare.hpp"
 #include "rabitqlib/fastscan/detail/simd.hpp"
 #include "rabitqlib/fastscan/fastscan.hpp"
 #include "rabitqlib/fastscan/highacc_fastscan.hpp"
+#include "rabitqlib/utils/space.hpp"
 #if defined(__aarch64__) || defined(_M_ARM64)
 #include "rabitqlib/utils/cpu_features.hpp"
 #endif
@@ -18,6 +21,8 @@ constexpr Backend kGenericBackend{
     "generic",
     HaccLutLayout::Pair32,
     simd::pack_lut_generic,
+    simd::pack_lut_range_generic,
+    rabitqlib::simd::scalar_quantize_uint8_generic,
     simd::transfer_lut_hacc_generic,
     simd::accumulate_unsupported,
     simd::accumulate_hacc_generic,
@@ -30,6 +35,8 @@ constexpr Backend kAvx2Backend{
     "avx2",
     HaccLutLayout::Pair32,
     simd::pack_lut_avx2,
+    simd::pack_lut_range_avx2,
+    rabitqlib::simd::scalar_quantize_nonnegative_uint8_avx2,
     simd::transfer_lut_hacc_avx2,
     simd::accumulate_avx2,
     simd::accumulate_hacc_avx2,
@@ -40,6 +47,8 @@ constexpr Backend kAvx512Backend{
     "avx512",
     HaccLutLayout::Quad64,
     simd::pack_lut_avx512,
+    simd::pack_lut_range_avx512,
+    rabitqlib::simd::scalar_quantize_nonnegative_uint8_avx512,
     simd::transfer_lut_hacc_avx512,
     simd::accumulate_avx512,
     simd::accumulate_hacc_avx512,
@@ -52,6 +61,8 @@ constexpr Backend kNeonBackend{
     "neon",
     HaccLutLayout::Pair32,
     simd::pack_lut_neon,
+    simd::pack_lut_range_neon,
+    rabitqlib::simd::scalar_quantize_uint8_neon,
     simd::transfer_lut_hacc_generic,
     simd::accumulate_neon,
     simd::accumulate_hacc_neon,
@@ -72,6 +83,26 @@ RABITQ_COLD_INIT const Backend& selected_backend() {
 #endif
     return *backend;
 }
+
+bool pack_lut_with_range(
+    size_t dim, const float* query, float* lut, float& minimum, float& maximum
+) {
+    static const auto kPackLutRangeFn = selected_backend().pack_lut_range;
+    kPackLutRangeFn(dim, query, lut, minimum, maximum);
+    // Raw kernels signal any NaN via minimum; infinities remain in the extrema.
+    const bool finite = std::isfinite(minimum) && std::isfinite(maximum);
+    if (!finite || minimum == 0.0F || maximum == 0.0F) {
+        data_range(lut, dim * 4, minimum, maximum);
+    }
+    return finite;
+}
+
+void quantize_lut_uint8(
+    uint8_t* result, const float* values, size_t count, float lo, float delta
+) {
+    static const auto kQuantizeLutFn = selected_backend().quantize_lut;
+    kQuantizeLutFn(result, values, count, lo, delta);
+}
 }  // namespace rabitqlib::fastscan::detail
 
 namespace rabitqlib::simd {