faiss 0.4.3 → 0.5.1

data/vendor/faiss/faiss/IndexIVFPQFastScan.cpp

@@ -7,6 +7,7 @@
 
 #include <faiss/IndexIVFPQFastScan.h>
 
+#include <array>
 #include <cassert>
 #include <cstdio>
 
@@ -14,6 +15,7 @@
 
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/utils/Heap.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/simdlib.h>
 
@@ -210,10 +212,11 @@ void IndexIVFPQFastScan::compute_LUT(
         const float* x,
         const CoarseQuantized& cq,
         AlignedTable<float>& dis_tables,
-        AlignedTable<float>& biases) const {
+        AlignedTable<float>& biases,
+        const FastScanDistancePostProcessing&) const {
     size_t dim12 = pq.ksub * pq.M;
     size_t d = pq.d;
-    size_t nprobe = this->nprobe;
+    size_t nprobe = cq.nprobe;
 
     if (by_residual) {
         if (metric_type == METRIC_L2) {
@@ -292,4 +295,133 @@ void IndexIVFPQFastScan::compute_LUT(
     }
 }
 
+/*********************************************************
+ * InvertedListScanner for IVFPQFS
+ *********************************************************/
+
+namespace {
+
+struct IVFPQFastScanScanner : InvertedListScanner {
+    static constexpr int impl = 10; // based on search_implem_10
+    static constexpr size_t nq = 1; // 1 query at a time.
+    const IndexIVFPQFastScan& index;
+    AlignedTable<uint8_t> dis_tables;
+    AlignedTable<uint16_t> biases;
+    std::array<float, 2> normalizers{};
+    const float* xi = nullptr;
+
+    IVFPQFastScanScanner(
+            const IndexIVFPQFastScan& index,
+            bool store_pairs,
+            const IDSelector* sel)
+            : InvertedListScanner(store_pairs, sel), index(index) {
+        this->keep_max = is_similarity_metric(index.metric_type);
+    }
+
+    void set_query(const float* query) override {
+        this->xi = query;
+    }
+
+    void set_list(idx_t list_no, float coarse_dis) override {
+        this->list_no = list_no;
+        IndexIVFFastScan::CoarseQuantized cq{
+                .nprobe = 1,        // 1 due to explicitly passing in list_no
+                .dis = &coarse_dis, // dis from query to list_no centroid.
+                .ids = &list_no,    // id of the current list we are scanning
+        };
+        FastScanDistancePostProcessing empty_context{};
+        index.compute_LUT_uint8(
+                1, xi, cq, dis_tables, biases, &normalizers[0], empty_context);
+    }
+
+    float distance_to_code(const uint8_t* /* code */) const override {
+        // It's not really possible to implement a distance_to_code since codes
+        // for 32 database vectors are intermixed.
+        FAISS_THROW_MSG("not implemented");
+    }
+
+    // Based on IVFFastScan search_implem_10, since it also deals with 1 query
+    // at a time.
+    size_t scan_codes(
+            size_t ntotal,
+            const uint8_t* codes,
+            const idx_t* ids,
+            float* distances,
+            idx_t* labels,
+            size_t k) const override {
+        // initialize the current iteration heap to the worst possible value of
+        // the prior loop
+        std::vector<float> curr_dists(k, distances[0]);
+        std::vector<idx_t> curr_labels(k, labels[0]);
+        FastScanDistancePostProcessing empty_context{};
+        std::unique_ptr<SIMDResultHandlerToFloat> handler(
+                index.make_knn_handler(
+                        !keep_max,
+                        impl,
+                        nq,
+                        k,
+                        curr_dists.data(),
+                        curr_labels.data(),
+                        sel,
+                        empty_context,
+                        &normalizers[0]));
+
+        // This does not quite match search_implem_10, but it is fine because
+        // the scanner operates on a single query at a time, and this value is
+        // used as the query index. For a single query, the value is always 0.
+        int qmap1[1] = {0};
+
+        handler->q_map = qmap1;
+        handler->begin(&normalizers[0]);
+
+        const uint8_t* LUT = dis_tables.get();
+        handler->dbias = biases.get();
+
+        handler->ntotal = ntotal;
+        handler->id_map = ids;
+
+        pq4_accumulate_loop(
+                1,
+                roundup(ntotal, index.bbs),
+                index.bbs,
+                static_cast<int>(index.M2),
+                codes,
+                LUT,
+                *handler,
+                nullptr);
+
+        // The handler is for the results of this iteration.
+        // Then we need a second heap to combine across iterations.
+        handler->end();
+        if (keep_max) {
+            minheap_addn(
+                    k,
+                    distances,
+                    labels,
+                    curr_dists.data(),
+                    curr_labels.data(),
+                    k);
+        } else {
+            maxheap_addn(
+                    k,
+                    distances,
+                    labels,
+                    curr_dists.data(),
+                    curr_labels.data(),
+                    k);
+        }
+
+        return handler->num_updates();
+    }
+};
+
+} // anonymous namespace
+
+InvertedListScanner* IndexIVFPQFastScan::get_InvertedListScanner(
+        bool store_pairs,
+        const IDSelector* sel,
+        const IVFSearchParameters*) const {
+    return new IVFPQFastScanScanner(*this, store_pairs, sel);
+}
+
 } // namespace faiss

data/vendor/faiss/faiss/IndexIVFPQFastScan.h

@@ -80,7 +80,13 @@ struct IndexIVFPQFastScan : IndexIVFFastScan {
             const float* x,
             const CoarseQuantized& cq,
             AlignedTable<float>& dis_tables,
-            AlignedTable<float>& biases) const override;
+            AlignedTable<float>& biases,
+            const FastScanDistancePostProcessing& context) const override;
+
+    InvertedListScanner* get_InvertedListScanner(
+            bool store_pairs,
+            const IDSelector* sel,
+            const IVFSearchParameters*) const override;
 };
 
 } // namespace faiss

