faiss 0.4.2 → 0.5.0

data/vendor/faiss/faiss/IndexFastScan.cpp

@@ -7,17 +7,20 @@
 
 #include <faiss/IndexFastScan.h>
 
-#include <cassert>
-#include <climits>
-#include <memory>
-
 #include <omp.h>
+#include <cstring>
+#include <memory>
 
+#include <faiss/impl/CodePacker.h>
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/FastScanDistancePostProcessing.h>
 #include <faiss/impl/IDSelector.h>
 #include <faiss/impl/LookupTableScaler.h>
-#include <faiss/impl/ResultHandler.h>
+#include <faiss/impl/RaBitQUtils.h>
+#include <faiss/impl/pq4_fast_scan.h>
+#include <faiss/impl/simd_result_handlers.h>
 #include <faiss/utils/hamming.h>
+#include <faiss/utils/utils.h>
 
 #include <faiss/impl/pq4_fast_scan.h>
 #include <faiss/impl/simd_result_handlers.h>
@@ -163,14 +166,14 @@ void estimators_from_tables_generic(
         size_t k,
         typename C::T* heap_dis,
         int64_t* heap_ids,
-        const NormTableScaler* scaler) {
+        const FastScanDistancePostProcessing& context) {
     using accu_t = typename C::T;
 
     for (size_t j = 0; j < ncodes; ++j) {
         BitstringReader bsr(codes + j * index.code_size, index.code_size);
         accu_t dis = 0;
         const dis_t* dt = dis_table;
-        int nscale = scaler ? scaler->nscale : 0;
+        int nscale = context.norm_scaler ? context.norm_scaler->nscale : 0;
 
         for (size_t m = 0; m < index.M - nscale; m++) {
             uint64_t c = bsr.read(index.nbits);
@@ -178,10 +181,10 @@ void estimators_from_tables_generic(
             dt += index.ksub;
         }
 
-        if (nscale) {
+        if (nscale && context.norm_scaler) {
             for (size_t m = 0; m < nscale; m++) {
                 uint64_t c = bsr.read(index.nbits);
-                dis += scaler->scale_one(dt[c]);
+                dis += context.norm_scaler->scale_one(dt[c]);
                 dt += index.ksub;
             }
         }
@@ -193,40 +196,58 @@ void estimators_from_tables_generic(
     }
 }
 
-template <class C>
-ResultHandlerCompare<C, false>* make_knn_handler(
+} // anonymous namespace
+
+// Default implementation of make_knn_handler with centralized fallback logic
+void* IndexFastScan::make_knn_handler(
+        bool is_max,
         int impl,
         idx_t n,
         idx_t k,
         size_t ntotal,
         float* distances,
         idx_t* labels,
-        const IDSelector* sel = nullptr) {
-    using HeapHC = HeapHandler<C, false>;
-    using ReservoirHC = ReservoirHandler<C, false>;
-    using SingleResultHC = SingleResultHandler<C, false>;
-
-    if (k == 1) {
-        return new SingleResultHC(n, ntotal, distances, labels, sel);
-    } else if (impl % 2 == 0) {
-        return new HeapHC(n, ntotal, k, distances, labels, sel);
-    } else /* if (impl % 2 == 1) */ {
-        return new ReservoirHC(n, ntotal, k, 2 * k, distances, labels, sel);
+        const IDSelector* sel,
+        const FastScanDistancePostProcessing&) const {
+    // Create default handlers based on k and impl
+    if (is_max) {
+        using HeapHC = HeapHandler<CMax<uint16_t, int>, false>;
+        using ReservoirHC = ReservoirHandler<CMax<uint16_t, int>, false>;
+        using SingleResultHC = SingleResultHandler<CMax<uint16_t, int>, false>;
+
+        if (k == 1) {
+            return new SingleResultHC(n, ntotal, distances, labels, sel);
+        } else if (impl % 2 == 0) {
+            return new HeapHC(n, ntotal, k, distances, labels, sel);
+        } else {
+            return new ReservoirHC(n, ntotal, k, 2 * k, distances, labels, sel);
+        }
+    } else {
+        using HeapHC = HeapHandler<CMin<uint16_t, int>, false>;
+        using ReservoirHC = ReservoirHandler<CMin<uint16_t, int>, false>;
+        using SingleResultHC = SingleResultHandler<CMin<uint16_t, int>, false>;
+
+        if (k == 1) {
+            return new SingleResultHC(n, ntotal, distances, labels, sel);
+        } else if (impl % 2 == 0) {
+            return new HeapHC(n, ntotal, k, distances, labels, sel);
+        } else {
+            return new ReservoirHC(n, ntotal, k, 2 * k, distances, labels, sel);
+        }
     }
 }
 
-} // anonymous namespace
-
 using namespace quantize_lut;
 
 void IndexFastScan::compute_quantized_LUT(
         idx_t n,
         const float* x,
         uint8_t* lut,
-        float* normalizers) const {
+        float* normalizers,
+        const FastScanDistancePostProcessing& context) const {
     size_t dim12 = ksub * M;
     std::unique_ptr<float[]> dis_tables(new float[n * dim12]);
-    compute_float_LUT(dis_tables.get(), n, x);
+    compute_float_LUT(dis_tables.get(), n, x, context);
 
     for (uint64_t i = 0; i < n; i++) {
         round_uint8_per_column(
@@ -263,10 +284,12 @@ void IndexFastScan::search(
             !params, "search params not supported for this index");
     FAISS_THROW_IF_NOT(k > 0);
 
+    FastScanDistancePostProcessing empty_context{};
     if (metric_type == METRIC_L2) {
-        search_dispatch_implem<true>(n, x, k, distances, labels, nullptr);
+        search_dispatch_implem<true>(n, x, k, distances, labels, empty_context);
     } else {
-        search_dispatch_implem<false>(n, x, k, distances, labels, nullptr);
+        search_dispatch_implem<false>(
+                n, x, k, distances, labels, empty_context);
     }
 }
 
@@ -277,7 +300,7 @@ void IndexFastScan::search_dispatch_implem(
         idx_t k,
         float* distances,
         idx_t* labels,
-        const NormTableScaler* scaler) const {
+        const FastScanDistancePostProcessing& context) const {
     using Cfloat = typename std::conditional<
             is_max,
             CMax<float, int64_t>,
@@ -308,15 +331,20 @@ void IndexFastScan::search_dispatch_implem(
         FAISS_THROW_MSG("not implemented");
     } else if (implem == 2 || implem == 3 || implem == 4) {
         FAISS_THROW_IF_NOT(orig_codes != nullptr);
-        search_implem_234<Cfloat>(n, x, k, distances, labels, scaler);
+        search_implem_234<Cfloat>(n, x, k, distances, labels, context);
     } else if (impl >= 12 && impl <= 15) {
         FAISS_THROW_IF_NOT(ntotal < INT_MAX);
         int nt = std::min(omp_get_max_threads(), int(n));
+        // Fall back to single-threaded implementations when parallelization not
+        // beneficial:
+        // - Single-core system (omp_get_max_threads() = 1)
+        // - Single query (n = 1)
+        // - OpenMP disabled (omp_get_max_threads() = 1)
         if (nt < 2) {
             if (impl == 12 || impl == 13) {
-                search_implem_12<C>(n, x, k, distances, labels, impl, scaler);
+                search_implem_12<C>(n, x, k, distances, labels, impl, context);
             } else {
-                search_implem_14<C>(n, x, k, distances, labels, impl, scaler);
+                search_implem_14<C>(n, x, k, distances, labels, impl, context);
             }
         } else {
             // explicitly slice over threads
@@ -324,14 +352,33 @@ void IndexFastScan::search_dispatch_implem(
             for (int slice = 0; slice < nt; slice++) {
                 idx_t i0 = n * slice / nt;
                 idx_t i1 = n * (slice + 1) / nt;
+
+                // Create per-thread context with adjusted query_factors pointer
+                FastScanDistancePostProcessing thread_context = context;
+                if (thread_context.query_factors != nullptr) {
+                    thread_context.query_factors += i0;
+                }
+
                 float* dis_i = distances + i0 * k;
                 idx_t* lab_i = labels + i0 * k;
                 if (impl == 12 || impl == 13) {
                     search_implem_12<C>(
-                            i1 - i0, x + i0 * d, k, dis_i, lab_i, impl, scaler);
+                            i1 - i0,
+                            x + i0 * d,
+                            k,
+                            dis_i,
+                            lab_i,
+                            impl,
+                            thread_context);
                 } else {
                     search_implem_14<C>(
-                            i1 - i0, x + i0 * d, k, dis_i, lab_i, impl, scaler);
+                            i1 - i0,
+                            x + i0 * d,
+                            k,
+                            dis_i,
+                            lab_i,
+                            impl,
+                            thread_context);
                 }
             }
         }
@@ -347,12 +394,12 @@ void IndexFastScan::search_implem_234(
         idx_t k,
         float* distances,
         idx_t* labels,
-        const NormTableScaler* scaler) const {
+        const FastScanDistancePostProcessing& context) const {
     FAISS_THROW_IF_NOT(implem == 2 || implem == 3 || implem == 4);
 
     const size_t dim12 = ksub * M;
     std::unique_ptr<float[]> dis_tables(new float[n * dim12]);
-    compute_float_LUT(dis_tables.get(), n, x);
+    compute_float_LUT(dis_tables.get(), n, x, context);
 
     std::vector<float> normalizers(n * 2);
 
@@ -384,7 +431,7 @@ void IndexFastScan::search_implem_234(
                 k,
                 heap_dis,
                 heap_ids,
-                scaler);
+                context);
 
         heap_reorder<Cfloat>(k, heap_dis, heap_ids);
 
@@ -407,7 +454,7 @@ void IndexFastScan::search_implem_12(
         float* distances,
         idx_t* labels,
         int impl,
-        const NormTableScaler* scaler) const {
+        const FastScanDistancePostProcessing& context) const {
     using RH = ResultHandlerCompare<C, false>;
     FAISS_THROW_IF_NOT(bbs == 32);
 
@@ -416,6 +463,11 @@ void IndexFastScan::search_implem_12(
     if (n > qbs2) {
         for (int64_t i0 = 0; i0 < n; i0 += qbs2) {
             int64_t i1 = std::min(i0 + qbs2, n);
+            // Create sub-context with adjusted query_factors pointer
+            FastScanDistancePostProcessing sub_context = context;
+            if (sub_context.query_factors != nullptr) {
+                sub_context.query_factors += i0;
+            }
             search_implem_12<C>(
                     i1 - i0,
                     x + d * i0,
@@ -423,7 +475,7 @@ void IndexFastScan::search_implem_12(
                     distances + i0 * k,
                     labels + i0 * k,
                     impl,
-                    scaler);
+                    sub_context);
         }
         return;
     }
@@ -436,7 +488,7 @@ void IndexFastScan::search_implem_12(
         quantized_dis_tables.clear();
     } else {
         compute_quantized_LUT(
-                n, x, quantized_dis_tables.get(), normalizers.get());
+                n, x, quantized_dis_tables.get(), normalizers.get(), context);
     }
 
     AlignedTable<uint8_t> LUT(n * dim12);
@@ -455,7 +507,17 @@ void IndexFastScan::search_implem_12(
     FAISS_THROW_IF_NOT(LUT_nq == n);
 
     std::unique_ptr<RH> handler(
-            make_knn_handler<C>(impl, n, k, ntotal, distances, labels));
+            static_cast<RH*>(make_knn_handler(
+                    C::is_max,
+                    impl,
+                    n,
+                    k,
+                    ntotal,
+                    distances,
+                    labels,
+                    nullptr,
+                    context)));
+
     handler->disable = bool(skip & 2);
     handler->normalizers = normalizers.get();
 
@@ -469,7 +531,7 @@ void IndexFastScan::search_implem_12(
                 codes.get(),
                 LUT.get(),
                 *handler.get(),
-                scaler);
+                context.norm_scaler);
     }
     if (!(skip & 8)) {
         handler->end();
@@ -486,7 +548,7 @@ void IndexFastScan::search_implem_14(
         float* distances,
         idx_t* labels,
         int impl,
-        const NormTableScaler* scaler) const {
+        const FastScanDistancePostProcessing& context) const {
     using RH = ResultHandlerCompare<C, false>;
     FAISS_THROW_IF_NOT(bbs % 32 == 0);
 
@@ -496,6 +558,11 @@ void IndexFastScan::search_implem_14(
     if (n > qbs2) {
         for (int64_t i0 = 0; i0 < n; i0 += qbs2) {
             int64_t i1 = std::min(i0 + qbs2, n);
+            // Create sub-context with adjusted query_factors pointer
+            FastScanDistancePostProcessing sub_context = context;
+            if (sub_context.query_factors != nullptr) {
+                sub_context.query_factors += i0;
+            }
             search_implem_14<C>(
                     i1 - i0,
                     x + d * i0,
@@ -503,7 +570,7 @@ void IndexFastScan::search_implem_14(
                     distances + i0 * k,
                     labels + i0 * k,
                     impl,
-                    scaler);
+                    sub_context);
         }
         return;
     }
@@ -516,14 +583,23 @@ void IndexFastScan::search_implem_14(
         quantized_dis_tables.clear();
     } else {
         compute_quantized_LUT(
-                n, x, quantized_dis_tables.get(), normalizers.get());
+                n, x, quantized_dis_tables.get(), normalizers.get(), context);
     }
 
     AlignedTable<uint8_t> LUT(n * dim12);
     pq4_pack_LUT(n, M2, quantized_dis_tables.get(), LUT.get());
 
     std::unique_ptr<RH> handler(
-            make_knn_handler<C>(impl, n, k, ntotal, distances, labels));
+            static_cast<RH*>(make_knn_handler(
+                    C::is_max,
+                    impl,
+                    n,
+                    k,
+                    ntotal,
+                    distances,
+                    labels,
+                    nullptr,
+                    context)));
     handler->disable = bool(skip & 2);
     handler->normalizers = normalizers.get();
 
@@ -538,7 +614,7 @@ void IndexFastScan::search_implem_14(
                 codes.get(),
                 LUT.get(),
                 *handler.get(),
-                scaler);
+                context.norm_scaler);
     }
     if (!(skip & 8)) {
         handler->end();
@@ -551,7 +627,7 @@ template void IndexFastScan::search_dispatch_implem<true>(
         idx_t k,
         float* distances,
         idx_t* labels,
-        const NormTableScaler* scaler) const;
+        const FastScanDistancePostProcessing& context) const;
 
 template void IndexFastScan::search_dispatch_implem<false>(
         idx_t n,
@@ -559,7 +635,7 @@ template void IndexFastScan::search_dispatch_implem<false>(
         idx_t k,
         float* distances,
         idx_t* labels,
-        const NormTableScaler* scaler) const;
+        const FastScanDistancePostProcessing& context) const;
 
 void IndexFastScan::reconstruct(idx_t key, float* recons) const {
     std::vector<uint8_t> code(code_size, 0);

data/vendor/faiss/faiss/IndexFastScan.h

@@ -8,6 +8,7 @@
 #pragma once
 
 #include <faiss/Index.h>
+#include <faiss/impl/FastScanDistancePostProcessing.h>
 #include <faiss/utils/AlignedTable.h>
 
 namespace faiss {
@@ -15,6 +16,13 @@ namespace faiss {
 struct CodePacker;
 struct NormTableScaler;
 
+// Forward declarations for result handlers
+namespace simd_result_handlers {
+template <class C, bool with_id_map>
+struct ResultHandlerCompare;
+}
+struct IDSelector;
+
 /** Fast scan version of IndexPQ and IndexAQ. Works for 4-bit PQ and AQ for now.
  *
  * The codes are not stored sequentially but grouped in blocks of size bbs.
@@ -54,6 +62,14 @@ struct IndexFastScan : Index {
     // (set when initialized by IndexPQ or IndexAQ)
     const uint8_t* orig_codes = nullptr;
 
+    /** Initialize the fast scan index
+     *
+     * @param d         dimensionality of vectors
+     * @param M         number of subquantizers
+     * @param nbits     number of bits per subquantizer
+     * @param metric    distance metric to use
+     * @param bbs       block size for SIMD processing
+     */
     void init_fastscan(
             int d,
             size_t M,
@@ -65,6 +81,15 @@ struct IndexFastScan : Index {
 
     void reset() override;
 
+    /** Search for k nearest neighbors
+     *
+     * @param n          number of query vectors
+     * @param x          query vectors (n * d)
+     * @param k          number of nearest neighbors to find
+     * @param distances  output distances (n * k)
+     * @param labels     output labels/indices (n * k)
+     * @param params     optional search parameters
+     */
     void search(
             idx_t n,
             const float* x,
@@ -73,20 +98,70 @@ struct IndexFastScan : Index {
             idx_t* labels,
             const SearchParameters* params = nullptr) const override;
 
+    /** Add vectors to the index
+     *
+     * @param n  number of vectors to add
+     * @param x  vectors to add (n * d)
+     */
     void add(idx_t n, const float* x) override;
 
+    /** Compute codes for vectors
+     *
+     * @param codes  output codes
+     * @param n      number of vectors to encode
+     * @param x      vectors to encode (n * d)
+     */
     virtual void compute_codes(uint8_t* codes, idx_t n, const float* x)
             const = 0;
 
-    virtual void compute_float_LUT(float* lut, idx_t n, const float* x)
-            const = 0;
+    /** Compute floating-point lookup table for distance computation
+     *
+     * @param lut          output lookup table
+     * @param n            number of query vectors
+     * @param x            query vectors (n * d)
+     * @param context      processing context containing all processors
+     */
+    virtual void compute_float_LUT(
+            float* lut,
+            idx_t n,
+            const float* x,
+            const FastScanDistancePostProcessing& context) const = 0;
+
+    /** Create a KNN handler for this index type
+     *
+     * This method can be overridden by derived classes to provide
+     * specialized handlers (e.g., RaBitQHeapHandler for RaBitQ indexes).
+     * Base implementation creates standard handlers based on k and impl.
+     *
+     * @param is_max       whether to use CMax comparator (true) or CMin (false)
+     * @param impl         implementation number
+     * @param n            number of queries
+     * @param k            number of neighbors to find
+     * @param ntotal       total number of vectors in database
+     * @param distances    output distances array
+     * @param labels       output labels array
+     * @param sel          optional ID selector
+     * @param query_offset query offset for batch processing
+     * @return             pointer to created handler (never returns nullptr)
+     */
+    virtual void* make_knn_handler(
+            bool is_max,
+            int impl,
+            idx_t n,
+            idx_t k,
+            size_t ntotal,
+            float* distances,
+            idx_t* labels,
+            const IDSelector* sel,
+            const FastScanDistancePostProcessing& context) const;
 
     // called by search function
     void compute_quantized_LUT(
             idx_t n,
             const float* x,
             uint8_t* lut,
-            float* normalizers) const;
+            float* normalizers,
+            const FastScanDistancePostProcessing& context) const;
 
     template <bool is_max>
     void search_dispatch_implem(
@@ -95,7 +170,7 @@ struct IndexFastScan : Index {
             idx_t k,
             float* distances,
             idx_t* labels,
-            const NormTableScaler* scaler) const;
+            const FastScanDistancePostProcessing& context) const;
 
     template <class Cfloat>
     void search_implem_234(
@@ -104,7 +179,7 @@ struct IndexFastScan : Index {
             idx_t k,
             float* distances,
             idx_t* labels,
-            const NormTableScaler* scaler) const;
+            const FastScanDistancePostProcessing& context) const;
 
     template <class C>
     void search_implem_12(
@@ -114,7 +189,7 @@ struct IndexFastScan : Index {
             float* distances,
             idx_t* labels,
             int impl,
-            const NormTableScaler* scaler) const;
+            const FastScanDistancePostProcessing& context) const;
 
     template <class C>
     void search_implem_14(
@@ -124,14 +199,39 @@ struct IndexFastScan : Index {
             float* distances,
             idx_t* labels,
             int impl,
-            const NormTableScaler* scaler) const;
+            const FastScanDistancePostProcessing& context) const;
 
+    /** Reconstruct a vector from its code
+     *
+     * @param key     index of vector to reconstruct
+     * @param recons  output reconstructed vector
+     */
     void reconstruct(idx_t key, float* recons) const override;
+
+    /** Remove vectors by ID selector
+     *
+     * @param sel  selector defining which vectors to remove
+     * @return     number of vectors removed
+     */
     size_t remove_ids(const IDSelector& sel) override;
 
+    /** Get the code packer for this index
+     *
+     * @return  pointer to the code packer
+     */
     CodePacker* get_CodePacker() const;
 
+    /** Merge another index into this one
+     *
+     * @param otherIndex  index to merge from
+     * @param add_id      ID offset to add to merged vectors
+     */
     void merge_from(Index& otherIndex, idx_t add_id = 0) override;
+
+    /** Check if another index is compatible for merging
+     *
+     * @param otherIndex  index to check compatibility with
+     */
     void check_compatible_for_merge(const Index& otherIndex) const override;
 
     /// standalone codes interface (but the codes are flattened)

data/vendor/faiss/faiss/IndexFlat.h

@@ -66,7 +66,7 @@ struct IndexFlat : IndexFlatCodes {
 
     FlatCodesDistanceComputer* get_FlatCodesDistanceComputer() const override;
 
-    /* The stanadlone codec interface (just memcopies in this case) */
+    /* The standalone codec interface (just memcopies in this case) */
     void sa_encode(idx_t n, const float* x, uint8_t* bytes) const override;
 
     void sa_decode(idx_t n, const uint8_t* bytes, float* x) const override;

data/vendor/faiss/faiss/IndexHNSW.cpp

@@ -450,7 +450,9 @@ void IndexHNSW::search_level_0(
             vt.advance();
         }
 #pragma omp critical
-        { hnsw_stats.combine(search_stats); }
+        {
+            hnsw_stats.combine(search_stats);
+        }
     }
     if (is_similarity_metric(this->metric_type)) {
 // we need to revert the negated distances

data/vendor/faiss/faiss/IndexHNSW.h

@@ -43,7 +43,7 @@ struct IndexHNSW : Index {
 
     // When set to true, all neighbors in level 0 are filled up
     // to the maximum size allowed (2 * M). This option is used by
-    // IndexHHNSWCagra to create a full base layer graph that is
+    // IndexHNSWCagra to create a full base layer graph that is
     // used when GpuIndexCagra::copyFrom(IndexHNSWCagra*) is invoked.
     bool keep_max_size_level0 = false;