faiss 0.5.3 → 0.6.1

data/vendor/faiss/faiss/impl/scalar_quantizer/training.h

@@ -0,0 +1,54 @@
+/*
+ * 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
+
+/*******************************************************************
+ * Quantizer range training for the scalar quantizer. This is independent of the
+ * searching code and needs not to be very optimized (scalar quantizer training
+ * is very efficient).
+ */
+
+#include <faiss/impl/ScalarQuantizer.h>
+
+namespace faiss {
+
+namespace scalar_quantizer {
+
+using RangeStat = ScalarQuantizer::RangeStat;
+
+void train_Uniform(
+        RangeStat rs,
+        float rs_arg,
+        idx_t n,
+        int k,
+        const float* x,
+        std::vector<float>& trained);
+
+void train_NonUniform(
+        RangeStat rs,
+        float rs_arg,
+        idx_t n,
+        int d,
+        int k,
+        const float* x,
+        std::vector<float>& trained);
+
+/** Build the TurboQuant MSE codebook using the beta-distribution-optimal
+ *  quantizer from the TurboQuant paper. The codebook is analytical
+ *  (depends only on d and nbits, no training data needed).
+ *
+ *  @param d         vector dimensionality (used for beta-distribution shape)
+ *  @param nbits     bits per component (1-8)
+ *  @param trained   output: [centroids (k floats), boundaries (k-1 floats)]
+ *                   where k = 2^nbits
+ */
+void train_TurboQuantMSE(size_t d, size_t nbits, std::vector<float>& trained);
+
+} // namespace scalar_quantizer
+
+} // namespace faiss

data/vendor/faiss/faiss/impl/simd_dispatch.h

@@ -0,0 +1,173 @@
+/*
+ * 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 simd_dispatch.h
+ * @brief Internal dispatch macros for SIMD level selection.
+ *
+ * This is a PRIVATE header - do not include in public APIs or user code.
+ * Only faiss internal .cpp files should include this header.
+ *
+ * For the public API (SIMDLevel enum, SIMDConfig class), use:
+ *   #include <faiss/utils/simd_levels.h>
+ */
+
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/utils/simd_levels.h>
+
+namespace faiss {
+
+/** Defining which SIMD levels are available for a given function is via a
+ * binary mask. Here we predefine the most common masks.
+ *  */
+
+constexpr int AVAILABLE_SIMD_LEVELS_NONE = (1 << int(SIMDLevel::NONE));
+
+constexpr int AVAILABLE_SIMD_LEVELS_AVX2_NEON = AVAILABLE_SIMD_LEVELS_NONE |
+        (1 << int(SIMDLevel::AVX2)) | (1 << int(SIMDLevel::ARM_NEON));
+
+// A0: same + AVX512 + RISCV_RVV
+constexpr int AVAILABLE_SIMD_LEVELS_A0 = AVAILABLE_SIMD_LEVELS_AVX2_NEON |
+        (1 << int(SIMDLevel::AVX512)) | (1 << int(SIMDLevel::RISCV_RVV));
+
+// A1: same + ARM_SVE (for functions with dedicated SVE implementations)
+constexpr int AVAILABLE_SIMD_LEVELS_A1 =
+        AVAILABLE_SIMD_LEVELS_A0 | (1 << int(SIMDLevel::ARM_SVE));
+
+// A2: NONE + AVX2 + ARM_SVE only (for functions with only these
+// implementations)
+constexpr int AVAILABLE_SIMD_LEVELS_A2 = AVAILABLE_SIMD_LEVELS_NONE |
+        (1 << int(SIMDLevel::AVX2)) | (1 << int(SIMDLevel::ARM_SVE));
+
+constexpr int AVAILABLE_SIMD_LEVELS_ALL = -1;
+
+/** The complete dispatching function. It takes into account:
+ * - the currently selected SIMD level
+ * - the compiled in SIMD levels (given by COMPILE_SIMD_XXX)
+ * - the available SIMD implementations for that particular function (given by
+ * available_levels)
+ */
+
+template <int available_levels, typename LambdaType>
+inline auto with_selected_simd_levels(LambdaType&& action) {
+#ifdef FAISS_ENABLE_DD
+    switch (SIMDConfig::level) {
+        // For x86 -- try from highest to lowest level
+
+#ifdef COMPILE_SIMD_AVX512_SPR
+        case SIMDLevel::AVX512_SPR:
+            if constexpr (
+                    available_levels & (1 << int(SIMDLevel::AVX512_SPR))) {
+                return action.template operator()<SIMDLevel::AVX512_SPR>();
+            }
+            [[fallthrough]];
+#endif
+
+#ifdef COMPILE_SIMD_AVX512
+        case SIMDLevel::AVX512:
+            if constexpr (available_levels & (1 << int(SIMDLevel::AVX512))) {
+                return action.template operator()<SIMDLevel::AVX512>();
+            }
+            [[fallthrough]];
+#endif
+
+#ifdef COMPILE_SIMD_AVX2
+        case SIMDLevel::AVX2:
+            if constexpr (available_levels & (1 << int(SIMDLevel::AVX2))) {
+                return action.template operator()<SIMDLevel::AVX2>();
+            }
+            [[fallthrough]];
+#endif
+
+            // For ARM, try from highest to lowest level
+#ifdef COMPILE_SIMD_ARM_SVE
+        case SIMDLevel::ARM_SVE:
+            if constexpr (available_levels & (1 << int(SIMDLevel::ARM_SVE))) {
+                return action.template operator()<SIMDLevel::ARM_SVE>();
+            }
+            [[fallthrough]];
+#endif
+
+#ifdef COMPILE_SIMD_ARM_NEON
+        case SIMDLevel::ARM_NEON:
+            if constexpr (available_levels & (1 << int(SIMDLevel::ARM_NEON))) {
+                return action.template operator()<SIMDLevel::ARM_NEON>();
+            }
+            [[fallthrough]];
+#endif
+
+#ifdef COMPILE_SIMD_RISCV_RVV
+        case SIMDLevel::RISCV_RVV:
+            if constexpr (available_levels & (1 << int(SIMDLevel::RISCV_RVV))) {
+                return action.template operator()<SIMDLevel::RISCV_RVV>();
+            }
+            [[fallthrough]];
+#endif
+        default:
+            return action.template operator()<SIMDLevel::NONE>();
+    }
+#else // static dispatch
+    // In static mode, SINGLE_SIMD_LEVEL is a constexpr resolved at compile
+    // time. If the compiled level is not in the available set, fall through
+    // to NONE (mirroring the DD fallthrough behavior). Only SINGLE_SIMD_LEVEL
+    // and NONE have compiled specializations.
+    if constexpr (available_levels & (1 << int(SINGLE_SIMD_LEVEL))) {
+        return action.template operator()<SINGLE_SIMD_LEVEL>();
+    } else {
+        return action.template operator()<SIMDLevel::NONE>();
+    }
+#endif
+}
+
+/**
+ * Dispatch to a lambda with SIMDLevel as a compile-time constant.
+ *
+ * This function calls the provided templated lambda with the current
+ * runtime SIMD level (from SIMDConfig::level) as a compile-time template
+ * argument. This enables SIMD-specialized code paths while keeping the
+ * dispatch logic centralized.
+ *
+ * The key benefit is that the SIMD dispatch happens once, outside any loops,
+ * so the loop body runs with the optimal SIMD implementation without
+ * per-iteration dispatch overhead.
+ *
+ * Example with a loop (the dispatch happens once, not per iteration):
+ *
+ *   std::vector<float> distances(n);
+ *   with_simd_level([&]<SIMDLevel level>() {
+ *       for (size_t i = 0; i < n; i++) {
+ *           distances[i] = fvec_L2sqr<level>(query, vectors + i * d, d);
+ *       }
+ *   });
+ *
+ * The lambda must be a generic lambda with a SIMDLevel template parameter.
+ * By default, the lambda uses levels AVX2 + AVX512 + NEON + RVV, since these
+ * are the most common cases.
+ *
+ * @param action A generic lambda with signature `template<SIMDLevel> T
+ * operator()()`
+ * @return The return value of the lambda
+ */
+template <typename LambdaType>
+inline auto with_simd_level(LambdaType&& action) {
+    return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A0>(
+            std::forward<LambdaType>(action));
+}
+
+/**
+ * Use for functions implemented with simdXintY (256-bit) operations
+ * that don't have dedicated AVX512 or SVE implementations.
+ */
+template <typename LambdaType>
+inline auto with_simd_level_256bit(LambdaType&& action) {
+    return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_AVX2_NEON>(
+            std::forward<LambdaType>(action));
+}
+
+} // namespace faiss

