SVE low-dim nearest fast path + SuperKMeans spherical support

Summary:
Two related optimizations for the SCANN coarse quantizer / PQ encoding
path, both targeting ARM SVE:

1. Low-dimensional L2sqr nearest fast path (d in {2,4,8})
   fvec_L2sqr_ny_nearest<ARM_SVE> previously wrote all ny distances to
   the scratch buffer then did a separate linear scan. For PQ encoding
   (compute_1_code) each call has d = dsub and ny = ksub, making this the
   hot path. The new path keeps one SVE lane per centroid via D svld1
   loads, tracks the min index in registers, and still fills the buffer
   (contract-preserving). ny stays generic.
   ProductQuantizer::compute_code switches to AVAILABLE_SIMD_LEVELS_A1 so
   the ARM_SVE implementation is reachable from PQ encoding (A0 does not
   include ARM_SVE).

2. SuperKMeans spherical (inner-product) support
   SuperKMeans previously only supported L2. With cp.spherical=true:
   - TrainState::R becomes std::unique_ptr<VectorTransform>; power-of-two
     d uses the fast HadamardRotation, L2 keeps RandomRotationMatrix.
   - update_centroids_and_split and Forgy init renormalize centroids to
     unit length, so minimizing L2 is equivalent to maximizing inner
     product.
   - HadamardRotation::reverse_transform fills the missing inverse.
   - ClusteringParameters::use_super_kmeans (default false) lets
     Level1Quantizer::train_q1 route coarse quantizer training through
     SuperKMeans when explicitly enabled.
   - block_l2<ARM_SVE> completes the SuperKMeans SIMD kernels.

Tests: low-dim nearest across SIMD levels (d in {2,4,8}, varied ny),
spherical objective vs vanilla spherical Clustering, unit-norm centroids,
use_super_kmeans field inheritance. Python stubs updated.

Measured on qwen 4096-dim, IP, nlist=1024, sub_dim=4, 1 thread, 50k rows:
- SCANN_DVR Train 67.73s -> 24.36s, Build 78.59s -> 31.52s (-60%)
- recall@10 unchanged vs Clustering baseline (diff <= 0.0002)

Author: marcelo-cjl

faiss/CMakeLists.txt

