faiss 0.6.2 → 0.6.4

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

@@ -8,9 +8,12 @@
 #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>
 
@@ -22,6 +25,8 @@
 
 namespace faiss {
 
+RaBitQStats rabitq_stats;
+
 // Import shared utilities from RaBitQUtils
 using rabitq_utils::ExtraBitsFactors;
 using rabitq_utils::QueryFactorsData;
@@ -183,8 +188,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;
@@ -220,6 +227,51 @@ void RaBitQuantizer::decode_core(
     }
 }
 
+template <SIMDLevel SL>
+float symmetric_dis_1bit(const RaBitQDistanceComputer& dc, idx_t i, idx_t j) {
+    FAISS_THROW_IF_NOT_MSG(
+            dc.metric_type == MetricType::METRIC_L2,
+            "RaBitQ symmetric distance supports only L2");
+    FAISS_ASSERT(i >= 0 && j >= 0);
+    FAISS_ASSERT(dc.codes != nullptr);
+
+    const size_t sign_bytes = (dc.d + 7) / 8;
+    const uint8_t* code_i = dc.codes + static_cast<size_t>(i) * dc.code_size;
+    const uint8_t* code_j = dc.codes + static_cast<size_t>(j) * dc.code_size;
+    const auto* factors_i =
+            reinterpret_cast<const SignBitFactors*>(code_i + sign_bytes);
+    const auto* factors_j =
+            reinterpret_cast<const SignBitFactors*>(code_j + sign_bytes);
+
+    const uint64_t xor_popcount =
+            rabitq::bitwise_xor_dot_product<SL>(code_i, code_j, sign_bytes, 1);
+    const float sign_dot =
+            static_cast<float>(dc.d) - 2.0f * static_cast<float>(xor_popcount);
+
+    // The L2-optimal reconstruction of residual r is alpha * sign(r), where
+    // alpha = ||r||_1 / d. The stored factors give
+    // alpha_i * alpha_j = ||r_i||^2 * ||r_j||^2 /
+    //     (d * dp_multiplier_i * dp_multiplier_j).
+    float cross_term = 0.0f;
+    if (factors_i->dp_multiplier != 0.0f && factors_j->dp_multiplier != 0.0f) {
+        // Dividing each norm first avoids overflowing the product of two
+        // squared norms even when the final distance is representable.
+        const float scaled_norm_i =
+                factors_i->or_minus_c_l2sqr / factors_i->dp_multiplier;
+        const float scaled_norm_j =
+                factors_j->or_minus_c_l2sqr / factors_j->dp_multiplier;
+        cross_term = (scaled_norm_i * (sign_dot / static_cast<float>(dc.d))) *
+                scaled_norm_j;
+    }
+    const float distance = factors_i->or_minus_c_l2sqr +
+            factors_j->or_minus_c_l2sqr - 2.0f * cross_term;
+    return std::max(0.0f, distance);
+}
+
+float RaBitQDistanceComputer::symmetric_dis(idx_t i, idx_t j) {
+    return symmetric_dis_1bit<SIMDLevel::NONE>(*this, i, j);
+}
+
 namespace {
 
 // Distance computers templatized on SIMDLevel to avoid per-call dynamic
@@ -228,7 +280,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
@@ -236,27 +288,14 @@ 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 symmetric_dis(idx_t i, idx_t j) final {
+        return symmetric_dis_1bit<SL>(*this, i, j);
+    }
 
+    // Compute distance using only 1-bit codes (fast)
+    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 +325,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 +333,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 +386,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 +427,68 @@ 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);
+
+        // Honor IDSelectorWithContext on the multibit path too, so a RaBitQ
+        // index does not silently lose the context hook once nb_bits >= 2 (the
+        // 1-bit path already routes through run_scan_codes1).
+        const IDSelectorContextDispatch sel_dispatch(sel, store_pairs);
+
+        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_dispatch.is_member(
+                            id, IDScanContext{ids, list_size, j})) {
+                    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
@@ -394,26 +507,15 @@ 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 symmetric_dis(idx_t i, idx_t j) final {
+        return symmetric_dis_1bit<SL>(*this, i, j);
+    }
 
+    // Compute distance using only 1-bit codes (fast)
+    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 +530,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 +550,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 +558,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 +609,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 +641,70 @@ 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);
+
+        // Honor IDSelectorWithContext on the multibit path too, so a RaBitQ
+        // index does not silently lose the context hook once nb_bits >= 2 (the
+        // 1-bit path already routes through run_scan_codes1).
+        const IDSelectorContextDispatch sel_dispatch(sel, store_pairs);
+
+        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_dispatch.is_member(
+                            id, IDScanContext{ids, list_size, j})) {
+                    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;
     }
 };
 