data/vendor/faiss/faiss/impl/simdlib/simdlib.h

@@ -0,0 +1,57 @@
+/*
+ * 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
+
+/** Abstractions for 256-bit and 512-bit SIMD registers.
+ *
+ * The objective is to separate the different interpretations of the same
+ * registers (as a vector of uint8, uint16 or uint32), to provide printing
+ * functions.
+ *
+ * The types are templatized on SIMDLevel. Each platform header provides
+ * explicit specializations for the appropriate level. Code without explicit
+ * SL context uses SINGLE_SIMD_LEVEL (see simd_levels.h).
+ */
+
+#include <faiss/utils/simd_levels.h>
+
+namespace faiss {
+
+// 256-bit primary templates
+template <SIMDLevel SL>
+struct simd256bit_tpl {};
+template <SIMDLevel SL>
+struct simd16uint16_tpl : simd256bit_tpl<SL> {};
+template <SIMDLevel SL>
+struct simd32uint8_tpl : simd256bit_tpl<SL> {};
+template <SIMDLevel SL>
+struct simd8uint32_tpl : simd256bit_tpl<SL> {};
+template <SIMDLevel SL>
+struct simd8float32_tpl : simd256bit_tpl<SL> {};
+
+// 512-bit primary templates
+template <SIMDLevel SL>
+struct simd512bit_tpl {};
+template <SIMDLevel SL>
+struct simd32uint16_tpl : simd512bit_tpl<SL> {};
+template <SIMDLevel SL>
+struct simd64uint8_tpl : simd512bit_tpl<SL> {};
+template <SIMDLevel SL>
+struct simd16float32_tpl : simd512bit_tpl<SL> {};
+
+} // namespace faiss
+
+// NONE specialization — always included.
+// Provides simd16uint16_tpl<NONE> etc. (scalar fallback).
+// On PPC64: uses PPC-optimized scalar code (hand-tuned loop unrolling).
+// Elsewhere: generic scalar implementation.
+#if defined(__PPC64__)
+#include <faiss/impl/simdlib/simdlib_ppc64.h>
+#else
+#include <faiss/impl/simdlib/simdlib_emulated.h>
+#endif