faiss 0.5.3 → 0.6.1

data/vendor/faiss/faiss/impl/binary_hamming/avx2.cpp

@@ -0,0 +1,26 @@
+/*
+ * 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_SIMD_LEVEL SIMDLevel::AVX2
+
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/utils/hamming_distance/hamming_computer-avx2.h>
+
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/impl/binary_hamming/IndexBinaryHNSW_impl.h>
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/impl/binary_hamming/IndexBinaryHash_impl.h>
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/impl/binary_hamming/IndexBinaryIVF_impl.h>
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/impl/binary_hamming/IndexIVFSpectralHash_impl.h>
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/impl/binary_hamming/IndexPQ_impl.h>
+
+#endif // COMPILE_SIMD_AVX2

data/vendor/faiss/faiss/impl/binary_hamming/avx512.cpp

@@ -0,0 +1,26 @@
+/*
+ * 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_SIMD_LEVEL SIMDLevel::AVX512
+
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/utils/hamming_distance/hamming_computer-avx512.h>
+
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/impl/binary_hamming/IndexBinaryHNSW_impl.h>
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/impl/binary_hamming/IndexBinaryHash_impl.h>
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/impl/binary_hamming/IndexBinaryIVF_impl.h>
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/impl/binary_hamming/IndexIVFSpectralHash_impl.h>
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/impl/binary_hamming/IndexPQ_impl.h>
+
+#endif // COMPILE_SIMD_AVX512

data/vendor/faiss/faiss/impl/binary_hamming/dispatch.h

@@ -0,0 +1,143 @@
+/*
+ * 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/utils/simd_levels.h>
+
+#include <cstddef>
+#include <cstdint>
+
+namespace faiss {
+
+// Forward declarations
+struct DistanceComputer;
+struct InvertedListScanner;
+struct BinaryInvertedListScanner;
+struct IndexBinaryFlat;
+struct IndexBinaryHash;
+struct IndexBinaryMultiHash;
+struct IndexBinaryIVF;
+struct IndexIVFSpectralHash;
+struct IndexPQ;
+struct SearchParametersIVF;
+using IVFSearchParameters = SearchParametersIVF;
+struct RangeQueryResult;
+using idx_t = int64_t;
+
+/** @name IndexBinaryHNSW dispatch
+ * @{ */
+template <SIMDLevel SL>
+DistanceComputer* make_binary_hnsw_distance_computer_fixSL(
+        int code_size,
+        IndexBinaryFlat* flat_storage);
+/** @} */
+
+/** @name IndexBinaryIVF dispatch
+ * @{ */
+template <SIMDLevel SL>
+BinaryInvertedListScanner* make_binary_ivf_scanner_fixSL(
+        size_t code_size,
+        bool store_pairs);
+
+template <SIMDLevel SL>
+void search_knn_hamming_per_invlist_fixSL(
+        int code_size,
+        const IndexBinaryIVF* ivf,
+        size_t n,
+        const uint8_t* x,
+        idx_t k,
+        const idx_t* keys_in,
+        const int32_t* coarse_dis,
+        int32_t* distances,
+        idx_t* labels,
+        bool store_pairs,
+        const IVFSearchParameters* params);
+
+template <SIMDLevel SL>
+void search_knn_hamming_count_fixSL(
+        int code_size,
+        bool store_pairs,
+        const IndexBinaryIVF* ivf,
+        size_t nx,
+        const uint8_t* x,
+        const idx_t* keys,
+        int k,
+        int32_t* distances,
+        idx_t* labels,
+        const IVFSearchParameters* params);
+/** @} */
+
+/** @name IndexBinaryHash dispatch
+ * @{ */
+template <SIMDLevel SL>
+void binary_hash_knn_search_fixSL(
+        const IndexBinaryHash& index,
+        const uint8_t* q,
+        idx_t k,
+        int32_t* heap_sim,
+        idx_t* heap_ids,
+        size_t& n0,
+        size_t& nlist,
+        size_t& ndis);
+
+template <SIMDLevel SL>
+void binary_hash_range_search_fixSL(
+        const IndexBinaryHash& index,
+        const uint8_t* q,
+        int radius,
+        RangeQueryResult& qres,
+        size_t& n0,
+        size_t& nlist,
+        size_t& ndis);
+
+template <SIMDLevel SL>
+void binary_multihash_knn_search_fixSL(
+        const IndexBinaryMultiHash& index,
+        const uint8_t* q,
+        idx_t k,
+        int32_t* heap_sim,
+        idx_t* heap_ids,
+        size_t& n0,
+        size_t& nlist,
+        size_t& ndis);
+
+template <SIMDLevel SL>
+void binary_multihash_range_search_fixSL(
+        const IndexBinaryMultiHash& index,
+        const uint8_t* q,
+        int radius,
+        RangeQueryResult& qres,
+        size_t& n0,
+        size_t& nlist,
+        size_t& ndis);
+/** @} */
+
+/** @name IndexIVFSpectralHash dispatch
+ * @{ */
+template <SIMDLevel SL>
+InvertedListScanner* make_spectral_hash_scanner_fixSL(
+        int code_size,
+        const IndexIVFSpectralHash* index,
+        bool store_pairs);
+/** @} */
+
+/** @name IndexPQ polysemous dispatch
+ * @{ */
+template <SIMDLevel SL>
+size_t polysemous_inner_loop_fixSL(
+        int code_size,
+        const IndexPQ* index,
+        const float* dis_table_qi,
+        const uint8_t* q_code,
+        size_t k,
+        float* heap_dis,
+        int64_t* heap_ids,
+        int ht);
+/** @} */
+
+} // namespace faiss

