faiss 0.6.0 → 0.6.1

data/vendor/faiss/faiss/impl/{pq4_fast_scan_search_qbs.cpp → fast_scan/kernels_simd512.h}

@@ -5,115 +5,20 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-#include <faiss/impl/pq4_fast_scan.h>
+#pragma once
 
-#include <faiss/impl/FaissAssert.h>
-#include <faiss/impl/LookupTableScaler.h>
-#include <faiss/impl/simd_result_handlers.h>
+#include <faiss/impl/platform_macros.h>
+#include <faiss/impl/simdlib/simdlib_dispatch.h>
 
-namespace faiss {
-
-// declared in simd_result_handlers.h
-bool simd_result_handlers_accept_virtual = true;
-
-using namespace simd_result_handlers;
-
-/************************************************************
- * Accumulation functions
- ************************************************************/
-
-namespace {
-
-/*
- * The computation kernel
- * It accumulates results for NQ queries and 2 * 16 database elements
- * writes results in a ResultHandler
- */
-
-#ifndef __AVX512F__
-
-template <int NQ, class ResultHandler, class Scaler>
-void kernel_accumulate_block(
-        int nsq,
-        const uint8_t* codes,
-        const uint8_t* LUT,
-        ResultHandler& res,
-        const Scaler& scaler) {
-    // dummy alloc to keep the windows compiler happy
-    constexpr int NQA = NQ > 0 ? NQ : 1;
-    // distance accumulators
-    // layout: accu[q][b]: distance accumulator for vectors 8*b..8*b+7
-    simd16uint16 accu[NQA][4];
-
-    for (int q = 0; q < NQ; q++) {
-        for (int b = 0; b < 4; b++) {
-            accu[q][b].clear();
-        }
-    }
+#ifdef __AVX512F__
 
-    // _mm_prefetch(codes + 768, 0);
-    for (int sq = 0; sq < nsq - scaler.nscale; sq += 2) {
-        simd32uint8 c;
-        c.loadu(codes);
-        codes += 32;
-
-        simd32uint8 mask(0xf);
-        // shift op does not exist for int8...
-        simd32uint8 chi = simd32uint8(simd16uint16(c) >> 4) & mask;
-        simd32uint8 clo = c & mask;
-
-        for (int q = 0; q < NQ; q++) {
-            // load LUTs for 2 quantizers
-            simd32uint8 lut(LUT);
-            LUT += 32;
-
-            simd32uint8 res0 = lut.lookup_2_lanes(clo);
-            simd32uint8 res1 = lut.lookup_2_lanes(chi);
-
-            accu[q][0] += simd16uint16(res0);
-            accu[q][1] += simd16uint16(res0) >> 8;
-
-            accu[q][2] += simd16uint16(res1);
-            accu[q][3] += simd16uint16(res1) >> 8;
-        }
-    }
-
-    for (int sq = 0; sq < scaler.nscale; sq += 2) {
-        simd32uint8 c;
-        c.loadu(codes);
-        codes += 32;
-
-        simd32uint8 mask(0xf);
-        // shift op does not exist for int8...
-        simd32uint8 chi = simd32uint8(simd16uint16(c) >> 4) & mask;
-        simd32uint8 clo = c & mask;
-
-        for (int q = 0; q < NQ; q++) {
-            // load LUTs for 2 quantizers
-            simd32uint8 lut(LUT);
-            LUT += 32;
-
-            simd32uint8 res0 = scaler.lookup(lut, clo);
-            accu[q][0] += scaler.scale_lo(res0); // handle vectors 0..7
-            accu[q][1] += scaler.scale_hi(res0); // handle vectors 8..15
-
-            simd32uint8 res1 = scaler.lookup(lut, chi);
-            accu[q][2] += scaler.scale_lo(res1); // handle vectors 16..23
-            accu[q][3] += scaler.scale_hi(res1); //  handle vectors 24..31
-        }
-    }
-
-    for (int q = 0; q < NQ; q++) {
-        accu[q][0] -= accu[q][1] << 8;
-        simd16uint16 dis0 = combine2x2(accu[q][0], accu[q][1]);
-        accu[q][2] -= accu[q][3] << 8;
-        simd16uint16 dis1 = combine2x2(accu[q][2], accu[q][3]);
-        res.handle(q, 0, dis0, dis1);
-    }
-}
+namespace faiss {
 
-#else
+// Explicit SIMD-level aliases for this file (no global bare aliases).
+using simd32uint16 = simd32uint16_tpl<SINGLE_SIMD_LEVEL_512>;
+using simd64uint8 = simd64uint8_tpl<SINGLE_SIMD_LEVEL_512>;
 
+// NQ=1 specialization: processes 512-bit chunks aggressively.
 // a special version for NQ=1.
 // Despite the function being large in the text form, it compiles to a very
 //    compact assembler code.
@@ -125,6 +30,13 @@ void kernel_accumulate_block_avx512_nq1(
         const uint8_t* LUT,
         ResultHandler& res,
         const Scaler& scaler) {
+    // Explicit SIMD levels for DD mode where bare aliases resolve to NONE
+    // (512-bit NONE types don't exist — empty primary templates).
+    using simd32uint16 = simd32uint16_tpl<SIMDLevel::AVX512>;
+    using simd64uint8 = simd64uint8_tpl<SIMDLevel::AVX512>;
+    using simd16uint16 = simd16uint16_tpl<SIMDLevel::AVX2>;
+    using simd32uint8 = simd32uint8_tpl<SIMDLevel::AVX2>;
+
     // NQ is kept in order to match the similarity to baseline function
     constexpr int NQ = 1;
     // distance accumulators. We can accept more for NQ=1
@@ -386,6 +298,12 @@ void kernel_accumulate_block_avx512_nqx(
         const uint8_t* LUT,
         ResultHandler& res,
         const Scaler& scaler) {
+    // Explicit SIMD levels for DD mode (see nq1 variant for explanation).
+    using simd32uint16 = simd32uint16_tpl<SIMDLevel::AVX512>;
+    using simd64uint8 = simd64uint8_tpl<SIMDLevel::AVX512>;
+    using simd16uint16 = simd16uint16_tpl<SIMDLevel::AVX2>;
+    using simd32uint8 = simd32uint8_tpl<SIMDLevel::AVX2>;
+
     // dummy alloc to keep the windows compiler happy
     constexpr int NQA = NQ > 0 ? NQ : 1;
     // distance accumulators
@@ -539,8 +457,9 @@ void kernel_accumulate_block_avx512_nqx(
     }
 }
 
+// Dispatcher: selects NQ=1 vs general case.
 template <int NQ, class ResultHandler, class Scaler>
-void kernel_accumulate_block(
+void pq4_kernel_qbs_512(
         int nsq,
         const uint8_t* codes,
         const uint8_t* LUT,
@@ -555,255 +474,6 @@ void kernel_accumulate_block(
     }
 }
 
-#endif
-
-// handle at most 4 blocks of queries
-template <int QBS, class ResultHandler, class Scaler>
-void accumulate_q_4step(
-        size_t ntotal2,
-        int nsq,
-        const uint8_t* codes,
-        const uint8_t* LUT0,
-        ResultHandler& res,
-        const Scaler& scaler,
-        size_t block_stride) {
-    constexpr int Q1 = QBS & 15;
-    constexpr int Q2 = (QBS >> 4) & 15;
-    constexpr int Q3 = (QBS >> 8) & 15;
-    constexpr int Q4 = (QBS >> 12) & 15;
-    constexpr int SQ = Q1 + Q2 + Q3 + Q4;
-
-    for (size_t j0 = 0; j0 < ntotal2; j0 += 32) {
-        FixedStorageHandler<SQ, 2> res2;
-        const uint8_t* LUT = LUT0;
-        kernel_accumulate_block<Q1>(nsq, codes, LUT, res2, scaler);
-        LUT += Q1 * nsq * 16;
-        if (Q2 > 0) {
-            res2.set_block_origin(Q1, 0);
-            kernel_accumulate_block<Q2>(nsq, codes, LUT, res2, scaler);
-            LUT += Q2 * nsq * 16;
-        }
-        if (Q3 > 0) {
-            res2.set_block_origin(Q1 + Q2, 0);
-            kernel_accumulate_block<Q3>(nsq, codes, LUT, res2, scaler);
-            LUT += Q3 * nsq * 16;
-        }
-        if (Q4 > 0) {
-            res2.set_block_origin(Q1 + Q2 + Q3, 0);
-            kernel_accumulate_block<Q4>(nsq, codes, LUT, res2, scaler);
-        }
-        res.set_block_origin(0, j0);
-        res2.to_other_handler(res);
-        codes += block_stride;
-    }
-}
-
-template <int NQ, class ResultHandler, class Scaler>
-void kernel_accumulate_block_loop(
-        size_t ntotal2,
-        int nsq,
-        const uint8_t* codes,
-        const uint8_t* LUT,
-        ResultHandler& res,
-        const Scaler& scaler,
-        size_t block_stride) {
-    for (size_t j0 = 0; j0 < ntotal2; j0 += 32) {
-        res.set_block_origin(0, j0);
-        kernel_accumulate_block<NQ, ResultHandler>(
-                nsq, codes, LUT, res, scaler);
-        codes += block_stride;
-    }
-}
-
-// non-template version of accumulate kernel -- dispatches dynamically
-template <class ResultHandler, class Scaler>
-void accumulate(
-        int nq,
-        size_t ntotal2,
-        int nsq,
-        const uint8_t* codes,
-        const uint8_t* LUT,
-        ResultHandler& res,
-        const Scaler& scaler,
-        size_t block_stride) {
-    assert(nsq % 2 == 0);
-    assert(is_aligned_pointer(LUT));
-
-#define DISPATCH(NQ)                                                  \
-    case NQ:                                                          \
-        kernel_accumulate_block_loop<NQ, ResultHandler>(              \
-                ntotal2, nsq, codes, LUT, res, scaler, block_stride); \
-        return
-
-    switch (nq) {
-        DISPATCH(1);
-        DISPATCH(2);
-        DISPATCH(3);
-        DISPATCH(4);
-    }
-    FAISS_THROW_FMT("accumulate nq=%d not instantiated", nq);
-
-#undef DISPATCH
-}
-
-template <class ResultHandler, class Scaler>
-void pq4_accumulate_loop_qbs_fixed_scaler(
-        int qbs,
-        size_t ntotal2,
-        int nsq,
-        const uint8_t* codes,
-        const uint8_t* LUT0,
-        ResultHandler& res,
-        const Scaler& scaler,
-        size_t block_stride = 0) {
-    assert(nsq % 2 == 0);
-    assert(is_aligned_pointer(codes));
-    assert(is_aligned_pointer(LUT0));
-
-    // try out optimized versions
-    switch (qbs) {
-#define DISPATCH(QBS)                                                  \
-    case QBS:                                                          \
-        accumulate_q_4step<QBS>(                                       \
-                ntotal2, nsq, codes, LUT0, res, scaler, block_stride); \
-        return;
-        DISPATCH(0x3333); // 12
-        DISPATCH(0x2333); // 11
-        DISPATCH(0x2233); // 10
-        DISPATCH(0x333);  // 9
-        DISPATCH(0x2223); // 9
-        DISPATCH(0x233);  // 8
-        DISPATCH(0x1223); // 8
-        DISPATCH(0x223);  // 7
-        DISPATCH(0x34);   // 7
-        DISPATCH(0x133);  // 7
-        DISPATCH(0x6);    // 6
-        DISPATCH(0x33);   // 6
-        DISPATCH(0x123);  // 6
-        DISPATCH(0x222);  // 6
-        DISPATCH(0x23);   // 5
-        DISPATCH(0x5);    // 5
-        DISPATCH(0x13);   // 4
-        DISPATCH(0x22);   // 4
-        DISPATCH(0x4);    // 4
-        DISPATCH(0x3);    // 3
-        DISPATCH(0x21);   // 3
-        DISPATCH(0x2);    // 2
-        DISPATCH(0x1);    // 1
-#undef DISPATCH
-    }
-
-    // default implementation where qbs is not known at compile time
-    for (size_t j0 = 0; j0 < ntotal2; j0 += 32) {
-        const uint8_t* LUT = LUT0;
-        int qi = qbs;
-        int i0 = 0;
-        while (qi) {
-            int nq = qi & 15;
-            qi >>= 4;
-            res.set_block_origin(i0, j0);
-#define DISPATCH(NQ)                                \
-    case NQ:                                        \
-        kernel_accumulate_block<NQ, ResultHandler>( \
-                nsq, codes, LUT, res, scaler);      \
-        break
-            switch (nq) {
-                DISPATCH(1);
-                DISPATCH(2);
-                DISPATCH(3);
-                DISPATCH(4);
-#undef DISPATCH
-                default:
-                    FAISS_THROW_FMT("accumulate nq=%d not instantiated", nq);
-            }
-            i0 += nq;
-            LUT += nq * nsq * 16;
-        }
-        codes += block_stride;
-    }
-}
-
-struct Run_pq4_accumulate_loop_qbs {
-    template <class ResultHandler>
-    void f(ResultHandler& res,
-           int qbs,
-           size_t nb,
-           int nsq,
-           const uint8_t* codes,
-           const uint8_t* LUT,
-           const NormTableScaler* scaler,
-           size_t block_stride) {
-        if (scaler) {
-            pq4_accumulate_loop_qbs_fixed_scaler(
-                    qbs, nb, nsq, codes, LUT, res, *scaler, block_stride);
-        } else {
-            DummyScaler dummy;
-            pq4_accumulate_loop_qbs_fixed_scaler(
-                    qbs, nb, nsq, codes, LUT, res, dummy, block_stride);
-        }
-    }
-};
-
-} // namespace
-
-void pq4_accumulate_loop_qbs(
-        int qbs,
-        size_t nb,
-        int nsq,
-        const uint8_t* codes,
-        const uint8_t* LUT,
-        SIMDResultHandler& res,
-        const NormTableScaler* scaler,
-        size_t block_stride) {
-    Run_pq4_accumulate_loop_qbs consumer;
-    dispatch_SIMDResultHandler(
-            res, consumer, qbs, nb, nsq, codes, LUT, scaler, block_stride);
-}
-
-/***************************************************************
- * Packing functions
- ***************************************************************/
-
-int pq4_qbs_to_nq(int qbs) {
-    int i0 = 0;
-    int qi = qbs;
-    while (qi) {
-        int nq = qi & 15;
-        qi >>= 4;
-        i0 += nq;
-    }
-    return i0;
-}
-
-void accumulate_to_mem(
-        int nq,
-        size_t ntotal2,
-        int nsq,
-        const uint8_t* codes,
-        const uint8_t* LUT,
-        uint16_t* accu) {
-    FAISS_THROW_IF_NOT(ntotal2 % 32 == 0);
-    StoreResultHandler handler(accu, ntotal2);
-    DummyScaler scaler;
-    accumulate(nq, ntotal2, nsq, codes, LUT, handler, scaler, 32 * nsq / 2);
-}
-
-int pq4_preferred_qbs(int n) {
-    // from timings in P141901742, P141902828
-    static int map[12] = {
-            0, 1, 2, 3, 0x13, 0x23, 0x33, 0x223, 0x233, 0x333, 0x2233, 0x2333};
-    if (n <= 11) {
-        return map[n];
-    } else if (n <= 24) {
-        // override qbs: all first stages with 3 steps
-        // then 1 stage with the rest
-        int nbit = 4 * (n / 3); // nbits with only 3s
-        int qbs = 0x33333333 & ((1 << nbit) - 1);
-        qbs |= (n % 3) << nbit;
-        return qbs;
-    } else {
-        FAISS_THROW_FMT("number of queries %d too large", n);
-    }
-}
-
 } // namespace faiss
+
+#endif // __AVX512F__

data/vendor/faiss/faiss/impl/fast_scan/rabitq_dispatching.h

@@ -0,0 +1,90 @@
+/*
+ * 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
+
+/**
+ * @file rabitq_dispatching.h
+ * @brief Per-SIMD TU dispatch for RaBitQ flat scanner.
+ *
+ * Included after dispatching.h in each per-SIMD TU, so that
+ * ScannerMixIn from dispatching.h is visible (same TU).
+ *
+ * Provides the rabitq_make_knn_scanner_impl<SL> specialization
+ * that wraps RaBitQHeapHandler in ScannerMixIn.
+ */
+
+#ifndef THE_LEVEL_TO_DISPATCH
+#error "Define THE_LEVEL_TO_DISPATCH before including this header"
+#endif
+
+#include <faiss/IndexIVFRaBitQFastScan.h>
+#include <faiss/IndexRaBitQFastScan.h>
+#include <faiss/impl/fast_scan/fast_scan.h>
+
+// ScannerMixIn is visible from dispatching.h (same TU)
+
+namespace faiss {
+
+template <>
+std::unique_ptr<FastScanCodeScanner> rabitq_make_knn_scanner_impl<
+        THE_LEVEL_TO_DISPATCH>(
+        const IndexRaBitQFastScan* index,
+        bool is_max,
+        size_t nq,
+        int64_t k,
+        float* distances,
+        int64_t* ids,
+        const IDSelector* sel,
+        const FastScanDistancePostProcessing& context,
+        bool is_multi_bit) {
+    if (is_max) {
+        using H = RaBitQHeapHandler<
+                CMax<uint16_t, int>,
+                false,
+                THE_LEVEL_TO_DISPATCH>;
+        return std::make_unique<ScannerMixIn<H>>(
+                index, nq, k, distances, ids, sel, &context, is_multi_bit);
+    } else {
+        using H = RaBitQHeapHandler<
+                CMin<uint16_t, int>,
+                false,
+                THE_LEVEL_TO_DISPATCH>;
+        return std::make_unique<ScannerMixIn<H>>(
+                index, nq, k, distances, ids, sel, &context, is_multi_bit);
+    }
+}
+
+// IVF RaBitQ scanner factory
+template <>
+std::unique_ptr<FastScanCodeScanner> rabitq_ivf_make_knn_scanner_impl<
+        THE_LEVEL_TO_DISPATCH>(
+        bool is_max,
+        const IndexIVFRaBitQFastScan* index,
+        size_t nq,
+        size_t k,
+        float* distances,
+        int64_t* ids,
+        const IDSelector* sel,
+        const FastScanDistancePostProcessing* context,
+        bool multi_bit) {
+    if (is_max) {
+        using C = CMax<uint16_t, int64_t>;
+        using H = simd_result_handlers::
+                IVFRaBitQHeapHandler<C, THE_LEVEL_TO_DISPATCH>;
+        return std::make_unique<ScannerMixIn<H>>(
+                index, nq, k, distances, ids, sel, context, multi_bit);
+    } else {
+        using C = CMin<uint16_t, int64_t>;
+        using H = simd_result_handlers::
+                IVFRaBitQHeapHandler<C, THE_LEVEL_TO_DISPATCH>;
+        return std::make_unique<ScannerMixIn<H>>(
+                index, nq, k, distances, ids, sel, context, multi_bit);
+    }
+}
+
+} // namespace faiss

data/vendor/faiss/faiss/impl/fast_scan/rabitq_result_handler.h

@@ -0,0 +1,108 @@
+/*
+ * 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 <memory>
+#include <vector>
+
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/RaBitQUtils.h>
+#include <faiss/impl/fast_scan/FastScanDistancePostProcessing.h>
+#include <faiss/impl/fast_scan/simd_result_handlers.h>
+#include <faiss/invlists/InvertedLists.h>
+#include <faiss/utils/Heap.h>
+
+namespace faiss {
+
+// Forward declaration — full definition in IndexIVFRaBitQFastScan.h
+struct IndexIVFRaBitQFastScan;
+
+namespace simd_result_handlers {
+
+// Import shared utilities from RaBitQUtils
+using rabitq_utils::ExtraBitsFactors;
+using rabitq_utils::SignBitFactors;
+using rabitq_utils::SignBitFactorsWithError;
+
+/** SIMD result handler for IndexIVFRaBitQFastScan that applies
+ * RaBitQ-specific distance corrections during batch processing.
+ *
+ * @tparam C Comparator type (CMin/CMax) for heap operations
+ * @tparam SL SIMD level for dynamic dispatch
+ */
+template <class C, SIMDLevel SL = SINGLE_SIMD_LEVEL_256>
+struct IVFRaBitQHeapHandler : ResultHandlerCompare<C, true, SL> {
+    using RHC = ResultHandlerCompare<C, true, SL>;
+    using typename RHC::simd16uint16;
+
+    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
+    size_t nup = 0;         // Number of heap updates
+
+    // Cached block-layout constants (invariant for handler lifetime)
+    const size_t storage_size;
+    const size_t packed_block_size;
+    const size_t full_block_size;
+    std::vector<uint8_t> unpack_buf; // sign bits scratch buffer
+
+    // Cached per-list values (set in set_list_context, avoid recomputing in
+    // handle)
+    size_t cached_nprobe = 0;
+    bool is_similarity = false; // metric == INNER_PRODUCT
+
+    // Use float-based comparator for heap operations
+    using Cfloat = typename std::conditional<
+            C::is_max,
+            CMax<float, int64_t>,
+            CMin<float, int64_t>>::type;
+
+    // Constructor and method bodies are defined at the bottom of
+    // IndexIVFRaBitQFastScan.h (after the full struct definition is
+    // available) to break the circular header dependency.
+    IVFRaBitQHeapHandler(
+            const IndexIVFRaBitQFastScan* idx,
+            size_t nq_val,
+            size_t k_val,
+            float* distances,
+            int64_t* labels,
+            const IDSelector* sel,
+            const FastScanDistancePostProcessing* ctx = nullptr,
+            bool multibit = false);
+
+    void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) final;
+
+    void set_list_context(size_t list_no, const std::vector<int>& probe_map)
+            override;
+
+    void begin(const float* norms) override;
+
+    void end() override;
+
+    size_t num_updates() override {
+        return nup;
+    }
+
+   private:
+    float compute_full_multibit_distance(
+            size_t local_q,
+            size_t global_q,
+            size_t local_offset,
+            const uint8_t* aux_ptr);
+};
+
+} // namespace simd_result_handlers
+
+} // namespace faiss