@@ -552,12 +721,10 @@ FlatCodesDistanceComputer* RaBitQuantizer::get_distance_computer(
     // call the SIMD-specialized rabitq functions directly (no per-call
     // with_simd_level overhead).
     //
-    // 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>(
+    // VPOPCNT rather than SPR: Ice Lake and Zen 4 have VPOPCNTDQ without the
+    // rest of the SPR feature set. Below it, dispatch falls through to
+    // rabitq_avx512.cpp.
+    return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_BASE_WITH_VPOPCNT>(
             [&]<SIMDLevel SL>() -> FlatCodesDistanceComputer* {
                 if (qb == 0) {
                     auto dc =

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

@@ -13,9 +13,40 @@
 #include <faiss/MetricType.h>
 #include <faiss/impl/DistanceComputer.h>
 #include <faiss/impl/Quantizer.h>
+#include <faiss/impl/RaBitQUtils.h>
+#include <faiss/impl/platform_macros.h>
 
 namespace faiss {
 
+struct IDSelector;
+template <typename T, typename TI>
+struct ResultHandlerUnordered;
+
+/** Statistics for staged RaBitQ distance evaluation.
+ *
+ * These counters belong to the quantizer rather than to a particular ANN
+ * graph. They are currently populated by HNSW's staged search path.
+ */
+struct RaBitQStats {
+    uint64_t n_1bit = 0;
+    uint64_t n_refine = 0;
+
+    void reset() {
+        n_1bit = n_refine = 0;
+    }
+
+    void add(const RaBitQStats& other) {
+        n_1bit += other.n_1bit;
+        n_refine += other.n_refine;
+    }
+
+    double refine_ratio() const {
+        return n_1bit ? double(n_refine) / double(n_1bit) : 0.0;
+    }
+};
+
+FAISS_API extern RaBitQStats rabitq_stats;
+
 // 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".
@@ -118,10 +149,10 @@ struct RaBitQDistanceComputer : FlatCodesDistanceComputer {
     // Used with f_error to compute error bounds for two-stage filtering
     float g_error = 0.0f;
 
-    float symmetric_dis(idx_t /*i*/, idx_t /*j*/) override {
-        // Not used for RaBitQ
-        FAISS_THROW_MSG("Not implemented");
-    }
+    /// Per-distance-computer counters, aggregated by the owning index.
+    RaBitQStats stats;
+
+    float symmetric_dis(idx_t i, idx_t j) override;
 
     // Compute 1-bit distance estimate (fast)
     virtual float distance_to_code_1bit(const uint8_t* code) = 0;
@@ -129,6 +160,37 @@ struct RaBitQDistanceComputer : FlatCodesDistanceComputer {
     // Compute full multi-bit distance (accurate)
     virtual float distance_to_code_full(const uint8_t* code) = 0;
 
+    /// Apply the RaBitQ error bound without exposing the packed code layout to
+    /// ANN consumers. Precondition: nb_bits >= 2.
+#ifndef SWIG
+    FAISS_ALWAYS_INLINE bool should_refine(
+            const uint8_t* code,
+            float estimate,
+            float threshold,
+            bool is_similarity) const {
+        const auto* factors =
+                reinterpret_cast<const rabitq_utils::SignBitFactorsWithError*>(
+                        code + (d + 7) / 8);
+        return rabitq_utils::should_refine_candidate(
+                estimate, factors->f_error, g_error, threshold, is_similarity);
+    }
+#endif
+
+    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 {

data/vendor/faiss/faiss/impl/RaBitQuantizerMultiBit.cpp

@@ -257,8 +257,8 @@ void quantize_ex_bits(
     const size_t ex_bits = nb_bits - 1;
     FAISS_THROW_IF_NOT_MSG(
             ex_bits >= 1 && ex_bits <= 8, "ex_bits must be in range [1, 8]");
-    FAISS_THROW_IF_NOT_MSG(residual != nullptr, "residual cannot be null");
-    FAISS_THROW_IF_NOT_MSG(ex_code != nullptr, "ex_code cannot be null");
+    FAISS_THROW_IF_MSG(residual == nullptr, "residual cannot be null");
+    FAISS_THROW_IF_MSG(ex_code == nullptr, "ex_code cannot be null");
 
     // Step 1: Compute L2 norm of residual
     float norm_sqr = fvec_norm_L2sqr(residual, d);
@@ -302,9 +302,9 @@ void quantize_ex_bits(
     }
 
     // Step 6: Handle negative dimensions (flip bits)
-    // For negative residuals, flip all bits: code' = ~code & max_code
+    // For negative or zero residuals, flip all bits: code' = ~code & max_code
     for (size_t i = 0; i < d; i++) {
-        if (residual[i] < 0) {
+        if (residual[i] <= 0) {
             tmp_code[i] = (~tmp_code[i]) & max_code;
         }
     }
@@ -312,15 +312,6 @@ void quantize_ex_bits(
     // Step 7: Pack codes into byte array
     pack_multibit_codes(tmp_code.data(), ex_code, d, nb_bits);
 
-    // Step 8: Compute factors for distance computation
-    // Reconstruct total_code for factor computation
-    std::vector<int> total_code(d);
-    for (size_t i = 0; i < d; i++) {
-        // Form total_code = (sign << ex_bits) + ex_code
-        bool sign_bit = (residual[i] >= 0);
-        total_code[i] = tmp_code[i] + ((sign_bit ? 1 : 0) << ex_bits);
-    }
-
     // Compute ex-factors; centroid is needed for IP metric correction
     compute_ex_factors(
             residual,

data/vendor/faiss/faiss/impl/ResultHandler.h

@@ -18,6 +18,7 @@
 #include <faiss/impl/InvertedListScannerStats.h>
 #include <faiss/utils/Heap.h>
 #include <faiss/utils/partitioning.h>
+#include <faiss/utils/simd_levels.h>
 #include <algorithm>
 #include <iostream>
 
@@ -122,14 +123,20 @@ struct SingleQueryBlockResultHandler : BlockResultHandler<C, use_sel> {
         using ResultHandlerT<C>::threshold;
 
         explicit SingleResultHandler(SingleQueryBlockResultHandler& hr)
-                : the_handler(hr.the_handler) {}
+                : the_handler(hr.the_handler) {
+            threshold = the_handler.threshold;
+        }
 
         /// begin results for query # i
-        void begin(const size_t /* qid */) {}
+        void begin(const size_t /* qid */) {
+            threshold = the_handler.threshold;
+        }
 
         /// add one result for query i
         bool add_result(T dis, TI idx) final {
-            return the_handler.add_result(dis, idx);
+            bool updated = the_handler.add_result(dis, idx);
+            threshold = the_handler.threshold;
+            return updated;
         }
 
         /// series of results for query i is done
@@ -234,23 +241,9 @@ struct Top1BlockResultHandler : TopkBlockResultHandler<C, use_sel> {
     }
 
     /// add results for query i0..i1 and j0..j1
-    void add_results(size_t j0, size_t j1, const T* dis_tab_2) final {
-        for (size_t i = i0; i < i1; i++) {
-            const T* dis_tab_i = dis_tab_2 + (j1 - j0) * (i - i0) - j0;
-
-            auto& min_distance = this->dis_tab[i];
-            auto& min_index = this->ids_tab[i];
-
-            for (size_t j = j0; j < j1; j++) {
-                const T distance = dis_tab_i[j];
-
-                if (C::cmp(min_distance, distance)) {
-                    min_distance = distance;
-                    min_index = j;
-                }
-            }
-        }
-    }
+    /// Implemented in result_handler.cpp; dispatches via
+    /// with_selected_simd_levels to top1_add_results_tpl<C, use_sel, SL>.
+    void add_results(size_t j0, size_t j1, const T* dis_tab_2) final;
 
     void add_result(const size_t i, const T dis, const TI idx) {
         auto& min_distance = this->dis_tab[i];
@@ -553,20 +546,9 @@ struct ReservoirBlockResultHandler : TopkBlockResultHandler<C, use_sel> {
     }
 
     /// add results for query i0..i1 and j0..j1
-    void add_results(size_t j0, size_t j1, const T* dis_in) {
-#pragma omp parallel for
-        for (int64_t i = static_cast<int64_t>(i0); i < static_cast<int64_t>(i1);
-             i++) {
-            ReservoirTopN<C>& reservoir =
-                    reservoirs[i - static_cast<int64_t>(i0)];
-            const T* dis_tab_i =
-                    dis_in + (j1 - j0) * (i - static_cast<int64_t>(i0)) - j0;
-            for (size_t j = j0; j < j1; j++) {
-                T dis = dis_tab_i[j];
-                reservoir.add_result(dis, j);
-            }
-        }
-    }
+    /// Implemented in result_handler.cpp; dispatches via
+    /// with_selected_simd_levels to reservoir_add_results_tpl<C, use_sel, SL>.
+    void add_results(size_t j0, size_t j1, const T* dis_in);
 
     /// series of results for queries i0..i1 is done
     void end_multiple() final {
@@ -799,4 +781,22 @@ typename Consumer::T dispatch_range_ResultHandler(
 #undef DISPATCH_C_SEL
 }
 
+// ------------------------------------------------------------------
+// SIMD-dispatch entry points for Top1 and Reservoir add_results.
+// ------------------------------------------------------------------
+
+template <class C, bool use_sel, SIMDLevel SL>
+void top1_add_results_tpl(
+        Top1BlockResultHandler<C, use_sel>* self,
+        size_t j0,
+        size_t j1,
+        const typename C::T* dis_tab);
+
+template <class C, bool use_sel, SIMDLevel SL>
+void reservoir_add_results_tpl(
+        ReservoirBlockResultHandler<C, use_sel>* self,
+        size_t j0,
+        size_t j1,
+        const typename C::T* dis_in);
+
 } // namespace faiss