faiss 0.6.0 → 0.6.1

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>
 
@@ -350,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

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

@@ -0,0 +1,341 @@
+/*
+ * 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 <cstdint>
+#include <cstdlib>
+#include <memory>
+
+#include <faiss/impl/CodePacker.h>
+#include <faiss/utils/simd_levels.h>
+
+/** PQ4 SIMD packing and accumulation functions
+ *
+ * The basic kernel accumulates nq query vectors with bbs = nb * 2 * 16 vectors
+ * and produces an output matrix for that. It is interesting for nq * nb <= 4,
+ * otherwise register spilling becomes too large.
+ *
+ * The implementation of these functions is spread over 3 cpp files to reduce
+ * parallel compile times. Templates are instantiated explicitly.
+ */
+
+namespace faiss {
+
+struct IDSelector;
+struct RangeSearchResult;
+struct RangeSearchPartialResult;
+struct SIMDResultHandlerToFloat;
+
+/** Pack codes for consumption by the SIMD kernels.
+ *  The unused bytes are set to 0.
+ *
+ * @param codes   input codes, size (ntotal, ceil(M / 2))
+ * @param ntotal  number of input codes
+ * @param nb      output number of codes (ntotal rounded up to a multiple of
+ *                bbs)
+ * @param nsq      number of sub-quantizers (=M rounded up to a multiple of 2)
+ * @param bbs     size of database blocks (multiple of 32)
+ * @param blocks  output array, size nb * nsq / 2.
+ * @param code_stride  optional stride between consecutive codes (0 = use
+default (M + 1) / 2)
+ */
+void pq4_pack_codes(
+        const uint8_t* codes,
+        size_t ntotal,
+        size_t M,
+        size_t nb,
+        size_t bbs,
+        size_t nsq,
+        uint8_t* blocks,
+        size_t code_stride = 0);
+
+/** Same as pack_codes but write in a given range of the output,
+ * leaving the rest untouched. Assumes allocated entries are 0 on input.
+ *
+ * @param codes   input codes, size (i1 - i0, ceil(M / 2))
+ * @param i0      first output code to write
+ * @param i1      last output code to write
+ * @param blocks  output array, size at least ceil(i1 / bbs) * bbs * nsq / 2
+ * @param code_stride  optional stride between consecutive codes (0 = use
+ * default (M + 1) / 2)
+ * @param block_stride  stride in bytes between consecutive blocks.
+ */
+void pq4_pack_codes_range(
+        const uint8_t* codes,
+        size_t M,
+        size_t i0,
+        size_t i1,
+        size_t bbs,
+        size_t nsq,
+        uint8_t* blocks,
+        size_t code_stride,
+        size_t block_stride);
+
+/** get a single element from a packed codes table
+ *
+ * @param vector_id        vector id
+ * @param sq       subquantizer (< nsq)
+ */
+uint8_t pq4_get_packed_element(
+        const uint8_t* data,
+        size_t bbs,
+        size_t nsq,
+        size_t vector_id,
+        size_t sq);
+
+/** set a single element "code" into a packed codes table
+ *
+ * @param vector_id       vector id
+ * @param sq       subquantizer (< nsq)
+ */
+void pq4_set_packed_element(
+        uint8_t* data,
+        uint8_t code,
+        size_t bbs,
+        size_t nsq,
+        size_t vector_id,
+        size_t sq);
+
+/** CodePacker API for the PQ4 fast-scan */
+struct CodePackerPQ4 : CodePacker {
+    size_t nsq;
+
+    CodePackerPQ4(size_t nsq, size_t bbs);
+
+    CodePacker* clone() const final;
+
+    void pack_1(const uint8_t* flat_code, size_t offset, uint8_t* block)
+            const final;
+    void unpack_1(const uint8_t* block, size_t offset, uint8_t* flat_code)
+            const final;
+};
+
+/** Pack Look-up table for consumption by the kernel.
+ *
+ * @param nq      number of queries
+ * @param nsq     number of sub-quantizers (multiple of 2)
+ * @param src     input array, size (nq, 16)
+ * @param dest    output array, size (nq, 16)
+ */
+void pq4_pack_LUT(int nq, int nsq, const uint8_t* src, uint8_t* dest);
+
+/* compute the number of queries from a base-16 decomposition */
+int pq4_qbs_to_nq(int qbs);
+
+/** return the preferred decomposition in blocks for a nb of queries. */
+int pq4_preferred_qbs(int nq);
+
+/** Pack Look-up table for consumption by the kernel.
+ *
+ * @param qbs     4-bit encoded number of query blocks, the total number of
+ *                queries handled (nq) is deduced from it
+ * @param nsq     number of sub-quantizers (multiple of 2)
+ * @param src     input array, size (nq, 16)
+ * @param dest    output array, size (nq, 16)
+ * @return nq
+ */
+int pq4_pack_LUT_qbs(int fqbs, int nsq, const uint8_t* src, uint8_t* dest);
+
+/** Same as pq4_pack_LUT_qbs, except the source vectors are remapped with q_map
+ */
+int pq4_pack_LUT_qbs_q_map(
+        int qbs,
+        int nsq,
+        const uint8_t* src,
+        const int* q_map,
+        uint8_t* dest);
+
+/** Wrapper using simple StoreResultHandler
+ *  and DummyScaler
+ *
+ * @param nq      number of queries
+ * @param ntotal2 number of database elements (multiple of 32)
+ * @param nsq     number of sub-quantizers (muliple of 2)
+ * @param codes   packed codes array
+ * @param LUT     packed look-up table
+ * @param accu    array to store the results
+ */
+void accumulate_to_mem(
+        int nq,
+        size_t ntotal2,
+        int nsq,
+        const uint8_t* codes,
+        const uint8_t* LUT,
+        uint16_t* accu);
+
+/// Per-SIMD specialization of accumulate_to_mem (defined in per-SIMD TUs)
+template <SIMDLevel SL>
+void accumulate_to_mem_impl(
+        int nq,
+        size_t ntotal2,
+        int nsq,
+        const uint8_t* codes,
+        const uint8_t* LUT,
+        uint16_t* accu);
+
+/***************************************************************
+ * FastScanCodeScanner: virtual base that bundles handler + kernel
+ * behind the SIMD dispatch boundary. Per-SIMD TUs instantiate this
+ * with the correct SIMDLevel so that handler and kernel share the
+ * same SIMD types.
+ ***************************************************************/
+
+struct FastScanCodeScanner {
+    virtual ~FastScanCodeScanner() = default;
+
+    /// Access the underlying result handler (for begin/end/normalizer calls)
+    virtual SIMDResultHandlerToFloat* handler() = 0;
+
+    /// Run the search_1 accumulation loop (bbs > 32, multi-BB kernel)
+    virtual 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) = 0;
+
+    /// Run the QBS accumulation loop (bbs == 32)
+    virtual 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) = 0;
+};
+
+/// Per-SIMD factory: explicitly specialized in each per-SIMD TU
+/// (impl-avx2.cpp, impl-avx512.cpp, impl-neon.cpp, fast_scan.cpp for NONE).
+/// Not called directly — use make_fast_scan_knn_scanner() instead.
+template <SIMDLevel SL>
+std::unique_ptr<FastScanCodeScanner> make_fast_scan_scanner_impl(
+        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);
+
+/// Runtime dispatch wrapper: selects the best available SIMD level
+/// (via with_simd_level) and delegates to the corresponding
+/// make_fast_scan_scanner_impl<SL> specialization.
+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 = false);
+
+/// Per-SIMD range scanner factories (defined in per-SIMD TUs via dispatching.h)
+template <SIMDLevel SL>
+std::unique_ptr<FastScanCodeScanner> make_range_scanner_impl(
+        bool is_max,
+        RangeSearchResult& rres,
+        float radius,
+        size_t ntotal,
+        const IDSelector* sel);
+
+template <SIMDLevel SL>
+std::unique_ptr<FastScanCodeScanner> make_partial_range_scanner_impl(
+        bool is_max,
+        RangeSearchPartialResult& pres,
+        float radius,
+        size_t ntotal,
+        size_t q0,
+        size_t q1,
+        const IDSelector* sel);
+
+/// Runtime dispatch: range search scanner.
+std::unique_ptr<FastScanCodeScanner> make_range_scanner(
+        bool is_max,
+        RangeSearchResult& rres,
+        float radius,
+        size_t ntotal,
+        const IDSelector* sel);
+
+/// Runtime dispatch: partial range search scanner (per-thread).
+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);
+
+/***************************************************************
+ * RaBitQ scanner factory: per-SIMD specializations live in
+ * rabitq_dispatching.h, included by each per-SIMD TU.
+ ***************************************************************/
+
+struct IndexRaBitQFastScan;
+struct IndexIVFRaBitQFastScan;
+struct FastScanDistancePostProcessing;
+
+/// Per-SIMD factory (primary template; specializations in rabitq_dispatching.h)
+template <SIMDLevel SL>
+std::unique_ptr<FastScanCodeScanner> rabitq_make_knn_scanner_impl(
+        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);
+
+/// Runtime dispatch wrapper for rabitq_make_knn_scanner_impl
+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);
+
+/// Per-SIMD IVF RaBitQ scanner factory.
+template <SIMDLevel SL>
+std::unique_ptr<FastScanCodeScanner> rabitq_ivf_make_knn_scanner_impl(
+        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);
+
+/// Runtime dispatch wrapper for IVF RaBitQ scanner.
+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);
+
+} // namespace faiss

