faiss 0.5.3 → 0.6.1

data/vendor/faiss/faiss/IndexIVFRaBitQ.cpp

@@ -15,21 +15,26 @@
 #include <vector>
 
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/RaBitQUtils.h>
 #include <faiss/impl/RaBitQuantizer.h>
+#include <faiss/impl/ResultHandler.h>
+#include <faiss/impl/expanded_scanners.h>
 
 namespace faiss {
 
+using rabitq_utils::SignBitFactorsWithError;
+
 IndexIVFRaBitQ::IndexIVFRaBitQ(
-        Index* quantizer,
-        const size_t d,
-        const size_t nlist,
+        Index* quantizer_in,
+        const size_t d_in,
+        const size_t nlist_in,
         MetricType metric,
-        bool own_invlists,
+        bool own_invlists_in,
         uint8_t nb_bits_in)
-        : IndexIVF(quantizer, d, nlist, 0, metric, own_invlists),
-          rabitq(d, metric, nb_bits_in) {
+        : IndexIVF(quantizer_in, d_in, nlist_in, 0, metric, own_invlists_in),
+          rabitq(d_in, metric, nb_bits_in) {
     code_size = rabitq.code_size;
-    if (own_invlists) {
+    if (own_invlists_in) {
         invlists->code_size = code_size;
     }
     is_trained = false;
@@ -44,7 +49,7 @@ IndexIVFRaBitQ::IndexIVFRaBitQ() {
 void IndexIVFRaBitQ::train_encoder(
         idx_t n,
         const float* x,
-        const idx_t* assign) {
+        const idx_t* /*assign*/) {
     rabitq.train(n, x);
 }
 
@@ -121,7 +126,7 @@ void IndexIVFRaBitQ::add_core(
         int rank = omp_get_thread_num();
 
         // each thread takes care of a subset of lists
-        for (size_t i = 0; i < n; i++) {
+        for (idx_t i = 0; i < n; i++) {
             int64_t list_no = precomputed_idx[i];
             if (list_no >= 0 && list_no % nt == rank) {
                 int64_t id = xids ? xids[i] : ntotal + i;
@@ -147,7 +152,10 @@ void IndexIVFRaBitQ::add_core(
     ntotal += n;
 }
 
+namespace {
+
 struct RaBitInvertedListScanner : InvertedListScanner {
+    using InvertedListScanner::scan_codes;
     const IndexIVFRaBitQ& ivf_rabitq;
 
     std::vector<float> reconstructed_centroid;
@@ -162,14 +170,14 @@ struct RaBitInvertedListScanner : InvertedListScanner {
 
     explicit RaBitInvertedListScanner(
             const IndexIVFRaBitQ& ivf_rabitq_in,
-            bool store_pairs = false,
-            const IDSelector* sel = nullptr,
+            bool store_pairs_in = false,
+            const IDSelector* sel_in = nullptr,
             uint8_t qb_in = 0,
-            bool centered = false)
-            : InvertedListScanner(store_pairs, sel),
+            bool centered_in = false)
+            : InvertedListScanner(store_pairs_in, sel_in),
               ivf_rabitq{ivf_rabitq_in},
               qb{qb_in},
-              centered(centered) {
+              centered(centered_in) {
         keep_max = is_similarity_metric(ivf_rabitq.metric_type);
         code_size = ivf_rabitq.code_size;
     }
@@ -182,46 +190,47 @@ struct RaBitInvertedListScanner : InvertedListScanner {
     }
 
     /// following codes come from this inverted list
-    void set_list(idx_t list_no, float coarse_dis) override {
-        this->list_no = list_no;
+    void set_list(idx_t list_no_in, float /*coarse_dis*/) override {
+        this->list_no = list_no_in;
 
         reconstructed_centroid.resize(ivf_rabitq.d);
         ivf_rabitq.quantizer->reconstruct(
-                list_no, reconstructed_centroid.data());
+                list_no_in, reconstructed_centroid.data());
 
         internal_try_setup_dc();
     }
 
     /// compute a single query-to-code distance
-    float distance_to_code(const uint8_t* code) const override {
+    float distance_to_code(const uint8_t* code) const final {
         return dc->distance_to_code(code);
     }
 
+    // redefiniing the scan_codes allows to inline the distance_to_code
+    // (this is unlikely to matter because it contains a virtual function call)
+    size_t scan_codes_1bit(
+            size_t list_size,
+            const uint8_t* codes,
+            const idx_t* ids,
+            ResultHandler& handler) const {
+        return run_scan_codes(*this, list_size, codes, ids, handler);
+    }
+
     /// Override scan_codes to implement adaptive filtering for multi-bit codes
     size_t scan_codes(
             size_t list_size,
             const uint8_t* codes,
             const idx_t* ids,
-            float* simi,
-            idx_t* idxi,
-            size_t k) const override {
+            ResultHandler& handler) const override {
         size_t ex_bits = ivf_rabitq.rabitq.nb_bits - 1;
 
         // For 1-bit codes, use default implementation
         if (ex_bits == 0 || rabitq_dc == nullptr) {
-            return InvertedListScanner::scan_codes(
-                    list_size, codes, ids, simi, idxi, k);
+            return scan_codes_1bit(list_size, codes, ids, handler);
         }
 
         // Multi-bit: Two-stage search with adaptive filtering
         size_t nup = 0;
 
-        // Stats tracking for multi-bit two-stage search
-        // n_1bit_evaluations: candidates evaluated using 1-bit lower bound
-        // n_multibit_evaluations: candidates requiring full multi-bit distance
-        size_t local_1bit_evaluations = 0;
-        size_t local_multibit_evaluations = 0;
-
         for (size_t j = 0; j < list_size; j++) {
             if (sel != nullptr) {
                 int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
@@ -231,47 +240,35 @@ struct RaBitInvertedListScanner : InvertedListScanner {
                 }
             }
 
-            local_1bit_evaluations++;
-
-            // Stage 1: Compute lower bound using 1-bit codes
-            float lower_bound = rabitq_dc->lower_bound_distance(codes);
+            float est_distance = rabitq_dc->distance_to_code_1bit(codes);
 
-            // Stage 2: Adaptive filtering
-            // L2 (min-heap): filter if lower_bound < simi[0]
-            // IP (max-heap): filter if lower_bound > simi[0]
-            // Note: Using simi[0] directly (not cached) enables more aggressive
-            // filtering as the heap is updated with better candidates
-            bool should_refine = keep_max ? (lower_bound > simi[0])
-                                          : (lower_bound < simi[0]);
+            size_t code_size_base = (ivf_rabitq.d + 7) / 8;
+            const rabitq_utils::SignBitFactorsWithError* base_fac =
+                    reinterpret_cast<
+                            const rabitq_utils::SignBitFactorsWithError*>(
+                            codes + code_size_base);
 
+            bool should_refine = rabitq_utils::should_refine_candidate(
+                    est_distance,
+                    base_fac->f_error,
+                    rabitq_dc->g_error,
+                    handler.threshold,
+                    keep_max);
             if (should_refine) {
-                local_multibit_evaluations++;
-                // Lower bound is promising, compute full distance
+                // Refining computes the full distance — counts as a
+                // post-filter "distance computed" for stats purposes.
+                handler.stats.scan_cnt++;
                 float dis = distance_to_code(codes);
+                int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
 
-                // Check if distance improves heap
-                bool improves_heap =
-                        keep_max ? (dis > simi[0]) : (dis < simi[0]);
-
-                if (improves_heap) {
-                    int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
-                    if (keep_max) {
-                        minheap_replace_top(k, simi, idxi, dis, id);
-                    } else {
-                        maxheap_replace_top(k, simi, idxi, dis, id);
-                    }
+                if (handler.add_result(dis, id)) {
+                    handler.stats.nheap_updates++;
                     nup++;
                 }
             }
             codes += code_size;
         }
 
-        // Update global stats atomically
-#pragma omp atomic
-        rabitq_stats.n_1bit_evaluations += local_1bit_evaluations;
-#pragma omp atomic
-        rabitq_stats.n_multibit_evaluations += local_multibit_evaluations;
-
         return nup;
     }
 
@@ -290,6 +287,8 @@ struct RaBitInvertedListScanner : InvertedListScanner {
     }
 };
 
+} // namespace
+
 InvertedListScanner* IndexIVFRaBitQ::get_InvertedListScanner(
         bool store_pairs,
         const IDSelector* sel,
@@ -380,7 +379,7 @@ float IVFRaBitDistanceComputer::operator()(idx_t i) {
     return distance;
 }
 
-float IVFRaBitDistanceComputer::symmetric_dis(idx_t i, idx_t j) {
+float IVFRaBitDistanceComputer::symmetric_dis(idx_t /*i*/, idx_t /*j*/) {
     FAISS_THROW_MSG("Not implemented");
 }
 

data/vendor/faiss/faiss/IndexIVFRaBitQ.h

@@ -13,13 +13,12 @@
 #include <faiss/Index.h>
 #include <faiss/IndexIVF.h>
 
-#include <faiss/impl/RaBitQStats.h>
 #include <faiss/impl/RaBitQuantizer.h>
 
 namespace faiss {
 
 struct IVFRaBitQSearchParameters : IVFSearchParameters {
-    uint8_t qb = 0;
+    uint8_t qb = 4;
     bool centered = false;
 };
 
@@ -29,7 +28,9 @@ struct IndexIVFRaBitQ : IndexIVF {
 
     // the default number of bits to quantize a query with.
     // use '0' to disable quantization and use raw fp32 values.
-    uint8_t qb = 0;
+    // Note: qb=0 is NOT supported by FastScan variants, which require
+    // quantized queries for SIMD lookup table construction.
+    uint8_t qb = 4;
 
     IndexIVFRaBitQ(
             Index* quantizer,