faiss 0.5.3 → 0.6.1

data/vendor/faiss/faiss/impl/scalar_quantizer/sq-rvv.cpp

@@ -0,0 +1,311 @@
+/*
+ * 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
+
+#include <faiss/impl/scalar_quantizer/codecs.h>
+#include <faiss/impl/scalar_quantizer/distance_computers.h>
+#include <faiss/impl/scalar_quantizer/quantizers.h>
+#include <faiss/impl/scalar_quantizer/scanners.h>
+#include <faiss/impl/scalar_quantizer/similarities.h>
+
+#include <riscv_vector.h>
+#include <cmath>
+
+namespace faiss {
+
+namespace scalar_quantizer {
+
+/*************************************************************************
+ * Marker specializations.
+ *
+ * Unlike x86/NEON sq-*.cpp files that expose a fixed 8-wide / 16-wide codec
+ * interface (reconstruct_8_components / reconstruct_16_components), RVV is
+ * variable-width: the native vector length is implementation-defined and
+ * queried at runtime via __riscv_vsetvl. Forcing RVV into a fixed-width
+ * codec would leave performance on the table on wider hardware.
+ *
+ * So the strategy here is: Codec / Quantizer / Similarity classes for
+ * RISCV_RVV act as opaque TAG TYPES — they only need to be complete types
+ * so that baseline's sq-dispatch.h can form template arguments like
+ * `DCTemplate<QuantizerTemplate<Codec4bit<RISCV_RVV>, UNIFORM, RISCV_RVV>,
+ *             SimilarityL2<RISCV_RVV>, RISCV_RVV>`.
+ *
+ * The real SIMD work lives in full DCTemplate specializations below.
+ * Unspecialized combinations fall through to scalar via the fallback
+ * `DCTemplate<Q, Sim, RISCV_RVV> : DCTemplate<Q, Sim, NONE>`.
+ ************************************************************************/
+
+template <>
+struct Codec8bit<SIMDLevel::RISCV_RVV> : Codec8bit<SIMDLevel::NONE> {};
+
+template <>
+struct Codec4bit<SIMDLevel::RISCV_RVV> : Codec4bit<SIMDLevel::NONE> {};
+
+template <>
+struct Codec6bit<SIMDLevel::RISCV_RVV> : Codec6bit<SIMDLevel::NONE> {};
+
+template <class Codec>
+struct QuantizerTemplate<
+        Codec,
+        QuantizerTemplateScaling::UNIFORM,
+        SIMDLevel::RISCV_RVV>
+        : QuantizerTemplate<
+                  Codec,
+                  QuantizerTemplateScaling::UNIFORM,
+                  SIMDLevel::NONE> {
+    QuantizerTemplate(size_t d, const std::vector<float>& trained)
+            : QuantizerTemplate<
+                      Codec,
+                      QuantizerTemplateScaling::UNIFORM,
+                      SIMDLevel::NONE>(d, trained) {}
+};
+
+template <class Codec>
+struct QuantizerTemplate<
+        Codec,
+        QuantizerTemplateScaling::NON_UNIFORM,
+        SIMDLevel::RISCV_RVV>
+        : QuantizerTemplate<
+                  Codec,
+                  QuantizerTemplateScaling::NON_UNIFORM,
+                  SIMDLevel::NONE> {
+    QuantizerTemplate(size_t d, const std::vector<float>& trained)
+            : QuantizerTemplate<
+                      Codec,
+                      QuantizerTemplateScaling::NON_UNIFORM,
+                      SIMDLevel::NONE>(d, trained) {}
+};
+
+template <>
+struct QuantizerFP16<SIMDLevel::RISCV_RVV> : QuantizerFP16<SIMDLevel::NONE> {
+    QuantizerFP16(size_t d, const std::vector<float>& trained)
+            : QuantizerFP16<SIMDLevel::NONE>(d, trained) {}
+};
+
+template <>
+struct QuantizerBF16<SIMDLevel::RISCV_RVV> : QuantizerBF16<SIMDLevel::NONE> {
+    QuantizerBF16(size_t d, const std::vector<float>& trained)
+            : QuantizerBF16<SIMDLevel::NONE>(d, trained) {}
+};
+
+template <>
+struct Quantizer8bitDirect<SIMDLevel::RISCV_RVV>
+        : Quantizer8bitDirect<SIMDLevel::NONE> {
+    Quantizer8bitDirect(size_t d, const std::vector<float>& trained)
+            : Quantizer8bitDirect<SIMDLevel::NONE>(d, trained) {}
+};
+
+template <>
+struct Quantizer8bitDirectSigned<SIMDLevel::RISCV_RVV>
+        : Quantizer8bitDirectSigned<SIMDLevel::NONE> {
+    Quantizer8bitDirectSigned(size_t d, const std::vector<float>& trained)
+            : Quantizer8bitDirectSigned<SIMDLevel::NONE>(d, trained) {}
+};
+
+template <>
+struct SimilarityL2<SIMDLevel::RISCV_RVV> : SimilarityL2<SIMDLevel::NONE> {
+    using SimilarityL2<SIMDLevel::NONE>::SimilarityL2;
+};
+
+template <>
+struct SimilarityIP<SIMDLevel::RISCV_RVV> : SimilarityIP<SIMDLevel::NONE> {
+    using SimilarityIP<SIMDLevel::NONE>::SimilarityIP;
+};
+
+/*************************************************************************
+ * Fallback DCTemplate / DistanceComputerByte for RISCV_RVV.
+ *
+ * Inheriting from the NONE specialization means every (Quantizer, Similarity)
+ * combination that does NOT have a hand-tuned RVV full specialization below
+ * falls through to scalar code. Callers and the dispatcher don't know or care.
+ ************************************************************************/
+
+template <class Quantizer, class Similarity>
+struct DCTemplate<Quantizer, Similarity, SIMDLevel::RISCV_RVV>
+        : DCTemplate<Quantizer, Similarity, SIMDLevel::NONE> {
+    using Base = DCTemplate<Quantizer, Similarity, SIMDLevel::NONE>;
+    using Base::Base;
+};
+
+template <class Similarity>
+struct DistanceComputerByte<Similarity, SIMDLevel::RISCV_RVV>
+        : DistanceComputerByte<Similarity, SIMDLevel::NONE> {
+    using Base = DistanceComputerByte<Similarity, SIMDLevel::NONE>;
+    using Base::Base;
+};
+
+/*************************************************************************
+ * Fast path — QT_4bit_uniform + L2
+ *
+ * 4-bit UNIFORM scaling: every component reconstructs as an affine function
+ * of the 4-bit code,
+ *     recon(c) = vmin + vdiff * (c + 0.5) / 15 = final_scale * c + bias
+ * where final_scale = vdiff / 15. L2 distance between two reconstructions
+ * therefore reduces to final_scale^2 * (q_c - c_c)^2 over integer codes,
+ * so we can stay in the int domain and pay one float multiply at the end.
+ *
+ * The RVV path pre-nibbles the query into q_lo / q_hi (even / odd lanes)
+ * once at set_query time and then processes native-VL-sized chunks of code
+ * without ever decoding to float.
+ ************************************************************************/
+
+template <>
+struct DCTemplate<
+        QuantizerTemplate<
+                Codec4bit<SIMDLevel::RISCV_RVV>,
+                QuantizerTemplateScaling::UNIFORM,
+                SIMDLevel::RISCV_RVV>,
+        SimilarityL2<SIMDLevel::RISCV_RVV>,
+        SIMDLevel::RISCV_RVV> : SQDistanceComputer {
+    using Sim = SimilarityL2<SIMDLevel::RISCV_RVV>;
+
+    size_t d;
+    float vmin;
+    float vdiff;
+    float final_scale_sq;
+    std::vector<uint8_t> q_lo;
+    std::vector<uint8_t> q_hi;
+
+    DCTemplate(size_t d_in, const std::vector<float>& trained)
+            : d(d_in),
+              vmin(trained[0]),
+              vdiff(trained[1]),
+              q_lo((d_in + 1) / 2, 0),
+              q_hi((d_in + 1) / 2, 0) {
+        const float final_scale = vdiff / 15.0f;
+        final_scale_sq = final_scale * final_scale;
+    }
+
+    void set_query(const float* x) final {
+        this->q = x;
+        const float inv_scale = (vdiff == 0.0f) ? 0.0f : 15.0f / vdiff;
+        for (size_t i = 0; i < d; i++) {
+            float val = (x[i] - vmin) * inv_scale;
+            int code = static_cast<int>(val);
+            if (code < 0) {
+                code = 0;
+            }
+            if (code > 15) {
+                code = 15;
+            }
+            if (i % 2 == 0) {
+                q_lo[i / 2] = static_cast<uint8_t>(code);
+            } else {
+                q_hi[i / 2] = static_cast<uint8_t>(code);
+            }
+        }
+    }
+
+    /// Squared integer-domain L2 between pre-nibbled q and packed 4-bit code.
+    /// Uses RVV's native VL; no fixed width assumptions. Returns the raw
+    /// integer sum — caller multiplies by final_scale_sq.
+    int64_t accumulate_int_l2(const uint8_t* code) const {
+        int64_t acc = 0;
+        size_t i = 0;
+        while (i < d) {
+            // Process up to vl codes per iteration. Each code byte packs two
+            // 4-bit codes, so we load (vl + 1) / 2 bytes; keep vl even to
+            // keep the nibble split aligned with the i % 2 split we used at
+            // set_query time.
+            size_t remaining = d - i;
+            size_t vl = __riscv_vsetvl_e8m1(remaining);
+            if (vl & 1) {
+                vl -= 1; // keep even; tail handled on next iter or scalar
+            }
+            if (vl == 0) {
+                break;
+            }
+            const size_t byte_vl = vl / 2;
+
+            vuint8m1_t packed = __riscv_vle8_v_u8m1(code + i / 2, byte_vl);
+            vuint8m1_t ql = __riscv_vle8_v_u8m1(q_lo.data() + i / 2, byte_vl);
+            vuint8m1_t qh = __riscv_vle8_v_u8m1(q_hi.data() + i / 2, byte_vl);
+
+            vuint8m1_t lo_nib = __riscv_vand_vx_u8m1(packed, 0x0F, byte_vl);
+            vuint8m1_t hi_nib = __riscv_vsrl_vx_u8m1(packed, 4, byte_vl);
+
+            // |ql - lo| and |qh - hi| fit in u8 (values are in [0, 15]).
+            vuint8m1_t d_lo = __riscv_vsub_vv_u8m1(
+                    __riscv_vmaxu_vv_u8m1(ql, lo_nib, byte_vl),
+                    __riscv_vminu_vv_u8m1(ql, lo_nib, byte_vl),
+                    byte_vl);
+            vuint8m1_t d_hi = __riscv_vsub_vv_u8m1(
+                    __riscv_vmaxu_vv_u8m1(qh, hi_nib, byte_vl),
+                    __riscv_vminu_vv_u8m1(qh, hi_nib, byte_vl),
+                    byte_vl);
+
+            // Square via widening multiply (each byte squared fits in u16,
+            // since max byte value is 15 -> 225).
+            vuint16m2_t sq_lo = __riscv_vwmulu_vv_u16m2(d_lo, d_lo, byte_vl);
+            vuint16m2_t sq_hi = __riscv_vwmulu_vv_u16m2(d_hi, d_hi, byte_vl);
+            vuint16m2_t sq_sum = __riscv_vadd_vv_u16m2(sq_lo, sq_hi, byte_vl);
+
+            // Reduce to a scalar u32 (safe: byte_vl * 450 fits in u32 for
+            // any realistic d).
+            vuint32m1_t zero = __riscv_vmv_v_x_u32m1(0, 1);
+            vuint32m1_t red =
+                    __riscv_vwredsumu_vs_u16m2_u32m1(sq_sum, zero, byte_vl);
+            acc += __riscv_vmv_x_s_u32m1_u32(red);
+
+            i += vl;
+        }
+        // Scalar tail: cover any leftover odd lane (at most one).
+        for (; i < d; i++) {
+            uint8_t c_code =
+                    (i % 2 == 0) ? (code[i / 2] & 0x0F) : (code[i / 2] >> 4);
+            uint8_t q_code = (i % 2 == 0) ? q_lo[i / 2] : q_hi[i / 2];
+            int diff = int(q_code) - int(c_code);
+            acc += diff * diff;
+        }
+        return acc;
+    }
+
+    float query_to_code(const uint8_t* code) const final {
+        return static_cast<float>(accumulate_int_l2(code)) * final_scale_sq;
+    }
+
+    float symmetric_dis(idx_t i, idx_t j) override {
+        // Not on the critical path for most workloads; reconstruct both
+        // codes into nibbles scalar-style and compute squared distance.
+        const uint8_t* c1 = codes + i * code_size;
+        const uint8_t* c2 = codes + j * code_size;
+        int64_t acc = 0;
+        for (size_t k = 0; k < d; k++) {
+            uint8_t a = (k % 2 == 0) ? (c1[k / 2] & 0x0F) : (c1[k / 2] >> 4);
+            uint8_t b = (k % 2 == 0) ? (c2[k / 2] & 0x0F) : (c2[k / 2] >> 4);
+            int diff = int(a) - int(b);
+            acc += diff * diff;
+        }
+        return static_cast<float>(acc) * final_scale_sq;
+    }
+
+    void query_to_codes_batch_4(
+            const uint8_t* code_0,
+            const uint8_t* code_1,
+            const uint8_t* code_2,
+            const uint8_t* code_3,
+            float& dis0,
+            float& dis1,
+            float& dis2,
+            float& dis3) const final {
+        // Simple 4x unroll of the single-code path; good enough as a first
+        // cut — gives ILP across the four independent accumulate loops.
+        dis0 = static_cast<float>(accumulate_int_l2(code_0)) * final_scale_sq;
+        dis1 = static_cast<float>(accumulate_int_l2(code_1)) * final_scale_sq;
+        dis2 = static_cast<float>(accumulate_int_l2(code_2)) * final_scale_sq;
+        dis3 = static_cast<float>(accumulate_int_l2(code_3)) * final_scale_sq;
+    }
+};
+
+} // namespace scalar_quantizer
+} // namespace faiss
+
+#define THE_LEVEL_TO_DISPATCH SIMDLevel::RISCV_RVV
+#include <faiss/impl/scalar_quantizer/sq-dispatch.h>
+
+#endif // COMPILE_SIMD_RISCV_RVV

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

