faiss 0.4.3 → 0.5.1

data/vendor/faiss/faiss/IndexIVFRaBitQFastScan.h

@@ -0,0 +1,252 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <vector>
+
+#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>
+#include <faiss/utils/AlignedTable.h>
+#include <faiss/utils/Heap.h>
+
+namespace faiss {
+
+// Forward declarations
+struct FastScanDistancePostProcessing;
+
+// Import shared utilities from RaBitQUtils
+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
+ * with RaBitQ's bit-level quantization and FastScan's batch processing.
+ *
+ * Key features:
+ * - Inherits from IndexIVFFastScan for IVF structure and search algorithms
+ * - Processes 32 database vectors at a time using SIMD
+ * - Separates factors from quantized bits for efficient processing
+ * - Supports both L2 and inner product metrics
+ * - Maintains compatibility with existing IVF search parameters
+ *
+ * Implementation details:
+ * - Batch size (bbs) is typically 32 for optimal SIMD performance
+ * - Factors are stored separately from packed codes for cache efficiency
+ * - Query factors are computed once per search and reused across lists
+ * - Uses specialized result handlers for RaBitQ distance corrections
+ */
+struct IndexIVFRaBitQFastScan : IndexIVFFastScan {
+    RaBitQuantizer rabitq;
+
+    /// Default number of bits to quantize a query with
+    uint8_t qb = 8;
+
+    /// Use zero-centered scalar quantizer for queries
+    bool centered = false;
+
+    /// 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
+
+    IndexIVFRaBitQFastScan();
+
+    IndexIVFRaBitQFastScan(
+            Index* quantizer,
+            size_t d,
+            size_t nlist,
+            MetricType metric = METRIC_L2,
+            int bbs = 32,
+            bool own_invlists = true,
+            uint8_t nb_bits = 1);
+
+    /// Build from an existing IndexIVFRaBitQ
+    explicit IndexIVFRaBitQFastScan(const IndexIVFRaBitQ& orig, int bbs = 32);
+
+    // Required overrides
+
+    void train_encoder(idx_t n, const float* x, const idx_t* assign) override;
+
+    void encode_vectors(
+            idx_t n,
+            const float* x,
+            const idx_t* list_nos,
+            uint8_t* codes,
+            bool include_listnos = false) const override;
+
+   protected:
+    /// Extract and store RaBitQ factors from encoded vectors
+    void preprocess_code_metadata(
+            idx_t n,
+            const uint8_t* flat_codes,
+            idx_t start_global_idx) override;
+
+    /// Return code_size as stride to skip embedded factor data during packing
+    size_t code_packing_stride() const override;
+
+   public:
+    /// Reconstruct a single vector from an inverted list
+    void reconstruct_from_offset(int64_t list_no, int64_t offset, float* recons)
+            const override;
+
+    /// Override sa_decode to handle RaBitQ reconstruction
+    void sa_decode(idx_t n, const uint8_t* bytes, float* x) const override;
+
+    /// 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(
+            const float* residual,
+            QueryFactorsData& query_factors,
+            float* lut_out,
+            const float* original_query = nullptr) const;
+
+    /// Decode FastScan code to RaBitQ residual vector with explicit
+    /// dp_multiplier
+    void decode_fastscan_to_residual(
+            const uint8_t* fastscan_code,
+            float* residual,
+            float dp_multiplier) const;
+
+   public:
+    /// Implementation methods for IVFRaBitQFastScan specialization
+    bool lookup_table_is_3d() const override;
+
+    void compute_LUT(
+            size_t n,
+            const float* x,
+            const CoarseQuantized& cq,
+            AlignedTable<float>& dis_tables,
+            AlignedTable<float>& biases,
+            const FastScanDistancePostProcessing& context) const override;
+
+    void search_preassigned(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            const idx_t* assign,
+            const float* centroid_dis,
+            float* distances,
+            idx_t* labels,
+            bool store_pairs,
+            const IVFSearchParameters* params = nullptr,
+            IndexIVFStats* stats = nullptr) const override;
+
+    /// Override to create RaBitQ-specific handlers
+    SIMDResultHandlerToFloat* make_knn_handler(
+            bool is_max,
+            int /* impl */,
+            idx_t n,
+            idx_t k,
+            float* distances,
+            idx_t* labels,
+            const IDSelector* sel,
+            const FastScanDistancePostProcessing& context,
+            const float* normalizers = nullptr) const override;
+
+    /** SIMD result handler for IndexIVFRaBitQFastScan that applies
+     * RaBitQ-specific distance corrections during batch processing.
+     *
+     * This handler processes batches of 32 distance computations from SIMD
+     * kernels, applies RaBitQ distance formula adjustments (factors and
+     * normalizers), and immediately updates result heaps. This eliminates the
+     * need for post-processing and provides significant performance benefits.
+     *
+     * Key optimizations:
+     * - Direct heap integration with no intermediate result storage
+     * - Batch-level computation of normalizers and query factors
+     * - 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
+     */
+    template <class C>
+    struct IVFRaBitQHeapHandler
+            : simd_result_handlers::ResultHandlerCompare<C, true> {
+        const IndexIVFRaBitQFastScan* index;
+        float* heap_distances; // [nq * k]
+        int64_t* heap_labels;  // [nq * k]
+        const size_t nq, k;
+        size_t current_list_no = 0;
+        std::vector<int>
+                probe_indices; // probe index for each query in current batch
+        const FastScanDistancePostProcessing*
+                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<
+                C::is_max,
+                CMax<float, int64_t>,
+                CMin<float, int64_t>>::type;
+
+        IVFRaBitQHeapHandler(
+                const IndexIVFRaBitQFastScan* idx,
+                size_t nq_val,
+                size_t k_val,
+                float* distances,
+                int64_t* labels,
+                const FastScanDistancePostProcessing* ctx = nullptr,
+                bool multibit = false);
+
+        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)
+                override;
+
+        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;
+    };
+};
+
+} // namespace faiss