data/vendor/faiss/faiss/impl/fast_scan/impl-avx2.cpp

@@ -0,0 +1,36 @@
+/*
+ * 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.
+ */
+
+#ifdef COMPILE_SIMD_AVX2
+
+#define THE_LEVEL_TO_DISPATCH SIMDLevel::AVX2
+#include <faiss/impl/fast_scan/dispatching.h>        // IWYU pragma: keep
+#include <faiss/impl/fast_scan/rabitq_dispatching.h> // IWYU pragma: keep
+
+#include <faiss/impl/fast_scan/decompose_qbs.h>
+
+namespace faiss {
+
+using namespace simd_result_handlers;
+
+template <>
+void accumulate_to_mem_impl<SIMDLevel::AVX2>(
+        int nq,
+        size_t ntotal2,
+        int nsq,
+        const uint8_t* codes,
+        const uint8_t* LUT,
+        uint16_t* accu) {
+    StoreResultHandler<SIMDLevel::AVX2> handler(accu, ntotal2);
+    DummyScaler<SIMDLevel::AVX2> scaler;
+    accumulate<SIMDLevel::AVX2>(
+            nq, ntotal2, nsq, codes, LUT, handler, scaler, 32 * nsq / 2);
+}
+
+} // namespace faiss
+
+#endif // COMPILE_SIMD_AVX2

