faiss 0.6.2 → 0.6.3

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;
     }
 };
 

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 {

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);
@@ -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>
 
@@ -234,23 +235,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 +540,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 +775,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