Optimize RaBitQ fastscan query setup + 80% qps (#5396)

Summary:
## What

Two hot-path fixes for RaBitQ FastScan search, from profiling Cohere 1M (d=1024, nprobe=64):

1. **LUT/query quantization**: `compute_LUT_uint8` and `compute_query_factors` call
   `std::roundf` per element (~330k libm calls per query, 16.9% of cycles in
   `__roundf`+PLT alone — half-away-from-zero rounding has no x86 instruction, so it
   never inlines). Replaced with an inline `x + 0.5f` + truncate helper, valid since
   these inputs are non-negative and range-bounded by construction. (recall impact was small enough to be unnoticed)

2. **Single-query pretransform**: `LinearTransform::apply_noalloc` always calls `sgemm`,
   so for n=1 OpenBLAS repacks the whole rotation matrix every query
   (`sgemm_incopy` was 12.2% of cycles vs 2.7% for the actual kernel). Added an
   `sgemv` path for n=1.

Note: rounding is not bit-exact with `roundf` for inputs within 1-2 ulp of a
`k + 0.5` boundary (±1 quantization code); ties agree. gemv also differs from gemm
in accumulation order. Neither affects recall (below), 4/10000 result ids changed.

<img width="1714" height="1884" alt="image" src="https://github.com/user-attachments/assets/9f391480-721d-4ab1-98a8-597c189c6142" />

## Results

Cohere 1M, IVF_RABITQ_FASTSCAN, k=100, nprobe=64, single thread:

- QPS 299.5 → 546.4 (median of 5)
- perf stat: cycles -36.7%, instructions -36.4%, L1D misses -62.2%, IPC unchanged

Recall vs exact ground truth (recomputed by brute force; the HDF5-provided
neighbors are invalid for IP/COSINE): recall@1/10/100 = 0.95 / 0.924 / 0.6554,
identical before and after.

Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5396

Reviewed By: mnorris11

Differential Revision: D112153740

Pulled By: alibeklfc

fbshipit-source-id: d4eae49d7dd751b542a12c4c1de94c12dc5cbf6f

faiss/IndexIVFRaBitQFastScan.cpp

@@ -28,6 +28,8 @@ namespace faiss {
 // Import shared utilities from RaBitQUtils
 using rabitq_utils::ExtraBitsFactors;
 using rabitq_utils::QueryFactorsData;
+using rabitq_utils::round_nonnegative_to_uint16;
+using rabitq_utils::round_nonnegative_to_uint8;
 using rabitq_utils::SignBitFactors;
 using rabitq_utils::SignBitFactorsWithError;
 
@@ -603,13 +605,12 @@ void IndexIVFRaBitQFastScan::compute_LUT_uint8(
                     float mn = all_mins[j2 * M + m];
                     uint8_t* out = out_base + j2 * dim12_2 + m * ksub;
                     for (size_t s = 0; s < ksub; s++) {
-                        out[s] = static_cast<uint8_t>(
-                                std::roundf(a * (tab[s] - mn)));
+                        out[s] = round_nonnegative_to_uint8(a * (tab[s] - mn));
                     }
                 }
                 memset(out_base + j2 * dim12_2 + M * ksub, 0, (M2 - M) * ksub);
-                bq[j2] = static_cast<uint16_t>(
-                        std::roundf(a * (probe_b[j2] - glob_b)));
+                bq[j2] =
+                        round_nonnegative_to_uint16(a * (probe_b[j2] - glob_b));
             }
             normalizers[2 * i] = a;
             normalizers[2 * i + 1] = glob_b;
@@ -851,8 +852,8 @@ struct IVFRaBitQFastScanScanner : InvertedListScanner {
         for (size_t m = 0; m < M; m++) {
             const float* tab = lut_float.get() + m * ksub;
             for (size_t s = 0; s < ksub; s++) {
-                out[m * ksub + s] = static_cast<uint8_t>(
-                        std::roundf(a * (tab[s] - mins[m])));
+                out[m * ksub + s] =
+                        round_nonnegative_to_uint8(a * (tab[s] - mins[m]));
             }
         }
         memset(out + M * ksub, 0, (M2 - M) * ksub);

faiss/VectorTransform.cpp

@@ -47,6 +47,19 @@ int sgemm_(
         float* c,
         FINTEGER* ldc);
 
+int sgemv_(
+        const char* trans,
+        FINTEGER* m,
+        FINTEGER* n,
+        float* alpha,
+        const float* a,
+        FINTEGER* lda,
+        const float* x,
+        FINTEGER* incx,
+        float* beta,
+        float* y,
+        FINTEGER* incy);
+
 int dgemm_(
         const char* transa,
         const char* transb,
@@ -192,6 +205,25 @@ void LinearTransform::apply_noalloc(idx_t n, const float* x, float* xt) const {
 
     float one = 1;
     FINTEGER nbiti = d_out, ni = static_cast<FINTEGER>(n), di = d_in;
+    if (n == 1) {
+        FINTEGER onei = 1;
+        // Avoid GEMM packing overhead for single-vector transforms. GEMV is
+        // mathematically equivalent but may not be bit-exact with GEMM because
+        // BLAS implementations can use different accumulation orders.
+        sgemv_("Transposed",
+               &di,
+               &nbiti,
+               &one,
+               A.data(),
+               &di,
+               x,
+               &onei,
+               &c_factor,
+               xt,
+               &onei);
+        return;
+    }
+
     sgemm_("Transposed",
            "Not transposed",
            &nbiti,

faiss/impl/RaBitQUtils.cpp

@@ -224,8 +224,8 @@ QueryFactorsData compute_query_factors(
     uint8_t* rqq = rotated_qq.data();
     for (size_t i = 0; i < d; i++) {
         const float v_q = rq[i];
-        const uint8_t v_qq = std::clamp<float>(
-                std::round((v_q - v_min) * inv_delta), 0, max_code);
+        const uint8_t v_qq =
+                round_clamped_to_uint8((v_q - v_min) * inv_delta, max_code);
         rqq[i] = v_qq;
         sum_qq += v_qq;
 

faiss/impl/RaBitQUtils.h

@@ -13,6 +13,7 @@
 #include <faiss/utils/rabitq_simd.h>
 #include <faiss/utils/simd_levels.h>
 #include <algorithm>
+#include <cassert>
 #include <cstddef>
 #include <cstdint>
 #include <cstring>
@@ -86,6 +87,50 @@ struct QueryFactorsData {
  */
 FAISS_API extern const float Z_MAX_BY_QB[8];
 
+/** Fast half-away rounding for non-negative RaBitQ quantization values.
+ *
+ * Contract: x must be finite, non-NaN, and in [0, 255.5). This is not a
+ * bit-exact replacement for roundf within a few ulps below k + 0.5, where the
+ * addition may round across the tie. It is intended only for RaBitQ LUT
+ * quantization paths where a +/- 1 code difference at those boundaries is
+ * acceptable.
+ */
+inline uint8_t round_nonnegative_to_uint8(float x) {
+    assert(x == x);
+    assert(x >= 0.0f);
+    assert(x < 255.5f);
+    return static_cast<uint8_t>(static_cast<int>(x + 0.5f));
+}
+
+/** Same as round_nonnegative_to_uint8 for uint16 RaBitQ bias values. */
+inline uint16_t round_nonnegative_to_uint16(float x) {
+    assert(x == x);
+    assert(x >= 0.0f);
+    assert(x < 65535.5f);
+    return static_cast<uint16_t>(static_cast<int>(x + 0.5f));
+}
+
+/** Fast clamped rounding for query-byte quantization.
+ *
+ * Contract: x must be non-NaN. Values outside [0, max_code] are clamped before
+ * applying the same non-bit-exact rounding used by
+ * round_nonnegative_to_uint8().
+ */
+inline uint8_t round_clamped_to_uint8(float x, uint8_t max_code) {
+    assert(x == x);
+
+    if (x <= 0.0f) {
+        return 0;
+    }
+
+    const float max_code_f = static_cast<float>(max_code);
+    if (x >= max_code_f) {
+        return max_code;
+    }
+
+    return static_cast<uint8_t>(static_cast<int>(x + 0.5f));
+}
+
 /** Compute factors for a single database vector using RaBitQ algorithm.
  * This function consolidates the mathematical logic that was duplicated
  * between IndexRaBitQ and IndexRaBitQFastScan.

tests/CMakeLists.txt

@@ -33,6 +33,7 @@ set(FAISS_TEST_SRC
   test_ivf_early_termination.cpp
   test_callback.cpp
   test_utils.cpp
+  test_rabitq_utils.cpp
   test_hamming.cpp
   test_zerocopy.cpp
   test_scalar_quantizer.cpp

tests/test_rabitq_utils.cpp

@@ -0,0 +1,131 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <gtest/gtest.h>
+
+#include <faiss/VectorTransform.h>
+#include <faiss/impl/RaBitQUtils.h>
+
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <random>
+#include <vector>
+
+namespace {
+
+uint8_t round_clamped_ref(float x, uint8_t max_code) {
+    float y = std::round(x);
+    y = std::max(0.0f, std::min(y, static_cast<float>(max_code)));
+    return static_cast<uint8_t>(y);
+}
+
+void check_linear_transform_batch_match(bool have_bias) {
+    constexpr int d_in = 7;
+    constexpr int d_out = 5;
+    faiss::LinearTransform transform(d_in, d_out, have_bias);
+    transform.is_trained = true;
+
+    transform.A.resize(d_in * d_out);
+    for (size_t i = 0; i < transform.A.size(); i++) {
+        transform.A[i] = std::sin(static_cast<float>(i) * 0.37f);
+    }
+
+    if (have_bias) {
+        transform.b.resize(d_out);
+        for (int i = 0; i < d_out; i++) {
+            transform.b[i] = 0.25f * i - 0.5f;
+        }
+    }
+
+    std::vector<float> x(d_in);
+    for (int i = 0; i < d_in; i++) {
+        x[i] = std::cos(static_cast<float>(i) * 0.19f);
+    }
+
+    std::vector<float> one_out(d_out);
+    transform.apply_noalloc(1, x.data(), one_out.data());
+
+    std::vector<float> batch_x(4 * d_in);
+    for (int i = 0; i < 4; i++) {
+        std::copy(x.begin(), x.end(), batch_x.begin() + i * d_in);
+    }
+
+    std::vector<float> batch_out(4 * d_out);
+    transform.apply_noalloc(4, batch_x.data(), batch_out.data());
+
+    for (int i = 0; i < d_out; i++) {
+        EXPECT_NEAR(one_out[i], batch_out[i], 1e-5f);
+    }
+}
+
+} // namespace
+
+TEST(RaBitQUtils, RoundNonnegativeToUint8) {
+    using faiss::rabitq_utils::round_nonnegative_to_uint8;
+
+    EXPECT_EQ(round_nonnegative_to_uint8(0.0f), 0);
+    EXPECT_EQ(round_nonnegative_to_uint8(0.49f), 0);
+    EXPECT_EQ(round_nonnegative_to_uint8(0.5f), 1);
+    EXPECT_EQ(round_nonnegative_to_uint8(1.5f), 2);
+    EXPECT_EQ(round_nonnegative_to_uint8(254.5f), 255);
+    EXPECT_EQ(round_nonnegative_to_uint8(255.0f), 255);
+
+    const float below_half = std::nextafter(0.5f, 0.0f);
+    EXPECT_LE(
+            std::abs(
+                    static_cast<int>(round_nonnegative_to_uint8(below_half)) -
+                    static_cast<int>(std::round(below_half))),
+            1);
+
+    std::mt19937 rng(123);
+    std::uniform_real_distribution<float> dist(0.0f, 255.49f);
+    for (int i = 0; i < 10000; i++) {
+        const float x = dist(rng);
+        EXPECT_LE(
+                std::abs(
+                        static_cast<int>(round_nonnegative_to_uint8(x)) -
+                        static_cast<int>(std::round(x))),
+                1);
+    }
+}
+
+TEST(RaBitQUtils, RoundNonnegativeToUint16) {
+    using faiss::rabitq_utils::round_nonnegative_to_uint16;
+
+    EXPECT_EQ(round_nonnegative_to_uint16(0.0f), 0);
+    EXPECT_EQ(round_nonnegative_to_uint16(0.5f), 1);
+    EXPECT_EQ(round_nonnegative_to_uint16(65534.5f), 65535);
+    EXPECT_EQ(round_nonnegative_to_uint16(65535.0f), 65535);
+}
+
+TEST(RaBitQUtils, RoundClampedToUint8) {
+    using faiss::rabitq_utils::round_clamped_to_uint8;
+
+    EXPECT_EQ(round_clamped_to_uint8(-1.0f, 15), 0);
+    EXPECT_EQ(round_clamped_to_uint8(0.49f, 15), 0);
+    EXPECT_EQ(round_clamped_to_uint8(0.5f, 15), 1);
+    EXPECT_EQ(round_clamped_to_uint8(14.5f, 15), 15);
+    EXPECT_EQ(round_clamped_to_uint8(15.4f, 15), 15);
+    EXPECT_EQ(round_clamped_to_uint8(16.0f, 15), 15);
+
+    std::mt19937 rng(456);
+    std::uniform_real_distribution<float> dist(-4.0f, 260.0f);
+    for (int i = 0; i < 10000; i++) {
+        const float x = dist(rng);
+        EXPECT_LE(
+                std::abs(
+                        static_cast<int>(round_clamped_to_uint8(x, 255)) -
+                        static_cast<int>(round_clamped_ref(x, 255))),
+                1);
+    }
+}
+
+TEST(LinearTransform, SingleVectorMatchesBatchPath) {
+    check_linear_transform_batch_match(false);
+    check_linear_transform_batch_match(true);
+}