@@ -64,6 +64,7 @@ set(FAISS_SIMD_NEON_SRC
 set(FAISS_SIMD_SVE_SRC
   impl/pq_code_distance/pq_code_distance-sve.cpp
   utils/simd_impl/distances_arm_sve.cpp
+  utils/simd_impl/super_kmeans_kernels_sve.cpp
 )
 set(FAISS_SIMD_RVV_SRC
   impl/fast_scan/impl-riscv.cpp

faiss/Clustering.h

@@ -75,6 +75,11 @@ struct ClusteringParameters {
     /// so the training process stops only if an error
     /// is unchanged from the previous iteration.
     double early_stop_threshold = 0.0;
+
+    /// Whether to use the SuperKMeans (super fast k-means) variant instead of
+    /// the vanilla Clustering implementation. Only honored by callers that
+    /// explicitly support it (e.g. IVF level-1 quantizer training).
+    bool use_super_kmeans = false;
 };
 
 struct ClusteringIterationStats {

faiss/IndexIVF.cpp

@@ -6,6 +6,7 @@
  */
 
 #include <faiss/IndexIVF.h>
+#include <faiss/SuperKMeans.h>
 
 #include <omp.h>
 #include <atomic>
@@ -78,13 +79,22 @@ void Level1Quantizer::train_q1(
             printf("Training level-1 quantizer on %zd vectors in %zdD\n", n, d);
         }
 
-        Clustering clus(static_cast<int>(d), static_cast<int>(nlist), cp);
         quantizer->reset();
-        if (clustering_index) {
-            clus.train(n, x, *clustering_index);
+        if (cp.use_super_kmeans && clustering_index == nullptr) {
+            SuperKMeansParameters super_cp;
+            static_cast<ClusteringParameters&>(super_cp) = cp;
+            SuperKMeans clus(
+                    static_cast<int>(d), static_cast<int>(nlist), super_cp);
+            clus.train(n, x);
             quantizer->add(nlist, clus.centroids.data());
         } else {
-            clus.train(n, x, *quantizer);
+            Clustering clus(static_cast<int>(d), static_cast<int>(nlist), cp);
+            if (clustering_index) {
+                clus.train(n, x, *clustering_index);
+                quantizer->add(nlist, clus.centroids.data());
+            } else {
+                clus.train(n, x, *quantizer);
+            }
         }
         quantizer->is_trained = true;
     } else if (quantizer_trains_alone == 2) {

faiss/SuperKMeans.cpp

@@ -57,9 +57,9 @@ namespace faiss {
 namespace {
 
 struct TrainState {
-    /// Orthogonal rotation. Train in rotated space (X_tilde = X * R);
+    /// Fast orthogonal rotation. Train in rotated space;
     /// un-rotate centroids before return.
-    faiss::RandomRotationMatrix R;
+    std::unique_ptr<faiss::VectorTransform> R;
 
     std::vector<float> X_tilde; // (n, d) row-major
     int n = 0;
@@ -77,7 +77,17 @@ struct TrainState {
     int low_pruning_streak = 0;
     bool low_pruning_warning_printed = false;
 
-    explicit TrainState(int d) : R(d, d) {}
+    explicit TrainState(int d, bool spherical)
+            : R([d, spherical]() -> std::unique_ptr<faiss::VectorTransform> {
+                  // Spherical (inner-product) clustering only: a power-of-two
+                  // dimension can use the fast Hadamard rotation instead of
+                  // the generic random rotation. L2 training keeps the
+                  // original RandomRotationMatrix path unchanged.
+                  if (spherical && d > 0 && (d & (d - 1)) == 0) {
+                      return std::make_unique<faiss::HadamardRotation>(d);
+                  }
+                  return std::make_unique<faiss::RandomRotationMatrix>(d, d);
+              }()) {}
 };
 
 /// PDX block layout for the trailing pruning sweep: block b covers original
@@ -198,7 +208,8 @@ int update_centroids_and_split(
         int k,
         TrainState& state,
         std::vector<int64_t>& labels64,
-        std::vector<float>& hassign) {
+        std::vector<float>& hassign,
+        bool spherical) {
     std::fill(hassign.begin(), hassign.end(), 0.0f);
     assert(!labels64.empty());
     assert(!state.assignments.empty());
@@ -216,16 +227,23 @@ int update_centroids_and_split(
             /*weights=*/nullptr,
             hassign.data(),
             state.Y_tilde.data());
+    if (spherical) {
+        fvec_renorm_L2(d, k, state.Y_tilde.data());
+    }
     if (state.n <= k) {
         return 0;
     }
-    return detail::split_clusters(
+    const int nsplit = detail::split_clusters(
             d,
             k,
             state.n,
             /*k_frozen=*/0,
             hassign.data(),
             state.Y_tilde.data());
+    if (spherical) {
+        fvec_renorm_L2(d, k, state.Y_tilde.data());
+    }
+    return nsplit;
 }
 
 /// Stay-in-band controller: nudge state.d_prime based on observed pruning
@@ -295,10 +313,17 @@ std::unique_ptr<uint8_t[]> setup_train_state(
             "SuperKMeans: training set size exceeds INT_MAX after sampling");
     state.n = static_cast<int>(nx);
 
-    state.R.init(cp.seed);
+    if (auto* R = dynamic_cast<HadamardRotation*>(state.R.get())) {
+        R->init(cp.seed);
+    } else {
+        auto* dense_rotation =
+                dynamic_cast<RandomRotationMatrix*>(state.R.get());
+        FAISS_ASSERT(dense_rotation != nullptr);
+        dense_rotation->init(cp.seed);
+    }
 
     state.X_tilde.resize(static_cast<size_t>(state.n) * d);
-    state.R.apply_noalloc(state.n, x_sampled, state.X_tilde.data());
+    state.R->apply_noalloc(state.n, x_sampled, state.X_tilde.data());
 
     // Forgy init: pick k random rows from the rotated pool as initial
     // centroids. These remain in rotated space; un-rotation happens
@@ -313,6 +338,9 @@ std::unique_ptr<uint8_t[]> setup_train_state(
                     state.X_tilde.data() + static_cast<size_t>(perm[j]) * d,
                     sizeof(float) * d);
         }
+        if (cp.spherical) {
+            fvec_renorm_L2(d, k, state.Y_tilde.data());
+        }
     }
 
     state.d_prime =
@@ -339,7 +367,7 @@ std::unique_ptr<uint8_t[]> setup_train_state(
 /// reverse_transform applies R^T = R^-1.
 void untransform_centroids(
         std::vector<float>& centroids,
-        const RandomRotationMatrix& R,
+        const VectorTransform& R,
         int d,
         int k,
         const float* Y_tilde) {
@@ -411,7 +439,7 @@ void SuperKMeans::train(idx_t n, const float* x) {
                static_cast<idx_t>(k) * cp.min_points_per_centroid);
     }
 
-    TrainState state(d);
+    TrainState state(d, cp.spherical);
     std::vector<int64_t> labels64;
     SuperKMeansAssignScratch assign_scratch;
     std::vector<float> hassign;
@@ -444,8 +472,8 @@ void SuperKMeans::train(idx_t n, const float* x) {
                     pruned_at_gemm);
         }
 
-        const int nsplit =
-                update_centroids_and_split(d, k, state, labels64, hassign);
+        const int nsplit = update_centroids_and_split(
+                d, k, state, labels64, hassign, cp.spherical);
         const float pruning_rate = (iter == 0)
                 ? 0.0f
                 : adapt_d_prime(d, cp, state, total_pairs, pruned_at_gemm);
@@ -492,7 +520,7 @@ void SuperKMeans::train(idx_t n, const float* x) {
                (getmillisecs() - t_train_start) / 1000.0);
     }
 
-    untransform_centroids(centroids, state.R, d, k, state.Y_tilde.data());
+    untransform_centroids(centroids, *state.R, d, k, state.Y_tilde.data());
 }
 
 void super_kmeans_assign_iteration(

faiss/SuperKMeans.h

@@ -13,8 +13,9 @@
 //   "A Super Fast K-means for Indexing Vector Embeddings."
 //   arXiv preprint arXiv:2603.20009.
 //
-// Use when: L2 metric, k >= 1024, d >= 128, dense float embeddings.
-// Do not use for: IP/cosine (use Clustering with cp.spherical=true), small k,
+// Use when: L2 metric, or spherical inner-product clustering with
+// cp.spherical=true, k >= 1024, d >= 128, dense float embeddings.
+// Do not use for: small k,
 // binary data (use IndexBinaryIVF), or near-unit-sphere embeddings with
 // k < 4096 (chi-squared assumption breaks down).
 //

faiss/VectorTransform.cpp

@@ -513,6 +513,41 @@ void HadamardRotation::apply_noalloc(idx_t n, const float* x, float* xt) const {
     }
 }
 
+void HadamardRotation::reverse_transform(idx_t n, const float* xt, float* x)
+        const {
+    FAISS_THROW_IF_NOT_MSG(is_trained, "Transformation not trained yet");
+    FAISS_THROW_IF_NOT_MSG(
+            d_in == d_out,
+            "HadamardRotation inverse requires equal input/output dimensions");
+
+    const size_t p = d_out;
+    const float inverse_scale = p * std::sqrt(static_cast<float>(p));
+
+#pragma omp parallel for schedule(dynamic)
+    for (idx_t i = 0; i < n; i++) {
+        const float* xi = xt + i * p;
+        float* xo = x + i * p;
+
+        // The inverse reverses the three sign-flip/Hadamard factors.
+        std::memcpy(xo, xi, p * sizeof(float));
+        fwht_inplace(xo, p);
+
+        for (size_t j = 0; j < p; j++) {
+            xo[j] *= signs3[j];
+        }
+        fwht_inplace(xo, p);
+
+        for (size_t j = 0; j < p; j++) {
+            xo[j] *= signs2[j];
+        }
+        fwht_inplace(xo, p);
+
+        for (size_t j = 0; j < p; j++) {
+            xo[j] *= signs1[j] * inverse_scale;
+        }
+    }
+}
+
 void HadamardRotation::check_identical(const VectorTransform& other) const {
     auto* hr = dynamic_cast<const HadamardRotation*>(&other);
     FAISS_THROW_IF_NOT_MSG(hr, "failed to cast to HadamardRotation");

faiss/VectorTransform.h

@@ -144,6 +144,9 @@ struct HadamardRotation : VectorTransform {
 
     void apply_noalloc(idx_t n, const float* x, float* xt) const override;
 
+    /// Apply the inverse transform when d_in == d_out.
+    void reverse_transform(idx_t n, const float* xt, float* x) const override;
+
     void check_identical(const VectorTransform& other) const override;
 
     HadamardRotation() {}

faiss/impl/ProductQuantizer.cpp

@@ -280,7 +280,9 @@ void compute_1_code(const ProductQuantizer& pq, const float* x, uint8_t* code) {
 } // namespace
 
 void ProductQuantizer::compute_code(const float* x, uint8_t* code) const {
-    with_simd_level([&]<SIMDLevel SL>() {
+    // A1 includes ARM_SVE so the low-dimensional SVE nearest kernel in
+    // fvec_L2sqr_ny_nearest is reachable from PQ encoding.
+    with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A1>([&]<SIMDLevel SL>() {
         switch (nbits) {
             case 8:
                 compute_1_code<PQEncoder8, SL>(*this, x, code);

faiss/python/__init__.pyi

@@ -2628,6 +2628,7 @@ class ClusteringParameters:
     init_method: ClusteringInitMethod
     afkmc2_chain_length: int  # chain length for AFK-MC² initialization
     early_stop_threshold: float  # early stop threshold [0, 1]
+    use_super_kmeans: bool  # route through SuperKMeans when supported
 
     def __init__(self) -> None: ...
 

faiss/python/extra_wrappers.py

@@ -656,8 +656,10 @@ class SuperKmeans(Kmeans):
     `iteration_stats`, and `assign()` surface; additionally exposes
     `gemm_pruning_rates`.
 
-    kwargs are forwarded to `SuperKMeansParameters`. Fields not present on it
-    (e.g. `spherical`, `int_centroids`, `nredo`, `frozen_centroids`,
+    kwargs are forwarded to `SuperKMeansParameters`. Fields inherited from
+    `ClusteringParameters` such as `spherical` are supported (unit-normalized
+    centroids enable inner-product clustering). Fields not present on it
+    (e.g. `int_centroids`, `nredo`, `frozen_centroids`,
     `init_method`, `update_index`, `early_stop_threshold`,
     `progressive_dim_steps`, `gpu`) raise `AttributeError`.
     """