diff --git a/cmake/libs/libfaiss.cmake b/cmake/libs/libfaiss.cmake index 59fc1f26d..be3097de9 100644 --- a/cmake/libs/libfaiss.cmake +++ b/cmake/libs/libfaiss.cmake @@ -187,6 +187,7 @@ knowhere_file_glob( FAISS_DD_SVE_SRCS thirdparty/faiss/faiss/impl/pq_code_distance/pq_code_distance-sve.cpp thirdparty/faiss/faiss/utils/simd_impl/distances_arm_sve.cpp + thirdparty/faiss/faiss/utils/simd_impl/super_kmeans_kernels_sve.cpp ) # combine files list(APPEND FAISS_SVE_SRCS ${FAISS_DD_SVE_SRCS}) @@ -548,6 +549,7 @@ if(__AARCH64) knowhere_utils) if(SVE_AVAILABLE) target_link_libraries(faiss PUBLIC faiss_sve) + target_compile_definitions(faiss PRIVATE COMPILE_SIMD_ARM_SVE) endif() target_compile_definitions(faiss PRIVATE FINTEGER=int FAISS_ENABLE_DD COMPILE_SIMD_ARM_NEON) endif() diff --git a/include/knowhere/comp/index_param.h b/include/knowhere/comp/index_param.h index 9d45d6cfb..429345df0 100644 --- a/include/knowhere/comp/index_param.h +++ b/include/knowhere/comp/index_param.h @@ -145,6 +145,7 @@ namespace indexparam { constexpr const char* NPROBE = "nprobe"; constexpr const char* NLIST = "nlist"; constexpr const char* USE_ELKAN = "use_elkan"; +constexpr const char* USE_SUPER_KMEANS = "use_super_kmeans"; constexpr const char* NBITS = "nbits"; // PQ/SQ constexpr const char* M = "m"; // PQ param for IVFPQ constexpr const char* IVF_SQ_TYPE = "sq_type"; // SQ param for IVFSQ diff --git a/src/index/ivf/ivf.cc b/src/index/ivf/ivf.cc index 047cc9283..8c6c3c76e 100644 --- a/src/index/ivf/ivf.cc +++ b/src/index/ivf/ivf.cc @@ -19,6 +19,7 @@ #include "faiss/IndexIVFRaBitQ.h" #include "faiss/IndexIVFRaBitQFastScan.h" #include "faiss/IndexRefine.h" +#include "faiss/SuperKMeans.h" #include "faiss/VectorTransform.h" #include "faiss/cppcontrib/knowhere/IndexBinaryFlat.h" #include "faiss/cppcontrib/knowhere/IndexBinaryIVF.h" @@ -599,6 +600,7 @@ IvfIndexNode::TrainInternal(const DataSetPtr dataset, std:: } // apply clustering config ApplyClusteringConfig(index->cp); + index->cp.use_super_kmeans = ivf_flat_cfg.use_super_kmeans.value(); // train index->train(rows, static_cast(data)); // transfer ownership of qzr to index @@ -621,6 +623,7 @@ IvfIndexNode::TrainInternal(const DataSetPtr dataset, std:: } // apply clustering config ApplyClusteringConfig(index->cp); + index->cp.use_super_kmeans = ivf_flat_cc_cfg.use_super_kmeans.value(); // train index->train(rows, static_cast(data)); // transfer ownership of qzr to index @@ -652,6 +655,7 @@ IvfIndexNode::TrainInternal(const DataSetPtr dataset, std:: // apply clustering config ApplyClusteringConfig(index->get_base_ivf_index()->cp); + index->get_base_ivf_index()->cp.use_super_kmeans = ivf_pq_cfg.use_super_kmeans.value(); // train index->train(rows, static_cast(data)); // transfer ownership of qzr to index @@ -678,6 +682,14 @@ IvfIndexNode::TrainInternal(const DataSetPtr dataset, std:: } // apply clustering config ApplyClusteringConfig(base_index->cp); + // SuperKMeans requires d >= 2 * d_prime_min (default 32). Fall back to + // Clustering below that hard limit; use_super_kmeans otherwise remains + // the user's choice, including for small nlist values. + bool use_super_kmeans = scann_cfg.use_super_kmeans.value(); + if (use_super_kmeans && dim < 2 * faiss::SuperKMeansParameters{}.d_prime_min) { + use_super_kmeans = false; + } + base_index->cp.use_super_kmeans = use_super_kmeans; // create scann index, which does not base_index by default, // but owns the refine index by default omg if (scann_cfg.with_raw_data.value()) { @@ -713,6 +725,7 @@ IvfIndexNode::TrainInternal(const DataSetPtr dataset, std:: // apply clustering config ApplyClusteringConfig(index->get_base_ivf_index()->cp); + index->get_base_ivf_index()->cp.use_super_kmeans = ivf_sq_cfg.use_super_kmeans.value(); // train index->train(rows, static_cast(data)); // transfer ownership of qzr to index @@ -730,6 +743,7 @@ IvfIndexNode::TrainInternal(const DataSetPtr dataset, std:: index = std::make_unique(qzr.get(), dim, nlist, metric.value()); // apply clustering config ApplyClusteringConfig(index->cp); + index->cp.use_super_kmeans = ivf_bin_cfg.use_super_kmeans.value(); // train index->train(rows, static_cast(data)); // transfer ownership of qzr to index @@ -760,6 +774,7 @@ IvfIndexNode::TrainInternal(const DataSetPtr dataset, std:: } // apply clustering config ApplyClusteringConfig(index->cp); + index->cp.use_super_kmeans = ivf_sq_cc_cfg.use_super_kmeans.value(); // train index->train(rows, static_cast(data)); // transfer ownership of qzr to index @@ -782,6 +797,7 @@ IvfIndexNode::TrainInternal(const DataSetPtr dataset, std:: index = std::move(result.value()); // apply clustering config ApplyClusteringConfig(index->get_ivfrabitq_index()->cp); + index->get_ivfrabitq_index()->cp.use_super_kmeans = ivf_rabitq_cfg.use_super_kmeans.value(); // train index->train(rows, static_cast(data)); } @@ -798,6 +814,7 @@ IvfIndexNode::TrainInternal(const DataSetPtr dataset, std:: auto* fs_idx = index->get_fastscan_index(); if (fs_idx) { ApplyClusteringConfig(fs_idx->cp); + fs_idx->cp.use_super_kmeans = fs_cfg.use_super_kmeans.value(); } index->train(rows, static_cast(data)); } diff --git a/src/index/ivf/ivf_config.h b/src/index/ivf/ivf_config.h index 4bb844afd..03125f74c 100644 --- a/src/index/ivf/ivf_config.h +++ b/src/index/ivf/ivf_config.h @@ -27,6 +27,7 @@ class IvfConfig : public BaseConfig { CFG_INT nlist; CFG_INT nprobe; CFG_BOOL use_elkan; + CFG_BOOL use_super_kmeans; CFG_BOOL ensure_topk_full; // internal config, used for temp index CFG_INT max_empty_result_buckets; KNOWHERE_DECLARE_CONFIG(IvfConfig) { @@ -46,6 +47,10 @@ class IvfConfig : public BaseConfig { .set_default(true) .description("whether to use elkan algorithm") .for_train(); + KNOWHERE_CONFIG_DECLARE_FIELD(use_super_kmeans) + .set_default(false) + .description("whether to use SuperKMeans for coarse quantizer training") + .for_train(); KNOWHERE_CONFIG_DECLARE_FIELD(ensure_topk_full) .set_default(true) .description("whether to make sure topk results full") @@ -196,6 +201,13 @@ class ScannConfig : public IvfFlatConfig { .set_default(false) .description("whether to make sure topk results full") .for_search(); + // SCANN defaults to SuperKMeans for coarse quantizer training: the + // super-fast k-means variant is recall-equivalent to Clustering on + // inner-product data but trains significantly faster for large nlist. + KNOWHERE_CONFIG_DECLARE_FIELD(use_super_kmeans) + .set_default(true) + .description("whether to use SuperKMeans for coarse quantizer training") + .for_train(); } Status diff --git a/tests/ut/test_cluster.cc b/tests/ut/test_cluster.cc index 1fcc32e50..db9068cac 100644 --- a/tests/ut/test_cluster.cc +++ b/tests/ut/test_cluster.cc @@ -18,6 +18,7 @@ #include "catch2/generators/catch_generators.hpp" #include "faiss/Clustering.h" #include "faiss/IndexFlat.h" +#include "faiss/SuperKMeans.h" #include "faiss/cppcontrib/knowhere/utils/binary_distances.h" #include "hnswlib/hnswalg.h" #include "knowhere/bitsetview.h" @@ -143,3 +144,68 @@ TEST_CASE("Test Kmeans With Float Vector", "[float metrics]") { REQUIRE(recall > kKnnRecallThreshold); } } + +// SuperKMeans spherical (inner-product) support: unit-normalized centroids +// make L2 assignment equivalent to IP argmax, so the final objective must +// track vanilla spherical Clustering and centroids must be unit norm. +TEST_CASE("Test SuperKMeans Spherical", "[cluster]") { + const int d = 64; + const int k = 16; + const size_t n = 2000; + + std::mt19937 rng(42); + std::normal_distribution dist(0.f, 1.f); + std::vector x(n * d); + for (auto& v : x) { + v = dist(rng); + } + + faiss::SuperKMeansParameters sp; + sp.seed = 42; + sp.niter = 10; + sp.spherical = true; + faiss::SuperKMeans sc(d, k, sp); + sc.train(n, x.data()); + + // Centroids must be unit norm under spherical clustering. + for (int j = 0; j < k; ++j) { + float norm = 0.f; + for (int i = 0; i < d; ++i) { + norm += sc.centroids[j * d + i] * sc.centroids[j * d + i]; + } + norm = std::sqrt(norm); + REQUIRE(norm == Catch::Approx(1.f).margin(1e-4)); + } + + // Final objective must track vanilla spherical Clustering. + const float sc_final = sc.iteration_stats.at(sc.iteration_stats.size() - 1).obj; + faiss::ClusteringParameters vp; + vp.seed = 42; + vp.niter = 10; + vp.spherical = true; + faiss::Clustering vanilla(d, k, vp); + faiss::IndexFlatL2 quantizer(d); + vanilla.train(n, x.data(), quantizer); + const float v_final = vanilla.iteration_stats.at(vanilla.iteration_stats.size() - 1).obj; + REQUIRE(std::abs(sc_final - v_final) / v_final < 0.05f); +} + +TEST_CASE("Test SuperKMeans with 256 centroids", "[cluster]") { + constexpr int d = 64; + constexpr int k = 256; + constexpr size_t n = 1024; + + std::mt19937 rng(43); + std::normal_distribution dist(0.f, 1.f); + std::vector x(n * d); + for (auto& v : x) { + v = dist(rng); + } + + faiss::SuperKMeansParameters sp; + sp.seed = 43; + sp.niter = 2; + faiss::SuperKMeans clustering(d, k, sp); + REQUIRE_NOTHROW(clustering.train(n, x.data())); + REQUIRE(clustering.centroids.size() == static_cast(k * d)); +} diff --git a/tests/ut/test_scann_superkmeans.cc b/tests/ut/test_scann_superkmeans.cc new file mode 100644 index 000000000..625a9a9e3 --- /dev/null +++ b/tests/ut/test_scann_superkmeans.cc @@ -0,0 +1,134 @@ +// Copyright (C) 2019-2023 Zilliz. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software distributed under the License +// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +// or implied. See the License for the specific language governing permissions and limitations under the License. + +#include +#include + +#include "catch2/catch_approx.hpp" +#include "catch2/catch_test_macros.hpp" +#include "knowhere/comp/brute_force.h" +#include "knowhere/comp/index_param.h" +#include "knowhere/index/index_factory.h" +#include "knowhere/version.h" +#include "utils.h" + +namespace { + +// Build a SCANN index over train_ds with the given use_super_kmeans value +// and return recall@k against brute-force ground truth. +float +BuildScannAndRecall(const knowhere::DataSetPtr& train_ds, const knowhere::DataSetPtr& query_ds, int64_t nlist, + int64_t nprobe, int64_t topk, bool use_super_kmeans) { + const auto version = knowhere::Version::GetCurrentVersion().VersionNumber(); + auto idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_FAISS_SCANN, version) + .value(); + + knowhere::Json cfg; + cfg[knowhere::meta::METRIC_TYPE] = knowhere::metric::IP; + cfg[knowhere::indexparam::NLIST] = nlist; + cfg[knowhere::indexparam::NPROBE] = nprobe; + cfg[knowhere::indexparam::SUB_DIM] = 4; + cfg[knowhere::indexparam::WITH_RAW_DATA] = false; + cfg[knowhere::indexparam::USE_SUPER_KMEANS] = use_super_kmeans; + + REQUIRE(idx.Build(train_ds, cfg) == knowhere::Status::success); + + knowhere::Json search_cfg; + search_cfg[knowhere::meta::METRIC_TYPE] = knowhere::metric::IP; + search_cfg[knowhere::meta::TOPK] = topk; + search_cfg[knowhere::indexparam::NPROBE] = nprobe; + auto results = idx.Search(query_ds, search_cfg, nullptr); + REQUIRE(results.has_value()); + + auto gt = knowhere::BruteForce::Search( + train_ds, query_ds, + knowhere::Json{{knowhere::meta::METRIC_TYPE, knowhere::metric::IP}, {knowhere::meta::TOPK, topk}}, nullptr); + REQUIRE(gt.has_value()); + + return GetKNNRecall(*gt.value(), *results.value()); +} + +} // namespace + +TEST_CASE("SCANN use_super_kmeans default matches Clustering recall", "[scann]") { + constexpr int64_t nb = 2000; + constexpr int64_t nq = 100; + constexpr int64_t dim = 64; + constexpr int64_t topk = 10; + constexpr int64_t nlist = 128; + constexpr int64_t nprobe = 8; + + const auto train_ds = GenDataSet(nb, dim, kSeed); + const auto query_ds = GenDataSet(nq, dim, kSeed); + + const float super_recall = BuildScannAndRecall(train_ds, query_ds, nlist, nprobe, topk, true); + const float cluster_recall = BuildScannAndRecall(train_ds, query_ds, nlist, nprobe, topk, false); + + CAPTURE(super_recall, cluster_recall); + // SuperKMeans coarse quantizer training is recall-equivalent to Clustering. + // MatchNlist shrinks nlist on this small synthetic set. The enabled build + // still honors SuperKMeans for the resulting small centroid count; recall + // equivalence between the two clustering implementations is the invariant. + REQUIRE(super_recall == Catch::Approx(cluster_recall).margin(0.02f)); +} + +TEST_CASE("SCANN use_super_kmeans field is honored", "[scann]") { + // Explicitly disabled must build and search successfully too. + constexpr int64_t nb = 1000; + constexpr int64_t nq = 50; + constexpr int64_t dim = 32; + constexpr int64_t topk = 10; + constexpr int64_t nlist = 64; + constexpr int64_t nprobe = 4; + + const auto train_ds = GenDataSet(nb, dim, kSeed + 1); + const auto query_ds = GenDataSet(nq, dim, kSeed + 1); + + const float cluster_recall = BuildScannAndRecall(train_ds, query_ds, nlist, nprobe, topk, false); + REQUIRE(cluster_recall > 0.0f); +} + +// SCANN with default use_super_kmeans=true must not fail the build on +// low-dimensional data (e.g. d=16 emb-list scenarios) where SuperKMeans is +// not applicable; it should fall back to Clustering. +TEST_CASE("SCANN low-dim build succeeds with default superkmeans", "[scann]") { + constexpr int64_t nb = 200; + constexpr int64_t nq = 20; + constexpr int64_t dim = 16; + constexpr int64_t topk = 5; + constexpr int64_t nlist = 16; + constexpr int64_t nprobe = 2; + + const auto train_ds = GenDataSet(nb, dim, kSeed + 2); + const auto query_ds = GenDataSet(nq, dim, kSeed + 2); + + const auto version = knowhere::Version::GetCurrentVersion().VersionNumber(); + auto idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_FAISS_SCANN, version) + .value(); + + knowhere::Json cfg; + cfg[knowhere::meta::METRIC_TYPE] = knowhere::metric::IP; + cfg[knowhere::indexparam::NLIST] = nlist; + cfg[knowhere::indexparam::NPROBE] = nprobe; + cfg[knowhere::indexparam::SUB_DIM] = 2; + cfg[knowhere::indexparam::WITH_RAW_DATA] = false; + // Default use_super_kmeans=true; must fall back to Clustering for d=16. + REQUIRE(idx.Build(train_ds, cfg) == knowhere::Status::success); + + knowhere::Json search_cfg; + search_cfg[knowhere::meta::METRIC_TYPE] = knowhere::metric::IP; + search_cfg[knowhere::meta::TOPK] = topk; + search_cfg[knowhere::indexparam::NPROBE] = nprobe; + auto results = idx.Search(query_ds, search_cfg, nullptr); + REQUIRE(results.has_value()); +} diff --git a/thirdparty/faiss/faiss/CMakeLists.txt b/thirdparty/faiss/faiss/CMakeLists.txt index 1a4a30a51..83e5a8753 100644 --- a/thirdparty/faiss/faiss/CMakeLists.txt +++ b/thirdparty/faiss/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 diff --git a/thirdparty/faiss/faiss/Clustering.h b/thirdparty/faiss/faiss/Clustering.h index acb501bce..3ed95b13a 100644 --- a/thirdparty/faiss/faiss/Clustering.h +++ b/thirdparty/faiss/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 { diff --git a/thirdparty/faiss/faiss/IndexIVF.cpp b/thirdparty/faiss/faiss/IndexIVF.cpp index 6fdba6785..122b2cce0 100644 --- a/thirdparty/faiss/faiss/IndexIVF.cpp +++ b/thirdparty/faiss/faiss/IndexIVF.cpp @@ -6,6 +6,7 @@ */ #include +#include #include #include @@ -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(d), static_cast(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(super_cp) = cp; + SuperKMeans clus( + static_cast(d), static_cast(nlist), super_cp); + clus.train(n, x); quantizer->add(nlist, clus.centroids.data()); } else { - clus.train(n, x, *quantizer); + Clustering clus(static_cast(d), static_cast(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) { diff --git a/thirdparty/faiss/faiss/SuperKMeans.cpp b/thirdparty/faiss/faiss/SuperKMeans.cpp index 9ca879c05..db9ad2b08 100644 --- a/thirdparty/faiss/faiss/SuperKMeans.cpp +++ b/thirdparty/faiss/faiss/SuperKMeans.cpp @@ -59,7 +59,7 @@ namespace { struct TrainState { /// Orthogonal rotation. Train in rotated space (X_tilde = X * R); /// un-rotate centroids before return. - faiss::RandomRotationMatrix R; + std::unique_ptr R; std::vector 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 { + // 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(d); + } + return std::make_unique(d, d); + }()) {} }; /// PDX block layout for the trailing pruning sweep: block b covers original @@ -295,10 +305,17 @@ std::unique_ptr setup_train_state( "SuperKMeans: training set size exceeds INT_MAX after sampling"); state.n = static_cast(nx); - state.R.init(cp.seed); + if (auto* R = dynamic_cast(state.R.get())) { + R->init(cp.seed); + } else { + auto* dense_rotation = + dynamic_cast(state.R.get()); + FAISS_ASSERT(dense_rotation != nullptr); + dense_rotation->init(cp.seed); + } state.X_tilde.resize(static_cast(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 +330,9 @@ std::unique_ptr setup_train_state( state.X_tilde.data() + static_cast(perm[j]) * d, sizeof(float) * d); } + if (cp.spherical) { + fvec_renorm_L2(d, k, state.Y_tilde.data()); + } } state.d_prime = @@ -339,7 +359,7 @@ std::unique_ptr setup_train_state( /// reverse_transform applies R^T = R^-1. void untransform_centroids( std::vector& centroids, - const RandomRotationMatrix& R, + const VectorTransform& R, int d, int k, const float* Y_tilde) { @@ -411,7 +431,7 @@ void SuperKMeans::train(idx_t n, const float* x) { static_cast(k) * cp.min_points_per_centroid); } - TrainState state(d); + TrainState state(d, cp.spherical); std::vector labels64; SuperKMeansAssignScratch assign_scratch; std::vector hassign; @@ -446,6 +466,9 @@ void SuperKMeans::train(idx_t n, const float* x) { const int nsplit = update_centroids_and_split(d, k, state, labels64, hassign); + if (cp.spherical) { + fvec_renorm_L2(d, k, state.Y_tilde.data()); + } const float pruning_rate = (iter == 0) ? 0.0f : adapt_d_prime(d, cp, state, total_pairs, pruned_at_gemm); @@ -492,7 +515,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( diff --git a/thirdparty/faiss/faiss/SuperKMeans.h b/thirdparty/faiss/faiss/SuperKMeans.h index 2ffdf8ee6..c91a9bb44 100644 --- a/thirdparty/faiss/faiss/SuperKMeans.h +++ b/thirdparty/faiss/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). // diff --git a/thirdparty/faiss/faiss/VectorTransform.cpp b/thirdparty/faiss/faiss/VectorTransform.cpp index 8516bd45c..61365149d 100644 --- a/thirdparty/faiss/faiss/VectorTransform.cpp +++ b/thirdparty/faiss/faiss/VectorTransform.cpp @@ -513,6 +513,44 @@ 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; + // Reverse of apply_noalloc: three unnormalized FWHT rounds scale norms + // by (sqrt(p))^3 = p*sqrt(p); the forward pass cancels this with + // total_scale = 1/(p*sqrt(p)), so the inverse applies the same factor. + const float inverse_scale = 1.0f / (p * std::sqrt(static_cast(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(&other); FAISS_THROW_IF_NOT_MSG(hr, "failed to cast to HadamardRotation"); diff --git a/thirdparty/faiss/faiss/VectorTransform.h b/thirdparty/faiss/faiss/VectorTransform.h index 6288f16f9..0dd6f26ba 100644 --- a/thirdparty/faiss/faiss/VectorTransform.h +++ b/thirdparty/faiss/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() {} diff --git a/thirdparty/faiss/faiss/impl/ProductQuantizer.cpp b/thirdparty/faiss/faiss/impl/ProductQuantizer.cpp index 15260ed35..3b917e4f5 100644 --- a/thirdparty/faiss/faiss/impl/ProductQuantizer.cpp +++ b/thirdparty/faiss/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([&]() { + // 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([&]() { switch (nbits) { case 8: compute_1_code(*this, x, code); diff --git a/thirdparty/faiss/faiss/utils/simd_impl/distances_arm_sve.cpp b/thirdparty/faiss/faiss/utils/simd_impl/distances_arm_sve.cpp index 0902db345..34b87b61c 100644 --- a/thirdparty/faiss/faiss/utils/simd_impl/distances_arm_sve.cpp +++ b/thirdparty/faiss/faiss/utils/simd_impl/distances_arm_sve.cpp @@ -7,6 +7,9 @@ #include +#include +#include + #include #include @@ -520,6 +523,100 @@ void fvec_L2sqr_ny( } } +namespace { + +/// Low-dimensional L2sqr nearest (D in {2,4,8}). The data is row-major +/// (centroid c is y[c*D .. c*D+D]). Each SVE lane tracks one centroid across +/// batches, keeping the lane-local minimum and index in registers. The scratch +/// buffer is not written, matching the AVX2/AVX512 D2/D4/D8 implementations. +template +size_t fvec_L2sqr_ny_nearest_lowdim( + float* /*distances_tmp_buffer*/, + const float* x, + const float* y, + size_t ny) { + const size_t lanes = svcntw(); + svfloat32_t global_mins = svdup_n_f32(HUGE_VALF); + svuint32_t global_ids = svdup_n_u32(0); + svuint32_t current_ids = svindex_u32(0, 1); + + for (size_t c = 0; c < ny; c += lanes) { + const svbool_t pg = svwhilelt_b32_u64(c, ny); + svfloat32_t distances = svdup_n_f32(0.0f); + for (uint32_t j = 0; j < D; ++j) { + const svuint32_t offsets = svindex_u32(j, D); + const svfloat32_t yv = + svld1_gather_u32index_f32(pg, y + c * D, offsets); + const svfloat32_t diff = svsub_n_f32_x(pg, yv, x[j]); + distances = svmla_f32_m(pg, distances, diff, diff); + } + + const svbool_t closer = svcmplt_f32(pg, distances, global_mins); + global_mins = svsel_f32(closer, distances, global_mins); + global_ids = svsel_u32(closer, current_ids, global_ids); + current_ids = svadd_n_u32_x( + pg, current_ids, static_cast(lanes)); + } + + const svbool_t all = svptrue_b32(); + const float global_min = svminv_f32(all, global_mins); + return svminv_u32( + svcmpeq_n_f32(all, global_mins, global_min), global_ids); +} + +} // namespace + +// Specializations for low-dimensional L2sqr nearest, mirroring the +// fvec_L2sqr_ny_nearest_D2/D4/D8 declarations used by the SSE/AVX +// implementations. One SVE lane covers one D-float centroid. +template +size_t fvec_L2sqr_ny_nearest_D2( + float* distances_tmp_buffer, + const float* x, + const float* y, + size_t ny); + +template +size_t fvec_L2sqr_ny_nearest_D4( + float* distances_tmp_buffer, + const float* x, + const float* y, + size_t ny); + +template +size_t fvec_L2sqr_ny_nearest_D8( + float* distances_tmp_buffer, + const float* x, + const float* y, + size_t ny); + +template <> +size_t fvec_L2sqr_ny_nearest_D2( + float* distances_tmp_buffer, + const float* x, + const float* y, + size_t ny) { + return fvec_L2sqr_ny_nearest_lowdim<2>(distances_tmp_buffer, x, y, ny); +} + +template <> +size_t fvec_L2sqr_ny_nearest_D4( + float* distances_tmp_buffer, + const float* x, + const float* y, + size_t ny) { + return fvec_L2sqr_ny_nearest_lowdim<4>(distances_tmp_buffer, x, y, ny); +} + +template <> +size_t fvec_L2sqr_ny_nearest_D8( + float* distances_tmp_buffer, + const float* x, + const float* y, + size_t ny) { + return fvec_L2sqr_ny_nearest_lowdim<8>(distances_tmp_buffer, x, y, ny); +} + template <> size_t fvec_L2sqr_ny_nearest( float* distances_tmp_buffer, @@ -527,6 +624,25 @@ size_t fvec_L2sqr_ny_nearest( const float* y, size_t d, size_t ny) { + // Low-dimensional L2sqr nearest (d in {2,4,8}): each SVE lane tracks one + // centroid while D gathers load its components. d is the SIMD layout + // condition; ny is the generic loop bound. The scratch buffer is not + // written (matching the x86 D2/D4/D8 implementations); minima and indices + // remain in registers. + if (d == 2 || d == 4 || d == 8) { + switch (d) { + case 2: + return fvec_L2sqr_ny_nearest_D2( + distances_tmp_buffer, x, y, ny); + case 4: + return fvec_L2sqr_ny_nearest_D4( + distances_tmp_buffer, x, y, ny); + default: + return fvec_L2sqr_ny_nearest_D8( + distances_tmp_buffer, x, y, ny); + } + } + fvec_L2sqr_ny(distances_tmp_buffer, x, y, d, ny); size_t nearest_idx = 0; diff --git a/thirdparty/faiss/faiss/utils/simd_impl/super_kmeans_kernels.h b/thirdparty/faiss/faiss/utils/simd_impl/super_kmeans_kernels.h index b36c05f5c..fbbd9daff 100644 --- a/thirdparty/faiss/faiss/utils/simd_impl/super_kmeans_kernels.h +++ b/thirdparty/faiss/faiss/utils/simd_impl/super_kmeans_kernels.h @@ -39,5 +39,10 @@ template <> float block_l2(const float* x, const float* y, int n); #endif +#ifdef COMPILE_SIMD_ARM_SVE +template <> +float block_l2(const float* x, const float* y, int n); +#endif + } // namespace detail } // namespace faiss diff --git a/thirdparty/faiss/faiss/utils/simd_impl/super_kmeans_kernels_sve.cpp b/thirdparty/faiss/faiss/utils/simd_impl/super_kmeans_kernels_sve.cpp new file mode 100644 index 000000000..e41ef5d2b --- /dev/null +++ b/thirdparty/faiss/faiss/utils/simd_impl/super_kmeans_kernels_sve.cpp @@ -0,0 +1,34 @@ +/* + * 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_ARM_SVE + +#include + +#include + +namespace faiss { +namespace detail { + +template <> +float block_l2(const float* x, const float* y, int n) { + svfloat32_t acc = svdup_n_f32(0.0f); + const int lanes = static_cast(svcntw()); + for (int m = 0; m < n; m += lanes) { + const svbool_t pg = svwhilelt_b32(m, n); + const svfloat32_t xv = svld1_f32(pg, x + m); + const svfloat32_t yv = svld1_f32(pg, y + m); + const svfloat32_t diff = svsub_f32_x(pg, xv, yv); + acc = svmla_f32_m(pg, acc, diff, diff); + } + return svaddv_f32(svptrue_b32(), acc); +} + +} // namespace detail +} // namespace faiss + +#endif // COMPILE_SIMD_ARM_SVE diff --git a/thirdparty/faiss/tests/test_distances_dispatch.cpp b/thirdparty/faiss/tests/test_distances_dispatch.cpp index 4217bfc64..861a6c818 100644 --- a/thirdparty/faiss/tests/test_distances_dispatch.cpp +++ b/thirdparty/faiss/tests/test_distances_dispatch.cpp @@ -32,6 +32,7 @@ #include +#include #include #include #include @@ -299,3 +300,52 @@ TEST(DistancesDispatch, FvecL2sqrNyNearest_AllLevels) { levels); } } + +// The low-dimensional nearest path (d in {2,4,8}, one SIMD lane per +// centroid) must agree with the generic implementation for any ny, including +// values around common SIMD widths and non-multiple tails. +TEST(DistancesDispatch, FvecL2sqrNyNearest_LowDim) { + SIMDLevelGuard guard; + auto levels = available_levels(); + SKIP_IF_SINGLE_LEVEL(levels); + constexpr size_t kLowDims[] = {2, 4, 8}; + constexpr size_t kNy[] = { + 1, 2, 3, 7, 8, 9, 15, 16, 17, + 23, 31, 32, 33, 63, 64, 65, 255, 256, 257}; + for (size_t d : kLowDims) { + for (size_t ny : kNy) { + auto x = rand_vec(d, 40); + auto y = rand_vec(d * ny, 41); + check_index_at_levels( + [&]() { + std::vector tmp(ny); + return fvec_L2sqr_ny_nearest( + tmp.data(), x.data(), y.data(), d, ny); + }, + levels); + } + } +} + +TEST(DistancesDispatch, FvecL2sqrNyNearest_LowDimKeepsFirstTie) { + SIMDLevelGuard guard; + auto levels = available_levels(); + SKIP_IF_SINGLE_LEVEL(levels); + constexpr size_t kLowDims[] = {2, 4, 8}; + constexpr size_t ny = 23; + constexpr size_t expected = 3; + for (size_t d : kLowDims) { + std::vector x(d, 0.0f); + std::vector y(d * ny, 1.0f); + std::fill_n(y.data() + expected * d, d, 0.0f); + std::fill_n(y.data() + 17 * d, d, 0.0f); + auto nearest = [&]() { + std::vector tmp(ny); + return fvec_L2sqr_ny_nearest( + tmp.data(), x.data(), y.data(), d, ny); + }; + SIMDConfig::set_level(SIMDLevel::NONE); + EXPECT_EQ(expected, nearest()); + check_index_at_levels(nearest, levels); + } +}