faiss 0.5.3 → 0.6.1

data/vendor/faiss/faiss/IndexFastScan.cpp

@@ -13,40 +13,34 @@
 
 #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/RaBitQUtils.h>
-#include <faiss/impl/pq4_fast_scan.h>
-#include <faiss/impl/simd_result_handlers.h>
+#include <faiss/impl/fast_scan/FastScanDistancePostProcessing.h>
+#include <faiss/impl/fast_scan/fast_scan.h>
+#include <faiss/impl/fast_scan/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>
 #include <faiss/utils/quantize_lut.h>
+#include <faiss/utils/utils.h>
 
 namespace faiss {
 
-using namespace simd_result_handlers;
-
 inline size_t roundup(size_t a, size_t b) {
     return (a + b - 1) / b * b;
 }
 
 void IndexFastScan::init_fastscan(
-        int d,
+        int d_,
         size_t M_init,
         size_t nbits_init,
         MetricType metric,
-        int bbs) {
+        int bbs_) {
     FAISS_THROW_IF_NOT(nbits_init == 4);
-    FAISS_THROW_IF_NOT(bbs % 32 == 0);
-    this->d = d;
+    FAISS_THROW_IF_NOT(bbs_ % 32 == 0);
+    this->d = d_;
     this->M = M_init;
     this->nbits = nbits_init;
     this->metric_type = metric;
-    this->bbs = bbs;
+    this->bbs = bbs_;
     ksub = (1 << nbits_init);
 
     code_size = (M_init * nbits_init + 7) / 8;
@@ -84,7 +78,8 @@ void IndexFastScan::add(idx_t n, const float* x) {
     compute_codes(tmp_codes.get(), n, x);
 
     ntotal2 = roundup(ntotal + n, bbs);
-    size_t new_size = ntotal2 * M2 / 2; // assume nbits = 4
+    size_t n_blocks = ntotal2 / bbs;
+    size_t new_size = n_blocks * get_block_stride();
     size_t old_size = codes.size();
     if (new_size > old_size) {
         codes.resize(new_size);
@@ -92,7 +87,15 @@ void IndexFastScan::add(idx_t n, const float* x) {
     }
 
     pq4_pack_codes_range(
-            tmp_codes.get(), M, ntotal, ntotal + n, bbs, M2, codes.get());
+            tmp_codes.get(),
+            M,
+            ntotal,
+            ntotal + n,
+            bbs,
+            M2,
+            codes.get(),
+            0,
+            get_block_stride());
 
     ntotal += n;
 }
@@ -101,17 +104,25 @@ CodePacker* IndexFastScan::get_CodePacker() const {
     return new CodePackerPQ4(M, bbs);
 }
 
+size_t IndexFastScan::get_block_stride() const {
+    std::unique_ptr<CodePacker> packer(get_CodePacker());
+    FAISS_THROW_IF_NOT_MSG(
+            packer->nvec == static_cast<size_t>(bbs),
+            "CodePacker must pack bbs vectors per block for fast-scan");
+    return packer->block_size;
+}
+
 size_t IndexFastScan::remove_ids(const IDSelector& sel) {
     idx_t j = 0;
     std::vector<uint8_t> buffer(code_size);
-    CodePackerPQ4 packer(M, bbs);
+    std::unique_ptr<CodePacker> packer(get_CodePacker());
     for (idx_t i = 0; i < ntotal; i++) {
         if (sel.is_member(i)) {
             // should be removed
         } else {
             if (i > j) {
-                packer.unpack_1(codes.data(), i, buffer.data());
-                packer.pack_1(buffer.data(), j, codes.data());
+                packer->unpack_1(codes.data(), i, buffer.data());
+                packer->pack_1(buffer.data(), j, codes.data());
             }
             j++;
         }
@@ -120,8 +131,7 @@ size_t IndexFastScan::remove_ids(const IDSelector& sel) {
     if (nremove > 0) {
         ntotal = j;
         ntotal2 = roundup(ntotal, bbs);
-        size_t new_size = ntotal2 * M2 / 2;
-        codes.resize(new_size);
+        codes.resize(ntotal2 / bbs * get_block_stride());
     }
     return nremove;
 }
@@ -139,17 +149,18 @@ void IndexFastScan::check_compatible_for_merge(const Index& otherIndex) const {
             "can only merge indexes of the same type");
 }
 
-void IndexFastScan::merge_from(Index& otherIndex, idx_t add_id) {
+void IndexFastScan::merge_from(Index& otherIndex, idx_t /*add_id*/) {
     check_compatible_for_merge(otherIndex);
     IndexFastScan* other = static_cast<IndexFastScan*>(&otherIndex);
     ntotal2 = roundup(ntotal + other->ntotal, bbs);
-    codes.resize(ntotal2 * M2 / 2);
+    codes.resize(ntotal2 / bbs * get_block_stride());
     std::vector<uint8_t> buffer(code_size);
-    CodePackerPQ4 packer(M, bbs);
+    std::unique_ptr<CodePacker> packer(get_CodePacker());
+    std::unique_ptr<CodePacker> other_packer(other->get_CodePacker());
 
-    for (int i = 0; i < other->ntotal; i++) {
-        packer.unpack_1(other->codes.data(), i, buffer.data());
-        packer.pack_1(buffer.data(), ntotal + i, codes.data());
+    for (idx_t i = 0; i < other->ntotal; i++) {
+        other_packer->unpack_1(other->codes.data(), i, buffer.data());
+        packer->pack_1(buffer.data(), ntotal + i, codes.data());
     }
     ntotal += other->ntotal;
     other->reset();
@@ -173,18 +184,18 @@ void estimators_from_tables_generic(
         BitstringReader bsr(codes + j * index.code_size, index.code_size);
         accu_t dis = 0;
         const dis_t* dt = dis_table;
-        int nscale = context.norm_scaler ? context.norm_scaler->nscale : 0;
+        int nscale = context.pq2x4_scale ? 2 : 0;
 
         for (size_t m = 0; m < index.M - nscale; m++) {
-            uint64_t c = bsr.read(index.nbits);
+            uint64_t c = bsr.read(static_cast<int>(index.nbits));
             dis += dt[c];
             dt += index.ksub;
         }
 
-        if (nscale && context.norm_scaler) {
+        if (nscale) {
             for (size_t m = 0; m < nscale; m++) {
-                uint64_t c = bsr.read(index.nbits);
-                dis += context.norm_scaler->scale_one(dt[c]);
+                uint64_t c = bsr.read(static_cast<int>(index.nbits));
+                dis += dt[c] * context.pq2x4_scale;
                 dt += index.ksub;
             }
         }
@@ -198,43 +209,18 @@ void estimators_from_tables_generic(
 
 } // anonymous namespace
 
-// Default implementation of make_knn_handler with centralized fallback logic
-SIMDResultHandlerToFloat* IndexFastScan::make_knn_handler(
+std::unique_ptr<FastScanCodeScanner> IndexFastScan::make_knn_scanner(
         bool is_max,
-        int impl,
         idx_t n,
         idx_t k,
-        size_t ntotal,
+        size_t ntotal_,
         float* distances,
         idx_t* labels,
         const IDSelector* sel,
+        int impl,
         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);
-        }
-    }
+    return make_fast_scan_knn_scanner(
+            is_max, impl, n, ntotal_, k, distances, labels, sel);
 }
 
 using namespace quantize_lut;
@@ -249,7 +235,7 @@ void IndexFastScan::compute_quantized_LUT(
     std::unique_ptr<float[]> dis_tables(new float[n * dim12]);
     compute_float_LUT(dis_tables.get(), n, x, context);
 
-    for (uint64_t i = 0; i < n; i++) {
+    for (idx_t i = 0; i < n; i++) {
         round_uint8_per_column(
                 dis_tables.get() + i * dim12,
                 M,
@@ -258,11 +244,11 @@ void IndexFastScan::compute_quantized_LUT(
                 &normalizers[2 * i + 1]);
     }
 
-    for (uint64_t i = 0; i < n; i++) {
+    for (idx_t i = 0; i < n; i++) {
         const float* t_in = dis_tables.get() + i * dim12;
         uint8_t* t_out = lut + i * M2 * ksub;
 
-        for (int j = 0; j < dim12; j++) {
+        for (size_t j = 0; j < dim12; j++) {
             t_out[j] = int(t_in[j]);
         }
         memset(t_out + dim12, 0, (M2 - M) * ksub);
@@ -406,7 +392,7 @@ void IndexFastScan::search_implem_234(
     if (implem == 2) {
         // default float
     } else if (implem == 3 || implem == 4) {
-        for (uint64_t i = 0; i < n; i++) {
+        for (idx_t i = 0; i < n; i++) {
             round_uint8_per_column(
                     dis_tables.get() + i * dim12,
                     M,
@@ -439,7 +425,7 @@ void IndexFastScan::search_implem_234(
             float a = normalizers[2 * i];
             float b = normalizers[2 * i + 1];
 
-            for (int j = 0; j < k; j++) {
+            for (idx_t j = 0; j < k; j++) {
                 heap_dis[j] = heap_dis[j] / a + b;
             }
         }
@@ -455,7 +441,6 @@ void IndexFastScan::search_implem_12(
         idx_t* labels,
         int impl,
         const FastScanDistancePostProcessing& context) const {
-    using RH = ResultHandlerCompare<C, false>;
     FAISS_THROW_IF_NOT(bbs == 32);
 
     // handle qbs2 blocking by recursive call
@@ -496,45 +481,37 @@ void IndexFastScan::search_implem_12(
     // block sizes are encoded in qbs, 4 bits at a time
 
     // caution: we override an object field
-    int qbs = this->qbs;
+    int qbs_ = this->qbs;
 
-    if (n != pq4_qbs_to_nq(qbs)) {
-        qbs = pq4_preferred_qbs(n);
+    if (n != pq4_qbs_to_nq(qbs_)) {
+        qbs_ = pq4_preferred_qbs(static_cast<int>(n));
     }
 
-    int LUT_nq =
-            pq4_pack_LUT_qbs(qbs, M2, quantized_dis_tables.get(), LUT.get());
+    int LUT_nq = pq4_pack_LUT_qbs(
+            qbs_, static_cast<int>(M2), quantized_dis_tables.get(), LUT.get());
     FAISS_THROW_IF_NOT(LUT_nq == n);
 
-    std::unique_ptr<RH> handler(
-            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();
+    auto scanner = make_knn_scanner(
+            C::is_max, n, k, ntotal, distances, labels, nullptr, impl, context);
+    auto* rh = scanner->handler();
+    rh->normalizers = normalizers.get();
+    // Note: skip & 2 previously set handler->disable (run kernel,
+    // discard results). Through the scanner path, skip & 2 now skips
+    // the kernel entirely (same as skip & 4), since disable is not
+    // accessible through the SIMDResultHandlerToFloat* interface.
+    if (!(skip & (2 | 4))) {
+        scanner->accumulate_loop_qbs(
+                qbs_,
 
-    if (skip & 4) {
-        // pass
-    } else {
-        pq4_accumulate_loop_qbs(
-                qbs,
                 ntotal2,
-                M2,
+                static_cast<int>(M2),
                 codes.get(),
                 LUT.get(),
-                *handler.get(),
-                context.norm_scaler);
+                context.pq2x4_scale,
+                get_block_stride());
     }
     if (!(skip & 8)) {
-        handler->end();
+        rh->end();
     }
 }
 
@@ -549,10 +526,16 @@ void IndexFastScan::search_implem_14(
         idx_t* labels,
         int impl,
         const FastScanDistancePostProcessing& context) const {
-    using RH = ResultHandlerCompare<C, false>;
     FAISS_THROW_IF_NOT(bbs % 32 == 0);
 
-    int qbs2 = qbs == 0 ? 4 : qbs;
+    // The accumulate loop dispatch table only instantiates certain
+    // (nq, BB) pairs where BB = bbs/32.  Cap the query batch size to
+    // the maximum nq instantiated for the current BB so the caller
+    // doesn't have to know about internal template constraints.
+    //   BB=1 → nq up to 4,  BB=2 → nq up to 2,  BB>=3 → nq=1
+    int BB = bbs / 32;
+    int max_qbs = BB <= 1 ? 4 : BB == 2 ? 2 : 1;
+    int qbs2 = std::min(qbs == 0 ? 4 : qbs, max_qbs);
 
     // handle qbs2 blocking by recursive call
     if (n > qbs2) {
@@ -587,37 +570,33 @@ void IndexFastScan::search_implem_14(
     }
 
     AlignedTable<uint8_t> LUT(n * dim12);
-    pq4_pack_LUT(n, M2, quantized_dis_tables.get(), LUT.get());
-
-    std::unique_ptr<RH> handler(
-            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();
-
-    if (skip & 4) {
-        // pass
-    } else {
-        pq4_accumulate_loop(
-                n,
+    pq4_pack_LUT(
+            static_cast<int>(n),
+            static_cast<int>(M2),
+            quantized_dis_tables.get(),
+            LUT.get());
+
+    auto scanner = make_knn_scanner(
+            C::is_max, n, k, ntotal, distances, labels, nullptr, impl, context);
+    auto* rh = scanner->handler();
+    rh->normalizers = normalizers.get();
+    // Note: skip & 2 previously set handler->disable (run kernel,
+    // discard results). Through the scanner path, skip & 2 now skips
+    // the kernel entirely (same as skip & 4), since disable is not
+    // accessible through the SIMDResultHandlerToFloat* interface.
+    if (!(skip & (2 | 4))) {
+        scanner->accumulate_loop(
+                static_cast<int>(n),
                 ntotal2,
                 bbs,
-                M2,
+                static_cast<int>(M2),
                 codes.get(),
                 LUT.get(),
-                *handler.get(),
-                context.norm_scaler);
+                context.pq2x4_scale,
+                get_block_stride());
     }
     if (!(skip & 8)) {
-        handler->end();
+        rh->end();
     }
 }
 
@@ -639,11 +618,8 @@ template void IndexFastScan::search_dispatch_implem<false>(
 
 void IndexFastScan::reconstruct(idx_t key, float* recons) const {
     std::vector<uint8_t> code(code_size, 0);
-    BitstringWriter bsw(code.data(), code_size);
-    for (size_t m = 0; m < M; m++) {
-        uint8_t c = pq4_get_packed_element(codes.data(), bbs, M2, key, m);
-        bsw.write(c, nbits);
-    }
+    std::unique_ptr<CodePacker> packer(get_CodePacker());
+    packer->unpack_1(codes.data(), key, code.data());
     sa_decode(1, code.data(), recons);
 }
 

data/vendor/faiss/faiss/IndexFastScan.h

@@ -7,16 +7,17 @@
 
 #pragma once
 
+#include <memory>
+
 #include <faiss/Index.h>
-#include <faiss/impl/FastScanDistancePostProcessing.h>
+#include <faiss/impl/fast_scan/FastScanDistancePostProcessing.h>
+#include <faiss/impl/fast_scan/fast_scan.h>
 #include <faiss/utils/AlignedTable.h>
 
 namespace faiss {
 
 struct CodePacker;
-struct NormTableScaler;
 struct IDSelector;
-struct SIMDResultHandlerToFloat;
 
 /** Fast scan version of IndexPQ and IndexAQ. Works for 4-bit PQ and AQ for now.
  *
@@ -43,8 +44,8 @@ struct IndexFastScan : Index {
 
     // vector quantizer
     size_t M;
-    size_t nbits;
-    size_t ksub;
+    size_t nbits = 0;
+    size_t ksub = 0;
     size_t code_size;
 
     // packed version of the codes
@@ -122,33 +123,22 @@ struct IndexFastScan : Index {
             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.
+    /** Create a SIMD-dispatched scanner for knn search.
      *
-     * @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 context      processing context for distance post-processing
-     * @return             pointer to created handler (never returns nullptr)
+     * Returns a FastScanCodeScanner that bundles handler + accumulation
+     * kernel behind the SIMD dispatch boundary.
+     * The scanner's accumulate methods dispatch to the optimal SIMD level.
      */
-    virtual SIMDResultHandlerToFloat* make_knn_handler(
+    virtual std::unique_ptr<FastScanCodeScanner> make_knn_scanner(
             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;
+            int impl = 0,
+            const FastScanDistancePostProcessing& context = {}) const;
 
     // called by search function
     void compute_quantized_LUT(
@@ -214,7 +204,16 @@ struct IndexFastScan : Index {
      *
      * @return  pointer to the code packer
      */
-    CodePacker* get_CodePacker() const;
+    virtual CodePacker* get_CodePacker() const;
+
+    /** Get stride in bytes between consecutive SIMD blocks.
+     *
+     * Derived from get_CodePacker()->block_size so that there is a
+     * single source of truth for the block layout.
+     *
+     * @return stride in bytes
+     */
+    size_t get_block_stride() const;
 
     /** Merge another index into this one
      *
@@ -237,6 +236,18 @@ struct IndexFastScan : Index {
     void sa_encode(idx_t n, const float* x, uint8_t* bytes) const override {
         compute_codes(bytes, n, x);
     }
+
+    /** Get the size of the code portion packed by pq4_pack_codes.
+     *
+     * Returns the number of bytes per vector that are interleaved into
+     * SIMD blocks by pq4_pack_codes, excluding any embedded metadata
+     * (e.g., RaBitQ factors). The meaning of these bytes depends on the
+     * quantizer: for PQ/AQ they are 4-bit sub-quantizer nibbles, for
+     * RaBitQ they are 1-bit-per-dimension sign bits packed into nibbles.
+     *
+     * Must be implemented by all derived classes.
+     */
+    virtual size_t fast_scan_code_size() const = 0;
 };
 
 struct FastScanStats {