faiss 0.6.0 → 0.6.2

data/vendor/faiss/faiss/IndexIVFRaBitQ.cpp

@@ -25,16 +25,16 @@ 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;
@@ -126,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;
@@ -152,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;
@@ -167,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;
     }
@@ -187,12 +190,12 @@ 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();
     }
@@ -228,12 +231,6 @@ struct RaBitInvertedListScanner : InvertedListScanner {
         // 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];
@@ -243,16 +240,8 @@ struct RaBitInvertedListScanner : InvertedListScanner {
                 }
             }
 
-            local_1bit_evaluations++;
-
-            // Stage 1: Compute distance bound using 1-bit codes
-            // For L2 (min-heap): use lower_bound to safely skip if it's
-            //                    already worse than heap worst
-            // For IP (max-heap): use upper_bound because with a lower bound,
-            //                    we can't safely skip any candidate
             float est_distance = rabitq_dc->distance_to_code_1bit(codes);
 
-            // Extract f_error and g_error for filtering
             size_t code_size_base = (ivf_rabitq.d + 7) / 8;
             const rabitq_utils::SignBitFactorsWithError* base_fac =
                     reinterpret_cast<
@@ -266,24 +255,20 @@ struct RaBitInvertedListScanner : InvertedListScanner {
                     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];
 
                 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;
     }
 
@@ -302,6 +287,8 @@ struct RaBitInvertedListScanner : InvertedListScanner {
     }
 };
 
+} // namespace
+
 InvertedListScanner* IndexIVFRaBitQ::get_InvertedListScanner(
         bool store_pairs,
         const IDSelector* sel,
@@ -322,12 +309,12 @@ void IndexIVFRaBitQ::reconstruct_from_offset(
         int64_t list_no,
         int64_t offset,
         float* recons) const {
-    const uint8_t* code = invlists->get_single_code(list_no, offset);
+    InvertedLists::ScopedCodes sc(invlists, list_no, offset);
 
     std::vector<float> centroid(d);
     quantizer->reconstruct(list_no, centroid.data());
 
-    rabitq.decode_core(code, recons, 1, centroid.data());
+    rabitq.decode_core(sc.get(), recons, 1, centroid.data());
 }
 
 void IndexIVFRaBitQ::sa_decode(idx_t n, const uint8_t* codes, float* x) const {
@@ -370,29 +357,20 @@ float IVFRaBitDistanceComputer::operator()(idx_t i) {
     uint64_t list_no = lo_listno(lo);
     uint64_t offset = lo_offset(lo);
 
-    const uint8_t* code = parent->invlists->get_single_code(list_no, offset);
+    InvertedLists::ScopedCodes sc(parent->invlists, list_no, offset);
 
     // ok, we know the appropriate cluster that we need
     std::vector<float> centroid(parent->d);
     parent->quantizer->reconstruct(list_no, centroid.data());
 
-    // compute the distance
-    float distance = 0;
-
     std::unique_ptr<FlatCodesDistanceComputer> dc(
             parent->rabitq.get_distance_computer(
                     parent->qb, centroid.data(), /*centered=*/false));
     dc->set_query(q);
-    distance = dc->distance_to_code(code);
-
-    // deallocate
-    parent->invlists->release_codes(list_no, code);
-
-    // done
-    return distance;
+    return dc->distance_to_code(sc.get());
 }
 
-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,7 +13,6 @@
 #include <faiss/Index.h>
 #include <faiss/IndexIVF.h>
 
-#include <faiss/impl/RaBitQStats.h>
 #include <faiss/impl/RaBitQuantizer.h>
 
 namespace faiss {
@@ -29,6 +28,8 @@ struct IndexIVFRaBitQ : IndexIVF {
 
     // the default number of bits to quantize a query with.
     // use '0' to disable quantization and use raw fp32 values.
+    // Note: qb=0 is NOT supported by FastScan variants, which require
+    // quantized queries for SIMD lookup table construction.
     uint8_t qb = 4;
 
     IndexIVFRaBitQ(