faiss 0.6.0 → 0.6.1

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

@@ -0,0 +1,120 @@
+/*
+ * 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_ARM_NEON
+
+#define THE_LEVEL_TO_DISPATCH SIMDLevel::ARM_NEON
+#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::ARM_NEON>(
+        int nq,
+        size_t ntotal2,
+        int nsq,
+        const uint8_t* codes,
+        const uint8_t* LUT,
+        uint16_t* accu) {
+    StoreResultHandler<SIMDLevel::ARM_NEON> handler(accu, ntotal2);
+    DummyScaler<SIMDLevel::ARM_NEON> scaler;
+    accumulate<SIMDLevel::ARM_NEON>(
+            nq, ntotal2, nsq, codes, LUT, handler, scaler, 32 * nsq / 2);
+}
+
+} // namespace faiss
+
+// ARM_SVE: forward to ARM_NEON implementation until a dedicated SVE
+// specialization is written (same pattern as scalar_quantizer/sq-neon.cpp).
+#ifdef COMPILE_SIMD_ARM_SVE
+
+namespace faiss {
+
+template <>
+std::unique_ptr<FastScanCodeScanner> make_fast_scan_scanner_impl<
+        SIMDLevel::ARM_SVE>(
+        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 make_fast_scan_scanner_impl<SIMDLevel::ARM_NEON>(
+            is_max, impl, nq, ntotal, k, distances, ids, sel, with_id_map);
+}
+
+template <>
+std::unique_ptr<FastScanCodeScanner> make_range_scanner_impl<
+        SIMDLevel::ARM_SVE>(
+        bool is_max,
+        RangeSearchResult& rres,
+        float radius,
+        size_t ntotal,
+        const IDSelector* sel) {
+    return make_range_scanner_impl<SIMDLevel::ARM_NEON>(
+            is_max, rres, radius, ntotal, sel);
+}
+
+template <>
+std::unique_ptr<FastScanCodeScanner> make_partial_range_scanner_impl<
+        SIMDLevel::ARM_SVE>(
+        bool is_max,
+        RangeSearchPartialResult& pres,
+        float radius,
+        size_t ntotal,
+        size_t q0,
+        size_t q1,
+        const IDSelector* sel) {
+    return make_partial_range_scanner_impl<SIMDLevel::ARM_NEON>(
+            is_max, pres, radius, ntotal, q0, q1, sel);
+}
+
+template <>
+std::unique_ptr<FastScanCodeScanner> rabitq_make_knn_scanner_impl<
+        SIMDLevel::ARM_SVE>(
+        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 rabitq_make_knn_scanner_impl<SIMDLevel::ARM_NEON>(
+            index, is_max, nq, k, distances, ids, sel, context, is_multi_bit);
+}
+
+template <>
+std::unique_ptr<FastScanCodeScanner> rabitq_ivf_make_knn_scanner_impl<
+        SIMDLevel::ARM_SVE>(
+        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 rabitq_ivf_make_knn_scanner_impl<SIMDLevel::ARM_NEON>(
+            is_max, index, nq, k, distances, ids, sel, context, multi_bit);
+}
+
+} // namespace faiss
+
+#endif // COMPILE_SIMD_ARM_SVE
+
+#endif // COMPILE_SIMD_ARM_NEON

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

@@ -0,0 +1,104 @@
+/*
+ * 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.
+ */
+
+// RISC-V RVV: forward all fast_scan specializations to NONE until
+// dedicated RVV implementations are written.
+
+#ifdef COMPILE_SIMD_RISCV_RVV
+
+#include <faiss/impl/fast_scan/fast_scan.h>
+
+namespace faiss {
+
+template <>
+void accumulate_to_mem_impl<SIMDLevel::RISCV_RVV>(
+        int nq,
+        size_t ntotal2,
+        int nsq,
+        const uint8_t* codes,
+        const uint8_t* LUT,
+        uint16_t* accu) {
+    accumulate_to_mem_impl<SIMDLevel::NONE>(nq, ntotal2, nsq, codes, LUT, accu);
+}
+
+template <>
+std::unique_ptr<FastScanCodeScanner> make_fast_scan_scanner_impl<
+        SIMDLevel::RISCV_RVV>(
+        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 make_fast_scan_scanner_impl<SIMDLevel::NONE>(
+            is_max, impl, nq, ntotal, k, distances, ids, sel, with_id_map);
+}
+
+template <>
+std::unique_ptr<FastScanCodeScanner> make_range_scanner_impl<
+        SIMDLevel::RISCV_RVV>(
+        bool is_max,
+        RangeSearchResult& rres,
+        float radius,
+        size_t ntotal,
+        const IDSelector* sel) {
+    return make_range_scanner_impl<SIMDLevel::NONE>(
+            is_max, rres, radius, ntotal, sel);
+}
+
+template <>
+std::unique_ptr<FastScanCodeScanner> make_partial_range_scanner_impl<
+        SIMDLevel::RISCV_RVV>(
+        bool is_max,
+        RangeSearchPartialResult& pres,
+        float radius,
+        size_t ntotal,
+        size_t q0,
+        size_t q1,
+        const IDSelector* sel) {
+    return make_partial_range_scanner_impl<SIMDLevel::NONE>(
+            is_max, pres, radius, ntotal, q0, q1, sel);
+}
+
+template <>
+std::unique_ptr<FastScanCodeScanner> rabitq_make_knn_scanner_impl<
+        SIMDLevel::RISCV_RVV>(
+        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 rabitq_make_knn_scanner_impl<SIMDLevel::NONE>(
+            index, is_max, nq, k, distances, ids, sel, context, is_multi_bit);
+}
+
+template <>
+std::unique_ptr<FastScanCodeScanner> rabitq_ivf_make_knn_scanner_impl<
+        SIMDLevel::RISCV_RVV>(
+        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 rabitq_ivf_make_knn_scanner_impl<SIMDLevel::NONE>(
+            is_max, index, nq, k, distances, ids, sel, context, multi_bit);
+}
+
+} // namespace faiss
+
+#endif // COMPILE_SIMD_RISCV_RVV

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