data/vendor/faiss/faiss/impl/binary_hamming/neon.cpp

@@ -0,0 +1,26 @@
+/*
+ * 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_SIMD_LEVEL SIMDLevel::ARM_NEON
+
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/utils/hamming_distance/hamming_computer-neon.h>
+
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/impl/binary_hamming/IndexBinaryHNSW_impl.h>
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/impl/binary_hamming/IndexBinaryHash_impl.h>
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/impl/binary_hamming/IndexBinaryIVF_impl.h>
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/impl/binary_hamming/IndexIVFSpectralHash_impl.h>
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/impl/binary_hamming/IndexPQ_impl.h>
+
+#endif // COMPILE_SIMD_ARM_NEON

data/vendor/faiss/faiss/impl/binary_hamming/rvv.cpp

@@ -0,0 +1,26 @@
+/*
+ * 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_RISCV_RVV
+
+#define THE_SIMD_LEVEL SIMDLevel::RISCV_RVV
+
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/utils/hamming_distance/hamming_computer-rvv.h>
+
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/impl/binary_hamming/IndexBinaryHNSW_impl.h>
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/impl/binary_hamming/IndexBinaryHash_impl.h>
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/impl/binary_hamming/IndexBinaryIVF_impl.h>
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/impl/binary_hamming/IndexIVFSpectralHash_impl.h>
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/impl/binary_hamming/IndexPQ_impl.h>
+
+#endif // COMPILE_SIMD_RISCV_RVV

data/vendor/faiss/faiss/impl/expanded_scanners.h

@@ -0,0 +1,163 @@
+/*
+ * 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 <cstdio>
+
+#include <faiss/IndexIVF.h>
+#include <faiss/impl/InvertedListScannerStats.h>
+#include <faiss/impl/ResultHandler.h>
+
+/* This is the inner loop of the inverted list scanners. The default version
+ * that is defined in IndexIVF.cpp works fine but it cannot inline the distance
+ * computation code by calling one or another of the run_scan_codes_* variants
+ * with the exact ScannerType and by setting distance_to_code to be a final
+ * function, the code can be inlined. The speed difference matters for very
+ * small distance computations (eg. SQ or Flat) */
+
+namespace faiss {
+
+namespace {
+
+template <class ScannerType, typename C, bool store_pairs, bool use_sel>
+size_t run_scan_codes1(
+        const ScannerType& scanner,
+        size_t list_size,
+        const uint8_t* codes,
+        const idx_t* ids,
+        ResultHandler& handler) {
+    size_t nup = 0;
+    size_t list_no = scanner.list_no;
+    size_t code_size = scanner.code_size;
+    const IDSelector* sel = scanner.sel;
+    float threshold = handler.threshold;
+    for (size_t j = 0; j < list_size; j++) {
+        if (use_sel) {
+            int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
+            // skip code without computing distance
+            if (!sel->is_member(id)) {
+                codes += code_size;
+                continue;
+            }
+        }
+
+        // post-IDSelector: distance is about to be computed for this code.
+        handler.stats.scan_cnt++;
+        float dis = scanner.distance_to_code(codes); // will be inlined if final
+        if (C::cmp(threshold, dis)) {
+            int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
+            if (handler.add_result(dis, id)) {
+                handler.stats.nheap_updates++;
+                nup++;
+                threshold = handler.threshold;
+            }
+        }
+        codes += code_size;
+    }
+
+    return nup;
+}
+
+/*****************************************************************************
+ * The following functions dispatch runtime parameters to templates, with
+ * possibly some already-fixed templates.
+ */
+
+template <bool store_pairs, bool use_sel, class ScannerType>
+size_t run_scan_codes_fix_store_pairs_fix_use_sel(
+        const ScannerType& scanner,
+        size_t list_size,
+        const uint8_t* codes,
+        const idx_t* ids,
+        ResultHandler& handler) {
+    if (scanner.keep_max) {
+        return run_scan_codes1<
+                ScannerType,
+                CMin<float, idx_t>,
+                store_pairs,
+                use_sel>(scanner, list_size, codes, ids, handler);
+    } else {
+        return run_scan_codes1<
+                ScannerType,
+                CMax<float, idx_t>,
+                store_pairs,
+                use_sel>(scanner, list_size, codes, ids, handler);
+    }
+}
+
+template <class C, bool use_sel, class ScannerType>
+size_t run_scan_codes_fix_C_fix_use_sel(
+        const ScannerType& scanner,
+        size_t list_size,
+        const uint8_t* codes,
+        const idx_t* ids,
+        ResultHandler& handler) {
+    if (scanner.store_pairs) {
+        return run_scan_codes1<ScannerType, C, true, use_sel>(
+                scanner, list_size, codes, ids, handler);
+    } else {
+        return run_scan_codes1<ScannerType, C, false, use_sel>(
+                scanner, list_size, codes, ids, handler);
+    }
+}
+
+template <class C, class ScannerType>
+size_t run_scan_codes_fix_C(
+        const ScannerType& scanner,
+        size_t list_size,
+        const uint8_t* codes,
+        const idx_t* ids,
+        ResultHandler& handler) {
+    if (scanner.sel) {
+        if (scanner.store_pairs) {
+            return run_scan_codes1<ScannerType, C, true, true>(
+                    scanner, list_size, codes, ids, handler);
+        } else {
+            return run_scan_codes1<ScannerType, C, false, true>(
+                    scanner, list_size, codes, ids, handler);
+        }
+    } else {
+        if (scanner.store_pairs) {
+            return run_scan_codes1<ScannerType, C, true, false>(
+                    scanner, list_size, codes, ids, handler);
+        } else {
+            return run_scan_codes1<ScannerType, C, false, false>(
+                    scanner, list_size, codes, ids, handler);
+        }
+    }
+}
+
+template <class ScannerType>
+size_t run_scan_codes(
+        const ScannerType& scanner,
+        size_t list_size,
+        const uint8_t* codes,
+        const idx_t* ids,
+        ResultHandler& handler) {
+    if (scanner.sel == nullptr) {
+        if (scanner.store_pairs) {
+            return run_scan_codes_fix_store_pairs_fix_use_sel<true, false>(
+                    scanner, list_size, codes, ids, handler);
+        } else {
+            return run_scan_codes_fix_store_pairs_fix_use_sel<false, false>(
+                    scanner, list_size, codes, ids, handler);
+        }
+    } else {
+        if (scanner.store_pairs) {
+            return run_scan_codes_fix_store_pairs_fix_use_sel<true, true>(
+                    scanner, list_size, codes, ids, handler);
+        } else {
+            return run_scan_codes_fix_store_pairs_fix_use_sel<false, true>(
+                    scanner, list_size, codes, ids, handler);
+        }
+    }
+}
+
+} // anonymous namespace
+
+} // namespace faiss

