diff --git a/benchs/bench_super_kmeans.py b/benchs/bench_super_kmeans.py new file mode 100644 index 000000000..d6ddbb4c3 --- /dev/null +++ b/benchs/bench_super_kmeans.py @@ -0,0 +1,112 @@ +# 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. + +"""Benchmark `faiss.SuperKMeans` against vanilla `faiss.Clustering` on the same +input. Inputs may be a .fvecs file or a synthetic Gaussian mixture. + +Examples: + # Synthetic data + python bench_super_kmeans.py --n 10000 --d 128 --k 64 --niter 10 + + # Real .fvecs (e.g., SIFT base) + python bench_super_kmeans.py --fvecs /path/to/sift_base.fvecs --k 1024 \\ + --niter 10 +""" + +import argparse +import sys +import time + +import faiss +import numpy as np +from datasets import fvecs_read + + +def gaussian_mixture(n, d, k, seed): + """k centers in [-1, 1]^d, ~n/k samples per center from N(center, 0.1).""" + rng = np.random.RandomState(seed) + centers = rng.uniform(-1.0, 1.0, size=(k, d)).astype("float32") + cluster = np.arange(n) % k + noise = rng.normal(0.0, 0.1, size=(n, d)).astype("float32") + return centers[cluster] + noise + + +def run_vanilla(x, d, k, niter, seed): + quant = faiss.IndexFlatL2(d) + cl = faiss.Clustering(d, k) + cl.seed = seed + cl.niter = niter + cl.verbose = False + t0 = time.perf_counter() + cl.train(x, quant) + dt = time.perf_counter() - t0 + final_obj = cl.iteration_stats.at(cl.iteration_stats.size() - 1).obj + return dt, final_obj, cl.iteration_stats.size() + + +def run_super(x, d, k, niter, seed): + p = faiss.SuperKMeansParameters() + p.seed = seed + p.niter = niter + p.verbose = False + sc = faiss.SuperKMeans(d, k, p) + t0 = time.perf_counter() + sc.train(x) + dt = time.perf_counter() - t0 + final_obj = sc.iteration_stats.at(sc.iteration_stats.size() - 1).obj + iter_sum = sum( + sc.iteration_stats.at(i).time for i in range(sc.iteration_stats.size()) + ) + setup = dt - iter_sum + rates = faiss.vector_to_array(sc.gemm_pruning_rates) + return dt, final_obj, sc.iteration_stats.size(), iter_sum, setup, rates + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + src = ap.add_mutually_exclusive_group() + src.add_argument("--fvecs", help="path to a .fvecs file") + ap.add_argument("--n", type=int, default=10000, help="synthetic n") + ap.add_argument("--d", type=int, default=128, help="synthetic d") + ap.add_argument("--k", type=int, default=64, help="number of centroids") + ap.add_argument("--niter", type=int, default=10, help="iterations") + ap.add_argument("--seed", type=int, default=42, help="RNG seed") + args = ap.parse_args() + + if args.fvecs: + x = fvecs_read(args.fvecs) + n, d = x.shape + print("=== bench_super_kmeans ===") + print( + f"fvecs={args.fvecs} n={n} d={d} k={args.k} " + f"niter={args.niter} seed={args.seed}" + ) + else: + n, d = args.n, args.d + x = gaussian_mixture(n, d, args.k, args.seed) + print("=== bench_super_kmeans ===") + print( + f"n={n} d={d} k={args.k} niter={args.niter} seed={args.seed}" + ) + + print("\n--- faiss.Clustering ---") + v_dt, v_obj, v_iters = run_vanilla(x, d, args.k, args.niter, args.seed) + print(f" wall_time={v_dt:.3f}s final_obj={v_obj:g} iters={v_iters}") + + print("\n--- faiss.SuperKMeans ---") + s_dt, s_obj, s_iters, iter_sum, setup, rates = run_super( + x, d, args.k, args.niter, args.seed + ) + print( + f" wall_time={s_dt:.3f}s final_obj={s_obj:g} iters={s_iters} " + f"iter_sum={iter_sum:.3f}s setup={setup:.3f}s" + ) + print(" gemm_pruning_rates:", " ".join(f"{r:.3f}" for r in rates)) + print() + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/faiss/CMakeLists.txt b/faiss/CMakeLists.txt index b5f406792..db97ceefd 100644 --- a/faiss/CMakeLists.txt +++ b/faiss/CMakeLists.txt @@ -73,6 +73,7 @@ endif() set(FAISS_SRC AutoTune.cpp Clustering.cpp + SuperKMeans.cpp IVFlib.cpp Index.cpp Index2Layer.cpp @@ -200,6 +201,7 @@ endif() set(FAISS_HEADERS AutoTune.h Clustering.h + SuperKMeans.h IVFlib.h Index.h Index2Layer.h diff --git a/faiss/SuperKMeans.cpp b/faiss/SuperKMeans.cpp new file mode 100644 index 000000000..ace16aa71 --- /dev/null +++ b/faiss/SuperKMeans.cpp @@ -0,0 +1,656 @@ +/* + * 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 +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#ifndef FINTEGER +#define FINTEGER long +#endif + +extern "C" { + +/* declare BLAS functions, see http://www.netlib.org/clapack/cblas/ */ + +int sgemm_( + const char* transa, + const char* transb, + FINTEGER* m, + FINTEGER* n, + FINTEGER* k, + const float* alpha, + const float* a, + FINTEGER* lda, + const float* b, + FINTEGER* ldb, + float* beta, + float* c, + FINTEGER* ldc); +} + +namespace faiss { + +namespace { + +struct TrainState { + /// Orthogonal rotation. Train in rotated space (X_tilde = X * R); + /// un-rotate centroids before return. + faiss::RandomRotationMatrix R; + + std::vector X_tilde; // (n, d) row-major + int n = 0; + std::vector Y_tilde; // (k, d) row-major + + std::vector assignments; // size n + std::vector best_dists; // size n; tau per vector + + /// ||X_tilde[i, 0:d_prime]||^2; recomputed when d_prime changes. + std::vector x_norms_partial; + + int d_prime = 0; + + /// ADSampling threshold table; size d+1. + std::vector ad_coeff; + + /// PDX block layout for the trailing pruning sweep: block b covers + /// original dims [true_block_end[b] - block_dim[b], true_block_end[b]). + /// Recomputed when d_prime changes. + std::vector block_dim; + std::vector true_block_end; + + /// Counter for the verbose-mode "low pruning" warning. + int low_pruning_streak = 0; + bool low_pruning_warning_printed = false; + + explicit TrainState(int d) : R(d, d) {} +}; + +/// Rebuild state.block_dim and state.true_block_end from the current +/// state.d_prime and pdx_block_size. Call after any change to d_prime. +void rebuild_pdx_block_layout(int d, int pdx_block_size, TrainState& state) { + const int dp = state.d_prime; + const int d_trail = d - dp; + const int n_full_blocks = d_trail / pdx_block_size; + const int tail = d_trail % pdx_block_size; + const int n_blocks = n_full_blocks + (tail > 0 ? 1 : 0); + state.block_dim.assign(n_blocks, pdx_block_size); + state.true_block_end.resize(n_blocks); + if (n_blocks > 0) { + assert(!state.block_dim.empty()); + assert(!state.true_block_end.empty()); + for (int b = 0; b < n_full_blocks; ++b) { + state.true_block_end[b] = dp + (b + 1) * pdx_block_size; + } + if (tail > 0) { + state.block_dim[n_full_blocks] = tail; + state.true_block_end[n_full_blocks] = d; + } + } +} + +struct IterScratch { + std::vector partial_ip; // (bx_max, by_max) for the GEMM tile + std::vector Y_pdx; // PDX-laid-out trailing block + std::vector Y_trail; // row-major (k, d_trail) input to pdxify + std::vector y_norms_partial; // ||Y_tilde[j, 0:dp]||^2 + std::vector labels64; // size n; widened state.assignments + int prev_d_trail = -1; +}; + +/// Iter 0: full GEMM via knn_L2sqr (vanilla Lloyd's). Fills +/// state.assignments and state.best_dists. Returns objective. +double run_iter0_full_gemm(int d, int k, TrainState& state) { + std::vector labels(state.n); + std::vector distances(state.n); + knn_L2sqr( + state.X_tilde.data(), + state.Y_tilde.data(), + d, + state.n, + k, + /*k=*/1, + distances.data(), + labels.data(), + /*y_norm2=*/nullptr); + + assert(!state.assignments.empty()); + assert(!state.best_dists.empty()); + double objective = 0.0; + for (int i = 0; i < state.n; ++i) { + state.assignments[i] = static_cast(labels[i]); + state.best_dists[i] = distances[i]; + objective += distances[i]; + } + return objective; +} + +/// Iter 1+: partial GEMM over [0, d_prime) + ADSampling progressive +/// pruning over the PDX-laid-out trailing block. Updates +/// state.assignments and state.best_dists. Writes total_pairs and +/// pruned_at_gemm. Returns objective. +double run_iter_pruned( + int d, + int k, + const SuperKMeansParameters& cp, + TrainState& state, + IterScratch& scratch, + int64_t& total_pairs, + int64_t& pruned_at_gemm) { + const int dp = state.d_prime; + assert(dp >= 1); + assert(!state.ad_coeff.empty()); + assert(!scratch.partial_ip.empty()); + assert(!scratch.y_norms_partial.empty()); + const int d_trail = d - dp; + const int n_train = state.n; + assert(static_cast(state.best_dists.size()) >= n_train); + assert(static_cast(state.x_norms_partial.size()) >= n_train); + assert(static_cast(state.assignments.size()) >= n_train); + + if (d_trail != scratch.prev_d_trail) { + scratch.Y_pdx.resize(static_cast(k) * d_trail); + scratch.Y_trail.resize(static_cast(k) * d_trail); + scratch.prev_d_trail = d_trail; + } + for (int j = 0; j < k; ++j) { + std::memcpy( + scratch.Y_trail.data() + static_cast(j) * d_trail, + state.Y_tilde.data() + static_cast(j) * d + dp, + d_trail * sizeof(float)); + } + detail::pdxify( + scratch.Y_trail.data(), + k, + d_trail, + cp.pdx_block_size, + scratch.Y_pdx.data()); + + detail::compute_partial_norms( + state.Y_tilde.data(), k, d, dp, scratch.y_norms_partial.data()); + + const int n_blocks = static_cast(state.block_dim.size()); + + for (int xi = 0; xi < n_train; xi += cp.x_batch) { + const int bx = std::min(cp.x_batch, n_train - xi); + + // Refresh tau: recompute full-d L2 distance to the previously + // assigned centroid. This is intentionally over all d dims (not + // just d_prime) because tau must be an exact distance for the + // chi-squared pruning bound to be valid. Cost is O(bx * d) per + // x-batch, amortized across the y-batch tiles that follow. +#pragma omp parallel for + for (int i = 0; i < bx; ++i) { + const int j_prev = state.assignments[xi + i]; + const float* xrow = + state.X_tilde.data() + static_cast(xi + i) * d; + const float* yrow = + state.Y_tilde.data() + static_cast(j_prev) * d; + float tau = 0.0f; + for (int m = 0; m < d; ++m) { + const float diff = xrow[m] - yrow[m]; + tau += diff * diff; + } + state.best_dists[xi + i] = tau; + } + + for (int yj = 0; yj < k; yj += cp.y_batch) { + const int by = std::min(cp.y_batch, k - yj); + + // GEMM phase: column-major sgemm computes + // partial_ip[i*by + j] = . + { + FINTEGER M = by; + FINTEGER N_ = bx; + FINTEGER K_ = dp; + float alpha = 1.0f; + float beta = 0.0f; + FINTEGER lda_y = d; + FINTEGER lda_x = d; + FINTEGER ldc = by; + sgemm_("Transpose", + "Not transpose", + &M, + &N_, + &K_, + &alpha, + state.Y_tilde.data() + static_cast(yj) * d, + &lda_y, + state.X_tilde.data() + static_cast(xi) * d, + &lda_x, + &beta, + scratch.partial_ip.data(), + &ldc); + } + + // One SIMD dispatch per (xi, yj) tile — block_l2 below is + // a direct call (no per-call switch on SIMDConfig::level). + with_simd_level([&]() { + [[maybe_unused]] const int omp_chunk_local = cp.omp_chunk; + int64_t total_pairs_local = 0; + int64_t pruned_at_gemm_local = 0; +#pragma omp parallel for schedule(dynamic, omp_chunk_local) \ + reduction(+ : total_pairs_local) reduction(+ : pruned_at_gemm_local) + for (int i = 0; i < bx; ++i) { + // tau is the best full-d distance found so far for this + // point; tightened as closer centroids are found. + float tau = state.best_dists[xi + i]; + int best_j = state.assignments[xi + i]; + const float xnp_i = state.x_norms_partial[xi + i]; + const float* xrow = state.X_tilde.data() + + static_cast(xi + i) * d; + + for (int j = 0; j < by; ++j) { + ++total_pairs_local; + + // L2-from-IP; clamp to handle catastrophic + // cancellation when the true distance is ~0. + float pd = xnp_i + scratch.y_norms_partial[yj + j] - + 2.0f * + scratch.partial_ip + [static_cast(i) * by + + j]; + if (pd < 0.0f) { + pd = 0.0f; + } + + if (pd > state.ad_coeff[dp] * tau) { + ++pruned_at_gemm_local; + continue; + } + + // double accumulator mitigates float drift over many + // block additions. + double dist = pd; + bool keep = true; + + // Progressive pruning across PDX blocks. Per block: + // stride = k * block_dim[b] floats, column-major + // across centroids. + size_t pdx_offset = 0; + for (int b = 0; b < n_blocks; ++b) { + const int n_in_block = state.block_dim.at(b); + const int true_end = state.true_block_end.at(b); + const float* xblk = xrow + (true_end - n_in_block); + const float* yblk = scratch.Y_pdx.data() + + pdx_offset + + static_cast(yj + j) * n_in_block; + dist += faiss::detail::block_l2( + xblk, yblk, n_in_block); + pdx_offset += static_cast(k) * n_in_block; + + if (dist > + static_cast(state.ad_coeff[true_end]) * + tau) { + keep = false; + break; + } + } + + if (keep && dist < tau) { + tau = static_cast(dist); + best_j = yj + j; + } + } + + state.best_dists[xi + i] = tau; + state.assignments[xi + i] = best_j; + } + total_pairs += total_pairs_local; + pruned_at_gemm += pruned_at_gemm_local; + }); + } + } + + double objective = 0.0; + for (int i = 0; i < n_train; ++i) { + objective += state.best_dists[i]; + } + return objective; +} + +/// Post-iteration: update centroids and split empties. Returns nsplit. +int update_centroids_and_split( + int d, + int k, + TrainState& state, + IterScratch& scratch, + std::vector& hassign) { + std::fill(hassign.begin(), hassign.end(), 0.0f); + assert(!scratch.labels64.empty()); + assert(!state.assignments.empty()); + for (int i = 0; i < state.n; ++i) { + scratch.labels64[i] = static_cast(state.assignments[i]); + } + detail::compute_centroids( + d, + k, + state.n, + /*k_frozen=*/0, + reinterpret_cast(state.X_tilde.data()), + /*codec=*/nullptr, + scratch.labels64.data(), + /*weights=*/nullptr, + hassign.data(), + state.Y_tilde.data()); + if (state.n <= k) { + return 0; + } + return detail::split_clusters( + d, + k, + state.n, + /*k_frozen=*/0, + hassign.data(), + state.Y_tilde.data()); +} + +/// Stay-in-band controller: nudge state.d_prime based on observed +/// pruning rate. Recomputes x_norms_partial if d_prime changed. Returns +/// the observed pruning rate (0 when there were no pairs). +float adapt_d_prime( + int d, + const SuperKMeansParameters& cp, + TrainState& state, + int64_t total_pairs, + int64_t pruned_at_gemm) { + if (total_pairs == 0) { + return 0.0f; + } + const float pruning_rate = static_cast(pruned_at_gemm) / + static_cast(total_pairs); + int new_dp = state.d_prime; + if (pruning_rate > cp.pruning_target_high) { + new_dp = static_cast( + std::lround(state.d_prime * (1.0f - cp.d_prime_adjust))); + } else if (pruning_rate < cp.pruning_target_low) { + new_dp = static_cast( + std::lround(state.d_prime * (1.0f + cp.d_prime_adjust))); + } + new_dp = std::max(cp.d_prime_min, new_dp); + new_dp = std::min(d / 2, new_dp); + if (new_dp != state.d_prime) { + state.d_prime = new_dp; + detail::compute_partial_norms( + state.X_tilde.data(), + state.n, + d, + state.d_prime, + state.x_norms_partial.data()); + rebuild_pdx_block_layout(d, cp.pdx_block_size, state); + } + return pruning_rate; +} + +/// Pre-loop setup: subsample, rotate, Forgy init, build ADSampling table, +/// allocate scratch. Returned `sampled_x_owner` keeps the subsampled buffer +/// alive when subsampling occurred (otherwise empty). +std::unique_ptr setup_train_state( + TrainState& state, + IterScratch& scratch, + std::vector& hassign, + const SuperKMeansParameters& cp, + int d, + int k, + idx_t n, + const float* x) { + const size_t line_size = sizeof(float) * static_cast(d); + idx_t nx = n; + const uint8_t* x_bytes = reinterpret_cast(x); + std::unique_ptr sampled_x_owner; + if (static_cast(nx) > + static_cast(k) * cp.max_points_per_centroid) { + Clustering tmp_clus(d, k, cp); + uint8_t* x_new = nullptr; + float* w_unused = nullptr; + nx = detail::subsample_training_set( + tmp_clus, + nx, + x_bytes, + line_size, + /*weights=*/nullptr, + &x_new, + &w_unused); + FAISS_ASSERT(x_new != nullptr); + sampled_x_owner.reset(x_new); + x_bytes = x_new; + } + const float* x_sampled = reinterpret_cast(x_bytes); + + FAISS_THROW_IF_NOT_MSG( + nx <= static_cast(std::numeric_limits::max()), + "SuperKMeans: training set size exceeds INT_MAX after sampling"); + state.n = static_cast(nx); + + state.R.init(cp.seed); + + state.X_tilde.resize(static_cast(state.n) * d); + 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 + // after the iteration loop. + state.Y_tilde.resize(static_cast(k) * d); + { + std::vector perm(state.n); + rand_perm(perm.data(), state.n, static_cast(cp.seed) + 1); + for (int j = 0; j < k; ++j) { + std::memcpy( + state.Y_tilde.data() + static_cast(j) * d, + state.X_tilde.data() + static_cast(perm[j]) * d, + sizeof(float) * d); + } + } + + state.d_prime = + std::max(cp.d_prime_min, static_cast(d * cp.d_prime_fraction)); + state.d_prime = std::min(state.d_prime, d / 2); + rebuild_pdx_block_layout(d, cp.pdx_block_size, state); + + // Iter 1+ uses L2-from-IP only over [0, d_prime), so full ||X[i]||^2 is + // never read; iter 0 routes through knn_L2sqr which carries its own. + state.x_norms_partial.resize(state.n); + detail::compute_partial_norms( + state.X_tilde.data(), + state.n, + d, + state.d_prime, + state.x_norms_partial.data()); + + const double epsilon = static_cast(cp.ad_epsilon_factor) / d; + state.ad_coeff = detail::precompute_ad_thresholds(d, epsilon); + FAISS_ASSERT_MSG( + state.ad_coeff.size() == static_cast(d + 1), + "ad_coeff size mismatch"); + + state.assignments.assign(state.n, 0); + state.best_dists.assign(state.n, std::numeric_limits::max()); + + hassign.assign(k, 0.0f); + + const int by_max = std::min(cp.y_batch, k); + const int bx_max = std::min(cp.x_batch, state.n); + scratch.partial_ip.resize(static_cast(bx_max) * by_max); + scratch.y_norms_partial.resize(k); + scratch.labels64.resize(state.n); + + return sampled_x_owner; +} + +/// Un-rotate centroids into output buffer. R orthogonal, so +/// reverse_transform applies R^T = R^-1. +void untransform_centroids( + std::vector& centroids, + const RandomRotationMatrix& R, + int d, + int k, + const float* Y_tilde) { + centroids.resize(static_cast(k) * d); + R.reverse_transform(k, Y_tilde, centroids.data()); +} + +} // namespace + +SuperKMeans::SuperKMeans(int d, int k, const SuperKMeansParameters& cp_in) + : cp(cp_in), d(d), k(k) { + FAISS_THROW_IF_NOT_MSG(d > 0, "SuperKMeans: d must be positive"); + FAISS_THROW_IF_NOT_MSG(k > 0, "SuperKMeans: k must be positive"); + FAISS_THROW_IF_NOT_MSG( + cp.d_prime_fraction > 0.0f && cp.d_prime_fraction <= 1.0f, + "SuperKMeans: d_prime_fraction must be in (0, 1]"); + FAISS_THROW_IF_NOT_MSG( + cp.d_prime_adjust >= 0.0f && cp.d_prime_adjust < 1.0f, + "SuperKMeans: d_prime_adjust must be in [0, 1)"); + // d >= 2 * d_prime_min keeps d_prime in the chi-squared validity + // range after both clamping steps (floor at d_prime_min, ceiling + // at d/2). See AdSampling.h on the p >= 16 contract. + FAISS_THROW_IF_NOT_FMT( + d >= 2 * cp.d_prime_min, + "SuperKMeans: d (%d) must be >= 2 * d_prime_min (%d)", + d, + cp.d_prime_min); + FAISS_THROW_IF_NOT_MSG( + cp.d_prime_min >= 16, + "SuperKMeans: d_prime_min must be >= 16 (chi-squared validity floor)"); + FAISS_THROW_IF_NOT_MSG( + cp.pdx_block_size > 0, "SuperKMeans: pdx_block_size must be > 0"); + FAISS_THROW_IF_NOT_MSG(cp.x_batch > 0, "SuperKMeans: x_batch must be > 0"); + FAISS_THROW_IF_NOT_MSG(cp.y_batch > 0, "SuperKMeans: y_batch must be > 0"); + FAISS_THROW_IF_NOT_MSG( + cp.pruning_target_low > 0.0f && + cp.pruning_target_low <= cp.pruning_target_high && + cp.pruning_target_high < 1.0f, + "SuperKMeans: require 0 < pruning_target_low <= pruning_target_high < 1"); + // epsilon = ad_epsilon_factor / d is the chi-squared significance + // level. precompute_ad_thresholds requires epsilon in (0, 1). + FAISS_THROW_IF_NOT_MSG( + cp.ad_epsilon_factor > 0.0f && + cp.ad_epsilon_factor < static_cast(d), + "SuperKMeans: ad_epsilon_factor must be in (0, d) " + "so epsilon = factor/d is in (0,1)"); + FAISS_THROW_IF_NOT_MSG( + cp.omp_chunk > 0, "SuperKMeans: omp_chunk must be > 0"); +} + +void SuperKMeans::train(idx_t n, const float* x) { + FAISS_THROW_IF_NOT_MSG(n > 0, "SuperKMeans: n must be positive"); + FAISS_THROW_IF_NOT_MSG(x != nullptr, "SuperKMeans: x must not be null"); + FAISS_THROW_IF_NOT_MSG( + n >= static_cast(k), "SuperKMeans: n must be >= k"); + if (cp.check_input_data_for_NaNs) { + for (size_t i = 0; i < static_cast(n) * d; i++) { + FAISS_THROW_IF_NOT_MSG( + std::isfinite(x[i]), + "SuperKMeans: input contains NaN's or Inf's"); + } + } + if (cp.verbose && n < static_cast(k) * cp.min_points_per_centroid) { + printf("WARNING: clustering %" PRId64 + " points to %d centroids: please provide at least " + "%" PRId64 " training points\n", + n, + k, + static_cast(k) * cp.min_points_per_centroid); + } + + TrainState state(d); + IterScratch scratch; + std::vector hassign; + [[maybe_unused]] auto sampled_x_owner = + setup_train_state(state, scratch, hassign, cp, d, k, n, x); + + iteration_stats.clear(); + iteration_stats.reserve(cp.niter); + gemm_pruning_rates.clear(); + gemm_pruning_rates.reserve(cp.niter); + + const double t_train_start = getmillisecs(); + + for (int iter = 0; iter < cp.niter; ++iter) { + const double t_iter_start = getmillisecs(); + double objective = 0.0; + int64_t total_pairs = 0; + int64_t pruned_at_gemm = 0; + + if (iter == 0) { + objective = run_iter0_full_gemm(d, k, state); + } else { + objective = run_iter_pruned( + d, k, cp, state, scratch, total_pairs, pruned_at_gemm); + } + + const int nsplit = + update_centroids_and_split(d, k, state, scratch, hassign); + const float pruning_rate = (iter == 0) + ? 0.0f + : adapt_d_prime(d, cp, state, total_pairs, pruned_at_gemm); + + ClusteringIterationStats stat{}; + stat.obj = static_cast(objective); + stat.time = (getmillisecs() - t_iter_start) / 1000.0; + stat.time_search = stat.time; + stat.imbalance_factor = std::numeric_limits::quiet_NaN(); + stat.nsplit = nsplit; + iteration_stats.push_back(stat); + gemm_pruning_rates.push_back(pruning_rate); + + if (iter > 0) { + if (pruning_rate < 0.85f) { + state.low_pruning_streak++; + } else { + state.low_pruning_streak = 0; + } + if (cp.verbose && state.low_pruning_streak >= 3 && + !state.low_pruning_warning_printed) { + fprintf(stderr, + "WARNING: SuperKMeans steady-state pruning < 0.85 for 3+ iters " + "(current=%.2f). Data may not be a good fit for ADSampling; " + "consider falling back to faiss::Clustering.\n", + pruning_rate); + state.low_pruning_warning_printed = true; + } + } + + if (cp.verbose) { + printf(" Iter %d: obj=%g time=%.3fs prune=%.4f dp=%d nsplit=%d\n", + iter, + stat.obj, + stat.time, + pruning_rate, + state.d_prime, + nsplit); + } + } + + if (cp.verbose) { + printf("Total training time: %.3fs\n", + (getmillisecs() - t_train_start) / 1000.0); + } + + untransform_centroids(centroids, state.R, d, k, state.Y_tilde.data()); +} + +} // namespace faiss diff --git a/faiss/SuperKMeans.h b/faiss/SuperKMeans.h new file mode 100644 index 000000000..033ec9bba --- /dev/null +++ b/faiss/SuperKMeans.h @@ -0,0 +1,97 @@ +/* + * 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. + */ + +// SuperKMeans — a faster k-means alternative to faiss::Clustering for IVF +// training and other large-k workloads. +// +// Based on: +// Kuffo, L., Hepkema, S., & Boncz, P. (2026). +// "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, +// binary data (use IndexBinaryIVF), or near-unit-sphere embeddings with +// k < 4096 (chi-squared assumption breaks down). +// +// `cp`, `d`, `k` are public mutable; mutating them post-construction bypasses +// validation and may produce undefined behavior. + +#pragma once + +#include + +#include +#include + +namespace faiss { + +struct SuperKMeansParameters : public ClusteringParameters { + /// Initial d_prime as a fraction of d (GEMM/PRUNING split). + float d_prime_fraction = 0.125f; + + /// Matches block_l2 SIMD width. + int pdx_block_size = 64; + + /// ADSampling significance: epsilon = ad_epsilon_factor / d. + float ad_epsilon_factor = 1.0f; + + /// Adaptive d_prime stay-in-band controller. Per iteration: + /// pruning > high → shrink d_prime (over-pruning, cheaper threshold) + /// pruning < low → grow d_prime (under-pruning, more GEMM work) + /// else: hold. + float pruning_target_low = 0.95f; + float pruning_target_high = 0.97f; + + /// Relative step size for d_prime adjustments. + float d_prime_adjust = 0.20f; + + /// Floor on d_prime; below this the chi-squared bound is unvalidated. + int d_prime_min = 16; + + int x_batch = 4096; + int y_batch = 1024; + + /// OpenMP dynamic-schedule chunk size for the pruning loop. + int omp_chunk = 8; +}; + +/** Drop-in faster k-means: same interface as faiss::Clustering. Per iteration: + * iter 0: full GEMM over all d dims (vanilla Lloyd's). + * iter 1..niter-1: GEMM over front d_prime dims, then ADSampling + * progressive pruning over PDX-laid trailing dims. + * + * Trains in a randomly-rotated space; centroids are un-rotated before return. + */ +struct SuperKMeans { + SuperKMeansParameters cp; + int d; + int k; + + /// Output: k * d floats, row-major, un-rotated. + std::vector centroids; + + /// Per-iter stats. `obj`, `time`, and `nsplit` are populated faithfully. + /// `time_search` mirrors `time` (phases not timed separately) and + /// `imbalance_factor` is set to NaN (not computed). + std::vector iteration_stats; + + /// Per-iter fraction of (vector, centroid) pairs pruned at the d_prime + /// GEMM-boundary chi-squared check. Per-PDX-block early-exits are NOT + /// counted. iter 0 uses full GEMM, so gemm_pruning_rates[0] == 0.0f. + std::vector gemm_pruning_rates; + + SuperKMeans(int d, int k, const SuperKMeansParameters& cp = {}); + + /// Train on `n` row-major vectors of dimension `d`. Honors the applicable + /// ClusteringParameters fields (niter, seed, verbose, + /// max_points_per_centroid, use_faster_subsampling, + /// check_input_data_for_NaNs). + void train(idx_t n, const float* x); +}; + +} // namespace faiss diff --git a/faiss/impl/ClusteringHelpers.cpp b/faiss/impl/ClusteringHelpers.cpp index b4704402b..669cdacf4 100644 --- a/faiss/impl/ClusteringHelpers.cpp +++ b/faiss/impl/ClusteringHelpers.cpp @@ -78,7 +78,7 @@ idx_t subsample_training_set( "subsample_training_set: perm size %zu < required nx %" PRId64, perm.size(), nx); - assert(!perm.empty()); // pattern for clang-tidy LocalUncheckedArrayBounds + assert(!perm.empty()); uint8_t* x_new = new uint8_t[nx * line_size]; *x_out = x_new; @@ -116,15 +116,6 @@ void compute_centroids( size_t line_size = codec ? codec->sa_code_size() : d * sizeof(float); - // NOTE: FAISS_THROW_IF_NOT inside this OMP region is formally UB per the - // OpenMP spec (exceptions cannot propagate out of a parallel region). - // In practice, most runtimes terminate the process, which is acceptable - // for a corrupted-assignment invariant violation. Inherited from - // Clustering.cpp; not worth refactoring without a broader FAISS effort. - // - // Thread safety of hassign[ci]: each thread owns the exclusive slice - // [c0, c1) and only writes hassign[ci] when ci falls in that slice - // (same guard as centroids[ci*d]). No two threads share an index. #pragma omp parallel { int nt = omp_get_num_threads(); diff --git a/faiss/python/__init__.py b/faiss/python/__init__.py index 05b376efc..f3aa4268a 100644 --- a/faiss/python/__init__.py +++ b/faiss/python/__init__.py @@ -22,8 +22,8 @@ from faiss.gpu_wrappers import * from faiss.array_conversions import * from faiss.extra_wrappers import kmin, kmax, pairwise_distances, rand, randint, \ lrand, randn, rand_smooth_vectors, eval_intersection, normalize_L2, \ - ResultHeap, knn, Kmeans, checksum, matrix_bucket_sort_inplace, bucket_sort, \ - merge_knn_results, MapInt64ToInt64, knn_hamming, \ + ResultHeap, knn, Kmeans, SuperKmeans, checksum, matrix_bucket_sort_inplace, \ + bucket_sort, merge_knn_results, MapInt64ToInt64, knn_hamming, \ pack_bitstrings, unpack_bitstrings @@ -34,6 +34,7 @@ __version__ = "%d.%d.%d" % (FAISS_VERSION_MAJOR, logger = logging.getLogger(__name__) class_wrappers.handle_Clustering(Clustering) +class_wrappers.handle_SuperKMeans(SuperKMeans) class_wrappers.handle_Clustering1D(Clustering1D) class_wrappers.handle_MatrixStats(MatrixStats) class_wrappers.handle_IOWriter(IOWriter) diff --git a/faiss/python/class_wrappers.py b/faiss/python/class_wrappers.py index 66552ee6c..a51e184d9 100644 --- a/faiss/python/class_wrappers.py +++ b/faiss/python/class_wrappers.py @@ -126,6 +126,24 @@ def handle_Clustering(the_class): replace_method(the_class, 'train_encoded', replacement_train_encoded) +def handle_SuperKMeans(the_class): + + def replacement_train(self, x): + """Perform SuperKMeans clustering on a set of vectors. + + Parameters + ---------- + x : array_like + Training vectors, shape (n, self.d). `dtype` must be float32. + """ + n, d = x.shape + assert d == self.d + x = np.ascontiguousarray(x, dtype='float32') + self.train_c(n, swig_ptr(x)) + + replace_method(the_class, 'train', replacement_train) + + def handle_Clustering1D(the_class): def replacement_train_exact(self, x): diff --git a/faiss/python/extra_wrappers.py b/faiss/python/extra_wrappers.py index b6fc1dbd4..77b3266b2 100644 --- a/faiss/python/extra_wrappers.py +++ b/faiss/python/extra_wrappers.py @@ -602,6 +602,56 @@ class Kmeans: return D.ravel(), I.ravel() +class SuperKmeans(Kmeans): + """Drop-in replacement for `Kmeans` that runs `SuperKMeans` (ADSampling + + PDX progressive pruning) instead of `Clustering`. Same `centroids`, `obj`, + `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`, + `init_method`, `update_index`, `early_stop_threshold`, + `progressive_dim_steps`, `gpu`) raise `AttributeError`. + """ + + def __init__(self, d, k, **kwargs): + self.d = d + self.reset(k) + self.gpu = False + self.cp = SuperKMeansParameters() + for key, v in kwargs.items(): + getattr(self.cp, key) + setattr(self.cp, key, v) + self.set_index() + + def set_index(self): + self.index = IndexFlatL2(self.d) + + def train(self, x, weights=None, init_centroids=None): + assert weights is None, "SuperKmeans does not support weights" + assert init_centroids is None, \ + "SuperKmeans does not support init_centroids" + x = np.ascontiguousarray(x, dtype='float32') + n, d = x.shape + assert d == self.d + + sc = SuperKMeans(d, self.k, self.cp) + sc.train(x) + + centroids = faiss.vector_to_array(sc.centroids) + self.centroids = centroids.reshape(self.k, d) + stats = sc.iteration_stats + stats = [stats.at(i) for i in range(stats.size())] + self.obj = np.array([st.obj for st in stats]) + stat_fields = 'obj time time_search imbalance_factor nsplit'.split() + self.iteration_stats = [ + {field: getattr(st, field) for field in stat_fields} + for st in stats + ] + self.gemm_pruning_rates = faiss.vector_to_array(sc.gemm_pruning_rates) + return self.obj[-1] if self.obj.size > 0 else 0.0 + + ########################################### # Packing and unpacking bitstrings ########################################### diff --git a/faiss/python/swigfaiss.swig b/faiss/python/swigfaiss.swig index 8ab89f5e9..74c9cea2d 100644 --- a/faiss/python/swigfaiss.swig +++ b/faiss/python/swigfaiss.swig @@ -188,6 +188,7 @@ typedef uint64_t size_t; #endif // !_MSC_VER #include +#include #include #include @@ -521,6 +522,7 @@ void gpu_sync_all_devices() %include %include %include +%include %include diff --git a/faiss/utils/simd_impl/super_kmeans_dispatch.h b/faiss/utils/simd_impl/super_kmeans_dispatch.h index 06d4072a7..9f6614ab8 100644 --- a/faiss/utils/simd_impl/super_kmeans_dispatch.h +++ b/faiss/utils/simd_impl/super_kmeans_dispatch.h @@ -13,7 +13,7 @@ // scalar primary template; adding NEON/SVE means just adding a new // specialization file alongside the AVX ones. // -// Known perf gap: aarch64 (NEON/SVE) specializations are not implemented in v1. +// Known perf gap: aarch64 (NEON/SVE) specializations are not implemented yet. // aarch64 falls through to the scalar primary template. Validating SVE requires // a Graviton-class host; deferred to a focused follow-up. diff --git a/faiss/utils/simd_impl/super_kmeans_kernels.h b/faiss/utils/simd_impl/super_kmeans_kernels.h index 0eec99b94..b36c05f5c 100644 --- a/faiss/utils/simd_impl/super_kmeans_kernels.h +++ b/faiss/utils/simd_impl/super_kmeans_kernels.h @@ -27,15 +27,17 @@ inline float block_l2(const float* x, const float* y, int n) { return s; } -// Mirrors the impl-file guards; turns off-arch direct calls into a -// compile-time error instead of a linker error. -#if defined(__x86_64__) +// COMPILE_SIMD_* is a build-system define (link-time promise that the +// specialization will be available). Mirrors the impl-file guards. +#ifdef COMPILE_SIMD_AVX2 template <> float block_l2(const float* x, const float* y, int n); +#endif +#ifdef COMPILE_SIMD_AVX512 template <> float block_l2(const float* x, const float* y, int n); -#endif // __x86_64__ +#endif } // namespace detail } // namespace faiss diff --git a/faiss/utils/simd_impl/super_kmeans_kernels_avx2.cpp b/faiss/utils/simd_impl/super_kmeans_kernels_avx2.cpp index 47cff327c..d5759a75a 100644 --- a/faiss/utils/simd_impl/super_kmeans_kernels_avx2.cpp +++ b/faiss/utils/simd_impl/super_kmeans_kernels_avx2.cpp @@ -5,7 +5,7 @@ * LICENSE file in the root directory of this source tree. */ -#if defined(__x86_64__) +#ifdef COMPILE_SIMD_AVX2 #include @@ -54,4 +54,4 @@ float block_l2(const float* x, const float* y, int n) { } // namespace detail } // namespace faiss -#endif // __x86_64__ +#endif // COMPILE_SIMD_AVX2 diff --git a/faiss/utils/simd_impl/super_kmeans_kernels_avx512.cpp b/faiss/utils/simd_impl/super_kmeans_kernels_avx512.cpp index 4fe67b1d1..dcb155bac 100644 --- a/faiss/utils/simd_impl/super_kmeans_kernels_avx512.cpp +++ b/faiss/utils/simd_impl/super_kmeans_kernels_avx512.cpp @@ -5,7 +5,7 @@ * LICENSE file in the root directory of this source tree. */ -#if defined(__x86_64__) && defined(__AVX512F__) +#ifdef COMPILE_SIMD_AVX512 #include @@ -42,4 +42,4 @@ float block_l2(const float* x, const float* y, int n) { } // namespace detail } // namespace faiss -#endif // __x86_64__ && __AVX512F__ +#endif // COMPILE_SIMD_AVX512 diff --git a/tests/test_super_kmeans.py b/tests/test_super_kmeans.py new file mode 100644 index 000000000..f0a339574 --- /dev/null +++ b/tests/test_super_kmeans.py @@ -0,0 +1,188 @@ +# 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. + +import unittest + +import faiss +import numpy as np +from faiss.contrib.datasets import SyntheticDataset + + +class SuperKMeansTest(unittest.TestCase): + def test_train_rejects_too_few_points(self): + sc = faiss.SuperKMeans(128, 16) + x = SyntheticDataset(128, 10, 0, 0).get_train() # n=10 < k=16 + with self.assertRaises(RuntimeError): + sc.train(x) + + def test_objective_decreases_monotonically(self): + d, k, n = 32, 8, 1000 + x = SyntheticDataset(d, n, 0, 0).get_train() + + p = faiss.SuperKMeansParameters() + p.seed = 42 + p.niter = 10 + sc = faiss.SuperKMeans(d, k, p) + sc.train(x) + + stats = sc.iteration_stats + objs = [stats.at(i).obj for i in range(stats.size())] + for i in range(1, len(objs)): + self.assertLessEqual(objs[i], objs[i - 1] * 1.001) + + def test_pruned_iter_assignments_close_to_vanilla(self): + d, k, n = 64, 16, 2000 + x = SyntheticDataset(d, n, 0, 0).get_train() + + p = faiss.SuperKMeansParameters() + p.seed = 42 + p.niter = 10 + sc = faiss.SuperKMeans(d, k, p) + sc.train(x) + + quant = faiss.IndexFlatL2(d) + vanilla = faiss.Clustering(d, k) + vanilla.seed = 42 + vanilla.niter = 10 + vanilla.train(x, quant) + + sc_final = sc.iteration_stats.at(sc.iteration_stats.size() - 1).obj + v_stats = vanilla.iteration_stats + v_final = v_stats.at(v_stats.size() - 1).obj + self.assertLess(abs(sc_final - v_final) / v_final, 0.05) + + def test_pruning_rate_in_expected_range(self): + # The 10-d intrinsic manifold of SyntheticDataset gives ADSampling + # the assigned-vs-other distance gap it needs to prune effectively. + d, k, n = 64, 64, 5000 + x = SyntheticDataset(d, n, 0, 0).get_train() + + p = faiss.SuperKMeansParameters() + p.seed = 42 + p.niter = 5 + sc = faiss.SuperKMeans(d, k, p) + sc.train(x) + + rates = faiss.vector_to_array(sc.gemm_pruning_rates) + self.assertEqual(rates.shape, (5,)) + self.assertEqual(rates[0], 0.0) # iter 0: no pruning + # Lower bound 0.30: looser than the 0.93 design target because + # synthetic d=64 data lacks the cluster gap real embeddings have at + # d=768; the assertion still confirms meaningful pruning occurs. + for i in range(2, 5): + self.assertGreaterEqual(rates[i], 0.30) + + def test_pdx_offset_correct_for_unaligned_y_batch(self): + # PDX block layout is sized for the full k centroids, not per y-batch + # tile. The y-batch tile boundary must be invisible to PDX-offset + # arithmetic. k=70 with y_batch=32 forces a 6-wide tail tile; assert + # the resulting centroids are identical to running with y_batch=70 + # (single tile, no tail). + d, k, n = 64, 70, 5000 + x = SyntheticDataset(d, n, 0, 0).get_train() + + def train_with_y_batch(y_batch): + p = faiss.SuperKMeansParameters() + p.seed = 1234 + p.niter = 5 + p.y_batch = y_batch + sc = faiss.SuperKMeans(d, k, p) + sc.train(x) + return faiss.vector_to_array(sc.centroids) + + tiled = train_with_y_batch(32) # k % y_batch = 6 + single = train_with_y_batch(k) # k % y_batch = 0 + np.testing.assert_array_equal(tiled, single) + + def test_determinism(self): + # Bit-exact reproducibility requires MKL_CBWR=COMPATIBLE and a fixed + # OMP thread count. The BUCK target sets these in `env`; for OSS + # ctest the user must export them before running. + d, k, n = 128, 64, 5000 + x = SyntheticDataset(d, n, 0, 0).get_train() + + p = faiss.SuperKMeansParameters() + p.niter = 5 + p.seed = 1234 + + sc1 = faiss.SuperKMeans(d, k, p) + sc1.train(x) + sc2 = faiss.SuperKMeans(d, k, p) + sc2.train(x) + + np.testing.assert_array_equal( + faiss.vector_to_array(sc1.centroids), + faiss.vector_to_array(sc2.centroids), + ) + np.testing.assert_array_equal( + faiss.vector_to_array(sc1.gemm_pruning_rates), + faiss.vector_to_array(sc2.gemm_pruning_rates), + ) + + def test_rejects_nan_input(self): + # check_input_data_for_NaNs defaults to True. + sc = faiss.SuperKMeans(128, 64) + x = SyntheticDataset(128, 5000, 0, 0).get_train() + x.flat[42] = float("nan") + with self.assertRaises(RuntimeError): + sc.train(x) + + def test_cluster_splitting_with_many_duplicates(self): + # 80% duplicates force empty clusters after the first assignment; + # split_clusters must fire and produce distinct centroids. + d, k, n = 64, 16, 1000 + n_dup = int(n * 0.8) + + ds = SyntheticDataset(d, (n - n_dup) + 1, 0, 0).get_train() + x = np.empty((n, d), dtype="float32") + x[:n_dup] = ds[0] + x[n_dup:] = ds[1:] + + p = faiss.SuperKMeansParameters() + p.seed = 42 + p.niter = 5 + sc = faiss.SuperKMeans(d, k, p) + sc.train(x) + + stats = sc.iteration_stats + total_splits = sum(stats.at(i).nsplit for i in range(stats.size())) + self.assertGreater(total_splits, 0) + + centroids = faiss.vector_to_array(sc.centroids).reshape(k, d) + unique_centroids = {centroids[j].tobytes() for j in range(k)} + self.assertEqual(len(unique_centroids), k) + + +class SuperKmeansWrapperTest(unittest.TestCase): + def test_train_populates_kmeans_surface(self): + d, k, n = 64, 16, 2000 + x = SyntheticDataset(d, n, 0, 0).get_train() + + km = faiss.SuperKmeans(d, k, niter=5, seed=42) + final_obj = km.train(x) + + self.assertEqual(km.centroids.shape, (k, d)) + self.assertEqual(km.obj.shape, (5,)) + self.assertEqual(len(km.iteration_stats), 5) + self.assertEqual(final_obj, km.obj[-1]) + self.assertEqual(km.gemm_pruning_rates.shape, (5,)) + + D, I = km.assign(x) + self.assertEqual(I.shape, (n,)) + self.assertEqual(D.shape, (n,)) + self.assertTrue((I >= 0).all() and (I < k).all()) + + def test_final_obj_matches_vanilla_kmeans(self): + d, k, n = 64, 16, 2000 + x = SyntheticDataset(d, n, 0, 0).get_train() + + vanilla = faiss.Kmeans(d, k, niter=10, seed=42) + vanilla.train(x) + + sk = faiss.SuperKmeans(d, k, niter=10, seed=42) + sk.train(x) + + rel_diff = abs(sk.obj[-1] - vanilla.obj[-1]) / vanilla.obj[-1] + self.assertLess(rel_diff, 0.05) diff --git a/tests/test_super_kmeans_foundations.cpp b/tests/test_super_kmeans_foundations.cpp index 2e87aa287..4936237cd 100644 --- a/tests/test_super_kmeans_foundations.cpp +++ b/tests/test_super_kmeans_foundations.cpp @@ -171,35 +171,6 @@ TEST(BlockL2, ScalarMatchesReference) { } } -TEST(BlockL2, Avx512MatchesScalar) { -#if defined(__x86_64__) && defined(__AVX512F__) - constexpr int N = 64; - std::mt19937 rng(43); - std::uniform_real_distribution dist(-1.0f, 1.0f); - std::vector x(N), y(N); - for (int m = 0; m < N; ++m) { - x[m] = dist(rng); - y[m] = dist(rng); - } - - const float scalar = faiss::detail::block_l2( - x.data(), y.data(), N); - const float avx512 = faiss::detail::block_l2( - x.data(), y.data(), N); - EXPECT_NEAR(avx512, scalar, 1e-4f); - - for (int n = 1; n < N; ++n) { - const float s = faiss::detail::block_l2( - x.data(), y.data(), n); - const float a = faiss::detail::block_l2( - x.data(), y.data(), n); - EXPECT_NEAR(a, s, 1e-4f) << "n=" << n; - } -#else - GTEST_SKIP() << "AVX-512 test requires __AVX512F__"; -#endif -} - TEST(BlockL2, DispatchMatchesScalar) { constexpr int N = 64; std::mt19937 rng(44);