@@ -0,0 +1,213 @@
+/*
+ * 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 <faiss/impl/simdlib/simdlib_dispatch.h>
+
+namespace faiss {
+
+/*
+ * Multi-BB variant: accumulates NQ queries x BB*32 database elements.
+ * Used by the search_1 path (bbs > 32).
+ *
+ * KernelSL selects the SIMD type width used for the inner accumulation
+ * loop.  In DD mode the caller passes THE_LEVEL_TO_DISPATCH so the
+ * kernel uses real AVX2/AVX512 types rather than the emulated-scalar
+ * fallback that SINGLE_SIMD_LEVEL_256 would give.
+ */
+template <
+        int NQ,
+        int BB,
+        SIMDLevel KernelSL = SINGLE_SIMD_LEVEL,
+        class ResultHandler,
+        class Scaler>
+void kernel_accumulate_block(
+        int nsq,
+        const uint8_t* codes,
+        const uint8_t* LUT,
+        ResultHandler& res,
+        const Scaler& scaler) {
+    constexpr SIMDLevel SL256 = simd256_level_selector<KernelSL>::value;
+    using simd16uint16 = simd16uint16_tpl<SL256>;
+    using simd32uint8 = simd32uint8_tpl<SL256>;
+
+    // distance accumulators
+    simd16uint16 accu[NQ][BB][4];
+
+    for (int q = 0; q < NQ; q++) {
+        for (int b = 0; b < BB; b++) {
+            accu[q][b][0].clear();
+            accu[q][b][1].clear();
+            accu[q][b][2].clear();
+            accu[q][b][3].clear();
+        }
+    }
+
+    for (int sq = 0; sq < nsq - scaler.nscale; sq += 2) {
+        simd32uint8 lut_cache[NQ];
+        for (int q = 0; q < NQ; q++) {
+            lut_cache[q] = simd32uint8(LUT);
+            LUT += 32;
+        }
+
+        for (int b = 0; b < BB; b++) {
+            simd32uint8 c = simd32uint8(codes);
+            codes += 32;
+            simd32uint8 mask(15);
+            simd32uint8 chi = simd32uint8(simd16uint16(c) >> 4) & mask;
+            simd32uint8 clo = c & mask;
+
+            for (int q = 0; q < NQ; q++) {
+                simd32uint8 lut = lut_cache[q];
+                simd32uint8 res0 = lut.lookup_2_lanes(clo);
+                simd32uint8 res1 = lut.lookup_2_lanes(chi);
+
+                accu[q][b][0] += simd16uint16(res0);
+                accu[q][b][1] += simd16uint16(res0) >> 8;
+
+                accu[q][b][2] += simd16uint16(res1);
+                accu[q][b][3] += simd16uint16(res1) >> 8;
+            }
+        }
+    }
+
+    for (int sq = 0; sq < scaler.nscale; sq += 2) {
+        simd32uint8 lut_cache[NQ];
+        for (int q = 0; q < NQ; q++) {
+            lut_cache[q] = simd32uint8(LUT);
+            LUT += 32;
+        }
+
+        for (int b = 0; b < BB; b++) {
+            simd32uint8 c = simd32uint8(codes);
+            codes += 32;
+            simd32uint8 mask(15);
+            simd32uint8 chi = simd32uint8(simd16uint16(c) >> 4) & mask;
+            simd32uint8 clo = c & mask;
+
+            for (int q = 0; q < NQ; q++) {
+                simd32uint8 lut = lut_cache[q];
+
+                simd32uint8 res0 = scaler.lookup(lut, clo);
+                accu[q][b][0] += scaler.scale_lo(res0); // handle vectors 0..7
+                accu[q][b][1] += scaler.scale_hi(res0); // handle vectors 8..15
+
+                simd32uint8 res1 = scaler.lookup(lut, chi);
+                accu[q][b][2] += scaler.scale_lo(res1); // handle vectors 16..23
+                accu[q][b][3] +=
+                        scaler.scale_hi(res1); //  handle vectors 24..31
+            }
+        }
+    }
+
+    for (int q = 0; q < NQ; q++) {
+        for (int b = 0; b < BB; b++) {
+            accu[q][b][0] -= accu[q][b][1] << 8;
+            simd16uint16 dis0 = combine2x2(accu[q][b][0], accu[q][b][1]);
+
+            accu[q][b][2] -= accu[q][b][3] << 8;
+            simd16uint16 dis1 = combine2x2(accu[q][b][2], accu[q][b][3]);
+
+            res.handle(q, b, dis0, dis1);
+        }
+    }
+}
+
+/*
+ * Single-BB QBS variant: accumulates NQ queries x 32 db elements (BB=1).
+ * Used by the decompose_qbs layer for non-AVX512 paths.
+ */
+template <
+        int NQ,
+        SIMDLevel KernelSL = SINGLE_SIMD_LEVEL,
+        class ResultHandler,
+        class Scaler>
+void pq4_kernel_qbs_256(
+        int nsq,
+        const uint8_t* codes,
+        const uint8_t* LUT,
+        ResultHandler& res,
+        const Scaler& scaler) {
+    constexpr SIMDLevel SL256 = simd256_level_selector<KernelSL>::value;
+    using simd16uint16 = simd16uint16_tpl<SL256>;
+    using simd32uint8 = simd32uint8_tpl<SL256>;
+
+    // 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();
+        }
+    }
+
+    // _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