data/vendor/faiss/faiss/impl/fast_scan/impl-avx512.cpp

@@ -0,0 +1,40 @@
+/*
+ * 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.
+ */
+
+#ifdef COMPILE_SIMD_AVX512
+
+#define THE_LEVEL_TO_DISPATCH SIMDLevel::AVX512
+#include <faiss/impl/fast_scan/dispatching.h>        // IWYU pragma: keep
+#include <faiss/impl/fast_scan/rabitq_dispatching.h> // IWYU pragma: keep
+
+#include <faiss/impl/fast_scan/decompose_qbs.h>
+
+namespace faiss {
+
+using namespace simd_result_handlers;
+
+template <>
+void accumulate_to_mem_impl<SIMDLevel::AVX512>(
+        int nq,
+        size_t ntotal2,
+        int nsq,
+        const uint8_t* codes,
+        const uint8_t* LUT,
+        uint16_t* accu) {
+    // Use AVX2-level handler (256-bit StoreResultHandler) since the 512-bit
+    // kernels reduce to AVX2-level simd16uint16 via FixedStorage512.
+    StoreResultHandler<SIMDLevel::AVX2> handler(accu, ntotal2);
+    DummyScaler<SIMDLevel::AVX512> scaler;
+    // kernel_accumulate_block in decompose_qbs.h selects pq4_kernel_qbs_512
+    // via #ifdef __AVX512F__ (which is set for this TU).
+    accumulate<SIMDLevel::AVX512>(
+            nq, ntotal2, nsq, codes, LUT, handler, scaler, 32 * nsq / 2);
+}
+
+} // namespace faiss
+
+#endif // COMPILE_SIMD_AVX512