faiss 0.4.2 → 0.5.0

data/vendor/faiss/faiss/IndexIVFRaBitQFastScan.h

@@ -0,0 +1,216 @@
+/*
+ * 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/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::FactorsData;
+using rabitq_utils::QueryFactorsData;
+
+/** 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;
+
+    /// Extracted factors storage for batch processing
+    /// Size: ntotal, stores factors separately from packed codes
+    std::vector<FactorsData> factors_storage;
+
+    // Constructors
+
+    IndexIVFRaBitQFastScan();
+
+    IndexIVFRaBitQFastScan(
+            Index* quantizer,
+            size_t d,
+            size_t nlist,
+            MetricType metric = METRIC_L2,
+            int bbs = 32,
+            bool own_invlists = true);
+
+    /// 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;
+
+   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 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
+    void decode_fastscan_to_residual(
+            const uint8_t* fastscan_code,
+            float* residual) 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
+     *
+     * @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
+
+        // 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);
+
+        void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) final;
+
+        /// 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;
+    };
+};
+
+} // 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.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/IndexRaBitQ.cpp

@@ -55,16 +55,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 +77,7 @@ struct Run_search_with_dc_res {
     using T = void;
 
     uint8_t qb = 0;
+    bool centered = false;
 
     template <class BlockResultHandler>
     void f(BlockResultHandler& res, const IndexRaBitQ* index, const float* xq) {
@@ -87,7 +89,7 @@ struct Run_search_with_dc_res {
 #pragma omp parallel // if (res.nq > 100)
         {
             std::unique_ptr<FlatCodesDistanceComputer> dc(
-                    index->get_quantized_distance_computer(qb));
+                    index->get_quantized_distance_computer(qb, centered));
             SingleResultHandler resi(res);
 #pragma omp for
             for (int64_t q = 0; q < res.nq; q++) {
@@ -114,14 +116,15 @@ void IndexRaBitQ::search(
         float* distances,
         idx_t* labels,
         const SearchParameters* params_in) const {
-    uint8_t used_qb = qb;
-    if (auto params = dynamic_cast<const RaBitQSearchParameters*>(params_in)) {
-        used_qb = params->qb;
-    }
-
     const IDSelector* sel = (params_in != nullptr) ? params_in->sel : nullptr;
     Run_search_with_dc_res r;
-    r.qb = used_qb;
+    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;
+    }
 
     dispatch_knn_ResultHandler(
             n, distances, labels, k, metric_type, sel, r, this, x);

data/vendor/faiss/faiss/IndexRaBitQ.h

@@ -14,6 +14,7 @@ namespace faiss {
 
 struct RaBitQSearchParameters : SearchParameters {
     uint8_t qb = 0;
+    bool centered = false;
 };
 
 struct IndexRaBitQ : IndexFlatCodes {
@@ -26,9 +27,12 @@ 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);
 
     void train(idx_t n, const float* x) override;
 
@@ -42,7 +46,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.