faiss 0.5.0 → 0.5.1

data/vendor/faiss/faiss/IndexIVFRaBitQFastScan.h

@@ -12,6 +12,7 @@
 #include <faiss/IndexIVFFastScan.h>
 #include <faiss/IndexIVFRaBitQ.h>
 #include <faiss/IndexRaBitQFastScan.h>
+#include <faiss/impl/RaBitQStats.h>
 #include <faiss/impl/RaBitQUtils.h>
 #include <faiss/impl/RaBitQuantizer.h>
 #include <faiss/impl/simd_result_handlers.h>
@@ -24,8 +25,9 @@ namespace faiss {
 struct FastScanDistancePostProcessing;
 
 // Import shared utilities from RaBitQUtils
-using rabitq_utils::FactorsData;
 using rabitq_utils::QueryFactorsData;
+using rabitq_utils::SignBitFactors;
+using rabitq_utils::SignBitFactorsWithError;
 
 /** Fast-scan version of IndexIVFRaBitQ that processes vectors in batches
  * using SIMD operations. Combines the inverted file structure of IVF
@@ -53,9 +55,16 @@ struct IndexIVFRaBitQFastScan : IndexIVFFastScan {
     /// Use zero-centered scalar quantizer for queries
     bool centered = false;
 
-    /// Extracted factors storage for batch processing
-    /// Size: ntotal, stores factors separately from packed codes
-    std::vector<FactorsData> factors_storage;
+    /// Per-vector auxiliary data (1-bit codes stored separately in `codes`)
+    ///
+    /// 1-bit codes (sign bits) are stored in the inherited `codes` array from
+    /// IndexFastScan in packed FastScan format for SIMD processing.
+    ///
+    /// This flat_storage holds per-vector factors and refinement-bit codes:
+    /// Layout for 1-bit: [SignBitFactors (8 bytes)]
+    /// Layout for multi-bit: [SignBitFactorsWithError
+    /// (12B)][ref_codes][ExtraBitsFactors (8B)]
+    std::vector<uint8_t> flat_storage;
 
     // Constructors
 
@@ -67,7 +76,8 @@ struct IndexIVFRaBitQFastScan : IndexIVFFastScan {
             size_t nlist,
             MetricType metric = METRIC_L2,
             int bbs = 32,
-            bool own_invlists = true);
+            bool own_invlists = true,
+            uint8_t nb_bits = 1);
 
     /// Build from an existing IndexIVFRaBitQ
     explicit IndexIVFRaBitQFastScan(const IndexIVFRaBitQ& orig, int bbs = 32);
@@ -101,13 +111,10 @@ struct IndexIVFRaBitQFastScan : IndexIVFFastScan {
     /// Override sa_decode to handle RaBitQ reconstruction
     void sa_decode(idx_t n, const uint8_t* bytes, float* x) const override;
 
-   private:
-    /// Encode a vector to FastScan format without computing factors
-    void encode_vector_to_fastscan(
-            const float* xi,
-            const float* centroid,
-            uint8_t* fastscan_code) const;
+    /// Compute storage size per vector in flat_storage based on nb_bits
+    size_t compute_per_vector_storage_size() const;
 
+   private:
     /// Compute query factors and lookup table for a residual vector
     /// (similar to IndexRaBitQFastScan::compute_float_LUT)
     void compute_residual_LUT(
@@ -116,10 +123,12 @@ struct IndexIVFRaBitQFastScan : IndexIVFFastScan {
             float* lut_out,
             const float* original_query = nullptr) const;
 
-    /// Decode FastScan code to RaBitQ residual vector
+    /// Decode FastScan code to RaBitQ residual vector with explicit
+    /// dp_multiplier
     void decode_fastscan_to_residual(
             const uint8_t* fastscan_code,
-            float* residual) const;
+            float* residual,
+            float dp_multiplier) const;
 
    public:
     /// Implementation methods for IVFRaBitQFastScan specialization
@@ -171,6 +180,7 @@ struct IndexIVFRaBitQFastScan : IndexIVFFastScan {
      * - Specialized handling for both centered and non-centered quantization
      * modes
      * - Efficient inner product metric corrections
+     * - Uses runtime boolean for multi-bit mode
      *
      * @tparam C Comparator type (CMin/CMax) for heap operations
      */
@@ -185,7 +195,8 @@ struct IndexIVFRaBitQFastScan : IndexIVFFastScan {
         std::vector<int>
                 probe_indices; // probe index for each query in current batch
         const FastScanDistancePostProcessing*
-                context; // Processing context with query factors
+                context;        // Processing context with query factors
+        const bool is_multibit; // Whether to use multi-bit two-stage search
 
         // Use float-based comparator for heap operations
         using Cfloat = typename std::conditional<
@@ -199,9 +210,11 @@ struct IndexIVFRaBitQFastScan : IndexIVFFastScan {
                 size_t k_val,
                 float* distances,
                 int64_t* labels,
-                const FastScanDistancePostProcessing* ctx = nullptr);
+                const FastScanDistancePostProcessing* ctx = nullptr,
+                bool multibit = false);
 
-        void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) final;
+        void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1)
+                override;
 
         /// Override base class virtual method to receive context information
         void set_list_context(size_t list_no, const std::vector<int>& probe_map)