data/vendor/faiss/faiss/IndexIVFSpectralHash.cpp

@@ -331,7 +331,7 @@ void IndexIVFSpectralHash::replace_vt(VectorTransform* vt_in, bool own) {
 /*
     Check that the encoder is a single vector transform followed by a LSH
     that just does thresholding.
-    If this is not the case, the linear transform + threhsolds of the IndexLSH
+    If this is not the case, the linear transform + thresholds of the IndexLSH
     should be merged into the VectorTransform (which is feasible).
 */
 

data/vendor/faiss/faiss/IndexIVFSpectralHash.h

@@ -79,7 +79,7 @@ struct IndexIVFSpectralHash : IndexIVF {
      */
     void replace_vt(VectorTransform* vt, bool own = false);
 
-    /** convenience function to get the VT from an index constucted by an
+    /** convenience function to get the VT from an index constructed by an
      * index_factory (should end in "LSH") */
     void replace_vt(IndexPreTransform* index, bool own = false);
 

data/vendor/faiss/faiss/IndexNNDescent.cpp

@@ -154,7 +154,7 @@ void IndexNNDescent::add(idx_t n, const float* x) {
 
     if (ntotal != 0) {
         fprintf(stderr,
-                "WARNING NNDescent doest not support dynamic insertions,"
+                "WARNING NNDescent does not support dynamic insertions,"
                 "multiple insertions would lead to re-building the index");
     }
 

data/vendor/faiss/faiss/IndexNSG.cpp

@@ -261,7 +261,7 @@ void IndexNSG::check_knn_graph(const idx_t* knn_graph, idx_t n, int K) const {
     }
     FAISS_THROW_IF_NOT_MSG(
             total_count < n / 10,
-            "There are too much invalid entries in the knn graph. "
+            "There are too many invalid entries in the knn graph. "
             "It may be an invalid knn graph.");
 }
 

data/vendor/faiss/faiss/IndexNeuralNetCodec.h

@@ -29,7 +29,7 @@ struct IndexNeuralNetCodec : IndexFlatCodes {
     void sa_encode(idx_t n, const float* x, uint8_t* codes) const override;
     void sa_decode(idx_t n, const uint8_t* codes, float* x) const override;
 
-    ~IndexNeuralNetCodec() {}
+    ~IndexNeuralNetCodec() override {}
 };
 
 struct IndexQINCo : IndexNeuralNetCodec {

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/IndexPQ.h

@@ -164,7 +164,7 @@ struct MultiIndexQuantizer : Index {
 // block size used in MultiIndexQuantizer::search
 FAISS_API extern int multi_index_quantizer_search_bs;
 
-/** MultiIndexQuantizer where the PQ assignmnet is performed by sub-indexes
+/** MultiIndexQuantizer where the PQ assignment is performed by sub-indexes
  */
 struct MultiIndexQuantizer2 : MultiIndexQuantizer {
     /// M Indexes on d / M dimensions

data/vendor/faiss/faiss/IndexPQFastScan.cpp

@@ -9,6 +9,7 @@
 
 #include <memory>
 
+#include <faiss/impl/FastScanDistancePostProcessing.h>
 #include <faiss/impl/pq4_fast_scan.h>
 #include <faiss/utils/utils.h>
 
@@ -53,8 +54,11 @@ void IndexPQFastScan::compute_codes(uint8_t* codes, idx_t n, const float* x)
     pq.compute_codes(x, codes, n);
 }
 
-void IndexPQFastScan::compute_float_LUT(float* lut, idx_t n, const float* x)
-        const {
+void IndexPQFastScan::compute_float_LUT(
+        float* lut,
+        idx_t n,
+        const float* x,
+        const FastScanDistancePostProcessing&) const {
     if (metric_type == METRIC_L2) {
         pq.compute_distance_tables(n, x, lut);
     } else {

data/vendor/faiss/faiss/IndexPQFastScan.h

@@ -45,7 +45,11 @@ struct IndexPQFastScan : IndexFastScan {
 
     void compute_codes(uint8_t* codes, idx_t n, const float* x) const override;
 
-    void compute_float_LUT(float* lut, idx_t n, const float* x) const override;
+    void compute_float_LUT(
+            float* lut,
+            idx_t n,
+            const float* x,
+            const FastScanDistancePostProcessing& context) const override;
 
     void sa_decode(idx_t n, const uint8_t* bytes, float* x) const override;
 };

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;
@@ -55,16 +60,17 @@ void IndexRaBitQ::sa_decode(idx_t n, const uint8_t* bytes, float* x) const {
 
 FlatCodesDistanceComputer* IndexRaBitQ::get_FlatCodesDistanceComputer() const {
     FlatCodesDistanceComputer* dc =
-            rabitq.get_distance_computer(qb, center.data());
+            rabitq.get_distance_computer(qb, center.data(), centered);
     dc->code_size = rabitq.code_size;
     dc->codes = codes.data();
     return dc;
 }
 
 FlatCodesDistanceComputer* IndexRaBitQ::get_quantized_distance_computer(
-        const uint8_t qb) const {
+        const uint8_t qb,
+        bool centered) const {
     FlatCodesDistanceComputer* dc =
-            rabitq.get_distance_computer(qb, center.data());
+            rabitq.get_distance_computer(qb, center.data(), centered);
     dc->code_size = rabitq.code_size;
     dc->codes = codes.data();
     return dc;
@@ -76,6 +82,8 @@ struct Run_search_with_dc_res {
     using T = void;
 
     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) {
@@ -83,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(
-                    index->get_quantized_distance_computer(qb));
+            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();
             }
         }
@@ -114,15 +187,25 @@ void IndexRaBitQ::search(
         float* distances,
         idx_t* labels,
         const SearchParameters* params_in) const {
+    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)) {
         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,12 +8,14 @@
 #pragma once
 
 #include <faiss/IndexFlatCodes.h>
+#include <faiss/impl/RaBitQStats.h>
 #include <faiss/impl/RaBitQuantizer.h>
 
 namespace faiss {
 
 struct RaBitQSearchParameters : SearchParameters {
     uint8_t qb = 0;
+    bool centered = false;
 };
 
 struct IndexRaBitQ : IndexFlatCodes {
@@ -26,9 +28,15 @@ struct IndexRaBitQ : IndexFlatCodes {
     // use '0' to disable quantization and use raw fp32 values.
     uint8_t qb = 0;
 
+    // quantize the query with a zero-centered scalar quantizer.
+    bool centered = false;
+
     IndexRaBitQ();
 
-    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;
 
@@ -42,7 +50,8 @@ struct IndexRaBitQ : IndexFlatCodes {
     // returns a quantized-to-qb bits DC if qb_in > 0
     // returns a default fp32-based DC if qb_in == 0
     FlatCodesDistanceComputer* get_quantized_distance_computer(
-            const uint8_t qb_in) const;
+            const uint8_t qb_in,
+            bool centered) const;
 
     // Don't rely on sa_decode(), bcz it is good for IP, but not for L2.
     //   As a result, use get_FlatCodesDistanceComputer() for the search.