faiss 0.6.1 → 0.6.3

data/vendor/faiss/faiss/impl/ProductQuantizer.cpp

@@ -13,12 +13,15 @@
 #include <cstdio>
 #include <cstring>
 #include <memory>
+#include <type_traits>
 
 #include <algorithm>
 
 #include <faiss/IndexFlat.h>
 #include <faiss/VectorTransform.h>
 #include <faiss/impl/FaissAssert.h>
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/impl/pq_code_distance/pq_code_distance-inl.h>
 #include <faiss/impl/simd_dispatch.h>
 #include <faiss/utils/distances.h>
 
@@ -719,8 +722,28 @@ void pq_knn_search_with_tables(
 
         switch (nbits) {
             case 8:
-                pq_estimators_from_tables<uint8_t, C>(
-                        pq, codes, ncodes, dis_table, k, heap_dis, heap_ids);
+                if (ksub == 256) {
+                    constexpr bool max_heap =
+                            std::is_same_v<C, CMax<float, int64_t>>;
+                    pq_code_distance::pq_scan_8bit(
+                            M,
+                            dis_table,
+                            codes,
+                            ncodes,
+                            k,
+                            heap_dis,
+                            heap_ids,
+                            max_heap);
+                } else {
+                    pq_estimators_from_tables<uint8_t, C>(
+                            pq,
+                            codes,
+                            ncodes,
+                            dis_table,
+                            k,
+                            heap_dis,
+                            heap_ids);
+                }
                 break;
 
             case 16:
@@ -839,32 +862,44 @@ void ProductQuantizer::search_sdc(
     size_t k = res->k;
     int64_t nq_signed = nq;
 
-#pragma omp parallel for
-    for (int64_t i = 0; i < nq_signed; i++) {
-        /* Compute distances and keep smallest values */
-        idx_t* heap_ids = res->ids + i * k;
-        float* heap_dis = res->val + i * k;
-        const uint8_t* qcode = qcodes + i * code_size;
-
-        if (init_finalize_heap)
-            maxheap_heapify(k, heap_dis, heap_ids);
-
-        const uint8_t* bcode = bcodes;
-        for (size_t j = 0; j < nb; j++) {
-            float dis = 0;
-            const float* tab = sdc_table.data();
+#pragma omp parallel
+    {
+        // One allocation per OMP thread instead of one per query.
+        std::vector<const float*> q_row(M);
+#pragma omp for
+        for (int64_t i = 0; i < nq_signed; i++) {
+            idx_t* heap_ids = res->ids + i * k;
+            float* heap_dis = res->val + i * k;
+            const uint8_t* qcode = qcodes + i * code_size;
+
+            if (init_finalize_heap)
+                maxheap_heapify(k, heap_dis, heap_ids);
+
+            // Precompute per-subquantizer row pointers: q_row[m] points to
+            // sdc_table[m*ksub^2 + qcode[m]*ksub], eliminating M
+            // multiplications and M pointer advances per database vector in the
+            // j-loop.
+            const float* sdc = sdc_table.data();
             for (size_t m = 0; m < M; m++) {
-                dis += tab[bcode[m] + qcode[m] * ksub];
-                tab += ksub * ksub;
+                q_row[m] = sdc + m * (size_t)(ksub * ksub) +
+                        (size_t)qcode[m] * ksub;
             }
-            if (dis < heap_dis[0]) {
-                maxheap_replace_top(k, heap_dis, heap_ids, dis, j);
+
+            const uint8_t* bcode = bcodes;
+            for (size_t j = 0; j < nb; j++) {
+                float dis = 0;
+                for (size_t m = 0; m < M; m++) {
+                    dis += q_row[m][bcode[m]];
+                }
+                if (dis < heap_dis[0]) {
+                    maxheap_replace_top(k, heap_dis, heap_ids, dis, j);
+                }
+                bcode += code_size;
             }
-            bcode += code_size;
-        }
 
-        if (init_finalize_heap)
-            maxheap_reorder(k, heap_dis, heap_ids);
+            if (init_finalize_heap)
+                maxheap_reorder(k, heap_dis, heap_ids);
+        }
     }
 }
 

data/vendor/faiss/faiss/impl/RaBitQUtils.cpp

@@ -188,46 +188,54 @@ QueryFactorsData compute_query_factors(
 
     const float inv_d_sqrt = 1.0f / std::sqrt(static_cast<float>(d));
 
-    // Compute quantization range
-    float v_min = std::numeric_limits<float>::max();
-    float v_max = std::numeric_limits<float>::lowest();
-
     const float* rq = rotated_q.data();
-    if (centered) {
-        float z_max = Z_MAX_BY_QB[qb - 1];
-        float v_radius = z_max * std::sqrt(query_factors.qr_to_c_L2sqr / d);
-        v_min = -v_radius;
-        v_max = v_radius;
-    } else {
-        for (size_t i = 0; i < d; i++) {
-            const float v_q = rq[i];
-            v_min = std::min(v_min, v_q);
-            v_max = std::max(v_max, v_q);
-        }
-    }
-
-    // Quantize the query
-    const uint8_t max_code = (1 << qb) - 1;
-    const float delta = (v_max - v_min) / max_code;
-    const float inv_delta = 1.0f / delta;
-
-    rotated_qq.resize(d);
+    float v_min;
+    float v_max;
+    float delta;
     size_t sum_qq = 0;
     int64_t sum2_signed_odd_int = 0;
-
+    const uint8_t max_code = (1 << qb) - 1;
+    rotated_qq.resize(d);
     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);
-        rqq[i] = v_qq;
-        sum_qq += v_qq;
-
-        if (centered) {
-            int64_t signed_odd_int = int64_t(v_qq) * 2 - max_code;
-            sum2_signed_odd_int += signed_odd_int * signed_odd_int;
-        }
-    }
+
+    // Select the SIMD implementation once for both range computation and
+    // quantization. This function runs once per query/probe pair.
+    with_selected_simd_levels<rabitq::RABITQ_QUANTIZATION_SIMD_LEVELS>(
+            [&]<SIMDLevel SL>() {
+                if (centered) {
+                    const float z_max = Z_MAX_BY_QB[qb - 1];
+                    const float v_radius =
+                            z_max * std::sqrt(query_factors.qr_to_c_L2sqr / d);
+                    v_min = -v_radius;
+                    v_max = v_radius;
+                } else {
+                    v_min = std::numeric_limits<float>::max();
+                    v_max = std::numeric_limits<float>::lowest();
+                    rabitq::minmax_values<SL>(rq, d, v_min, v_max);
+                }
+
+                delta = (v_max - v_min) / max_code;
+                // A constant (or zero-norm) query has delta == 0. Preserve the
+                // scalar path's centered correction terms while avoiding
+                // 0 * inf during quantization.
+                if (delta <= 0.0f) {
+                    memset(rqq, 0, d * sizeof(uint8_t));
+                    if (centered) {
+                        sum2_signed_odd_int = int64_t(d) * max_code * max_code;
+                    }
+                } else {
+                    rabitq::quantize_query_values<SL>(
+                            rq,
+                            d,
+                            v_min,
+                            1.0f / delta,
+                            max_code,
+                            centered,
+                            rqq,
+                            sum_qq,
+                            sum2_signed_odd_int);
+                }
+            });
 
     // Compute query factors
     query_factors.c1 = 2 * delta * inv_d_sqrt;
@@ -321,7 +329,7 @@ float compute_full_multibit_distance(
         size_t d,
         size_t ex_bits,
         MetricType metric_type) {
-    return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A0>(
+    return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A0_SPR>(
             [&]<SIMDLevel SL>() {
                 return compute_full_multibit_distance<SL>(
                         sign_bits,

data/vendor/faiss/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,40 @@ 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 rabitq::round_nonnegative_byte_scalar(x);
+}
+
+/** 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);
+    return rabitq::round_clamped_byte_scalar(x, max_code);
+}
+
 /** Compute factors for a single database vector using RaBitQ algorithm.
  * This function consolidates the mathematical logic that was duplicated
  * between IndexRaBitQ and IndexRaBitQFastScan.

data/vendor/faiss/faiss/impl/RaBitQuantizer.cpp

@@ -8,13 +8,15 @@
 #include <faiss/impl/RaBitQuantizer.h>
 
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/IDSelector.h>
 #include <faiss/impl/RaBitQUtils.h>
 #include <faiss/impl/RaBitQuantizerMultiBit.h>
+#include <faiss/impl/ResultHandler.h>
 #include <faiss/impl/simd_dispatch.h>
+#include <faiss/invlists/DirectMap.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/rabitq_simd.h>
 
-#include <algorithm>
 #include <cmath>
 #include <cstring>
 #include <memory>
@@ -183,8 +185,10 @@ void RaBitQuantizer::decode_core(
         float* x,
         size_t n,
         const float* centroid_in) const {
-    FAISS_ASSERT(codes != nullptr);
-    FAISS_ASSERT(x != nullptr);
+    FAISS_THROW_IF_MSG(
+            codes == nullptr, "RaBitQuantizer::decode_core: null codes buffer");
+    FAISS_THROW_IF_MSG(
+            x == nullptr, "RaBitQuantizer::decode_core: null output buffer");
 
     const float inv_d_sqrt = (d == 0) ? 1.0f : (1.0f / std::sqrt((float)d));
     const size_t ex_bits = nb_bits - 1;
@@ -228,7 +232,7 @@ namespace {
 // directly to the SIMD-specialized code.
 
 template <SIMDLevel SL>
-struct RaBitQDistanceComputerNotQ : RaBitQDistanceComputer {
+struct RaBitQDistanceComputerNotQ final : RaBitQDistanceComputer {
     // the rotated query (qr - c)
     std::vector<float> rotated_q;
     // some additional numbers for the query
@@ -237,26 +241,9 @@ struct RaBitQDistanceComputerNotQ : RaBitQDistanceComputer {
     RaBitQDistanceComputerNotQ() = default;
 
     // Compute distance using only 1-bit codes (fast)
-    float distance_to_code_1bit(const uint8_t* code) override {
-        FAISS_ASSERT(code != nullptr);
-        FAISS_ASSERT(
-                (metric_type == MetricType::METRIC_L2 ||
-                 metric_type == MetricType::METRIC_INNER_PRODUCT));
-        FAISS_ASSERT(rotated_q.size() == d);
-
-        // split the code into parts
-        const uint8_t* binary_data = code;
-
-        // Cast to appropriate type based on nb_bits
-        // For 1-bit: use SignBitFactors (8 bytes)
-        // For multi-bit: use SignBitFactorsWithError (12 bytes) which includes
-        // f_error
-        size_t ex_bits = nb_bits - 1;
-        const SignBitFactors* base_fac = (ex_bits == 0)
-                ? reinterpret_cast<const SignBitFactors*>(code + (d + 7) / 8)
-                : reinterpret_cast<const SignBitFactorsWithError*>(
-                          code + (d + 7) / 8);
-
+    float distance_to_code_1bit_impl(
+            const uint8_t* binary_data,
+            const SignBitFactors* base_fac) const {
         // this is the baseline code
         //
         // compute <q,o> using floats
@@ -286,7 +273,7 @@ struct RaBitQDistanceComputerNotQ : RaBitQDistanceComputer {
 
         if (metric_type == MetricType::METRIC_L2) {
             // ||or - q||^ 2
-            return pre_dist;
+            return std::max(0.0f, pre_dist);
         } else {
             // metric == MetricType::METRIC_INNER_PRODUCT
             // 2 * (or, q) = (||or - q||^2 - ||q||^2 - ||or||^2)
@@ -294,8 +281,24 @@ struct RaBitQDistanceComputerNotQ : RaBitQDistanceComputer {
         }
     }
 
+    float distance_to_code_1bit(const uint8_t* code) final {
+        FAISS_ASSERT(code != nullptr);
+        FAISS_ASSERT(
+                (metric_type == MetricType::METRIC_L2 ||
+                 metric_type == MetricType::METRIC_INNER_PRODUCT));
+        FAISS_ASSERT(rotated_q.size() == d);
+
+        const size_t code_size_base = (d + 7) / 8;
+        const size_t ex_bits = nb_bits - 1;
+        const SignBitFactors* base_fac = (ex_bits == 0)
+                ? reinterpret_cast<const SignBitFactors*>(code + code_size_base)
+                : reinterpret_cast<const SignBitFactorsWithError*>(
+                          code + code_size_base);
+        return distance_to_code_1bit_impl(code, base_fac);
+    }
+
     // Compute full distance using 1-bit + ex-bits (accurate)
-    float distance_to_code_full(const uint8_t* code) override {
+    float distance_to_code_full(const uint8_t* code) final {
         FAISS_ASSERT(code != nullptr);
         FAISS_ASSERT(
                 (metric_type == MetricType::METRIC_L2 ||
@@ -331,7 +334,7 @@ struct RaBitQDistanceComputerNotQ : RaBitQDistanceComputer {
                 metric_type);
     }
 
-    void set_query(const float* x) override {
+    void set_query(const float* x) final {
         q = x;
         FAISS_ASSERT(x != nullptr);
         FAISS_ASSERT(
@@ -372,10 +375,62 @@ struct RaBitQDistanceComputerNotQ : RaBitQDistanceComputer {
                     centroid ? fvec_inner_product(x, centroid, d) : 0.0f;
         }
     }
+
+    size_t scan_codes_multibit(
+            size_t list_size,
+            const uint8_t* codes,
+            const idx_t* ids,
+            size_t code_size,
+            idx_t list_no,
+            bool store_pairs,
+            const IDSelector* sel,
+            bool keep_max,
+            ResultHandler& handler) final {
+        const size_t code_size_base = (d + 7) / 8;
+        const size_t ex_bits = nb_bits - 1;
+        FAISS_ASSERT(ex_bits > 0);
+
+        size_t nup = 0;
+        for (size_t j = 0; j < list_size; j++) {
+            if (sel != nullptr) {
+                idx_t id = store_pairs ? lo_build(list_no, j) : ids[j];
+                if (!sel->is_member(id)) {
+                    codes += code_size;
+                    continue;
+                }
+            }
+
+            const auto* base_fac =
+                    reinterpret_cast<const SignBitFactorsWithError*>(
+                            codes + code_size_base);
+            const float est_distance =
+                    distance_to_code_1bit_impl(codes, base_fac);
+
+            const bool should_refine = rabitq_utils::should_refine_candidate(
+                    est_distance,
+                    base_fac->f_error,
+                    g_error,
+                    handler.threshold,
+                    keep_max);
+            if (should_refine) {
+                handler.stats.scan_cnt++;
+                const float dis = distance_to_code_full(codes);
+                idx_t id = store_pairs ? lo_build(list_no, j) : ids[j];
+
+                if (handler.add_result(dis, id)) {
+                    handler.stats.nheap_updates++;
+                    nup++;
+                }
+            }
+            codes += code_size;
+        }
+
+        return nup;
+    }
 };
 
 template <SIMDLevel SL>
-struct RaBitQDistanceComputerQ : RaBitQDistanceComputer {
+struct RaBitQDistanceComputerQ final : RaBitQDistanceComputer {
     // the rotated and quantized query (qr - c)
     std::vector<float> rotated_q;
     // the rotated and quantized query (qr - c) for fast 1-bit computation
@@ -395,25 +450,10 @@ struct RaBitQDistanceComputerQ : RaBitQDistanceComputer {
     RaBitQDistanceComputerQ() = default;
 
     // Compute distance using only 1-bit codes (fast)
-    float distance_to_code_1bit(const uint8_t* code) override {
-        FAISS_ASSERT(code != nullptr);
-        FAISS_ASSERT(
-                (metric_type == MetricType::METRIC_L2 ||
-                 metric_type == MetricType::METRIC_INNER_PRODUCT));
-
-        // split the code into parts
-        size_t size = (d + 7) / 8;
-        const uint8_t* binary_data = code;
-
-        // Cast to appropriate type based on nb_bits
-        // For 1-bit: use SignBitFactors (8 bytes)
-        // For multi-bit: use SignBitFactorsWithError (12 bytes) which
-        // includes f_error
-        size_t ex_bits = nb_bits - 1;
-        const SignBitFactors* base_fac = (ex_bits == 0)
-                ? reinterpret_cast<const SignBitFactors*>(code + size)
-                : reinterpret_cast<const SignBitFactorsWithError*>(code + size);
-
+    float distance_to_code_1bit_impl(
+            const uint8_t* binary_data,
+            const SignBitFactors* base_fac,
+            size_t size) const {
         // this is ||or - c||^2 - (IP ? ||or||^2 : 0)
         float final_dot = 0;
         if (centered) {
@@ -428,17 +468,16 @@ struct RaBitQDistanceComputerQ : RaBitQDistanceComputer {
                                qb);
             final_dot += int_dot * query_fac.int_dot_scale;
         } else {
-            auto dot_qo = rabitq::bitwise_and_dot_product<SL>(
-                    rearranged_rotated_qq.data(), binary_data, size, qb);
-            // It was a willful decision (after the discussion) to not to
-            // pre-cache the sum of all bits, just in order to reduce the
-            // overhead per vector.
-            // process 64-bit popcounts
-            auto sum_q = rabitq::popcount<SL>(binary_data, size);
+            auto bitwise_result =
+                    rabitq::bitwise_and_dot_product_with_popcount<SL>(
+                            rearranged_rotated_qq.data(),
+                            binary_data,
+                            size,
+                            qb);
             // dot-product itself
-            final_dot += query_fac.c1 * dot_qo;
+            final_dot += query_fac.c1 * bitwise_result.dot_product;
             // normalizer coefficients
-            final_dot += query_fac.c2 * sum_q;
+            final_dot += query_fac.c2 * bitwise_result.popcount;
             // normalizer coefficients
             final_dot -= query_fac.c34;
         }
@@ -449,7 +488,7 @@ struct RaBitQDistanceComputerQ : RaBitQDistanceComputer {
 
         if (metric_type == MetricType::METRIC_L2) {
             // ||or - q||^ 2
-            return pre_dist;
+            return std::max(0.0f, pre_dist);
         } else {
             // metric == MetricType::METRIC_INNER_PRODUCT
             // 2 * (or, q) = (||or - q||^2 - ||q||^2 - ||or||^2)
@@ -457,8 +496,22 @@ struct RaBitQDistanceComputerQ : RaBitQDistanceComputer {
         }
     }
 
+    float distance_to_code_1bit(const uint8_t* code) final {
+        FAISS_ASSERT(code != nullptr);
+        FAISS_ASSERT(
+                (metric_type == MetricType::METRIC_L2 ||
+                 metric_type == MetricType::METRIC_INNER_PRODUCT));
+
+        const size_t size = (d + 7) / 8;
+        const size_t ex_bits = nb_bits - 1;
+        const SignBitFactors* base_fac = (ex_bits == 0)
+                ? reinterpret_cast<const SignBitFactors*>(code + size)
+                : reinterpret_cast<const SignBitFactorsWithError*>(code + size);
+        return distance_to_code_1bit_impl(code, base_fac, size);
+    }
+
     // Compute full distance using 1-bit + ex-bits (accurate)
-    float distance_to_code_full(const uint8_t* code) override {
+    float distance_to_code_full(const uint8_t* code) final {
         FAISS_ASSERT(code != nullptr);
         FAISS_ASSERT(
                 (metric_type == MetricType::METRIC_L2 ||
@@ -494,7 +547,7 @@ struct RaBitQDistanceComputerQ : RaBitQDistanceComputer {
                 metric_type);
     }
 
-    void set_query(const float* x) override {
+    void set_query(const float* x) final {
         q = x;
         FAISS_ASSERT(x != nullptr);
         FAISS_ASSERT(
@@ -526,16 +579,64 @@ struct RaBitQDistanceComputerQ : RaBitQDistanceComputer {
         size_t offset = (d + 7) / 8;
 
         rearranged_rotated_qq.resize(offset * qb);
-        std::fill(
-                rearranged_rotated_qq.begin(), rearranged_rotated_qq.end(), 0);
-
-        for (size_t idim = 0; idim < d; idim++) {
-            for (size_t iv = 0; iv < qb; iv++) {
-                const bool bit = ((rotated_qq[idim] & (1 << iv)) != 0);
-                rearranged_rotated_qq[iv * offset + idim / 8] |=
-                        bit ? (1 << (idim % 8)) : 0;
+        with_selected_simd_levels<
+                AVAILABLE_SIMD_LEVELS_NONE | (1 << int(SIMDLevel::AVX2)) |
+                (1 << int(SIMDLevel::AVX512))>([&]<SIMDLevel RSL>() {
+            rabitq::rearrange_bit_planes<RSL>(
+                    rotated_qq.data(), d, qb, rearranged_rotated_qq.data());
+        });
+    }
+
+    size_t scan_codes_multibit(
+            size_t list_size,
+            const uint8_t* codes,
+            const idx_t* ids,
+            size_t code_size,
+            idx_t list_no,
+            bool store_pairs,
+            const IDSelector* sel,
+            bool keep_max,
+            ResultHandler& handler) final {
+        const size_t code_size_base = (d + 7) / 8;
+        const size_t ex_bits = nb_bits - 1;
+        FAISS_ASSERT(ex_bits > 0);
+
+        size_t nup = 0;
+        for (size_t j = 0; j < list_size; j++) {
+            if (sel != nullptr) {
+                idx_t id = store_pairs ? lo_build(list_no, j) : ids[j];
+                if (!sel->is_member(id)) {
+                    codes += code_size;
+                    continue;
+                }
             }
+
+            const auto* base_fac =
+                    reinterpret_cast<const SignBitFactorsWithError*>(
+                            codes + code_size_base);
+            const float est_distance =
+                    distance_to_code_1bit_impl(codes, base_fac, code_size_base);
+
+            const bool should_refine = rabitq_utils::should_refine_candidate(
+                    est_distance,
+                    base_fac->f_error,
+                    g_error,
+                    handler.threshold,
+                    keep_max);
+            if (should_refine) {
+                handler.stats.scan_cnt++;
+                const float dis = distance_to_code_full(codes);
+                idx_t id = store_pairs ? lo_build(list_no, j) : ids[j];
+
+                if (handler.add_result(dis, id)) {
+                    handler.stats.nheap_updates++;
+                    nup++;
+                }
+            }
+            codes += code_size;
         }
+
+        return nup;
     }
 };
 
@@ -551,7 +652,13 @@ FlatCodesDistanceComputer* RaBitQuantizer::get_distance_computer(
     // Dispatch on SIMDLevel once here so the distance computer methods
     // call the SIMD-specialized rabitq functions directly (no per-call
     // with_simd_level overhead).
-    return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A0>(
+    //
+    // Use A0_SPR (which includes AVX512_SPR) so that on Sapphire Rapids
+    // and later x86 microarchitectures the VPOPCNTDQ-based RaBitQ
+    // specialization in rabitq_avx512_spr.cpp is selected. On AVX-512
+    // CPUs without VPOPCNTDQ, dispatch falls through to the AVX512
+    // specialization in rabitq_avx512.cpp.
+    return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A0_SPR>(
             [&]<SIMDLevel SL>() -> FlatCodesDistanceComputer* {
                 if (qb == 0) {
                     auto dc =

data/vendor/faiss/faiss/impl/RaBitQuantizer.h

@@ -16,6 +16,10 @@
 
 namespace faiss {
 
+struct IDSelector;
+template <typename T, typename TI>
+struct ResultHandlerUnordered;
+
 // the reference implementation of the https://arxiv.org/pdf/2405.12497
 //   Jianyang Gao, Cheng Long, "RaBitQ: Quantizing High-Dimensional Vectors
 //   with a Theoretical Error Bound for Approximate Nearest Neighbor Search".
@@ -129,6 +133,21 @@ struct RaBitQDistanceComputer : FlatCodesDistanceComputer {
     // Compute full multi-bit distance (accurate)
     virtual float distance_to_code_full(const uint8_t* code) = 0;
 
+    virtual void set_centroid(const float* centroid_in) {
+        centroid = centroid_in;
+    }
+
+    virtual size_t scan_codes_multibit(
+            size_t list_size,
+            const uint8_t* codes,
+            const idx_t* ids,
+            size_t code_size,
+            idx_t list_no,
+            bool store_pairs,
+            const IDSelector* sel,
+            bool keep_max,
+            ResultHandlerUnordered<float, idx_t>& handler) = 0;
+
     // Override from FlatCodesDistanceComputer
     // Delegates to distance_to_code_full() for multi-bit distance computation
     float distance_to_code(const uint8_t* code) final {