@@ -210,6 +223,29 @@ struct IndexIVFRaBitQFastScan : IndexIVFFastScan {
         void begin(const float* norms) override;
 
         void end() override;
+
+       private:
+        /// Compute full multi-bit distance for a candidate vector (multi-bit
+        /// only)
+        /// @param db_idx Global database vector index
+        /// @param local_q Batch-local query index (for probe_indices access)
+        /// @param global_q Global query index (for storage indexing)
+        /// @param local_offset Offset within the current inverted list
+        float compute_full_multibit_distance(
+                size_t db_idx,
+                size_t local_q,
+                size_t global_q,
+                size_t local_offset) const;
+
+        /// Compute lower bound using 1-bit distance and error bound (multi-bit
+        /// only)
+        /// @param local_q Batch-local query index (for probe_indices access)
+        /// @param global_q Global query index (for storage indexing)
+        float compute_lower_bound(
+                float dist_1bit,
+                size_t db_idx,
+                size_t local_q,
+                size_t global_q) const;
     };
 };
 

data/vendor/faiss/faiss/IndexPQ.cpp

@@ -81,6 +81,7 @@ struct PQDistanceComputer : FlatCodesDistanceComputer {
     const float* sdc;
     std::vector<float> precomputed_table;
     size_t ndis;
+    const float* q;
 
     float distance_to_code(const uint8_t* code) final {
         ndis++;
@@ -109,7 +110,8 @@ struct PQDistanceComputer : FlatCodesDistanceComputer {
             : FlatCodesDistanceComputer(
                       storage.codes.data(),
                       storage.code_size),
-              pq(storage.pq) {
+              pq(storage.pq),
+              q(nullptr) {
         precomputed_table.resize(pq.M * pq.ksub);
         nb = storage.ntotal;
         d = storage.d;
@@ -123,6 +125,7 @@ struct PQDistanceComputer : FlatCodesDistanceComputer {
     }
 
     void set_query(const float* x) override {
+        q = x;
         if (metric == METRIC_L2) {
             pq.compute_distance_table(x, precomputed_table.data());
         } else {

data/vendor/faiss/faiss/IndexPreTransform.cpp

@@ -197,6 +197,20 @@ void IndexPreTransform::range_search(
             n, tv.x, radius, result, extract_index_search_params(params));
 }
 
+void IndexPreTransform::search_subset(
+        idx_t n,
+        const float* x,
+        idx_t k_base,
+        const idx_t* base_labels,
+        idx_t k,
+        float* distances,
+        idx_t* labels) const {
+    FAISS_THROW_IF_NOT(k > 0);
+    FAISS_THROW_IF_NOT(is_trained);
+    TransformedVectors tv(x, apply_chain(n, x));
+    index->search_subset(n, tv.x, k_base, base_labels, k, distances, labels);
+}
+
 void IndexPreTransform::reset() {
     index->reset();
     ntotal = 0;

data/vendor/faiss/faiss/IndexPreTransform.h

@@ -57,6 +57,15 @@ struct IndexPreTransform : Index {
             idx_t* labels,
             const SearchParameters* params = nullptr) const override;
 
+    void search_subset(
+            idx_t n,
+            const float* x,
+            idx_t k_base,
+            const idx_t* base_labels,
+            idx_t k,
+            float* distances,
+            idx_t* labels) const override;
+
     /* range search, no attempt is done to change the radius */
     void range_search(
             idx_t n,

data/vendor/faiss/faiss/IndexRaBitQ.cpp

@@ -9,13 +9,18 @@
 
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/ResultHandler.h>
+#include <memory>
 
 namespace faiss {
 
+// Forward declaration from RaBitQuantizer.cpp
+struct RaBitQDistanceComputer;
+
 IndexRaBitQ::IndexRaBitQ() = default;
 
-IndexRaBitQ::IndexRaBitQ(idx_t d, MetricType metric)
-        : IndexFlatCodes(0, d, metric), rabitq(d, metric) {
+IndexRaBitQ::IndexRaBitQ(idx_t d, MetricType metric, uint8_t nb_bits_in)
+        : IndexFlatCodes(0, d, metric), rabitq(d, metric, nb_bits_in) {
+    // Update code size based on nb_bits
     code_size = rabitq.code_size;
 
     is_trained = false;
@@ -78,6 +83,7 @@ struct Run_search_with_dc_res {
 
     uint8_t qb = 0;
     bool centered = false;
+    uint8_t nb_bits = 1; // Number of bits per dimension
 
     template <class BlockResultHandler>
     void f(BlockResultHandler& res, const IndexRaBitQ* index, const float* xq) {
@@ -85,22 +91,87 @@ struct Run_search_with_dc_res {
         using SingleResultHandler =
                 typename BlockResultHandler::SingleResultHandler;
         const int d = index->d;
+        size_t ex_bits = nb_bits - 1;
 
-#pragma omp parallel // if (res.nq > 100)
+#pragma omp parallel
         {
-            std::unique_ptr<FlatCodesDistanceComputer> dc(
+            std::unique_ptr<FlatCodesDistanceComputer> dc_base(
                     index->get_quantized_distance_computer(qb, centered));
             SingleResultHandler resi(res);
 #pragma omp for
             for (int64_t q = 0; q < res.nq; q++) {
                 resi.begin(q);
-                dc->set_query(xq + d * q);
-                for (size_t i = 0; i < ntotal; i++) {
-                    if (res.is_in_selection(i)) {
-                        float dis = (*dc)(i);
-                        resi.add_result(dis, i);
+                dc_base->set_query(xq + d * q);
+
+                // Stats tracking for multi-bit two-stage search only
+                // 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;
+
+                if (ex_bits == 0) {
+                    // 1-bit: Standard single-stage search (no stats tracking)
+                    for (size_t i = 0; i < ntotal; i++) {
+                        if (res.is_in_selection(i)) {
+                            float dis = (*dc_base)(i);
+                            resi.add_result(dis, i);
+                        }
+                    }
+                } else {
+                    // Multi-bit: Two-stage search with adaptive filtering
+                    // Note: Even with query quantization (qb > 0), ex-bits
+                    // distance computation uses the float query to maintain
+                    // consistency with encoding-time factor computation. See
+                    // RaBitQuantizer.cpp for details.
+                    auto* dc = dynamic_cast<RaBitQDistanceComputer*>(
+                            dc_base.get());
+                    FAISS_THROW_IF_NOT_MSG(
+                            dc != nullptr,
+                            "Failed to cast to RaBitQDistanceComputer for two-stage search");
+
+                    // Use appropriate comparison based on metric type
+                    bool is_similarity =
+                            is_similarity_metric(index->metric_type);
+
+                    for (size_t i = 0; i < ntotal; i++) {
+                        if (res.is_in_selection(i)) {
+                            const uint8_t* code =
+                                    index->codes.data() + i * index->code_size;
+
+                            local_1bit_evaluations++;
+
+                            // Stage 1: Compute 1-bit lower bound
+                            float lower_bound = dc->lower_bound_distance(code);
+
+                            // Stage 2: Adaptive filtering using threshold
+                            // For L2 (min-heap): filter if lower_bound <
+                            // resi.threshold For IP (max-heap): filter if
+                            // lower_bound > resi.threshold Note: Using
+                            // resi.threshold directly (not cached) enables more
+                            // aggressive filtering as the heap is updated
+                            bool should_refine = is_similarity
+                                    ? (lower_bound > resi.threshold)
+                                    : (lower_bound < resi.threshold);
+
+                            if (should_refine) {
+                                local_multibit_evaluations++;
+                                // Compute full multi-bit distance
+                                float dist_full =
+                                        dc->distance_to_code_full(code);
+                                resi.add_result(dist_full, i);
+                            }
+                        }
                     }
                 }
+
+                // 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;
+
                 resi.end();
             }
         }
@@ -116,16 +187,25 @@ void IndexRaBitQ::search(
         float* distances,
         idx_t* labels,
         const SearchParameters* params_in) const {
-    const IDSelector* sel = (params_in != nullptr) ? params_in->sel : nullptr;
-    Run_search_with_dc_res r;
+    FAISS_THROW_IF_NOT(is_trained);
+
+    // Extract search parameters
+    uint8_t used_qb = qb;
+    bool used_centered = centered;
     if (auto params = dynamic_cast<const RaBitQSearchParameters*>(params_in)) {
-        r.qb = params->qb;
-        r.centered = params->centered;
-    } else {
-        r.qb = this->qb;
-        r.centered = this->centered;
+        used_qb = params->qb;
+        used_centered = params->centered;
     }
 
+    const IDSelector* sel = (params_in != nullptr) ? params_in->sel : nullptr;
+
+    // Set up functor with all necessary parameters
+    Run_search_with_dc_res r;
+    r.qb = used_qb;
+    r.centered = used_centered;
+    r.nb_bits = rabitq.nb_bits; // Pass multi-bit info to functor
+
+    // Use Faiss framework for all cases (single-stage and two-stage)
     dispatch_knn_ResultHandler(
             n, distances, labels, k, metric_type, sel, r, this, x);
 }

data/vendor/faiss/faiss/IndexRaBitQ.h

@@ -8,6 +8,7 @@
 #pragma once
 
 #include <faiss/IndexFlatCodes.h>
+#include <faiss/impl/RaBitQStats.h>
 #include <faiss/impl/RaBitQuantizer.h>
 
 namespace faiss {
@@ -32,7 +33,10 @@ struct IndexRaBitQ : IndexFlatCodes {
 
     IndexRaBitQ();
 
-    explicit IndexRaBitQ(idx_t d, MetricType metric = METRIC_L2);
+    explicit IndexRaBitQ(
+            idx_t d,
+            MetricType metric = METRIC_L2,
+            uint8_t nb_bits = 1);
 
     void train(idx_t n, const float* x) override;