@@ -0,0 +1,387 @@
+/*
+ * 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.
+ */
+
+#include <faiss/impl/scalar_quantizer/training.h>
+
+#include <faiss/impl/FaissAssert.h>
+#include <algorithm>
+#include <cmath>
+
+namespace faiss {
+
+namespace scalar_quantizer {
+/*******************************************************************
+ * Quantizer range training
+ */
+
+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,
+        idx_t n,
+        int k,
+        const float* x,
+        std::vector<float>& trained) {
+    FAISS_THROW_IF_NOT(n > 0);
+    trained.resize(2);
+    float& vmin = trained[0];
+    float& vmax = trained[1];
+
+    if (rs == ScalarQuantizer::RS_minmax) {
+        vmin = HUGE_VAL;
+        vmax = -HUGE_VAL;
+        for (idx_t i = 0; i < n; i++) {
+            if (x[i] < vmin) {
+                vmin = x[i];
+            }
+            if (x[i] > vmax) {
+                vmax = x[i];
+            }
+        }
+        float vexp = (vmax - vmin) * rs_arg;
+        vmin -= vexp;
+        vmax += vexp;
+    } else if (rs == ScalarQuantizer::RS_meanstd) {
+        double sum = 0, sum2 = 0;
+        for (idx_t i = 0; i < n; i++) {
+            sum += x[i];
+            sum2 += x[i] * x[i];
+        }
+        float mean = sum / n;
+        float var = sum2 / n - mean * mean;
+        float std = var <= 0 ? 1.0 : std::sqrt(var);
+
+        vmin = mean - std * rs_arg;
+        vmax = mean + std * rs_arg;
+    } else if (rs == ScalarQuantizer::RS_quantiles) {
+        std::vector<float> x_copy(n);
+        memcpy(x_copy.data(), x, n * sizeof(*x));
+        idx_t o = static_cast<idx_t>(rs_arg * n);
+        if (o < 0) {
+            o = 0;
+        }
+        if (o > n - o) {
+            o = n / 2;
+        }
+        std::nth_element(x_copy.begin(), x_copy.begin() + o, x_copy.end());
+        vmin = x_copy[o];
+        std::nth_element(
+                x_copy.begin(), x_copy.begin() + (n - 1 - o), x_copy.end());
+        vmax = x_copy[n - 1 - o];
+
+    } else if (rs == ScalarQuantizer::RS_optim) {
+        float a, b;
+        float sx = 0;
+        {
+            vmin = HUGE_VAL, vmax = -HUGE_VAL;
+            for (idx_t i = 0; i < n; i++) {
+                if (x[i] < vmin) {
+                    vmin = x[i];
+                }
+                if (x[i] > vmax) {
+                    vmax = x[i];
+                }
+                sx += x[i];
+            }
+            b = vmin;
+            a = (vmax - vmin) / (k - 1);
+        }
+        int verbose = false;
+        int niter = 2000;
+        float last_err = -1;
+        int iter_last_err = 0;
+        for (int it = 0; it < niter; it++) {
+            float sn = 0, sn2 = 0, sxn = 0, err1 = 0;
+
+            for (idx_t i = 0; i < n; i++) {
+                float xi = x[i];
+                float ni = floor((xi - b) / a + 0.5);
+                if (ni < 0) {
+                    ni = 0;
+                }
+                if (ni >= k) {
+                    ni = k - 1;
+                }
+                err1 += sqr(xi - (ni * a + b));
+                sn += ni;
+                sn2 += ni * ni;
+                sxn += ni * xi;
+            }
+
+            if (err1 == last_err) {
+                iter_last_err++;
+                if (iter_last_err == 16) {
+                    break;
+                }
+            } else {
+                last_err = err1;
+                iter_last_err = 0;
+            }
+
+            float det = sqr(sn) - sn2 * n;
+
+            b = (sn * sxn - sn2 * sx) / det;
+            a = (sn * sx - n * sxn) / det;
+            if (verbose) {
+                printf("it %d, err1=%g            \r", it, err1);
+                fflush(stdout);
+            }
+        }
+        if (verbose) {
+            printf("\n");
+        }
+
+        vmin = b;
+        vmax = b + a * (k - 1);
+
+    } else {
+        FAISS_THROW_MSG("Invalid qtype");
+    }
+    vmax -= vmin;
+}
+
+void train_NonUniform(
+        RangeStat rs,
+        float rs_arg,
+        idx_t n,
+        int d,
+        int k,
+        const float* x,
+        std::vector<float>& trained) {
+    trained.resize(static_cast<size_t>(2) * d);
+    float* vmin = trained.data();
+    float* vmax = trained.data() + d;
+    if (rs == ScalarQuantizer::RS_minmax) {
+        memcpy(vmin, x, sizeof(*x) * d);
+        memcpy(vmax, x, sizeof(*x) * d);
+        for (idx_t i = 1; i < n; i++) {
+            const float* xi = x + i * d;
+            for (int j = 0; j < d; j++) {
+                if (xi[j] < vmin[j]) {
+                    vmin[j] = xi[j];
+                }
+                if (xi[j] > vmax[j]) {
+                    vmax[j] = xi[j];
+                }
+            }
+        }
+        float* vdiff = vmax;
+        for (int j = 0; j < d; j++) {
+            float vexp = (vmax[j] - vmin[j]) * rs_arg;
+            vmin[j] -= vexp;
+            vmax[j] += vexp;
+            vdiff[j] = vmax[j] - vmin[j];
+        }
+    } else {
+        // transpose
+        std::vector<float> xt(n * d);
+        for (idx_t i = 1; i < n; i++) {
+            const float* xi = x + i * d;
+            for (int j = 0; j < d; j++) {
+                xt[j * n + i] = xi[j];
+            }
+        }
+        std::vector<float> trained_d(2);
+#pragma omp parallel for
+        for (int j = 0; j < d; j++) {
+            train_Uniform(rs, rs_arg, n, k, xt.data() + j * n, trained_d);
+            vmin[j] = trained_d[0];
+            vmax[j] = trained_d[1];
+        }
+    }
+}
+
+} // namespace scalar_quantizer
+
+} // namespace faiss