data/vendor/faiss/faiss/impl/{FastScanDistancePostProcessing.h → fast_scan/FastScanDistancePostProcessing.h}

@@ -11,9 +11,6 @@
 
 namespace faiss {
 
-// Forward declarations
-struct NormTableScaler;
-
 namespace rabitq_utils {
 struct QueryFactorsData;
 }
@@ -22,8 +19,10 @@ struct QueryFactorsData;
  * Simple context object that holds processors for FastScan operations.
  * */
 struct FastScanDistancePostProcessing {
-    /// Norm scaling processor for Additive Quantizers (nullptr if not needed)
-    const NormTableScaler* norm_scaler = nullptr;
+    /// Norm scaling processor for Additive Quantizers.
+    /// The scale is encoded in a 2x4 bit PQ table, then scaled by this int.
+    /// Set to 0 if unused.
+    int pq2x4_scale = 0;
 
     /// Query factors data pointer for RaBitQ (nullptr if not needed)
     /// This pointer should point to the beginning of the relevant
@@ -36,12 +35,20 @@ struct FastScanDistancePostProcessing {
     /// Set to 0 to use index->nprobe as fallback.
     size_t nprobe = 0;
 
+    /// RaBitQ query quantization bits override.
+    /// Set to 0 to use the index default (index->qb).
+    uint8_t qb = 0;
+
+    /// RaBitQ centered scalar quantizer override.
+    /// Only used when qb > 0 (i.e., when params are overridden).
+    bool centered = false;
+
     /// Default constructor - no processing
     FastScanDistancePostProcessing() = default;
 
     /// Check if norm scaling is enabled
     bool has_norm_scaling() const {
-        return norm_scaler != nullptr;
+        return pq2x4_scale != 0;
     }
 
     /// Check if query factors processing is enabled

data/vendor/faiss/faiss/impl/{LookupTableScaler.h → fast_scan/LookupTableScaler.h}

@@ -10,7 +10,7 @@
 #include <cstdint>
 #include <cstdlib>
 
-#include <faiss/utils/simdlib.h>
+#include <faiss/impl/simdlib/simdlib_dispatch.h>
 
 /*******************************************
  * The Scaler objects are used to specialize the handling of the
@@ -20,8 +20,15 @@
 namespace faiss {
 
 /// no-op handler
+template <SIMDLevel SL = SINGLE_SIMD_LEVEL>
 struct DummyScaler {
     static constexpr int nscale = 0;
+    static constexpr SIMDLevel SL256 = simd256_level_selector<SL>::value;
+    static constexpr SIMDLevel SL512 = simd512_level_selector<SL>::value;
+    using simd32uint8 = simd32uint8_tpl<SL256>;
+    using simd16uint16 = simd16uint16_tpl<SL256>;
+    using simd64uint8 = simd64uint8_tpl<SL512>;
+    using simd32uint16 = simd32uint16_tpl<SL512>;
 
     inline simd32uint8 lookup(const simd32uint8&, const simd32uint8&) const {
         FAISS_THROW_MSG("DummyScaler::lookup should not be called.");
@@ -38,7 +45,6 @@ struct DummyScaler {
         return simd16uint16(0);
     }
 
-#ifdef __AVX512F__
     inline simd64uint8 lookup(const simd64uint8&, const simd64uint8&) const {
         FAISS_THROW_MSG("DummyScaler::lookup should not be called.");
         return simd64uint8(0);
@@ -53,7 +59,6 @@ struct DummyScaler {
         FAISS_THROW_MSG("DummyScaler::scale_hi should not be called.");
         return simd32uint16(0);
     }
-#endif
 
     template <class dist_t>
     inline dist_t scale_one(const dist_t&) const {
@@ -64,8 +69,16 @@ struct DummyScaler {
 
 /// consumes 2x4 bits to encode a norm as a scalar additive quantizer
 /// the norm is scaled because its range is larger than other components
+template <SIMDLevel SL = SINGLE_SIMD_LEVEL>
 struct NormTableScaler {
     static constexpr int nscale = 2;
+    static constexpr SIMDLevel SL256 = simd256_level_selector<SL>::value;
+    static constexpr SIMDLevel SL512 = simd512_level_selector<SL>::value;
+    using simd32uint8 = simd32uint8_tpl<SL256>;
+    using simd16uint16 = simd16uint16_tpl<SL256>;
+    using simd64uint8 = simd64uint8_tpl<SL512>;
+    using simd32uint16 = simd32uint16_tpl<SL512>;
+
     int scale_int;
     simd16uint16 scale_simd;
 
@@ -84,7 +97,6 @@ struct NormTableScaler {
         return (simd16uint16(res) >> 8) * scale_simd;
     }
 
-#ifdef __AVX512F__
     inline simd64uint8 lookup(const simd64uint8& lut, const simd64uint8& c)
             const {
         return lut.lookup_4_lanes(c);
@@ -99,7 +111,6 @@ struct NormTableScaler {
         auto scale_simd_wide = simd32uint16(scale_simd, scale_simd);
         return (simd32uint16(res) >> 8) * scale_simd_wide;
     }
-#endif
 
     // for non-SIMD implem 2, 3, 4
     template <class dist_t>