data/vendor/faiss/faiss/IndexIVFRaBitQ.cpp

@@ -24,9 +24,10 @@ IndexIVFRaBitQ::IndexIVFRaBitQ(
         const size_t d,
         const size_t nlist,
         MetricType metric,
-        bool own_invlists)
+        bool own_invlists,
+        uint8_t nb_bits_in)
         : IndexIVF(quantizer, d, nlist, 0, metric, own_invlists),
-          rabitq(d, metric) {
+          rabitq(d, metric, nb_bits_in) {
     code_size = rabitq.code_size;
     if (own_invlists) {
         invlists->code_size = code_size;
@@ -153,17 +154,22 @@ struct RaBitInvertedListScanner : InvertedListScanner {
     std::vector<float> query_vector;
 
     std::unique_ptr<FlatCodesDistanceComputer> dc;
+    RaBitQDistanceComputer* rabitq_dc =
+            nullptr; // For multi-bit adaptive filtering
 
     uint8_t qb = 0;
+    bool centered = false;
 
-    RaBitInvertedListScanner(
+    explicit RaBitInvertedListScanner(
             const IndexIVFRaBitQ& ivf_rabitq_in,
             bool store_pairs = false,
             const IDSelector* sel = nullptr,
-            uint8_t qb_in = 0)
+            uint8_t qb_in = 0,
+            bool centered = false)
             : InvertedListScanner(store_pairs, sel),
               ivf_rabitq{ivf_rabitq_in},
-              qb{qb_in} {
+              qb{qb_in},
+              centered(centered) {
         keep_max = is_similarity_metric(ivf_rabitq.metric_type);
         code_size = ivf_rabitq.code_size;
     }
@@ -191,14 +197,95 @@ struct RaBitInvertedListScanner : InvertedListScanner {
         return dc->distance_to_code(code);
     }
 
+    /// 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 {
+        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);
+        }
+
+        // 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];
+                if (!sel->is_member(id)) {
+                    codes += code_size;
+                    continue;
+                }
+            }
+
+            local_1bit_evaluations++;
+
+            // Stage 1: Compute lower bound using 1-bit codes
+            float lower_bound = rabitq_dc->lower_bound_distance(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]);
+
+            if (should_refine) {
+                local_multibit_evaluations++;
+                // Lower bound is promising, compute full distance
+                float dis = distance_to_code(codes);
+
+                // 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);
+                    }
+                    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;
+    }
+
     void internal_try_setup_dc() {
         if (!query_vector.empty() && !reconstructed_centroid.empty()) {
             // both query_vector and centroid are available!
             // set up DistanceComputer
             dc.reset(ivf_rabitq.rabitq.get_distance_computer(
-                    qb, reconstructed_centroid.data()));
+                    qb, reconstructed_centroid.data(), centered));
 
             dc->set_query(query_vector.data());
+
+            // Try to cast to RaBitQDistanceComputer for multi-bit support
+            rabitq_dc = dynamic_cast<RaBitQDistanceComputer*>(dc.get());
         }
     }
 };
@@ -208,12 +295,15 @@ InvertedListScanner* IndexIVFRaBitQ::get_InvertedListScanner(
         const IDSelector* sel,
         const IVFSearchParameters* search_params_in) const {
     uint8_t used_qb = qb;
+    bool centered = false;
     if (auto params = dynamic_cast<const IVFRaBitQSearchParameters*>(
                 search_params_in)) {
         used_qb = params->qb;
+        centered = params->centered;
     }
 
-    return new RaBitInvertedListScanner(*this, store_pairs, sel, used_qb);
+    return new RaBitInvertedListScanner(
+            *this, store_pairs, sel, used_qb, centered);
 }
 
 void IndexIVFRaBitQ::reconstruct_from_offset(
@@ -278,7 +368,8 @@ float IVFRaBitDistanceComputer::operator()(idx_t i) {
     float distance = 0;
 
     std::unique_ptr<FlatCodesDistanceComputer> dc(
-            parent->rabitq.get_distance_computer(parent->qb, centroid.data()));
+            parent->rabitq.get_distance_computer(
+                    parent->qb, centroid.data(), /*centered=*/false));
     dc->set_query(q);
     distance = dc->distance_to_code(code);
 

data/vendor/faiss/faiss/IndexIVFRaBitQ.h

@@ -13,12 +13,14 @@
 #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;
+    bool centered = false;
 };
 
 // * by_residual is true, just by design
@@ -34,7 +36,8 @@ struct IndexIVFRaBitQ : IndexIVF {
             const size_t d,
             const size_t nlist,
             MetricType metric = METRIC_L2,
-            bool own_invlists = true);
+            bool own_invlists = true,
+            uint8_t nb_bits = 1);
 
     IndexIVFRaBitQ();