faiss 0.5.3 → 0.6.1

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

@@ -0,0 +1,268 @@
+/*
+ * 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 dispatching.h
+ * @brief Per-SIMD TU dispatch template for fast scan.
+ *
+ * This header is included once per SIMD TU with THE_LEVEL_TO_DISPATCH
+ * set to the desired SIMDLevel. It provides:
+ *   - ScannerMixIn: wraps a handler + calls kernel at the TU's SIMD level
+ *   - make_fast_scan_scanner_impl<SL>: factory specialization
+ *
+ * Usage (in a per-SIMD .cpp file):
+ *   #define THE_LEVEL_TO_DISPATCH SIMDLevel::AVX2
+ *   #include <faiss/impl/fast_scan/dispatching.h>
+ *
+ * Kernel helpers come from accumulate_loops.h (search_1 multi-BB path
+ * and QBS 256-bit path) and accumulate_loops_512.h (QBS 512-bit path,
+ * AVX512 TU only).
+ */
+
+#ifndef THE_LEVEL_TO_DISPATCH
+#error "Define THE_LEVEL_TO_DISPATCH before including this header"
+#endif
+
+#include <memory>
+
+#include <faiss/impl/fast_scan/accumulate_loops.h>
+#include <faiss/impl/fast_scan/fast_scan.h>
+
+#if defined(COMPILE_SIMD_AVX512) && defined(__AVX512F__)
+#include <faiss/impl/fast_scan/accumulate_loops_512.h>
+#endif
+
+namespace faiss {
+
+using namespace simd_result_handlers;
+
+/***************************************************************
+ * ScannerMixIn: wraps a concrete handler + calls accumulation
+ * kernels. Lives behind the virtual FastScanCodeScanner interface
+ * so callers don't need to know the handler type.
+ ***************************************************************/
+
+template <class Handler>
+struct ScannerMixIn : FastScanCodeScanner {
+    Handler handler_;
+
+    template <typename... Args>
+    explicit ScannerMixIn(Args&&... args)
+            : handler_(std::forward<Args>(args)...) {}
+
+    SIMDResultHandlerToFloat* handler() override {
+        return &handler_;
+    }
+
+    void accumulate_loop(
+            int nq,
+            size_t nb,
+            int bbs,
+            int nsq,
+            const uint8_t* codes,
+            const uint8_t* LUT,
+            int pq2x4_scale,
+            size_t block_stride) override {
+        if (pq2x4_scale) {
+            NormTableScaler<THE_LEVEL_TO_DISPATCH> scaler(pq2x4_scale);
+            pq4_accumulate_loop_fixed_scaler<THE_LEVEL_TO_DISPATCH>(
+                    nq,
+                    nb,
+                    bbs,
+                    nsq,
+                    codes,
+                    LUT,
+                    handler_,
+                    scaler,
+                    block_stride);
+        } else {
+            DummyScaler<THE_LEVEL_TO_DISPATCH> dummy;
+            pq4_accumulate_loop_fixed_scaler<THE_LEVEL_TO_DISPATCH>(
+                    nq,
+                    nb,
+                    bbs,
+                    nsq,
+                    codes,
+                    LUT,
+                    handler_,
+                    dummy,
+                    block_stride);
+        }
+    }
+
+    void accumulate_loop_qbs(
+            int qbs,
+            size_t nb,
+            int nsq,
+            const uint8_t* codes,
+            const uint8_t* LUT,
+            int pq2x4_scale,
+            size_t block_stride) override {
+#if defined(COMPILE_SIMD_AVX512) && defined(__AVX512F__)
+        constexpr bool use_avx512_qbs =
+                (THE_LEVEL_TO_DISPATCH == SIMDLevel::AVX512 ||
+                 THE_LEVEL_TO_DISPATCH == SIMDLevel::AVX512_SPR);
+#else
+        constexpr bool use_avx512_qbs = false;
+#endif
+        if constexpr (use_avx512_qbs) {
+            // Use 512-bit QBS kernels with properly-leveled scalers.
+            if (pq2x4_scale) {
+                NormTableScaler<THE_LEVEL_TO_DISPATCH> scaler(pq2x4_scale);
+                pq4_accumulate_loop_qbs_fixed_scaler_512(
+                        qbs,
+                        nb,
+                        nsq,
+                        codes,
+                        LUT,
+                        handler_,
+                        scaler,
+                        block_stride);
+            } else {
+                DummyScaler<THE_LEVEL_TO_DISPATCH> dummy;
+                pq4_accumulate_loop_qbs_fixed_scaler_512(
+                        qbs,
+                        nb,
+                        nsq,
+                        codes,
+                        LUT,
+                        handler_,
+                        dummy,
+                        block_stride);
+            }
+        } else {
+            if (pq2x4_scale) {
+                NormTableScaler<THE_LEVEL_TO_DISPATCH> scaler(pq2x4_scale);
+                pq4_accumulate_loop_qbs_fixed_scaler_256<THE_LEVEL_TO_DISPATCH>(
+                        qbs,
+                        nb,
+                        nsq,
+                        codes,
+                        LUT,
+                        handler_,
+                        scaler,
+                        block_stride);
+            } else {
+                DummyScaler<THE_LEVEL_TO_DISPATCH> dummy;
+                pq4_accumulate_loop_qbs_fixed_scaler_256<THE_LEVEL_TO_DISPATCH>(
+                        qbs,
+                        nb,
+                        nsq,
+                        codes,
+                        LUT,
+                        handler_,
+                        dummy,
+                        block_stride);
+            }
+        }
+    }
+};
+
+/***************************************************************
+ * Factory specialization for this SIMD level.
+ *
+ * Combinatorial dispatch: is_max × with_id_map × handler type
+ *   k == 1:  SingleResultHandler
+ *   impl even: HeapHandler
+ *   impl odd:  ReservoirHandler (capacity = 2*k)
+ ***************************************************************/
+
+template <>
+std::unique_ptr<FastScanCodeScanner> make_fast_scan_scanner_impl<
+        THE_LEVEL_TO_DISPATCH>(
+        bool is_max,
+        int impl,
+        size_t nq,
+        size_t ntotal,
+        int64_t k,
+        float* distances,
+        int64_t* ids,
+        const IDSelector* sel,
+        bool with_id_map) {
+    // Helper lambda: given comparator C and with_id_map W, select handler
+    auto make = [&]<class C, bool W>() -> std::unique_ptr<FastScanCodeScanner> {
+        if (k == 1) {
+            using H = SingleResultHandler<C, W, THE_LEVEL_TO_DISPATCH>;
+            return std::make_unique<ScannerMixIn<H>>(
+                    nq, ntotal, distances, ids, sel);
+        } else if (impl % 2 == 0) {
+            using H = HeapHandler<C, W, THE_LEVEL_TO_DISPATCH>;
+            return std::make_unique<ScannerMixIn<H>>(
+                    nq, ntotal, k, distances, ids, sel);
+        } else {
+            using H = ReservoirHandler<C, W, THE_LEVEL_TO_DISPATCH>;
+            return std::make_unique<ScannerMixIn<H>>(
+                    nq, ntotal, size_t(k), size_t(2 * k), distances, ids, sel);
+        }
+    };
+
+    if (is_max) {
+        if (with_id_map) {
+            return make.template operator()<CMax<uint16_t, int64_t>, true>();
+        } else {
+            return make.template operator()<CMax<uint16_t, int>, false>();
+        }
+    } else {
+        if (with_id_map) {
+            return make.template operator()<CMin<uint16_t, int64_t>, true>();
+        } else {
+            return make.template operator()<CMin<uint16_t, int>, false>();
+        }
+    }
+}
+
+/***************************************************************
+ * Range search scanner factories.
+ ***************************************************************/
+
+template <>
+std::unique_ptr<FastScanCodeScanner> make_range_scanner_impl<
+        THE_LEVEL_TO_DISPATCH>(
+        bool is_max,
+        RangeSearchResult& rres,
+        float radius,
+        size_t ntotal,
+        const IDSelector* sel) {
+    if (is_max) {
+        using C = CMax<uint16_t, int64_t>;
+        return std::make_unique<
+                ScannerMixIn<RangeHandler<C, true, THE_LEVEL_TO_DISPATCH>>>(
+                rres, radius, ntotal, sel);
+    } else {
+        using C = CMin<uint16_t, int64_t>;
+        return std::make_unique<
+                ScannerMixIn<RangeHandler<C, true, THE_LEVEL_TO_DISPATCH>>>(
+                rres, radius, ntotal, sel);
+    }
+}
+
+template <>
+std::unique_ptr<FastScanCodeScanner> make_partial_range_scanner_impl<
+        THE_LEVEL_TO_DISPATCH>(
+        bool is_max,
+        RangeSearchPartialResult& pres,
+        float radius,
+        size_t ntotal,
+        size_t q0,
+        size_t q1,
+        const IDSelector* sel) {
+    if (is_max) {
+        using C = CMax<uint16_t, int64_t>;
+        return std::make_unique<ScannerMixIn<
+                PartialRangeHandler<C, true, THE_LEVEL_TO_DISPATCH>>>(
+                pres, radius, ntotal, q0, q1, sel);
+    } else {
+        using C = CMin<uint16_t, int64_t>;
+        return std::make_unique<ScannerMixIn<
+                PartialRangeHandler<C, true, THE_LEVEL_TO_DISPATCH>>>(
+                pres, radius, ntotal, q0, q1, sel);
+    }
+}
+
+} // namespace faiss

