faiss 0.6.0 → 0.6.1

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

@@ -22,6 +22,190 @@ static float sqr(float x) {
     return x * x;
 }
 
+constexpr size_t kTurboQuantMaxBits = 8;
+// TurboQuant builds a 1-D optimal scalar quantizer analytically. We approximate
+// the target density on a uniform grid over [-1, 1]; the grid is kept dense
+// enough both in absolute terms and per output centroid.
+constexpr size_t kTurboQuantGridMin = 1 << 15;
+constexpr size_t kTurboQuantGridPerCentroid = 512;
+constexpr int kTurboQuantMaxIter = 100;
+constexpr double kTurboQuantTol = 1e-8;
+
+void build_TurboQuantMSECodebook(
+        size_t d,
+        size_t nbits,
+        std::vector<float>& centroids,
+        std::vector<float>& boundaries) {
+    FAISS_THROW_IF_NOT_FMT(
+            nbits <= kTurboQuantMaxBits,
+            "invalid TurboQuant nbits %zu (must be in [0, %zu])",
+            nbits,
+            kTurboQuantMaxBits);
+
+    if (nbits == 0) {
+        centroids.clear();
+        boundaries.clear();
+        return;
+    }
+
+    const size_t k = size_t(1) << nbits;
+
+    if (d == 1) {
+        // In 1-D, a unit vector can only be -1 or +1, so the marginal
+        // distribution collapses to two atoms. The TurboQuant codebook is
+        // therefore a repeated pair of endpoint centroids.
+        centroids.resize(k);
+        for (size_t i = 0; i < k; i++) {
+            centroids[i] = i < k / 2 ? -1.0f : 1.0f;
+        }
+        boundaries.resize(k - 1);
+        for (size_t i = 0; i + 1 < k; i++) {
+            boundaries[i] = 0.5f * (centroids[i] + centroids[i + 1]);
+        }
+        return;
+    }
+
+    // For d > 1, TurboQuant uses the marginal distribution of one coordinate of
+    // a random unit vector in R^d. On [-1, 1], this density is proportional to
+    // (1 - x^2)^((d - 3) / 2), which is a symmetric beta-law after a change of
+    // variables. The code below discretizes that density.
+    const size_t ngrid =
+            std::max(kTurboQuantGridMin, k * kTurboQuantGridPerCentroid);
+    const double step = 2.0 / ngrid;
+    const double alpha = 0.5 * (double(d) - 3.0);
+
+    std::vector<double> xs(ngrid);
+    // prefix_w stores the cumulative mass of the discretized density and
+    // prefix_wx stores its cumulative first moment, so interval means can be
+    // recovered in O(1).
+    std::vector<double> prefix_w(ngrid + 1, 0.0);
+    std::vector<double> prefix_wx(ngrid + 1, 0.0);
+
+    for (size_t i = 0; i < ngrid; i++) {
+        const double x = -1.0 + (i + 0.5) * step;
+        const double one_minus_x2 = std::max(0.0, 1.0 - x * x);
+        double w;
+        if (alpha == 0.0) { // when d == 3
+            w = 1.0;
+        } else {
+            // (1-x^2)^((d-3)/2)
+            w = std::pow(one_minus_x2, alpha);
+        }
+        if (!std::isfinite(w) || w < 0.0) {
+            w = 0.0;
+        }
+        xs[i] = x;
+        prefix_w[i + 1] = prefix_w[i] + w;
+        prefix_wx[i + 1] = prefix_wx[i] + w * x;
+    }
+
+    auto range_mean = [&](size_t i0, size_t i1, double fallback) {
+        const double w = prefix_w[i1] - prefix_w[i0];
+        if (w <= 0.0) {
+            return fallback;
+        }
+        return (prefix_wx[i1] - prefix_wx[i0]) / w;
+    };
+
+    const double total_w = prefix_w.back();
+    std::vector<size_t> cuts(k + 1, 0);
+    cuts[k] = ngrid;
+
+    // Initialize with k equal-mass cells under the target density. This gives
+    // a stable starting point before the Lloyd refinements below.
+    for (size_t i = 1; i < k; i++) {
+        const double target = total_w * i / k;
+        cuts[i] = std::lower_bound(prefix_w.begin(), prefix_w.end(), target) -
+                prefix_w.begin();
+        cuts[i] = std::min(cuts[i], ngrid);
+    }
+
+    std::vector<double> centroids_d(k);
+    for (size_t i = 0; i < k; i++) {
+        const double left = -1.0 + 2.0 * i / k;
+        const double right = -1.0 + 2.0 * (i + 1) / k;
+        // First estimate of each centroid: the conditional mean of its initial
+        // equal-mass cell, with a uniform-cell midpoint as a fallback.
+        centroids_d[i] = range_mean(cuts[i], cuts[i + 1], 0.5 * (left + right));
+    }
+
+    std::vector<double> boundaries_d(k > 0 ? k - 1 : 0);
+
+    // Refine the 1-D codebook with a weighted Lloyd iteration over the
+    // discretized marginal density on [-1, 1]:
+    // 1. boundaries_d are the Voronoi separators implied by neighboring
+    //    centroids.
+    // 2. cuts map each boundary interval back to a contiguous range of the
+    //    integration grid xs[].
+    // 3. each centroid becomes the weighted mean of the samples currently in
+    //    its cell, clipped to stay within its neighboring boundaries.
+    //
+    // The loop stops once the largest centroid update is below kTurboQuantTol.
+    for (int iter = 0; iter < kTurboQuantMaxIter; iter++) {
+        // Midpoints between adjacent centroids define the current Voronoi
+        // partition of [-1, 1].
+        for (size_t i = 0; i + 1 < k; i++) {
+            boundaries_d[i] = 0.5 * (centroids_d[i] + centroids_d[i + 1]);
+        }
+
+        cuts[0] = 0;
+        cuts[k] = ngrid;
+        // Reassign the discretized density samples to the Voronoi cell induced
+        // by each boundary. Because xs is sorted, the reassignment reduces to
+        // finding the first grid point strictly greater than each boundary.
+        for (size_t i = 1; i < k; i++) {
+            cuts[i] = std::upper_bound(
+                              xs.begin(), xs.end(), boundaries_d[i - 1]) -
+                    xs.begin();
+        }
+
+        double max_delta = 0.0;
+        for (size_t i = 0; i < k; i++) {
+            const double left = i == 0 ? -1.0 : boundaries_d[i - 1];
+            const double right = i + 1 == k ? 1.0 : boundaries_d[i];
+            // Lloyd update: replace the centroid with the weighted average of
+            // the mass assigned to its cell. Empty cells fall back to the cell
+            // midpoint, and we clamp to [left, right] to preserve ordering.
+            double c = range_mean(cuts[i], cuts[i + 1], 0.5 * (left + right));
+            c = std::min(std::max(c, left), right);
+            max_delta = std::max(max_delta, std::abs(c - centroids_d[i]));
+            centroids_d[i] = c;
+        }
+
+        if (max_delta < kTurboQuantTol) {
+            break;
+        }
+    }
+
+    std::sort(centroids_d.begin(), centroids_d.end());
+
+    centroids.resize(k);
+    boundaries.resize(k - 1);
+    for (size_t i = 0; i < k; i++) {
+        centroids[i] = centroids_d[i];
+    }
+    for (size_t i = 0; i + 1 < k; i++) {
+        boundaries[i] = 0.5f * (centroids[i] + centroids[i + 1]);
+    }
+}
+
+void train_TurboQuantMSE(size_t d, size_t nbits, std::vector<float>& trained) {
+    FAISS_THROW_IF_NOT_FMT(
+            nbits > 0, "invalid TurboQuant SQ nbits %zu (must be > 0)", nbits);
+    std::vector<float> centroids;
+    std::vector<float> boundaries;
+    build_TurboQuantMSECodebook(d, nbits, centroids, boundaries);
+    const size_t k = centroids.size();
+
+    trained.resize(k + (k - 1));
+    for (size_t i = 0; i < k; i++) {
+        trained[i] = centroids[i];
+    }
+    for (size_t i = 0; i + 1 < k; i++) {
+        trained[k + i] = boundaries[i];
+    }
+}
+
 void train_Uniform(
         RangeStat rs,
         float rs_arg,
@@ -37,7 +221,7 @@ void train_Uniform(
     if (rs == ScalarQuantizer::RS_minmax) {
         vmin = HUGE_VAL;
         vmax = -HUGE_VAL;
-        for (size_t i = 0; i < n; i++) {
+        for (idx_t i = 0; i < n; i++) {
             if (x[i] < vmin) {
                 vmin = x[i];
             }
@@ -50,7 +234,7 @@ void train_Uniform(
         vmax += vexp;
     } else if (rs == ScalarQuantizer::RS_meanstd) {
         double sum = 0, sum2 = 0;
-        for (size_t i = 0; i < n; i++) {
+        for (idx_t i = 0; i < n; i++) {
             sum += x[i];
             sum2 += x[i] * x[i];
         }
@@ -81,7 +265,7 @@ void train_Uniform(
         float sx = 0;
         {
             vmin = HUGE_VAL, vmax = -HUGE_VAL;
-            for (size_t i = 0; i < n; i++) {
+            for (idx_t i = 0; i < n; i++) {
                 if (x[i] < vmin) {
                     vmin = x[i];
                 }
@@ -161,9 +345,9 @@ void train_NonUniform(
     if (rs == ScalarQuantizer::RS_minmax) {
         memcpy(vmin, x, sizeof(*x) * d);
         memcpy(vmax, x, sizeof(*x) * d);
-        for (size_t i = 1; i < n; i++) {
+        for (idx_t i = 1; i < n; i++) {
             const float* xi = x + i * d;
-            for (size_t j = 0; j < d; j++) {
+            for (int j = 0; j < d; j++) {
                 if (xi[j] < vmin[j]) {
                     vmin[j] = xi[j];
                 }
@@ -173,7 +357,7 @@ void train_NonUniform(
             }
         }
         float* vdiff = vmax;
-        for (size_t j = 0; j < d; j++) {
+        for (int j = 0; j < d; j++) {
             float vexp = (vmax[j] - vmin[j]) * rs_arg;
             vmin[j] -= vexp;
             vmax[j] += vexp;
@@ -182,9 +366,9 @@ void train_NonUniform(
     } else {
         // transpose
         std::vector<float> xt(n * d);
-        for (size_t i = 1; i < n; i++) {
+        for (idx_t i = 1; i < n; i++) {
             const float* xi = x + i * d;
-            for (size_t j = 0; j < d; j++) {
+            for (int j = 0; j < d; j++) {
                 xt[j * n + i] = xi[j];
             }
         }

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

@@ -37,6 +37,18 @@ void train_NonUniform(
         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

@@ -23,86 +23,107 @@
 
 namespace faiss {
 
-/*********************** x86 SIMD dispatch cases */
+/** Defining which SIMD levels are available for a given function is via a
+ * binary mask. Here we predefine the most common masks.
+ *  */
 
-#ifdef COMPILE_SIMD_AVX2
-#define DISPATCH_SIMDLevel_AVX2(f, ...) \
-    case SIMDLevel::AVX2:               \
-        return f<SIMDLevel::AVX2>(__VA_ARGS__)
-#else
-#define DISPATCH_SIMDLevel_AVX2(f, ...)
-#endif
+constexpr int AVAILABLE_SIMD_LEVELS_NONE = (1 << int(SIMDLevel::NONE));
 
-#ifdef COMPILE_SIMD_AVX512
-#define DISPATCH_SIMDLevel_AVX512(f, ...) \
-    case SIMDLevel::AVX512:               \
-        return f<SIMDLevel::AVX512>(__VA_ARGS__)
-#else
-#define DISPATCH_SIMDLevel_AVX512(f, ...)
-#endif
+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
-#define DISPATCH_SIMDLevel_AVX512_SPR(f, ...) \
-    case SIMDLevel::AVX512_SPR:               \
-        return f<SIMDLevel::AVX512_SPR>(__VA_ARGS__)
-#else
-#define DISPATCH_SIMDLevel_AVX512_SPR(f, ...)
+        case SIMDLevel::AVX512_SPR:
+            if constexpr (
+                    available_levels & (1 << int(SIMDLevel::AVX512_SPR))) {
+                return action.template operator()<SIMDLevel::AVX512_SPR>();
+            }
+            [[fallthrough]];
 #endif
 
-/*********************** ARM SIMD dispatch cases */
+#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_ARM_NEON
-#define DISPATCH_SIMDLevel_ARM_NEON(f, ...) \
-    case SIMDLevel::ARM_NEON:               \
-        return f<SIMDLevel::ARM_NEON>(__VA_ARGS__)
-#else
-#define DISPATCH_SIMDLevel_ARM_NEON(f, ...)
+#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
-#define DISPATCH_SIMDLevel_ARM_SVE(f, ...) \
-    case SIMDLevel::ARM_SVE:               \
-        return f<SIMDLevel::ARM_SVE>(__VA_ARGS__)
-#else
-#define DISPATCH_SIMDLevel_ARM_SVE(f, ...)
+        case SIMDLevel::ARM_SVE:
+            if constexpr (available_levels & (1 << int(SIMDLevel::ARM_SVE))) {
+                return action.template operator()<SIMDLevel::ARM_SVE>();
+            }
+            [[fallthrough]];
 #endif
 
-/*********************** Main dispatch macro */
-
-#ifdef FAISS_ENABLE_DD
+#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
 
-// DD mode: runtime dispatch based on SIMDConfig::level
-#define DISPATCH_SIMDLevel(f, ...)                         \
-    switch (SIMDConfig::level) {                           \
-        case SIMDLevel::NONE:                              \
-            return f<SIMDLevel::NONE>(__VA_ARGS__);        \
-            DISPATCH_SIMDLevel_AVX2(f, __VA_ARGS__);       \
-            DISPATCH_SIMDLevel_AVX512(f, __VA_ARGS__);     \
-            DISPATCH_SIMDLevel_AVX512_SPR(f, __VA_ARGS__); \
-            DISPATCH_SIMDLevel_ARM_NEON(f, __VA_ARGS__);   \
-            DISPATCH_SIMDLevel_ARM_SVE(f, __VA_ARGS__);    \
-        default:                                           \
-            FAISS_THROW_MSG("Invalid SIMD level");         \
+#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>();
     }
-
-#else // Static mode
-
-// Static mode: direct call to compiled-in SIMD level (no runtime switch)
-#if defined(COMPILE_SIMD_AVX512_SPR)
-#define DISPATCH_SIMDLevel(f, ...) return f<SIMDLevel::AVX512_SPR>(__VA_ARGS__)
-#elif defined(COMPILE_SIMD_AVX512)
-#define DISPATCH_SIMDLevel(f, ...) return f<SIMDLevel::AVX512>(__VA_ARGS__)
-#elif defined(COMPILE_SIMD_AVX2)
-#define DISPATCH_SIMDLevel(f, ...) return f<SIMDLevel::AVX2>(__VA_ARGS__)
-#elif defined(COMPILE_SIMD_ARM_SVE)
-#define DISPATCH_SIMDLevel(f, ...) return f<SIMDLevel::ARM_SVE>(__VA_ARGS__)
-#elif defined(COMPILE_SIMD_ARM_NEON)
-#define DISPATCH_SIMDLevel(f, ...) return f<SIMDLevel::ARM_NEON>(__VA_ARGS__)
-#else
-#define DISPATCH_SIMDLevel(f, ...) return f<SIMDLevel::NONE>(__VA_ARGS__)
 #endif
-
-#endif // FAISS_ENABLE_DD
+}
 
 /**
  * Dispatch to a lambda with SIMDLevel as a compile-time constant.
@@ -126,6 +147,8 @@ namespace faiss {
  *   });
  *
  * 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()()`
@@ -133,7 +156,18 @@ namespace faiss {
  */
 template <typename LambdaType>
 inline auto with_simd_level(LambdaType&& action) {
-    DISPATCH_SIMDLevel(action.template operator());
+    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