data/vendor/faiss/faiss/impl/{pq4_fast_scan.cpp → fast_scan/fast_scan.cpp}

@@ -6,8 +6,11 @@
  */
 
 #include <faiss/impl/FaissAssert.h>
-#include <faiss/impl/pq4_fast_scan.h>
-#include <faiss/impl/simd_result_handlers.h>
+#include <faiss/impl/fast_scan/LookupTableScaler.h>
+#include <faiss/impl/fast_scan/decompose_qbs.h>
+#include <faiss/impl/fast_scan/fast_scan.h>
+#include <faiss/impl/fast_scan/simd_result_handlers.h>
+#include <faiss/impl/simd_dispatch.h>
 
 #include <array>
 
@@ -111,7 +114,8 @@ void pq4_pack_codes_range(
         size_t bbs,
         size_t nsq,
         uint8_t* blocks,
-        size_t code_stride) {
+        size_t code_stride,
+        size_t block_stride) {
     // Determine stride: use custom if provided, otherwise use legacy
     // calculation
     size_t actual_stride = (code_stride == 0) ? (M + 1) / 2 : code_stride;
@@ -136,7 +140,7 @@ void pq4_pack_codes_range(
     size_t block1 = ((i1 - 1) / bbs) + 1;
 
     for (size_t b = block0; b < block1; b++) {
-        uint8_t* codes2 = blocks + b * bbs * nsq / 2;
+        uint8_t* codes2 = blocks + b * block_stride;
         int64_t i_base = b * bbs - i0;
         for (int sq = 0; sq < nsq; sq += 2) {
             for (size_t i = 0; i < bbs; i += 32) {
@@ -272,6 +276,10 @@ void CodePackerPQ4::unpack_1(
     }
 }
 
+CodePacker* CodePackerPQ4::clone() const {
+    return new CodePackerPQ4(*this);
+}
+
 /***************************************************************
  * Packing functions for Look-Up Tables (LUT)
  ***************************************************************/
@@ -345,4 +353,168 @@ int pq4_pack_LUT_qbs_q_map(
     return i0;
 }
 
+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;
+}
+
+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
+
+/***************************************************************
+ * FastScanCodeScanner: NONE specialization + dispatch wrapper.
+ *
+ * The NONE specialization provides the scalar fallback.
+ * Per-SIMD specializations (AVX2, AVX512, ARM_NEON) are in
+ * impl-avx2.cpp, impl-avx512.cpp, impl-neon.cpp respectively.
+ ***************************************************************/
+
+#define THE_LEVEL_TO_DISPATCH SIMDLevel::NONE
+#include <faiss/impl/fast_scan/dispatching.h>        // IWYU pragma: keep
+#include <faiss/impl/fast_scan/rabitq_dispatching.h> // IWYU pragma: keep
+#undef THE_LEVEL_TO_DISPATCH
+
+namespace faiss {
+
+using namespace simd_result_handlers;
+
+/***************************************************************
+ * accumulate_to_mem: NONE specialization + runtime dispatch.
+ ***************************************************************/
+
+template <>
+void accumulate_to_mem_impl<SIMDLevel::NONE>(
+        int nq,
+        size_t ntotal2,
+        int nsq,
+        const uint8_t* codes,
+        const uint8_t* LUT,
+        uint16_t* accu) {
+    StoreResultHandler<SIMDLevel::NONE> handler(accu, ntotal2);
+    DummyScaler<SIMDLevel::NONE> scaler;
+    accumulate<SIMDLevel::NONE>(
+            nq, ntotal2, nsq, codes, LUT, handler, scaler, 32 * nsq / 2);
+}
+
+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);
+    with_simd_level([&]<SIMDLevel SL>() {
+        accumulate_to_mem_impl<SL>(nq, ntotal2, nsq, codes, LUT, accu);
+    });
+}
+
+} // namespace faiss
+
+namespace faiss {
+
+std::unique_ptr<FastScanCodeScanner> make_fast_scan_knn_scanner(
+        bool is_max,
+        int impl,
+        size_t nq,
+        size_t ntotal,
+        int64_t k,
+        float* distances,
+        int64_t* ids,
+        const IDSelector* sel,
+        bool with_id_map) {
+    return with_simd_level([&]<SIMDLevel SL>() {
+        return make_fast_scan_scanner_impl<SL>(
+                is_max, impl, nq, ntotal, k, distances, ids, sel, with_id_map);
+    });
+}
+
+std::unique_ptr<FastScanCodeScanner> make_range_scanner(
+        bool is_max,
+        RangeSearchResult& rres,
+        float radius,
+        size_t ntotal,
+        const IDSelector* sel) {
+    return with_simd_level([&]<SIMDLevel SL>() {
+        return make_range_scanner_impl<SL>(is_max, rres, radius, ntotal, sel);
+    });
+}
+
+std::unique_ptr<FastScanCodeScanner> make_partial_range_scanner(
+        bool is_max,
+        RangeSearchPartialResult& pres,
+        float radius,
+        size_t ntotal,
+        size_t q0,
+        size_t q1,
+        const IDSelector* sel) {
+    return with_simd_level([&]<SIMDLevel SL>() {
+        return make_partial_range_scanner_impl<SL>(
+                is_max, pres, radius, ntotal, q0, q1, sel);
+    });
+}
+
+std::unique_ptr<FastScanCodeScanner> rabitq_make_knn_scanner(
+        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) {
+    return with_simd_level([&]<SIMDLevel SL>() {
+        return rabitq_make_knn_scanner_impl<SL>(
+                index,
+                is_max,
+                nq,
+                k,
+                distances,
+                ids,
+                sel,
+                context,
+                is_multi_bit);
+    });
+}
+
+std::unique_ptr<FastScanCodeScanner> rabitq_ivf_make_knn_scanner(
+        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) {
+    return with_simd_level([&]<SIMDLevel SL>() {
+        return rabitq_ivf_make_knn_scanner_impl<SL>(
+                is_max, index, nq, k, distances, ids, sel, context, multi_bit);
+    });
+}
+
 } // namespace faiss