From 430a41ee4c52cd070470e180c6d0f57888588d6b Mon Sep 17 00:00:00 2001 From: Jaeyoon Jeong Date: Thu, 1 Oct 2026 04:06:09 -0700 Subject: [PATCH] Native fp16 k-means training for Faiss IVF (#5671) Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5671 Train IVF coarse quantizers from packed fp16 (IEEE binary16) data without an fp32 copy of the training set, and let the k-means assignment step consume fp16 natively. **k-means: `Clustering::train_ex(n, x, NumericType, index, weights)`** Float16 runs k-means directly on the fp16 rows, without a codec: - NaN/Inf check on the fp16 bit patterns - subsampling and random initialization copy fp16 rows - the centroid update accumulates fp16 rows into the fp32 centroids with F16C/NEON (`compute_centroids_fp16`) - the assignment step passes blocks of at most `decode_block_size` fp16 rows to `Index::search_ex(Float16)` Only k-means++ / AFK-MC2 initialization still widens the (subsampled) training set. **Assignment: `Index::search_ex(Float16)`** - Default for every index: widen bounded blocks and call `search()`. - `IndexFlat` (L2 / IP): the BLAS kernels take fp16 queries (`knn_L2sqr_fp16`, `knn_inner_product_fp16`). Queries are widened inside the kernel in chunks of at most 16 MiB, each reused for all database tiles; results are bit-identical to the fp32 path on the rounded queries. The AVX2 / AVX512 / SVE top-1 kernels share one templated body. Selectors, small inputs, d <= 32 and database-parallel searches widen chunks and use the fp32 path. `IndexFlat1D` and `IndexFlatPanorama` keep their own search. - GPU: every `GpuIndex` accepts Float16 (queries are copied to the device as fp16 and widened there); `GpuIndexFlat` with `useFloat16` storage searches fp16 queries directly with the fp16 GEMM. **IVF** - `IndexIVF::train_ex(Float16)` and `Level1Quantizer::train_q1_ex` use the native path; only the `train_encoder_num_vectors()` rows sampled for the encoder are widened. `IndexIVFFlat` trains only the coarse quantizer; `IndexIVFFlatDedup` rejects Float16. - `train_encoded` (IVF and Clustering) remains the generic codec path, e.g. for decode-time normalization. - `faiss::fp16_to_fp32` (faiss/utils/utils.h) is the shared bulk widening helper. **Performance** (opt build, AMD Genoa: F16C conversion only, no fp16 arithmetic) - top-1 knn, d=256, 65536 queries x 1024 centroids: fp32 22.2 ms, fp16 25.3 ms - k-means, 1024 centroids, 10 iterations: fp32 270 ms, fp16 319 ms, identical centroids On CPUs without fp16 arithmetic the GEMM still runs in fp32 on widened tiles, so fp16 is not faster than fp32 there. The gains are memory (fp16 data plus a bounded fp32 tile instead of a full fp32 copy) and a single hook (`search_ex`) for fp16-capable backends; GPU `useFloat16` already runs the fp16 GEMM directly. Reviewed By: mnorris11 Differential Revision: D121848262 fbshipit-source-id: 1400d3eeb7c713359895923dcb86c544570cba69 --- faiss/CMakeLists.txt | 3 + faiss/Clustering.cpp | 167 +++- faiss/Clustering.h | 29 + faiss/Index.cpp | 36 + faiss/Index.h | 19 +- faiss/IndexFlat.cpp | 56 ++ faiss/IndexFlat.h | 31 + faiss/IndexIVF.cpp | 212 ++++- faiss/IndexIVF.h | 46 +- faiss/IndexIVFFlat.cpp | 33 + faiss/IndexIVFFlat.h | 16 + faiss/gpu/GpuIndex.cu | 28 + faiss/gpu/GpuIndex.h | 17 +- faiss/gpu/GpuIndexFlat.cu | 35 + faiss/gpu/GpuIndexFlat.h | 11 + faiss/gpu/test/TestGpuIndexFlat.cpp | 176 +++++ faiss/impl/ClusteringHelpers.cpp | 181 +++-- faiss/impl/ClusteringHelpers.h | 27 + faiss/utils/distances.cpp | 232 +++++- faiss/utils/distances.h | 30 + faiss/utils/simd_impl/distances_arm_sve.cpp | 42 +- faiss/utils/simd_impl/distances_avx2.cpp | 42 +- faiss/utils/simd_impl/distances_avx512.cpp | 42 +- .../simd_impl/exhaustive_L2sqr_blas_cmax.h | 93 +++ faiss/utils/simd_impl/fp16_kernels.h | 61 ++ faiss/utils/simd_impl/fp16_kernels_avx2.cpp | 56 ++ faiss/utils/simd_impl/fp16_kernels_neon.cpp | 52 ++ faiss/utils/utils.cpp | 22 + faiss/utils/utils.h | 8 + tests/test_ivf_index.cpp | 741 ++++++++++++++++++ 30 files changed, 2394 insertions(+), 150 deletions(-) create mode 100644 faiss/utils/simd_impl/fp16_kernels.h create mode 100644 faiss/utils/simd_impl/fp16_kernels_avx2.cpp create mode 100644 faiss/utils/simd_impl/fp16_kernels_neon.cpp diff --git a/faiss/CMakeLists.txt b/faiss/CMakeLists.txt index 601d7027b..b709bd0b1 100644 --- a/faiss/CMakeLists.txt +++ b/faiss/CMakeLists.txt @@ -17,6 +17,7 @@ set(FAISS_SIMD_AVX2_SRC impl/approx_topk/avx2.cpp impl/binary_hamming/avx2.cpp utils/simd_impl/distances_avx2.cpp + utils/simd_impl/fp16_kernels_avx2.cpp utils/hamming_distance/hamming_avx2.cpp utils/simd_impl/partitioning_avx2.cpp utils/distances_fused/simdlib_based.cpp @@ -57,6 +58,7 @@ set(FAISS_SIMD_NEON_SRC impl/binary_hamming/neon.cpp impl/pq_code_distance/neon.cpp utils/simd_impl/distances_aarch64.cpp + utils/simd_impl/fp16_kernels_neon.cpp utils/hamming_distance/hamming_neon.cpp utils/simd_impl/partitioning_neon.cpp utils/distances_fused/simdlib_based_neon.cpp @@ -418,6 +420,7 @@ set(FAISS_HEADERS utils/simd_impl/distances_simdlib256.h utils/hamming_distance/hamming_impl.h utils/simd_impl/exhaustive_L2sqr_blas_cmax.h + utils/simd_impl/fp16_kernels.h utils/simd_impl/super_kmeans_dispatch.h utils/simd_impl/super_kmeans_kernels.h utils/simd_impl/IVFFlatScanner-inl.h diff --git a/faiss/Clustering.cpp b/faiss/Clustering.cpp index b82b1be40..38a2d1d00 100644 --- a/faiss/Clustering.cpp +++ b/faiss/Clustering.cpp @@ -11,6 +11,7 @@ #include #include +#include #include #include #include @@ -63,6 +64,47 @@ void Clustering::train_encoded( const Index* codec, Index& index, const float* weights) { + train_impl(nx, x_in, codec, NumericType::Float32, index, weights); +} + +void Clustering::train_ex( + idx_t nx, + const void* x_in, + NumericType numeric_type, + Index& index, + const float* weights) { + if (numeric_type == NumericType::Float32) { + train(nx, static_cast(x_in), index, weights); + return; + } + FAISS_THROW_IF_NOT_MSG( + numeric_type == NumericType::Float16, + "Clustering::train_ex: unsupported numeric type"); + train_impl( + nx, + static_cast(x_in), + nullptr, + numeric_type, + index, + weights); +} + +void Clustering::train_impl( + idx_t nx, + const uint8_t* x_in, + const Index* codec, + NumericType numeric_type, + Index& index, + const float* weights) { + FAISS_THROW_IF_NOT_MSG( + codec || numeric_type == NumericType::Float32 || + numeric_type == NumericType::Float16, + "unsupported numeric type for clustering"); + // packed IEEE binary16 rows, consumed without a codec + const bool fp16 = !codec && numeric_type == NumericType::Float16; + // rows that are not fp32 are widened by blocks for the assignment index + const bool encoded = codec || fp16; + FAISS_THROW_IF_NOT_FMT( nx >= static_cast(k), "Number of training points (%" PRId64 @@ -70,6 +112,7 @@ void Clustering::train_encoded( "as large as number of clusters (%zd)", nx, k); + FAISS_THROW_IF_NOT_MSG(x_in, "training data must not be null"); FAISS_THROW_IF_NOT_FMT( (!codec || static_cast(codec->d) == d), @@ -83,23 +126,46 @@ void Clustering::train_encoded( int(index.d), int(d)); + FAISS_THROW_IF_NOT_MSG( + !encoded || decode_block_size > 0, + "decode_block_size must be positive"); + double t0 = getmillisecs(); if (!codec && check_input_data_for_NaNs) { // Check for NaNs in input data. Normally it is the user's - // responsibility, but it may spare us some hard-to-debug - // reports. - const float* x = reinterpret_cast(x_in); - for (size_t i = 0; i < nx * d; i++) { + // responsibility, but it may spare us some hard-to-debug reports. + if (fp16) { FAISS_THROW_IF_NOT_MSG( - std::isfinite(x[i]), "input contains NaN's or Inf's"); + detail::fp16_all_finite( + nx * d, reinterpret_cast(x_in)), + "input contains NaN's or Inf's"); + } else { + const float* x = reinterpret_cast(x_in); + for (size_t i = 0; i < nx * d; i++) { + FAISS_THROW_IF_NOT_MSG( + std::isfinite(x[i]), "input contains NaN's or Inf's"); + } } } const uint8_t* x = x_in; std::unique_ptr del1; std::unique_ptr del3; - size_t line_size = codec ? codec->sa_code_size() : sizeof(float) * d; + size_t line_size = codec ? codec->sa_code_size() + : fp16 ? sizeof(uint16_t) * d + : sizeof(float) * d; + + // widen n consecutive training rows to fp32 + auto decode_rows = [&](size_t n, const uint8_t* rows, float* out) { + if (codec) { + codec->sa_decode(n, rows, out); + } else if (fp16) { + fp16_to_fp32(n * d, reinterpret_cast(rows), out); + } else { + memcpy(out, rows, n * line_size); + } + }; if (static_cast(nx) > k * max_points_per_centroid) { uint8_t* x_new; @@ -129,11 +195,7 @@ void Clustering::train_encoded( nx); } centroids.resize(d * k); - if (!codec) { - memcpy(centroids.data(), x_in, sizeof(float) * d * k); - } else { - codec->sa_decode(nx, x_in, centroids.data()); - } + decode_rows(k, x_in, centroids.data()); // one fake iteration... ClusteringIterationStats stats = {0.0, 0.0, 0.0, 1.0, 0}; @@ -153,9 +215,8 @@ void Clustering::train_encoded( k, nredo, niter); - if (codec) { - printf("Input data encoded in %zd bytes per vector\n", - codec->sa_code_size()); + if (encoded) { + printf("Input data encoded in %zd bytes per vector\n", line_size); } } @@ -191,8 +252,10 @@ void Clustering::train_encoded( // initialize seed const uint64_t actual_seed = detail::get_actual_rng_seed(seed); - // temporary buffer to decode vectors during the optimization - std::vector decode_buffer(codec ? d * decode_block_size : 0); + // temporary buffer to decode codec-encoded vectors during the optimization + std::vector decode_buffer( + codec ? d * std::min(static_cast(nx), decode_block_size) + : 0); for (int redo = 0; redo < nredo; redo++) { if (verbose && nredo > 1) { @@ -210,28 +273,22 @@ void Clustering::train_encoded( std::vector perm(nx); rand_perm(perm.data(), nx, actual_seed + 1 + redo * 15486557L); for (size_t i = 0; i < k_to_init; i++) { - if (!codec) { - memcpy(centroids.data() + (n_input_centroids + i) * d, - x + perm[n_input_centroids + i] * line_size, - line_size); - } else { - codec->sa_decode( - 1, - x + perm[n_input_centroids + i] * line_size, - centroids.data() + (n_input_centroids + i) * d); - } + decode_rows( + 1, + x + perm[n_input_centroids + i] * line_size, + centroids.data() + (n_input_centroids + i) * d); } } else { // For k-means++ and AFK-MC², we need all vectors decoded const float* x_float = nullptr; std::vector x_decoded; - if (!codec) { + if (!encoded) { x_float = reinterpret_cast(x); } else { // Decode all vectors for initialization x_decoded.resize(nx * d); - codec->sa_decode(nx, x, x_decoded.data()); + decode_rows(nx, x, x_decoded.data()); x_float = x_decoded.data(); } @@ -268,7 +325,7 @@ void Clustering::train_encoded( for (int i = 0; i < niter; i++) { double t0s = getmillisecs(); - if (!codec) { + if (!encoded) { index.search( nx, reinterpret_cast(x), @@ -277,15 +334,26 @@ void Clustering::train_encoded( assign.get()); } else { // search by blocks of decode_block_size vectors - size_t code_size = codec->sa_code_size(); for (size_t i0 = 0; i0 < static_cast(nx); i0 += decode_block_size) { size_t i1 = i0 + decode_block_size; if (i1 > static_cast(nx)) { i1 = nx; } - codec->sa_decode( - i1 - i0, x + code_size * i0, decode_buffer.data()); + if (fp16) { + // the index consumes fp16 queries itself (natively + // if it can, see Index::search_ex) + index.search_ex( + i1 - i0, + x + line_size * i0, + NumericType::Float16, + 1, + dis.get() + i0, + assign.get() + i0); + continue; + } + decode_rows( + i1 - i0, x + line_size * i0, decode_buffer.data()); index.search( i1 - i0, decode_buffer.data(), @@ -308,17 +376,30 @@ void Clustering::train_encoded( std::vector hassign(k); size_t k_frozen = frozen_centroids ? n_input_centroids : 0; - detail::compute_centroids( - d, - k, - nx, - k_frozen, - x, - codec, - assign.get(), - weights, - hassign.data(), - centroids.data()); + if (fp16) { + detail::compute_centroids_fp16( + d, + k, + nx, + k_frozen, + reinterpret_cast(x), + assign.get(), + weights, + hassign.data(), + centroids.data()); + } else { + detail::compute_centroids( + d, + k, + nx, + k_frozen, + x, + codec, + assign.get(), + weights, + hassign.data(), + centroids.data()); + } int nsplit = detail::split_clusters( d, k, nx, k_frozen, hassign.data(), centroids.data()); diff --git a/faiss/Clustering.h b/faiss/Clustering.h index abb124681..3185fbfdb 100644 --- a/faiss/Clustering.h +++ b/faiss/Clustering.h @@ -148,11 +148,40 @@ struct Clustering : ClusteringParameters { Index& index, const float* weights = nullptr); + /** run k-means on vectors of the given numeric type + * + * Float32 is the same as train(). Float16 takes packed IEEE binary16 + * vectors (n * d uint16_t) and runs k-means natively on them: the NaN + * check, subsampling, initialization and centroid updates read the fp16 + * rows directly, and the assignment step passes blocks of at most + * decode_block_size fp16 rows to index.search_ex(Float16). Indexes with a + * native fp16 search consume them directly; the Index::search_ex default + * widens them block by block. There is no fp32 copy of the training set, + * unless init_method needs one. Centroids are fp32. + */ + void train_ex( + idx_t n, + const void* x, + NumericType numeric_type, + Index& index, + const float* weights = nullptr); + /// Post-process the centroids after each centroid update. /// includes optional L2 normalization and nearest integer rounding void post_process_centroids(); virtual ~Clustering() {} + + protected: + /// k-means on rows encoded by `codec` or, if codec is null, stored as + /// `numeric_type` (Float32 or Float16) + void train_impl( + idx_t nx, + const uint8_t* x_in, + const Index* codec, + NumericType numeric_type, + Index& index, + const float* weights); }; /** Exact 1D clustering algorithm diff --git a/faiss/Index.cpp b/faiss/Index.cpp index a9a8413b0..ee4d3801c 100644 --- a/faiss/Index.cpp +++ b/faiss/Index.cpp @@ -13,8 +13,11 @@ #include #include #include +#include +#include #include +#include namespace faiss { @@ -123,6 +126,39 @@ void Index::search_subset( FAISS_THROW_MSG("search_subset not implemented for this type of index"); } +void Index::search_ex( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params) const { + if (numeric_type == NumericType::Float32) { + search(n, static_cast(x), k, distances, labels, params); + return; + } + FAISS_THROW_IF_NOT_MSG( + numeric_type == NumericType::Float16, + "Index::search: unsupported numeric type"); + FAISS_THROW_IF_NOT(d > 0); + // widen blocks of at most 64 MiB of fp32 queries + const idx_t bs = std::max( + 1, idx_t(64) * 1024 * 1024 / (sizeof(float) * idx_t(d))); + const uint16_t* x16 = static_cast(x); + std::vector buffer(size_t(std::min(n, bs)) * d); + for (idx_t i0 = 0; i0 < n; i0 += bs) { + const idx_t i1 = std::min(n, i0 + bs); + fp16_to_fp32(size_t(i1 - i0) * d, x16 + i0 * d, buffer.data()); + search(i1 - i0, + buffer.data(), + k, + distances + i0 * k, + labels + i0 * k, + params); + } +} + void Index::search1(const float*, ResultHandler&, SearchParameters*) const { FAISS_THROW_MSG("search1 not implemented for this type of index"); } diff --git a/faiss/Index.h b/faiss/Index.h index 97edfe644..e40fbe13d 100644 --- a/faiss/Index.h +++ b/faiss/Index.h @@ -212,6 +212,12 @@ struct Index { idx_t* labels, const SearchParameters* params = nullptr) const = 0; + /** search() for queries of the given numeric type + * + * Float32 calls search(). By default, Float16 queries (packed IEEE + * binary16, size n * d) are widened to fp32 in bounded blocks and passed + * to search(); indexes with a native fp16 search override this. + */ virtual void search_ex( idx_t n, const void* x, @@ -219,18 +225,7 @@ struct Index { idx_t k, float* distances, idx_t* labels, - const SearchParameters* params = nullptr) const { - if (numeric_type == NumericType::Float32) { - search(n, - static_cast(x), - k, - distances, - labels, - params); - } else { - FAISS_THROW_MSG("Index::search: unsupported numeric type"); - } - } + const SearchParameters* params = nullptr) const; /** search one vector with a custom result handler */ virtual void search1( diff --git a/faiss/IndexFlat.cpp b/faiss/IndexFlat.cpp index 7204bc748..f4fdc103f 100644 --- a/faiss/IndexFlat.cpp +++ b/faiss/IndexFlat.cpp @@ -59,6 +59,40 @@ void IndexFlat::search( } } +void IndexFlat::search_ex( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params) const { + if (numeric_type != NumericType::Float16 || + (metric_type != METRIC_L2 && metric_type != METRIC_INNER_PRODUCT)) { + Index::search_ex(n, x, numeric_type, k, distances, labels, params); + return; + } + FAISS_THROW_IF_NOT(k > 0); + const IDSelector* sel = params ? params->sel : nullptr; + const uint16_t* x16 = static_cast(x); + if (metric_type == METRIC_INNER_PRODUCT) { + knn_inner_product_fp16( + x16, get_xb(), d, n, ntotal, k, distances, labels, sel); + } else { + knn_L2sqr_fp16( + x16, + get_xb(), + d, + n, + ntotal, + k, + distances, + labels, + nullptr, + sel); + } +} + void IndexFlat::range_search( idx_t n, const float* x, @@ -446,6 +480,17 @@ void IndexFlat1D::reset() { perm.clear(); } +void IndexFlat1D::search_ex( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params) const { + Index::search_ex(n, x, numeric_type, k, distances, labels, params); +} + void IndexFlat1D::search( idx_t n, const float* x, @@ -670,6 +715,17 @@ void IndexFlatPanorama::add(idx_t n, const float* x) { pano.compute_cumulative_sums(cum_sums.data(), offset, n, x); } +void IndexFlatPanorama::search_ex( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params) const { + Index::search_ex(n, x, numeric_type, k, distances, labels, params); +} + void IndexFlatPanorama::search( idx_t n, const float* x, diff --git a/faiss/IndexFlat.h b/faiss/IndexFlat.h index 7c8fb889f..52547ef1e 100644 --- a/faiss/IndexFlat.h +++ b/faiss/IndexFlat.h @@ -31,6 +31,17 @@ struct IndexFlat : IndexFlatCodes { idx_t* labels, const SearchParameters* params = nullptr) const override; + /// Float16 queries with METRIC_L2 / METRIC_INNER_PRODUCT are searched + /// natively (knn_L2sqr_fp16 / knn_inner_product_fp16). + void search_ex( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params = nullptr) const override; + void range_search( idx_t n, const float* x, @@ -135,6 +146,16 @@ struct IndexFlatPanorama : IndexFlat { idx_t* labels, const SearchParameters* params = nullptr) const override; + /// widens Float16 queries for search() (Index::search_ex default) + void search_ex( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params = nullptr) const override; + void range_search( idx_t n, const float* x, @@ -221,6 +242,16 @@ struct IndexFlat1D : IndexFlatL2 { float* distances, idx_t* labels, const SearchParameters* params = nullptr) const override; + + /// widens Float16 queries for search() (Index::search_ex default) + void search_ex( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params = nullptr) const override; }; } // namespace faiss diff --git a/faiss/IndexIVF.cpp b/faiss/IndexIVF.cpp index cda998889..53939aeb6 100644 --- a/faiss/IndexIVF.cpp +++ b/faiss/IndexIVF.cpp @@ -18,6 +18,7 @@ #include #include +#include #include #include @@ -53,28 +54,59 @@ Level1Quantizer::~Level1Quantizer() { } } -void Level1Quantizer::train_q1( +namespace { + +/** Trains `q1` on rows encoded by `codec` or, if codec is null, stored as + * `numeric_type`. */ +void train_q1_impl( + Level1Quantizer& q1, size_t n, - const float* x, + const uint8_t* x, + NumericType numeric_type, + const Index* codec, bool verbose, MetricType metric_type) { + Index* quantizer = q1.quantizer; + const size_t nlist = q1.nlist; + const ClusteringParameters& cp = q1.cp; + Index* clustering_index = q1.clustering_index; + const bool float32 = !codec && numeric_type == NumericType::Float32; + + // k-means on the training rows, with `assigner` as assignment index + auto train_clustering = [&](Clustering& clus, Index& assigner) { + if (codec) { + clus.train_encoded(n, x, codec, assigner); + } else { + clus.train_ex(n, x, numeric_type, assigner); + } + }; + FAISS_THROW_IF_NOT_MSG(quantizer, "IVF quantizer must not be null"); size_t d = quantizer->d; + FAISS_THROW_IF_NOT_FMT( + !codec || static_cast(codec->d) == d, + "Codec dimension %d not the same as quantizer dimension %d", + codec ? codec->d : 0, + int(d)); if (quantizer->is_trained && (static_cast(quantizer->ntotal) == nlist)) { if (verbose) { printf("IVF quantizer does not need training.\n"); } - } else if (quantizer_trains_alone == 1) { + } else if (q1.quantizer_trains_alone == 1) { + FAISS_THROW_IF_MSG( + codec, + "encoded training is not supported when the IVF quantizer " + "trains alone"); if (verbose) { printf("IVF quantizer trains alone...\n"); } quantizer->verbose = verbose; - quantizer->train(n, x); + quantizer->train_ex(n, x, numeric_type); FAISS_THROW_IF_NOT_MSG( static_cast(quantizer->ntotal) == nlist, "nlist not consistent with quantizer size"); - } else if (quantizer_trains_alone == 0) { + } else if (q1.quantizer_trains_alone == 0) { if (verbose) { printf("Training level-1 quantizer on %zd vectors in %zdD\n", n, d); } @@ -83,6 +115,9 @@ void Level1Quantizer::train_q1( cp.use_super_kmeans && clustering_index, "cp.use_super_kmeans is incompatible with a user-provided " "clustering_index: SuperKMeans assigns with its own index"); + FAISS_THROW_IF_MSG( + cp.use_super_kmeans && !float32, + "SuperKMeans requires fp32 training data"); quantizer->reset(); if (cp.use_super_kmeans) { @@ -90,19 +125,19 @@ void Level1Quantizer::train_q1( static_cast(super_cp) = cp; SuperKMeans clus( static_cast(d), static_cast(nlist), super_cp); - clus.train(n, x); + clus.train(n, reinterpret_cast(x)); quantizer->add(nlist, clus.centroids.data()); } else { Clustering clus(static_cast(d), static_cast(nlist), cp); if (clustering_index) { - clus.train(n, x, *clustering_index); + train_clustering(clus, *clustering_index); quantizer->add(nlist, clus.centroids.data()); } else { - clus.train(n, x, *quantizer); + train_clustering(clus, *quantizer); } } quantizer->is_trained = true; - } else if (quantizer_trains_alone == 2) { + } else if (q1.quantizer_trains_alone == 2) { if (verbose) { printf("Training L2 quantizer on %zd vectors in %zdD%s\n", n, @@ -118,9 +153,9 @@ void Level1Quantizer::train_q1( Clustering clus(static_cast(d), static_cast(nlist), cp); if (!clustering_index) { IndexFlatL2 assigner(d); - clus.train(n, x, assigner); + train_clustering(clus, assigner); } else { - clus.train(n, x, *clustering_index); + train_clustering(clus, *clustering_index); } if (verbose) { printf("Adding centroids to quantizer\n"); @@ -135,6 +170,42 @@ void Level1Quantizer::train_q1( } } +} // namespace + +void Level1Quantizer::train_q1( + size_t n, + const float* x, + bool verbose, + MetricType metric_type) { + train_q1_ex(n, x, NumericType::Float32, verbose, metric_type); +} + +void Level1Quantizer::train_q1_ex( + size_t n, + const void* x, + NumericType numeric_type, + bool verbose, + MetricType metric_type) { + train_q1_impl( + *this, + n, + static_cast(x), + numeric_type, + nullptr, + verbose, + metric_type); +} + +void Level1Quantizer::train_q1_encoded( + size_t n, + const uint8_t* x, + const Index* codec, + bool verbose, + MetricType metric_type) { + train_q1_impl( + *this, n, x, NumericType::Float32, codec, verbose, metric_type); +} + size_t Level1Quantizer::coarse_code_size() const { size_t nl = nlist - 1; size_t nbyte = 0; @@ -1308,6 +1379,125 @@ void IndexIVF::update_vectors(int n, const idx_t* new_ids, const float* x) { invlists, n, new_ids, assign.data(), flat_codes.data()); } +namespace { + +/** Trains the encoder of `ivf` on training rows that are not fp32. + * + * Only the train_encoder_num_vectors() rows sampled for the encoder are + * widened to fp32. `decode_rows(n, rows, out)` widens n consecutive rows of + * `row_size` bytes. + */ +template +void train_encoder_on_encoded_rows( + IndexIVF& ivf, + idx_t n, + const uint8_t* x, + size_t row_size, + const DecodeRows& decode_rows) { + const size_t d = ivf.d; + idx_t max_nt = ivf.train_encoder_num_vectors(); + if (max_nt <= 0 || n <= max_nt) { + max_nt = n; + } + + std::vector decoded(static_cast(max_nt) * d); + if (max_nt == n) { + decode_rows(n, x, decoded.data()); + } else { + FAISS_THROW_IF_NOT_FMT( + n <= static_cast(std::numeric_limits::max()), + "Dataset too large (%" PRId64 ") for standard subsampling", + n); + if (ivf.verbose) { + printf(" Input training set too big (max size is %" PRId64 + "), sampling %" PRId64 " / %" PRId64 " vectors\n", + max_nt, + max_nt, + n); + } + std::vector subset(static_cast(n)); + rand_perm(subset.data(), n, 1234); + for (idx_t i = 0; i < max_nt; ++i) { + decode_rows( + 1, + x + static_cast(subset[i]) * row_size, + decoded.data() + static_cast(i) * d); + } + } + n = max_nt; + + if (ivf.by_residual) { + FAISS_THROW_IF_NOT_MSG(ivf.quantizer, "IVF quantizer must not be null"); + std::vector assign(n); + ivf.quantizer->assign(n, decoded.data(), assign.data()); + + std::vector residuals(static_cast(n) * d); + ivf.quantizer->compute_residual_n( + n, decoded.data(), residuals.data(), assign.data()); + + ivf.train_encoder(n, residuals.data(), assign.data()); + } else { + ivf.train_encoder(n, decoded.data(), nullptr); + } +} + +} // namespace + +void IndexIVF::train_ex(idx_t n, const void* x, NumericType numeric_type) { + if (numeric_type != NumericType::Float16) { + Index::train_ex(n, x, numeric_type); + return; + } + FAISS_THROW_IF_NOT_MSG(x, "training data must not be null"); + if (verbose) { + printf("Training level-1 quantizer on fp16 data\n"); + } + + train_q1_ex(n, x, numeric_type, verbose, metric_type); + + if (verbose) { + printf("Training IVF residual\n"); + } + + const size_t dim = d; + train_encoder_on_encoded_rows( + *this, + n, + static_cast(x), + sizeof(uint16_t) * dim, + [dim](idx_t ni, const uint8_t* rows, float* out) { + fp16_to_fp32( + ni * dim, reinterpret_cast(rows), out); + }); + + is_trained = true; +} + +void IndexIVF::train_encoded(idx_t n, const uint8_t* x, const Index* codec) { + FAISS_THROW_IF_NOT_MSG(codec, "encoded training requires a codec"); + FAISS_THROW_IF_NOT_MSG(x, "encoded training data must not be null"); + if (verbose) { + printf("Training level-1 quantizer\n"); + } + + train_q1_encoded(n, x, codec, verbose, metric_type); + + if (verbose) { + printf("Training IVF residual\n"); + } + + train_encoder_on_encoded_rows( + *this, + n, + x, + codec->sa_code_size(), + [codec](idx_t ni, const uint8_t* rows, float* out) { + codec->sa_decode(ni, rows, out); + }); + + is_trained = true; +} + void IndexIVF::train(idx_t n, const float* x) { if (verbose) { printf("Training level-1 quantizer\n"); diff --git a/faiss/IndexIVF.h b/faiss/IndexIVF.h index a58ecf765..e6f2e6485 100644 --- a/faiss/IndexIVF.h +++ b/faiss/IndexIVF.h @@ -46,13 +46,40 @@ struct Level1Quantizer { /// to override index used during clustering Index* clustering_index = nullptr; - /// Trains the quantizer and calls train_residual to train sub-quantizers + /// Trains the level-1 quantizer from fp32 vectors. void train_q1( size_t n, const float* x, bool verbose, MetricType metric_type); + /** Trains the level-1 quantizer from Float32 or packed Float16 vectors. + * + * Float16 runs k-means natively on the fp16 rows (see + * Clustering::train_ex). With quantizer_trains_alone == 1 the data is + * passed on to quantizer->train_ex(). SuperKMeans requires Float32. + */ + void train_q1_ex( + size_t n, + const void* x, + NumericType numeric_type, + bool verbose, + MetricType metric_type); + + /** Trains the level-1 quantizer from vectors encoded by `codec`. + * + * The codec must remain alive for the duration of this call and decode + * each input row to `quantizer->d` floats. Encoded training is supported + * when clustering trains the quantizer (quantizer_trains_alone 0 or 2) + * without SuperKMeans. + */ + void train_q1_encoded( + size_t n, + const uint8_t* x, + const Index* codec, + bool verbose, + MetricType metric_type); + /// compute the number of bytes required to store list ids size_t coarse_code_size() const; void encode_listno(idx_t list_no, uint8_t* code) const; @@ -236,6 +263,23 @@ struct IndexIVF : Index, IndexIVFInterface { /// Trains the quantizer and calls train_encoder to train sub-quantizers void train(idx_t n, const float* x) override; + /** Train from vectors decoded on demand by `codec`. + * + * Coarse quantizer training decodes in bounded batches. Subclasses with a + * trainable encoder decode at most train_encoder_num_vectors() rows into + * fp32; IndexIVFFlat overrides this method and needs no fp32 corpus buffer. + */ + virtual void train_encoded(idx_t n, const uint8_t* x, const Index* codec); + + /** Supports Float32 and packed IEEE binary16 (Float16) inputs. + * + * Float16 trains natively without a codec: the coarse quantizer runs + * fp16 k-means (Clustering::train_ex) and only the + * train_encoder_num_vectors() rows sampled for the encoder are widened + * to fp32. IndexIVFFlat overrides this and widens no rows at all. + */ + void train_ex(idx_t n, const void* x, NumericType numeric_type) override; + /// Calls add_with_ids with NULL ids void add(idx_t n, const float* x) override; diff --git a/faiss/IndexIVFFlat.cpp b/faiss/IndexIVFFlat.cpp index ea37e6d0d..93992d00e 100644 --- a/faiss/IndexIVFFlat.cpp +++ b/faiss/IndexIVFFlat.cpp @@ -57,6 +57,25 @@ IndexIVFFlat::IndexIVFFlat() { by_residual = false; } +void IndexIVFFlat::train_encoded( + idx_t n, + const uint8_t* x, + const Index* codec) { + FAISS_THROW_IF_NOT_MSG(codec, "encoded training requires a codec"); + train_q1_encoded(n, x, codec, verbose, metric_type); + is_trained = true; +} + +void IndexIVFFlat::train_ex(idx_t n, const void* x, NumericType numeric_type) { + if (numeric_type != NumericType::Float16) { + IndexIVF::train_ex(n, x, numeric_type); + return; + } + FAISS_THROW_IF_NOT_MSG(x, "training data must not be null"); + train_q1_ex(n, x, numeric_type, verbose, metric_type); + is_trained = true; +} + void IndexIVFFlat::add_core( idx_t n, const float* x, @@ -186,6 +205,20 @@ IndexIVFFlatDedup::IndexIVFFlatDedup( metric_type_in, own_invlists_in) {} +void IndexIVFFlatDedup::train_encoded(idx_t, const uint8_t*, const Index*) { + FAISS_THROW_MSG("encoded training is not supported by IndexIVFFlatDedup"); +} + +void IndexIVFFlatDedup::train_ex( + idx_t n, + const void* x, + NumericType numeric_type) { + FAISS_THROW_IF_NOT_MSG( + numeric_type == NumericType::Float32, + "only fp32 training is supported by IndexIVFFlatDedup"); + train(n, static_cast(x)); +} + void IndexIVFFlatDedup::train(idx_t n, const float* x) { std::unordered_map map; std::unique_ptr x2(new float[n * d]); diff --git a/faiss/IndexIVFFlat.h b/faiss/IndexIVFFlat.h index 1665a06e6..530fdd471 100644 --- a/faiss/IndexIVFFlat.h +++ b/faiss/IndexIVFFlat.h @@ -29,6 +29,16 @@ struct IndexIVFFlat : IndexIVF { MetricType = METRIC_L2, bool own_invlists = true); + /** Train the coarse quantizer from vectors decoded on demand by `codec`. + * + * IndexIVFFlat has no trainable residual encoder, so this avoids a + * corpus-sized decoded buffer. The codec must remain alive for this call. + */ + void train_encoded(idx_t n, const uint8_t* x, const Index* codec) override; + + /// Float16 trains only the coarse quantizer, natively on the fp16 rows. + void train_ex(idx_t n, const void* x, NumericType numeric_type) override; + void add_core( idx_t n, const float* x, @@ -110,6 +120,12 @@ struct IndexIVFFlatDedup : IndexIVFFlat { /// also dedups the training set void train(idx_t n, const float* x) override; + /// Encoded training cannot preserve the byte-exact deduplication contract. + void train_encoded(idx_t n, const uint8_t* x, const Index* codec) override; + + /// Non-fp32 input is rejected for the same reason as train_encoded. + void train_ex(idx_t n, const void* x, NumericType numeric_type) override; + /// implemented for all IndexIVF* classes void add_with_ids(idx_t n, const float* x, const idx_t* xids) override; diff --git a/faiss/gpu/GpuIndex.cu b/faiss/gpu/GpuIndex.cu index a2e41cd34..2fd8e612e 100644 --- a/faiss/gpu/GpuIndex.cu +++ b/faiss/gpu/GpuIndex.cu @@ -15,6 +15,7 @@ #include #include #include +#include #include #include @@ -475,6 +476,33 @@ void GpuIndex::search_and_reconstruct( reconstruct_batch(n * k, labels, recons); } +void GpuIndex::searchImpl_ex_( + idx_t n, + const void* x, + NumericType numeric_type, + int k, + float* distances, + idx_t* labels, + const SearchParameters* params) const { + if (numeric_type == NumericType::Float32) { + searchImpl_( + n, static_cast(x), k, distances, labels, params); + } else if (numeric_type == NumericType::Float16) { + // The fp16 queries were copied to the device at half the cost of + // float32; widen them there so that any index can search them + auto stream = resources_->getDefaultStream(config_.device); + + Tensor queries( + const_cast(static_cast(x)), {n, this->d}); + auto queriesFloat = convertTensorTemporary( + resources_.get(), stream, queries); + + searchImpl_(n, queriesFloat.data(), k, distances, labels, params); + } else { + FAISS_THROW_MSG("GpuIndex::searchImpl_: unsupported numeric type"); + } +} + void GpuIndex::searchNonPaged_ex_( idx_t n, const void* x, diff --git a/faiss/gpu/GpuIndex.h b/faiss/gpu/GpuIndex.h index b688d6270..f0e621c83 100644 --- a/faiss/gpu/GpuIndex.h +++ b/faiss/gpu/GpuIndex.h @@ -213,6 +213,9 @@ class GpuIndex : public faiss::Index { idx_t* labels, const SearchParameters* params) const = 0; + /// Search with queries of the given numeric type; all data is guaranteed + /// to be resident on our device. The default handles Float32 and widens + /// Float16 queries to float32 on the device before calling searchImpl_ virtual void searchImpl_ex_( idx_t n, const void* x, @@ -220,19 +223,7 @@ class GpuIndex : public faiss::Index { int k, float* distances, idx_t* labels, - const SearchParameters* params) const { - if (numeric_type == NumericType::Float32) { - searchImpl_( - n, - static_cast(x), - k, - distances, - labels, - params); - } else { - FAISS_THROW_MSG("GpuIndex::searchImpl_: unsupported numeric type"); - } - } + const SearchParameters* params) const; private: /// Calls addImpl_ for a single page of GPU-resident data diff --git a/faiss/gpu/GpuIndexFlat.cu b/faiss/gpu/GpuIndexFlat.cu index 0a57e1dc9..001946d20 100644 --- a/faiss/gpu/GpuIndexFlat.cu +++ b/faiss/gpu/GpuIndexFlat.cu @@ -241,6 +241,41 @@ void GpuIndexFlat::searchImpl_( sel); } +void GpuIndexFlat::searchImpl_ex_( + idx_t n, + const void* x, + NumericType numeric_type, + int k, + float* distances, + idx_t* labels, + const SearchParameters* params) const { + if (numeric_type != NumericType::Float16 || !flatConfig_.useFloat16) { + // float32 storage: fp16 queries are widened on the device + GpuIndex::searchImpl_ex_( + n, x, numeric_type, k, distances, labels, params); + return; + } + + // current device already set + // n/k already validated + // Input and output data are already resident on the GPU + Tensor queries( + const_cast(static_cast(x)), {n, this->d}); + Tensor outDistances(distances, {n, k}); + Tensor outLabels(labels, {n, k}); + + const IDSelector* sel = params ? params->sel : nullptr; + data_->query( + queries, + k, + metric_type, + metric_arg, + outDistances, + outLabels, + true, + sel); +} + void GpuIndexFlat::reconstruct(idx_t key, float* out) const { DeviceScope scope(config_.device); diff --git a/faiss/gpu/GpuIndexFlat.h b/faiss/gpu/GpuIndexFlat.h index ee4c14466..9912fa9ed 100644 --- a/faiss/gpu/GpuIndexFlat.h +++ b/faiss/gpu/GpuIndexFlat.h @@ -131,6 +131,17 @@ class GpuIndexFlat : public GpuIndex { idx_t* labels, const SearchParameters* params) const override; + /// Called from GpuIndex for search; Float16 queries are searched without + /// conversion against float16 storage + void searchImpl_ex_( + idx_t n, + const void* x, + NumericType numeric_type, + int k, + float* distances, + idx_t* labels, + const SearchParameters* params) const override; + protected: /// Our configuration options const GpuIndexFlatConfig flatConfig_; diff --git a/faiss/gpu/test/TestGpuIndexFlat.cpp b/faiss/gpu/test/TestGpuIndexFlat.cpp index 3b2ca1c19..fdb8ce1d3 100644 --- a/faiss/gpu/test/TestGpuIndexFlat.cpp +++ b/faiss/gpu/test/TestGpuIndexFlat.cpp @@ -5,14 +5,19 @@ * 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 @@ -807,6 +812,177 @@ TEST(TestCuvsGpuIndexFlat, IDSelector_IP) { } #endif +// Packs float vectors as IEEE binary16, the Float16 input of search_ex +std::vector toFp16(const std::vector& v) { + std::vector out(v.size()); + std::transform(v.begin(), v.end(), out.begin(), [](float f) { + return faiss::encode_fp16(f); + }); + return out; +} + +// A copy of host data on a GPU device +struct DeviceCopy { + DeviceCopy(int device, const void* src, size_t bytes) { + faiss::gpu::DeviceScope scope(device); + CUDA_VERIFY(cudaMalloc(&data, bytes)); + CUDA_VERIFY(cudaMemcpy(data, src, bytes, cudaMemcpyHostToDevice)); + } + + DeviceCopy(const DeviceCopy&) = delete; + DeviceCopy& operator=(const DeviceCopy&) = delete; + + ~DeviceCopy() { + CUDA_VERIFY(cudaFree(data)); + } + + void* data = nullptr; +}; + +// search_ex on Float16 queries must return exactly what search returns on the +// float32 queries rounded to fp16, as the rounded values are what the index +// computes with, whether the queries are on the host (paged or not) or on the +// device +void testSearchExFloat16(faiss::gpu::GpuIndex& index) { + int numQuery = faiss::gpu::randVal(1, 512); + int k = std::min(faiss::gpu::randVal(1, 50), int(index.ntotal)); + + auto queries = faiss::gpu::randVecs(numQuery, index.d); + auto queries16 = toFp16(queries); + auto roundedQueries = faiss::gpu::roundToHalf(queries); + + std::vector refDistances(numQuery * k); + std::vector refLabels(numQuery * k); + index.search( + numQuery, + roundedQueries.data(), + k, + refDistances.data(), + refLabels.data()); + + auto expectSearchExMatches = [&](const void* x, const std::string& where) { + std::vector distances(numQuery * k, 0); + std::vector labels(numQuery * k, -1); + index.search_ex( + numQuery, + x, + faiss::NumericType::Float16, + k, + distances.data(), + labels.data()); + + EXPECT_EQ(labels, refLabels) << where; + EXPECT_EQ(distances, refDistances) << where; + }; + + expectSearchExMatches(queries16.data(), "host queries"); + + auto minPagingSize = index.getMinPagingSize(); + index.setMinPagingSize(0); + expectSearchExMatches(queries16.data(), "paged host queries"); + index.setMinPagingSize(minPagingSize); + + DeviceCopy queries16Device( + index.getDevice(), + queries16.data(), + queries16.size() * sizeof(uint16_t)); + expectSearchExMatches(queries16Device.data, "device queries"); +} + +// Covers both float32 storage (queries widened on the device) and float16 +// storage (queries searched natively) +template +void testFlatSearchExFloat16() { + faiss::gpu::StandardGpuResources res; + res.noTempMemory(); + + int dim = faiss::gpu::randVal(50, 400); + int numVecs = faiss::gpu::randVal(1000, 5000); + auto vecs = faiss::gpu::randVecs(numVecs, dim); + + for (bool useFloat16 : {false, true}) { + faiss::gpu::GpuIndexFlatConfig config; + config.device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1); + config.useFloat16 = useFloat16; + + GpuIndexFlatT gpuIndex(&res, dim, config); + gpuIndex.add(numVecs, vecs.data()); + + SCOPED_TRACE(useFloat16 ? "useFloat16" : "float32 storage"); + testSearchExFloat16(gpuIndex); + } +} + +TEST(TestGpuIndexFlat, SearchExFloat16_L2) { + testFlatSearchExFloat16(); +} + +TEST(TestGpuIndexFlat, SearchExFloat16_IP) { + testFlatSearchExFloat16(); +} + +// Indexes without a native fp16 search widen the queries on the device +TEST(TestGpuIndexFlat, SearchExFloat16_IVFFlatDefault) { + faiss::gpu::StandardGpuResources res; + res.noTempMemory(); + + int dim = 64; + int numVecs = 5000; + auto vecs = faiss::gpu::randVecs(numVecs, dim); + + faiss::gpu::GpuIndexIVFFlatConfig config; + config.device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1); + + int numLists = 32; + faiss::gpu::GpuIndexIVFFlat gpuIndex( + &res, dim, numLists, faiss::METRIC_L2, config); + gpuIndex.nprobe = 4; + gpuIndex.train(numVecs, vecs.data()); + gpuIndex.add(numVecs, vecs.data()); + + testSearchExFloat16(gpuIndex); +} + +// k-means on fp16 data with a GPU assignment index must match k-means on the +// float32 rounding of that data with the same kind of index +TEST(TestGpuIndexFlat, ClusteringTrainExFloat16) { + faiss::gpu::StandardGpuResources res; + res.noTempMemory(); + + int dim = 32; + int numCentroids = 20; + // Fits in a single decode block, so that both runs search the same batches + int numVecs = 2000; + auto vecs = faiss::gpu::randVecs(numVecs, dim); + auto vecs16 = toFp16(vecs); + auto roundedVecs = faiss::gpu::roundToHalf(vecs); + + faiss::ClusteringParameters cp; + cp.niter = 10; + + for (bool useFloat16 : {false, true}) { + faiss::gpu::GpuIndexFlatConfig config; + config.device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1); + config.useFloat16 = useFloat16; + + faiss::gpu::GpuIndexFlatL2 refIndex(&res, dim, config); + faiss::Clustering refClustering(dim, numCentroids, cp); + refClustering.train(numVecs, roundedVecs.data(), refIndex); + + faiss::gpu::GpuIndexFlatL2 gpuIndex(&res, dim, config); + faiss::Clustering clustering(dim, numCentroids, cp); + clustering.train_ex( + numVecs, vecs16.data(), faiss::NumericType::Float16, gpuIndex); + + ASSERT_EQ(clustering.centroids.size(), refClustering.centroids.size()); + for (size_t i = 0; i < clustering.centroids.size(); ++i) { + EXPECT_NEAR( + clustering.centroids[i], refClustering.centroids[i], 1e-5f) + << "useFloat16 " << useFloat16 << " element " << i; + } + } +} + int main(int argc, char** argv) { testing::InitGoogleTest(&argc, argv); diff --git a/faiss/impl/ClusteringHelpers.cpp b/faiss/impl/ClusteringHelpers.cpp index 669cdacf4..a4accebd5 100644 --- a/faiss/impl/ClusteringHelpers.cpp +++ b/faiss/impl/ClusteringHelpers.cpp @@ -7,6 +7,8 @@ #include +#include +#include #include #include #include @@ -20,11 +22,86 @@ #include #include +#include #include +#include namespace faiss { namespace detail { +namespace { + +/// Below this many values, the fp16 finiteness check stays on one thread. +constexpr size_t fp16_parallel_threshold = size_t(1) << 16; + +/** Shared body of compute_centroids and compute_centroids_fp16. + * + * `add_row(i, w, c, scratch)` adds training vector i to the fp32 vector c, + * scaled by *w when w is not null. `scratch` is a per-thread buffer of d + * floats. + */ +template +void compute_centroids_impl( + size_t d, + size_t k, + size_t n, + size_t k_frozen, + const int64_t* assign, + const float* weights, + float* hassign, + float* centroids, + const AddRow& add_row) { + k -= k_frozen; + centroids += k_frozen * d; + + memset(centroids, 0, sizeof(*centroids) * d * k); + + std::atomic invalid_assignment{false}; + const int64_t num_centroids = static_cast(k + k_frozen); +#pragma omp parallel + { + int nt = omp_get_num_threads(); + int rank = omp_get_thread_num(); + + // this thread is taking care of centroids c0:c1 + size_t c0 = (k * rank) / nt; + size_t c1 = (k * (rank + 1)) / nt; + std::vector scratch(d); + + for (size_t i = 0; i < n; i++) { + int64_t ci = assign[i]; + if (ci < 0 || ci >= num_centroids) { + invalid_assignment.store(true, std::memory_order_relaxed); + continue; + } + ci -= k_frozen; + if (ci >= static_cast(c0) && + ci < static_cast(c1)) { + const float* w = weights ? weights + i : nullptr; + hassign[ci] += w ? *w : 1.0f; + add_row(i, w, centroids + ci * d, scratch.data()); + } + } + } + FAISS_THROW_IF_MSG( + invalid_assignment.load(std::memory_order_relaxed), + "invalid cluster assignment"); + +#pragma omp parallel for + for (idx_t ci = 0; ci < static_cast(k); ci++) { + if (hassign[ci] == 0) { + continue; + } + float norm = 1 / hassign[ci]; + float* c = centroids + ci * d; + for (size_t j = 0; j < d; j++) { + c[j] *= norm; + } + } +} + +} // namespace + uint64_t get_actual_rng_seed(const int seed) { return (seed >= 0) ? seed @@ -109,66 +186,80 @@ void compute_centroids( const float* weights, float* hassign, float* centroids) { - k -= k_frozen; - centroids += k_frozen * d; - - memset(centroids, 0, sizeof(*centroids) * d * k); - - size_t line_size = codec ? codec->sa_code_size() : d * sizeof(float); - -#pragma omp parallel - { - int nt = omp_get_num_threads(); - int rank = omp_get_thread_num(); - - // this thread is taking care of centroids c0:c1 - size_t c0 = (k * rank) / nt; - size_t c1 = (k * (rank + 1)) / nt; - std::vector decode_buffer(d); - - for (size_t i = 0; i < n; i++) { - int64_t ci = assign[i]; - FAISS_THROW_IF_NOT_MSG( - ci >= 0 && ci < k + k_frozen, "invalid cluster assignment"); - ci -= k_frozen; - if (ci >= static_cast(c0) && - ci < static_cast(c1)) { - float* c = centroids + ci * d; + const size_t line_size = codec ? codec->sa_code_size() : d * sizeof(float); + compute_centroids_impl( + d, + k, + n, + k_frozen, + assign, + weights, + hassign, + centroids, + [&](size_t i, const float* w, float* c, float* scratch) { const float* xi; if (!codec) { xi = reinterpret_cast(x + i * line_size); } else { - float* xif = decode_buffer.data(); - codec->sa_decode(1, x + i * line_size, xif); - xi = xif; + codec->sa_decode(1, x + i * line_size, scratch); + xi = scratch; } - if (weights) { - float w = weights[i]; - hassign[ci] += w; + if (w) { + const float wi = *w; for (size_t j = 0; j < d; j++) { - c[j] += xi[j] * w; + c[j] += xi[j] * wi; } } else { - hassign[ci] += 1.0; for (size_t j = 0; j < d; j++) { c[j] += xi[j]; } } - } - } - } + }); +} -#pragma omp parallel for - for (idx_t ci = 0; ci < static_cast(k); ci++) { - if (hassign[ci] == 0) { - continue; - } - float norm = 1 / hassign[ci]; - float* c = centroids + ci * d; - for (size_t j = 0; j < d; j++) { - c[j] *= norm; +void compute_centroids_fp16( + size_t d, + size_t k, + size_t n, + size_t k_frozen, + const uint16_t* x, + const int64_t* assign, + const float* weights, + float* hassign, + float* centroids) { + with_selected_simd_levels( + [&]() { + compute_centroids_impl( + d, + k, + n, + k_frozen, + assign, + weights, + hassign, + centroids, + [&](size_t i, + const float* w, + float* c, + float* /*scratch*/) { + // x * 1.0f is exact, so unweighted sums match the + // fp32 path. + fp16_madd(d, x + i * d, w ? *w : 1.0f, c); + }); + }); +} + +bool fp16_all_finite(size_t n, const uint16_t* x) { + // binary16 NaNs and infinities are the values with all exponent bits set + constexpr uint16_t exponent_mask = 0x7c00; + std::atomic all_finite{true}; +#pragma omp parallel for if (n > fp16_parallel_threshold) + for (int64_t i = 0; i < static_cast(n); i++) { + if ((x[i] & exponent_mask) == exponent_mask) { + all_finite.store(false, std::memory_order_relaxed); } } + return all_finite.load(std::memory_order_relaxed); } // a bit above machine epsilon for float16 diff --git a/faiss/impl/ClusteringHelpers.h b/faiss/impl/ClusteringHelpers.h index a9e67cb56..75a8dde6e 100644 --- a/faiss/impl/ClusteringHelpers.h +++ b/faiss/impl/ClusteringHelpers.h @@ -75,6 +75,33 @@ void compute_centroids( float* hassign, float* centroids); +/** Same as compute_centroids for training vectors stored as packed IEEE + * binary16 values. + * + * Rows are widened in registers and accumulated directly into the fp32 + * centroids: there is no per-row decode call and no decode buffer. + * + * @param x training vectors, size n * d + */ +void compute_centroids_fp16( + size_t d, + size_t k, + size_t n, + size_t k_frozen, + const uint16_t* x, + const int64_t* assign, + const float* weights, + float* hassign, + float* centroids); + +/** Check packed IEEE binary16 values for NaNs and infinities. + * + * The check runs on the bit patterns, without widening to fp32. + * + * @return true iff none of the `n` values is a NaN or an infinity + */ +bool fp16_all_finite(size_t n, const uint16_t* x); + /** Handle empty clusters by splitting larger ones. * * It works by slightly changing the centroids to make 2 clusters from diff --git a/faiss/utils/distances.cpp b/faiss/utils/distances.cpp index 2702b8ffc..5f4f5d893 100644 --- a/faiss/utils/distances.cpp +++ b/faiss/utils/distances.cpp @@ -26,6 +26,7 @@ #include #include #include +#include #ifndef FINTEGER #define FINTEGER long @@ -361,10 +362,11 @@ void exhaustive_L2sqr_seq( } } -/** Find the nearest neighbors for nx queries in a set of ny vectors */ -template -void exhaustive_inner_product_blas( - const float* x, +/** Find the nearest neighbors for nx queries (see Fp32QueryTiles) in a set + * of ny vectors */ +template +void exhaustive_inner_product_blas_tiles( + QueryTiles& queries, const float* y, size_t d, size_t nx, @@ -386,6 +388,8 @@ void exhaustive_inner_product_blas( i1 = nx; } + const float* xt = queries.tile(i0, i1, nullptr); + res.begin_multiple(i0, i1); for (size_t j0 = 0; j0 < ny; j0 += bs_y) { @@ -405,7 +409,7 @@ void exhaustive_inner_product_blas( &one, y + j0 * d, &di, - x + i0 * d, + xt, &di, &zero, ip_block.get(), @@ -421,9 +425,9 @@ void exhaustive_inner_product_blas( // distance correction is an operator that can be applied to transform // the distances -template -void exhaustive_L2sqr_blas_default_impl( - const float* x, +template +void exhaustive_L2sqr_blas_tiles( + QueryTiles& queries, const float* y, size_t d, size_t nx, @@ -443,7 +447,7 @@ void exhaustive_L2sqr_blas_default_impl( std::unique_ptr x_norms(new float[nx]); std::unique_ptr del2; - fvec_norms_L2sqr(x_norms.get(), x, d, nx); + queries.prepare_norms(nx, x_norms.get()); if (!y_norms) { float* y_norms2 = new float[ny]; @@ -458,6 +462,8 @@ void exhaustive_L2sqr_blas_default_impl( i1 = nx; } + const float* xt = queries.tile(i0, i1, x_norms.get()); + res.begin_multiple(i0, i1); for (size_t j0 = 0; j0 < ny; j0 += bs_y) { @@ -477,7 +483,7 @@ void exhaustive_L2sqr_blas_default_impl( &one, y + j0 * d, &di, - x + i0 * d, + xt, &di, &zero, ip_block.get(), @@ -510,6 +516,31 @@ void exhaustive_L2sqr_blas_default_impl( } } +template +void exhaustive_inner_product_blas( + const float* x, + const float* y, + size_t d, + size_t nx, + size_t ny, + BlockResultHandler& res) { + Fp32QueryTiles queries{x, d}; + exhaustive_inner_product_blas_tiles(queries, y, d, nx, ny, res); +} + +template +void exhaustive_L2sqr_blas_default_impl( + const float* x, + const float* y, + size_t d, + size_t nx, + size_t ny, + BlockResultHandler& res, + const float* y_norms = nullptr) { + Fp32QueryTiles queries{x, d}; + exhaustive_L2sqr_blas_tiles(queries, y, d, nx, ny, res, y_norms); +} + template void exhaustive_L2sqr_blas( const float* x, @@ -556,6 +587,44 @@ void exhaustive_L2sqr_blas>>( }); } +/// exhaustive_L2sqr_blas for packed IEEE binary16 queries +template +void exhaustive_L2sqr_blas_fp16( + const uint16_t* x, + const float* y, + size_t d, + size_t nx, + size_t ny, + BlockResultHandler& res, + const float* y_norms) { + Fp16QueryTiles queries(x, d, nx); + exhaustive_L2sqr_blas_tiles(queries, y, d, nx, ny, res, y_norms); +} + +template <> +void exhaustive_L2sqr_blas_fp16>>( + const uint16_t* x, + const float* y, + size_t d, + size_t nx, + size_t ny, + Top1BlockResultHandler>& res, + const float* y_norms) { + with_selected_simd_levels( + [&]() { + if constexpr ( + SL == SIMDLevel::AVX2 || SL == SIMDLevel::AVX512 || + SL == SIMDLevel::ARM_SVE) { + exhaustive_L2sqr_blas_cmax_fp16( + x, y, d, nx, ny, res, y_norms); + } else { + Fp16QueryTiles queries(x, d, nx); + exhaustive_L2sqr_blas_tiles( + queries, y, d, nx, ny, res, y_norms); + } + }); +} + struct Run_search_inner_product { using T = void; template @@ -598,6 +667,35 @@ struct Run_search_L2sqr { } }; +// fp16 consumers only run the BLAS kernels, see use_fp16_blas_kernels +struct Run_search_inner_product_fp16 { + using T = void; + template + void f(BlockResultHandler& res, + const uint16_t* x, + const float* y, + size_t d, + size_t nx, + size_t ny) { + Fp16QueryTiles queries(x, d, nx); + exhaustive_inner_product_blas_tiles(queries, y, d, nx, ny, res); + } +}; + +struct Run_search_L2sqr_fp16 { + using T = void; + template + void f(BlockResultHandler& res, + const uint16_t* x, + const float* y, + size_t d, + size_t nx, + size_t ny, + const float* y_norm2) { + exhaustive_L2sqr_blas_fp16(x, y, d, nx, ny, res, y_norm2); + } +}; + } // anonymous namespace /******************************************************* @@ -894,6 +992,120 @@ void knn_L2sqr( knn_L2sqr(x, y, d, nx, ny, res->k, res->val, res->ids, y_norm2, sel); } +namespace { + +/** Whether fp16 queries go to the BLAS kernels directly. + * + * Other cases widen chunks of queries and run the fp32 search: selectors, + * sequential and database-parallel searches, and small d, where the fused + * top-1 kernels need fp32 rows. + */ +bool use_fp16_blas_kernels( + size_t d, + size_t nx, + size_t ny, + const IDSelector* sel) { + return !sel && ny > 0 && d > 32 && + nx * d >= static_cast(distance_compute_blas_threshold) && + !should_use_db_parallel(nx, ny, sel); +} + +/// widens chunks of queries and calls search(x_chunk, n_chunk, i0) on them +template +void search_fp16_by_chunks( + const uint16_t* x, + size_t d, + size_t nx, + const Fp32Search& search) { + // a multiple of the query tile, so the fp32 kernels see the same tiles + const size_t bs_x = std::max(distance_compute_blas_query_bs, 1); + const size_t chunk = bs_x * + std::max(1, (size_t(8) << 20) / (bs_x * d * sizeof(float))); + std::vector buffer(std::min(nx, chunk) * d); + for (size_t i0 = 0; i0 < nx; i0 += chunk) { + const size_t i1 = std::min(nx, i0 + chunk); + fp16_to_fp32((i1 - i0) * d, x + i0 * d, buffer.data()); + search(buffer.data(), i1 - i0, i0); + } +} + +} // namespace + +void knn_inner_product_fp16( + const uint16_t* x, + const float* y, + size_t d, + size_t nx, + size_t ny, + size_t k, + float* vals, + int64_t* ids, + const IDSelector* sel) { + if (!use_fp16_blas_kernels(d, nx, ny, sel)) { + search_fp16_by_chunks( + x, d, nx, [&](const float* xc, size_t nc, size_t i0) { + knn_inner_product( + xc, + y, + d, + nc, + ny, + k, + vals + i0 * k, + ids + i0 * k, + sel); + }); + return; + } + Run_search_inner_product_fp16 r; + dispatch_knn_ResultHandler( + nx, + vals, + ids, + k, + METRIC_INNER_PRODUCT, + nullptr, + r, + x, + y, + d, + nx, + ny); +} + +void knn_L2sqr_fp16( + const uint16_t* x, + const float* y, + size_t d, + size_t nx, + size_t ny, + size_t k, + float* vals, + int64_t* ids, + const float* y_norm2, + const IDSelector* sel) { + if (!use_fp16_blas_kernels(d, nx, ny, sel)) { + search_fp16_by_chunks( + x, d, nx, [&](const float* xc, size_t nc, size_t i0) { + knn_L2sqr( + xc, + y, + d, + nc, + ny, + k, + vals + i0 * k, + ids + i0 * k, + y_norm2, + sel); + }); + return; + } + Run_search_L2sqr_fp16 r; + dispatch_knn_ResultHandler( + nx, vals, ids, k, METRIC_L2, nullptr, r, x, y, d, nx, ny, y_norm2); +} + /*************************************************************************** * Range search ***************************************************************************/ diff --git a/faiss/utils/distances.h b/faiss/utils/distances.h index f2d4a80a7..253a9329b 100644 --- a/faiss/utils/distances.h +++ b/faiss/utils/distances.h @@ -455,6 +455,36 @@ void knn_L2sqr( const float* y_norm2 = nullptr, const IDSelector* sel = nullptr); +/** knn_inner_product for packed IEEE binary16 queries (x: nx * d uint16_t). + * + * The BLAS kernels widen each query tile once and reuse it for all database + * tiles, so there is no fp32 copy of the queries. Results are the same as + * knn_inner_product on the widened queries. + */ +void knn_inner_product_fp16( + const uint16_t* x, + const float* y, + size_t d, + size_t nx, + size_t ny, + size_t k, + float* distances, + int64_t* indexes, + const IDSelector* sel = nullptr); + +/// knn_L2sqr for packed IEEE binary16 queries, see knn_inner_product_fp16 +void knn_L2sqr_fp16( + const uint16_t* x, + const float* y, + size_t d, + size_t nx, + size_t ny, + size_t k, + float* distances, + int64_t* indexes, + const float* y_norm2 = nullptr, + const IDSelector* sel = nullptr); + /** Find the max inner product neighbors for nx queries in a set of ny vectors * indexed by ids. May be useful for re-ranking a pre-selected vector list * diff --git a/faiss/utils/simd_impl/distances_arm_sve.cpp b/faiss/utils/simd_impl/distances_arm_sve.cpp index ad020cef4..0fed90d8d 100644 --- a/faiss/utils/simd_impl/distances_arm_sve.cpp +++ b/faiss/utils/simd_impl/distances_arm_sve.cpp @@ -758,9 +758,11 @@ size_t fvec_L2sqr_ny_nearest_y_transposed( return current_min_idx; } -template <> -void exhaustive_L2sqr_blas_cmax( - const float* x, +namespace { + +template +void exhaustive_L2sqr_blas_cmax_sve( + QueryTiles& queries, const float* y, size_t d, size_t nx, @@ -779,7 +781,7 @@ void exhaustive_L2sqr_blas_cmax( std::unique_ptr x_norms(new float[nx]); std::unique_ptr del2; - fvec_norms_L2sqr(x_norms.get(), x, d, nx); + queries.prepare_norms(nx, x_norms.get()); const size_t lanes = svcntw(); @@ -795,6 +797,8 @@ void exhaustive_L2sqr_blas_cmax( if (i1 > nx) i1 = nx; + const float* xt = queries.tile(i0, i1, x_norms.get()); + res.begin_multiple(i0, i1); for (size_t j0 = 0; j0 < ny; j0 += bs_y) { @@ -813,7 +817,7 @@ void exhaustive_L2sqr_blas_cmax( &one, y + j0 * d, &di, - x + i0 * d, + xt, &di, &zero, ip_block.get(), @@ -938,4 +942,32 @@ void exhaustive_L2sqr_blas_cmax( } } +} // namespace + +template <> +void exhaustive_L2sqr_blas_cmax( + const float* x, + const float* y, + size_t d, + size_t nx, + size_t ny, + Top1BlockResultHandler>& res, + const float* y_norms) { + Fp32QueryTiles queries{x, d}; + exhaustive_L2sqr_blas_cmax_sve(queries, y, d, nx, ny, res, y_norms); +} + +template <> +void exhaustive_L2sqr_blas_cmax_fp16( + const uint16_t* x, + const float* y, + size_t d, + size_t nx, + size_t ny, + Top1BlockResultHandler>& res, + const float* y_norms) { + Fp16QueryTiles queries(x, d, nx); + exhaustive_L2sqr_blas_cmax_sve(queries, y, d, nx, ny, res, y_norms); +} + } // namespace faiss diff --git a/faiss/utils/simd_impl/distances_avx2.cpp b/faiss/utils/simd_impl/distances_avx2.cpp index f7d9711ba..1276e57f2 100644 --- a/faiss/utils/simd_impl/distances_avx2.cpp +++ b/faiss/utils/simd_impl/distances_avx2.cpp @@ -1213,9 +1213,11 @@ int fvec_madd_and_argmin( return fvec_madd_and_argmin_sse(n, a, bf, b, c); } -template <> -void exhaustive_L2sqr_blas_cmax( - const float* x, +namespace { + +template +void exhaustive_L2sqr_blas_cmax_avx2( + QueryTiles& queries, const float* y, size_t d, size_t nx, @@ -1235,7 +1237,7 @@ void exhaustive_L2sqr_blas_cmax( std::unique_ptr x_norms(new float[nx]); std::unique_ptr del2; - fvec_norms_L2sqr(x_norms.get(), x, d, nx); + queries.prepare_norms(nx, x_norms.get()); if (!y_norms) { float* y_norms2 = new float[ny]; @@ -1250,6 +1252,8 @@ void exhaustive_L2sqr_blas_cmax( i1 = nx; } + const float* xt = queries.tile(i0, i1, x_norms.get()); + res.begin_multiple(i0, i1); for (size_t j0 = 0; j0 < ny; j0 += bs_y) { @@ -1269,7 +1273,7 @@ void exhaustive_L2sqr_blas_cmax( &one, y + j0 * d, &di, - x + i0 * d, + xt, &di, &zero, ip_block.get(), @@ -1427,4 +1431,32 @@ void exhaustive_L2sqr_blas_cmax( } } +} // namespace + +template <> +void exhaustive_L2sqr_blas_cmax( + const float* x, + const float* y, + size_t d, + size_t nx, + size_t ny, + Top1BlockResultHandler>& res, + const float* y_norms) { + Fp32QueryTiles queries{x, d}; + exhaustive_L2sqr_blas_cmax_avx2(queries, y, d, nx, ny, res, y_norms); +} + +template <> +void exhaustive_L2sqr_blas_cmax_fp16( + const uint16_t* x, + const float* y, + size_t d, + size_t nx, + size_t ny, + Top1BlockResultHandler>& res, + const float* y_norms) { + Fp16QueryTiles queries(x, d, nx); + exhaustive_L2sqr_blas_cmax_avx2(queries, y, d, nx, ny, res, y_norms); +} + } // namespace faiss diff --git a/faiss/utils/simd_impl/distances_avx512.cpp b/faiss/utils/simd_impl/distances_avx512.cpp index 13dd4197d..284a43c51 100644 --- a/faiss/utils/simd_impl/distances_avx512.cpp +++ b/faiss/utils/simd_impl/distances_avx512.cpp @@ -1133,9 +1133,11 @@ int fvec_madd_and_argmin( return fvec_madd_and_argmin_sse(n, a, bf, b, c); } -template <> -void exhaustive_L2sqr_blas_cmax( - const float* x, +namespace { + +template +void exhaustive_L2sqr_blas_cmax_avx512( + QueryTiles& queries, const float* y, size_t d, size_t nx, @@ -1154,7 +1156,7 @@ void exhaustive_L2sqr_blas_cmax( std::unique_ptr x_norms(new float[nx]); std::unique_ptr del2; - fvec_norms_L2sqr(x_norms.get(), x, d, nx); + queries.prepare_norms(nx, x_norms.get()); if (!y_norms) { float* y_norms2 = new float[ny]; @@ -1169,6 +1171,8 @@ void exhaustive_L2sqr_blas_cmax( i1 = nx; } + const float* xt = queries.tile(i0, i1, x_norms.get()); + res.begin_multiple(i0, i1); for (size_t j0 = 0; j0 < ny; j0 += bs_y) { @@ -1188,7 +1192,7 @@ void exhaustive_L2sqr_blas_cmax( &one, y + j0 * d, &di, - x + i0 * d, + xt, &di, &zero, ip_block.get(), @@ -1340,4 +1344,32 @@ void exhaustive_L2sqr_blas_cmax( } } +} // namespace + +template <> +void exhaustive_L2sqr_blas_cmax( + const float* x, + const float* y, + size_t d, + size_t nx, + size_t ny, + Top1BlockResultHandler>& res, + const float* y_norms) { + Fp32QueryTiles queries{x, d}; + exhaustive_L2sqr_blas_cmax_avx512(queries, y, d, nx, ny, res, y_norms); +} + +template <> +void exhaustive_L2sqr_blas_cmax_fp16( + const uint16_t* x, + const float* y, + size_t d, + size_t nx, + size_t ny, + Top1BlockResultHandler>& res, + const float* y_norms) { + Fp16QueryTiles queries(x, d, nx); + exhaustive_L2sqr_blas_cmax_avx512(queries, y, d, nx, ny, res, y_norms); +} + } // namespace faiss diff --git a/faiss/utils/simd_impl/exhaustive_L2sqr_blas_cmax.h b/faiss/utils/simd_impl/exhaustive_L2sqr_blas_cmax.h index 021d50133..e84269f78 100644 --- a/faiss/utils/simd_impl/exhaustive_L2sqr_blas_cmax.h +++ b/faiss/utils/simd_impl/exhaustive_L2sqr_blas_cmax.h @@ -7,11 +7,93 @@ #pragma once +#include +#include +#include +#include + #include +#include #include +#include namespace faiss { +/** fp32 query tiles for the BLAS exhaustive-search kernels. + * + * A kernel calls prepare_norms() once, then tile(i0, i1, x_norms) for each + * query tile; the returned rows stay valid until the next tile() call. + */ +struct Fp32QueryTiles { + const float* x; + size_t d; + + /// squared L2 norms of all nx queries + void prepare_norms(size_t nx, float* x_norms) const { + fvec_norms_L2sqr(x_norms, x, d, nx); + } + + const float* tile(size_t i0, size_t /*i1*/, float* /*x_norms*/) const { + return x + i0 * d; + } +}; + +/** Packed IEEE binary16 query tiles for the BLAS exhaustive-search kernels. + * + * Queries are widened in chunks of whole tiles (at most chunk_bytes of fp32, + * at least one tile) into a buffer reused across chunks: the fp32 copy stays + * bounded and each widened tile serves all database tiles. Norms are computed + * from the widened rows, so results match the fp32 kernels on the widened + * queries. + */ +struct Fp16QueryTiles { + static constexpr size_t chunk_bytes = size_t(16) << 20; + + const uint16_t* x; + size_t d; + size_t nx; + std::unique_ptr buffer; + size_t chunk_rows = 0; + size_t chunk_begin = 0; + size_t chunk_end = 0; + + Fp16QueryTiles(const uint16_t* x_in, size_t d_in, size_t nx_in) + : x(x_in), d(d_in), nx(nx_in) {} + + void prepare_norms(size_t /*nx*/, float* /*x_norms*/) const {} + + /** widened rows [i0, i1), and their norms if x_norms is not null + * + * Tiles are requested in order, all of the same size except the last. + */ + const float* tile(size_t i0, size_t i1, float* x_norms) { + if (i1 > chunk_end) { + if (!buffer) { + const size_t tile_rows = i1 - i0; + chunk_rows = + tile_rows * + std::max( + 1, + chunk_bytes / (tile_rows * d * sizeof(float))); + chunk_rows = std::min(chunk_rows, nx - i0); + buffer.reset(new float[chunk_rows * d]); + } + chunk_begin = i0; + chunk_end = std::min(nx, i0 + chunk_rows); + const int64_t n = static_cast(chunk_end - chunk_begin); +#pragma omp parallel for if (n * static_cast(d) > (1 << 16)) + for (int64_t i = 0; i < n; i++) { + float* row = buffer.get() + i * d; + fp16_to_fp32(d, x + (chunk_begin + i) * d, row); + if (x_norms) { + x_norms[chunk_begin + i] = fvec_norm_L2sqr(row, d); + } + } + } + return buffer.get() + (i0 - chunk_begin) * d; + } +}; + /// BLAS-accelerated exhaustive L2 search for the k=1 (top-1) case. /// Specializations live in the per-SIMD translation units under simd_impl/. template @@ -24,4 +106,15 @@ void exhaustive_L2sqr_blas_cmax( Top1BlockResultHandler>& res, const float* y_norms); +/// Same as exhaustive_L2sqr_blas_cmax for packed IEEE binary16 queries. +template +void exhaustive_L2sqr_blas_cmax_fp16( + const uint16_t* x, + const float* y, + size_t d, + size_t nx, + size_t ny, + Top1BlockResultHandler>& res, + const float* y_norms); + } // namespace faiss diff --git a/faiss/utils/simd_impl/fp16_kernels.h b/faiss/utils/simd_impl/fp16_kernels.h new file mode 100644 index 000000000..4b7f707e6 --- /dev/null +++ b/faiss/utils/simd_impl/fp16_kernels.h @@ -0,0 +1,61 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include +#include + +#include +#include + +namespace faiss { +namespace detail { + +// c[j] += fp32(x[j]) * w for j < d +template +inline void fp16_madd(size_t d, const uint16_t* x, float w, float* c) { + for (size_t j = 0; j < d; ++j) { + c[j] += decode_fp16(x[j]) * w; + } +} + +template +inline void fp16_to_fp32_kernel(size_t n, const uint16_t* x, float* out) { + for (size_t j = 0; j < n; ++j) { + out[j] = decode_fp16(x[j]); + } +} + +#ifdef COMPILE_SIMD_AVX2 +template <> +void fp16_madd(size_t d, const uint16_t* x, float w, float* c); + +template <> +void fp16_to_fp32_kernel( + size_t n, + const uint16_t* x, + float* out); +#endif + +#ifdef COMPILE_SIMD_ARM_NEON +template <> +void fp16_madd( + size_t d, + const uint16_t* x, + float w, + float* c); + +template <> +void fp16_to_fp32_kernel( + size_t n, + const uint16_t* x, + float* out); +#endif + +} // namespace detail +} // namespace faiss diff --git a/faiss/utils/simd_impl/fp16_kernels_avx2.cpp b/faiss/utils/simd_impl/fp16_kernels_avx2.cpp new file mode 100644 index 000000000..66d340631 --- /dev/null +++ b/faiss/utils/simd_impl/fp16_kernels_avx2.cpp @@ -0,0 +1,56 @@ +/* + * 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_AVX2 + +#include + +#include + +namespace faiss { +namespace detail { + +template <> +void fp16_madd( + size_t d, + const uint16_t* x, + float w, + float* c) { + const __m256 wv = _mm256_set1_ps(w); + size_t j = 0; + for (; j + 8 <= d; j += 8) { + const __m256 xv = _mm256_cvtph_ps( + _mm_loadu_si128(reinterpret_cast(x + j))); + _mm256_storeu_ps( + c + j, + _mm256_add_ps(_mm256_loadu_ps(c + j), _mm256_mul_ps(xv, wv))); + } + for (; j < d; ++j) { + c[j] += decode_fp16(x[j]) * w; + } +} + +template <> +void fp16_to_fp32_kernel( + size_t n, + const uint16_t* x, + float* out) { + size_t j = 0; + for (; j + 8 <= n; j += 8) { + const __m128i h = + _mm_loadu_si128(reinterpret_cast(x + j)); + _mm256_storeu_ps(out + j, _mm256_cvtph_ps(h)); + } + for (; j < n; ++j) { + out[j] = decode_fp16(x[j]); + } +} + +} // namespace detail +} // namespace faiss + +#endif // COMPILE_SIMD_AVX2 diff --git a/faiss/utils/simd_impl/fp16_kernels_neon.cpp b/faiss/utils/simd_impl/fp16_kernels_neon.cpp new file mode 100644 index 000000000..2dcb89a98 --- /dev/null +++ b/faiss/utils/simd_impl/fp16_kernels_neon.cpp @@ -0,0 +1,52 @@ +/* + * 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_NEON + +#include + +#include + +namespace faiss { +namespace detail { + +template <> +void fp16_madd( + size_t d, + const uint16_t* x, + float w, + float* c) { + const float32x4_t wv = vdupq_n_f32(w); + size_t j = 0; + for (; j + 4 <= d; j += 4) { + const float32x4_t xv = + vcvt_f32_f16(vreinterpret_f16_u16(vld1_u16(x + j))); + vst1q_f32(c + j, vaddq_f32(vld1q_f32(c + j), vmulq_f32(xv, wv))); + } + for (; j < d; ++j) { + c[j] += decode_fp16(x[j]) * w; + } +} + +template <> +void fp16_to_fp32_kernel( + size_t n, + const uint16_t* x, + float* out) { + size_t j = 0; + for (; j + 4 <= n; j += 4) { + vst1q_f32(out + j, vcvt_f32_f16(vreinterpret_f16_u16(vld1_u16(x + j)))); + } + for (; j < n; ++j) { + out[j] = decode_fp16(x[j]); + } +} + +} // namespace detail +} // namespace faiss + +#endif // COMPILE_SIMD_ARM_NEON diff --git a/faiss/utils/utils.cpp b/faiss/utils/utils.cpp index abd4fc756..e39a7dda5 100644 --- a/faiss/utils/utils.cpp +++ b/faiss/utils/utils.cpp @@ -36,7 +36,9 @@ #include #include +#include #include +#include #ifndef FINTEGER #define FINTEGER long @@ -639,6 +641,26 @@ void CombinerRangeKNN::write_result(T* D_res, int64_t* I_res) { template struct CombinerRangeKNN; template struct CombinerRangeKNN; +void fp16_to_fp32(size_t n, const uint16_t* x, float* out) { + // below this many values, stay on the calling thread + constexpr size_t block = size_t(1) << 16; + with_selected_simd_levels( + [&]() { + if (n <= block) { + detail::fp16_to_fp32_kernel(n, x, out); + return; + } + const int64_t nblock = + static_cast((n + block - 1) / block); +#pragma omp parallel for + for (int64_t b = 0; b < nblock; b++) { + const size_t j0 = static_cast(b) * block; + detail::fp16_to_fp32_kernel( + std::min(block, n - j0), x + j0, out + j0); + } + }); +} + void CodeSet::insert(size_t n, const uint8_t* codes, bool* inserted) { for (size_t i = 0; i < n; i++) { auto res = s.insert( diff --git a/faiss/utils/utils.h b/faiss/utils/utils.h index e387c6263..f3403bf51 100644 --- a/faiss/utils/utils.h +++ b/faiss/utils/utils.h @@ -164,6 +164,14 @@ void binary_to_real(size_t d, const uint8_t* x_in, float* x_out); */ void real_to_binary(size_t d, const float* x_in, uint8_t* x_out); +/** Convert packed IEEE binary16 values to fp32 (exact). + * + * @param n number of values + * @param x input fp16 values, size n + * @param out output fp32 values, size n + */ +void fp16_to_fp32(size_t n, const uint16_t* x, float* out); + /** A reasonable hashing function */ uint64_t hash_bytes(const uint8_t* bytes, int64_t n); diff --git a/tests/test_ivf_index.cpp b/tests/test_ivf_index.cpp index 27fe907f9..9e1c3a9e5 100644 --- a/tests/test_ivf_index.cpp +++ b/tests/test_ivf_index.cpp @@ -7,6 +7,7 @@ #include #include +#include #include #include #include @@ -17,9 +18,13 @@ #include #include +#include #include #include +#include #include +#include +#include namespace { @@ -511,3 +516,739 @@ TEST(IVF, search_callbacks) { EXPECT_GE(distance_count, heap_count) << "not every distance computation leads to a heap change"; } + +namespace { + +class LimitedEncoderIndex : public faiss::IndexIVFScalarQuantizer { + public: + LimitedEncoderIndex(faiss::Index* quantizer, int d, int nlist) + : faiss::IndexIVFScalarQuantizer( + quantizer, + d, + nlist, + faiss::ScalarQuantizer::QT_8bit) {} + + faiss::idx_t train_encoder_num_vectors() const override { + return 7; + } + + void train_encoder( + faiss::idx_t n, + const float* x, + const faiss::idx_t* assign) override { + encoder_input.assign(x, x + n * d); + encoder_assignments.assign(assign, assign + n); + } + + std::vector encoder_input; + std::vector encoder_assignments; +}; + +class TrackingFp16Codec : public faiss::IndexScalarQuantizer { + public: + explicit TrackingFp16Codec(int d) + : faiss::IndexScalarQuantizer( + d, + faiss::ScalarQuantizer::QuantizerType::QT_fp16) {} + + void sa_decode(faiss::idx_t n, const uint8_t* bytes, float* x) + const override { + decode_calls.fetch_add(1, std::memory_order_relaxed); + max_decode_rows = std::max(max_decode_rows, static_cast(n)); + faiss::IndexScalarQuantizer::sa_decode(n, bytes, x); + } + + mutable std::atomic decode_calls{0}; + mutable size_t max_decode_rows = 0; +}; + +std::vector encode_fp16(const std::vector& values) { + std::vector encoded(values.size()); + for (size_t i = 0; i < values.size(); ++i) { + encoded[i] = faiss::encode_fp16(values[i]); + } + return encoded; +} + +std::vector decode_fp16(const std::vector& values) { + std::vector decoded(values.size()); + for (size_t i = 0; i < values.size(); ++i) { + decoded[i] = faiss::decode_fp16(values[i]); + } + return decoded; +} + +} // namespace + +TEST(IVF, train_float16_matches_float32_on_rounded_input) { + constexpr int d = 4; + constexpr int n = 64; + constexpr int nlist = 4; + + std::vector input(n * d); + for (size_t i = 0; i < input.size(); ++i) { + input[i] = static_cast((i * 17) % 101) / 13.0f; + } + auto encoded = encode_fp16(input); + auto rounded = decode_fp16(encoded); + + faiss::IndexFlatL2 float_quantizer(d); + faiss::IndexIVFFlat float_index(&float_quantizer, d, nlist); + float_index.cp.seed = 1234; + float_index.cp.niter = 4; + float_index.cp.min_points_per_centroid = 1; + float_index.train(n, rounded.data()); + + faiss::IndexFlatL2 half_quantizer(d); + faiss::IndexIVFFlat half_index(&half_quantizer, d, nlist); + half_index.cp.seed = 1234; + half_index.cp.niter = 4; + half_index.cp.min_points_per_centroid = 1; + half_index.train_ex(n, encoded.data(), faiss::NumericType::Float16); + + ASSERT_TRUE(half_index.is_trained); + ASSERT_EQ(half_quantizer.ntotal, nlist); + std::vector float_centroids(nlist * d); + std::vector half_centroids(nlist * d); + float_quantizer.reconstruct_n(0, nlist, float_centroids.data()); + half_quantizer.reconstruct_n(0, nlist, half_centroids.data()); + EXPECT_EQ(half_centroids, float_centroids); +} + +TEST(IVF, encoded_training_rejects_non_finite_values) { + constexpr int d = 2; + constexpr int n = 4; + constexpr int nlist = 2; + auto encoded = encode_fp16( + {0.0f, + 1.0f, + 2.0f, + 3.0f, + std::numeric_limits::infinity(), + 5.0f, + 6.0f, + 7.0f}); + + faiss::IndexFlatL2 quantizer(d); + faiss::IndexIVFFlat index(&quantizer, d, nlist); + EXPECT_THROW( + index.train_ex(n, encoded.data(), faiss::NumericType::Float16), + faiss::FaissException); +} + +TEST(IVF, encoded_training_decodes_in_bounded_batches) { + constexpr int d = 4; + constexpr int n = 40; + constexpr int nlist = 2; + constexpr size_t decode_block_size = 5; + + std::vector input(n * d); + for (size_t i = 0; i < input.size(); ++i) { + input[i] = static_cast((i * 11) % 97) / 17.0f; + } + auto encoded = encode_fp16(input); + TrackingFp16Codec codec(d); + faiss::IndexFlatL2 quantizer(d); + faiss::IndexIVFFlat index(&quantizer, d, nlist); + index.cp.decode_block_size = decode_block_size; + index.cp.niter = 2; + index.cp.min_points_per_centroid = 1; + index.train_encoded( + n, reinterpret_cast(encoded.data()), &codec); + + EXPECT_TRUE(index.is_trained); + EXPECT_GT(codec.max_decode_rows, 0); + EXPECT_LE(codec.max_decode_rows, decode_block_size); +} + +TEST(IVF, encoded_training_skips_nan_scan) { + constexpr int d = 2; + constexpr int n = 2; + auto encoded = encode_fp16({0.0f, 1.0f, 2.0f, 3.0f}); + TrackingFp16Codec codec(d); + faiss::Clustering clustering(d, n); + faiss::IndexFlatL2 index(d); + + clustering.train_encoded( + n, reinterpret_cast(encoded.data()), &codec, index); + + EXPECT_EQ(codec.decode_calls.load(std::memory_order_relaxed), 1); +} + +TEST(IVF, train_float16_trains_scalar_quantizer_encoder) { + constexpr int d = 4; + constexpr int n = 128; + constexpr int nlist = 4; + + std::vector input(n * d); + for (size_t i = 0; i < input.size(); ++i) { + input[i] = static_cast((i * 19) % 113) / 23.0f; + } + auto encoded = encode_fp16(input); + + faiss::IndexFlatL2 quantizer(d); + faiss::IndexIVFScalarQuantizer index( + &quantizer, + d, + nlist, + faiss::ScalarQuantizer::QuantizerType::QT_8bit); + index.cp.seed = 1234; + index.cp.niter = 4; + index.cp.min_points_per_centroid = 1; + index.train_ex(n, encoded.data(), faiss::NumericType::Float16); + + EXPECT_TRUE(index.is_trained); + EXPECT_EQ(quantizer.ntotal, nlist); + EXPECT_FALSE(index.sq.trained.empty()); +} + +TEST(IVF, encoded_training_validates_codec) { + constexpr int d = 4; + constexpr int n = 4; + constexpr int nlist = 2; + std::vector encoded(n * (d + 1), faiss::encode_fp16(1.0f)); + + faiss::IndexFlatL2 quantizer(d); + std::vector centroids(nlist * d, 0.0f); + quantizer.add(nlist, centroids.data()); + faiss::IndexIVFFlat index(&quantizer, d, nlist); + TrackingFp16Codec wrong_dimension_codec(d + 1); + + EXPECT_THROW( + index.train_encoded( + n, + reinterpret_cast(encoded.data()), + &wrong_dimension_codec), + faiss::FaissException); + EXPECT_THROW( + index.train_encoded( + n, + reinterpret_cast(encoded.data()), + nullptr), + faiss::FaissException); +} + +TEST(IVF, encoded_encoder_subsampling_matches_float32) { + constexpr int d = 4; + constexpr int n = 40; + constexpr int nlist = 2; + + std::vector input(n * d); + for (size_t i = 0; i < input.size(); ++i) { + input[i] = static_cast((i * 23) % 127) / 29.0f; + } + auto encoded = encode_fp16(input); + auto rounded = decode_fp16(encoded); + + faiss::IndexFlatL2 float_quantizer(d); + LimitedEncoderIndex float_index(&float_quantizer, d, nlist); + float_index.cp.seed = 1234; + float_index.cp.niter = 2; + float_index.cp.min_points_per_centroid = 1; + float_index.train(n, rounded.data()); + + faiss::IndexFlatL2 half_quantizer(d); + LimitedEncoderIndex half_index(&half_quantizer, d, nlist); + half_index.cp.seed = 1234; + half_index.cp.niter = 2; + half_index.cp.min_points_per_centroid = 1; + half_index.train_ex(n, encoded.data(), faiss::NumericType::Float16); + + EXPECT_EQ(half_index.encoder_input, float_index.encoder_input); + EXPECT_EQ(half_index.encoder_assignments, float_index.encoder_assignments); +} + +namespace { + +class TrackingFlatL2 : public faiss::IndexFlatL2 { + public: + explicit TrackingFlatL2(faiss::idx_t d) : faiss::IndexFlatL2(d) {} + + void search( + faiss::idx_t n, + const float* x, + faiss::idx_t k, + float* distances, + faiss::idx_t* labels, + const faiss::SearchParameters* params = nullptr) const override { + max_search_rows = std::max(max_search_rows, static_cast(n)); + faiss::IndexFlatL2::search(n, x, k, distances, labels, params); + } + + void search_ex( + faiss::idx_t n, + const void* x, + faiss::NumericType numeric_type, + faiss::idx_t k, + float* distances, + faiss::idx_t* labels, + const faiss::SearchParameters* params = nullptr) const override { + max_search_rows = std::max(max_search_rows, static_cast(n)); + if (numeric_type == faiss::NumericType::Float16) { + fp16_search_calls++; + } + faiss::IndexFlatL2::search_ex( + n, x, numeric_type, k, distances, labels, params); + } + + mutable size_t max_search_rows = 0; + mutable size_t fp16_search_calls = 0; +}; + +class InvalidAssignmentIndex : public faiss::IndexFlatL2 { + public: + explicit InvalidAssignmentIndex(faiss::idx_t d) : faiss::IndexFlatL2(d) {} + + void search_ex( + faiss::idx_t n, + const void* /*x*/, + faiss::NumericType /*numeric_type*/, + faiss::idx_t k, + float* distances, + faiss::idx_t* labels, + const faiss::SearchParameters* /*params*/ = + nullptr) const override { + const size_t result_size = + static_cast(n) * static_cast(k); + std::fill_n(distances, result_size, 0.0f); + std::fill_n(labels, result_size, faiss::idx_t{-1}); + } +}; + +std::vector make_training_data(size_t n, size_t d) { + std::vector x(n * d); + for (size_t i = 0; i < x.size(); ++i) { + x[i] = static_cast((i * 29) % 131) / 17.0f - 3.0f; + } + return x; +} + +faiss::ClusteringParameters small_clustering_params(int niter) { + faiss::ClusteringParameters cp; + cp.niter = niter; + cp.seed = 1234; + cp.min_points_per_centroid = 1; + return cp; +} + +} // namespace + +TEST(Clustering, train_ex_float16_matches_float32_on_rounded_input) { + constexpr size_t n = 200; + constexpr size_t k = 5; + // 16 exercises only the SIMD kernels, 13 and 37 also the scalar tails + for (int d : {13, 16, 37}) { + auto encoded = encode_fp16(make_training_data(n, d)); + auto rounded = decode_fp16(encoded); + auto cp = small_clustering_params(6); + + faiss::Clustering full(d, k, cp); + faiss::IndexFlatL2 full_index(d); + full.train(n, rounded.data(), full_index); + + faiss::Clustering half(d, k, cp); + faiss::IndexFlatL2 half_index(d); + half.train_ex( + n, encoded.data(), faiss::NumericType::Float16, half_index); + + EXPECT_EQ(half.centroids, full.centroids) << "d=" << d; + ASSERT_EQ(half.iteration_stats.size(), full.iteration_stats.size()); + for (size_t i = 0; i < full.iteration_stats.size(); ++i) { + EXPECT_EQ(half.iteration_stats[i].obj, full.iteration_stats[i].obj); + } + EXPECT_EQ(half_index.ntotal, static_cast(k)); + } +} + +TEST(Clustering, train_ex_float16_weighted_matches_float32) { + constexpr size_t n = 150; + constexpr size_t d = 20; + constexpr size_t k = 4; + auto encoded = encode_fp16(make_training_data(n, d)); + auto rounded = decode_fp16(encoded); + std::vector weights(n); + for (size_t i = 0; i < n; ++i) { + weights[i] = 0.25f + static_cast(i % 7); + } + auto cp = small_clustering_params(1); + + faiss::Clustering full(d, k, cp); + faiss::IndexFlatL2 full_index(d); + full.train(n, rounded.data(), full_index, weights.data()); + + faiss::Clustering half(d, k, cp); + faiss::IndexFlatL2 half_index(d); + half.train_ex( + n, + encoded.data(), + faiss::NumericType::Float16, + half_index, + weights.data()); + + ASSERT_EQ(half.centroids.size(), full.centroids.size()); + for (size_t i = 0; i < full.centroids.size(); ++i) { + EXPECT_NEAR(half.centroids[i], full.centroids[i], 1e-4); + } +} + +TEST(Clustering, train_ex_float16_subsampling_and_kmeanspp_match_float32) { + constexpr size_t n = 100; + constexpr size_t d = 12; + constexpr size_t k = 4; + auto encoded = encode_fp16(make_training_data(n, d)); + auto rounded = decode_fp16(encoded); + auto cp = small_clustering_params(4); + cp.max_points_per_centroid = 8; + cp.init_method = faiss::ClusteringInitMethod::KMEANS_PLUS_PLUS; + + faiss::Clustering full(d, k, cp); + faiss::IndexFlatL2 full_index(d); + full.train(n, rounded.data(), full_index); + + faiss::Clustering half(d, k, cp); + faiss::IndexFlatL2 half_index(d); + half.train_ex(n, encoded.data(), faiss::NumericType::Float16, half_index); + + EXPECT_EQ(half.centroids, full.centroids); +} + +TEST(Clustering, train_ex_float16_assigns_in_bounded_batches) { + constexpr size_t n = 50; + constexpr size_t d = 8; + constexpr size_t k = 3; + auto encoded = encode_fp16(make_training_data(n, d)); + auto cp = small_clustering_params(3); + cp.decode_block_size = 7; + + faiss::Clustering clus(d, k, cp); + TrackingFlatL2 index(d); + clus.train_ex(n, encoded.data(), faiss::NumericType::Float16, index); + + EXPECT_GT(index.max_search_rows, 0); + EXPECT_LE(index.max_search_rows, cp.decode_block_size); + // fp16 blocks reach the index through search_ex, not widened beforehand + EXPECT_GT(index.fp16_search_calls, 0); +} + +TEST(Clustering, train_ex_invalid_assignment_throws) { + constexpr int n = 8; + constexpr int d = 3; + constexpr int k = 2; + auto encoded = encode_fp16(make_training_data(n, d)); + faiss::Clustering clustering(d, k, small_clustering_params(1)); + InvalidAssignmentIndex index(d); + + EXPECT_THROW( + clustering.train_ex( + n, encoded.data(), faiss::NumericType::Float16, index), + faiss::FaissException); +} + +TEST(Index, search_ex_float16_default_matches_float32) { + constexpr int d = 12; + constexpr int nb = 300; + constexpr int nq = 40; + constexpr int k = 5; + auto base = decode_fp16(encode_fp16(make_training_data(nb, d))); + std::vector queries_in(nq * d); + for (size_t i = 0; i < queries_in.size(); ++i) { + queries_in[i] = static_cast((i * 37) % 113) / 19.0f - 3.0f; + } + auto queries = encode_fp16(queries_in); + auto rounded = decode_fp16(queries); + + // IndexIVFFlat does not override search_ex: exercises the default + faiss::IndexFlatL2 quantizer(d); + faiss::IndexIVFFlat index(&quantizer, d, 4); + index.cp = small_clustering_params(3); + index.train(nb, base.data()); + index.add(nb, base.data()); + index.nprobe = 2; + + std::vector float_distances(nq * k), half_distances(nq * k); + std::vector float_labels(nq * k), half_labels(nq * k); + index.search( + nq, rounded.data(), k, float_distances.data(), float_labels.data()); + index.search_ex( + nq, + queries.data(), + faiss::NumericType::Float16, + k, + half_distances.data(), + half_labels.data()); + + EXPECT_EQ(half_labels, float_labels); + EXPECT_EQ(half_distances, float_distances); + EXPECT_THROW( + index.search_ex( + nq, + queries.data(), + faiss::NumericType::UInt8, + k, + half_distances.data(), + half_labels.data()), + faiss::FaissException); +} + +TEST(Clustering, train_ex_float16_rejects_non_finite_values) { + constexpr size_t n = 8; + constexpr size_t d = 3; + constexpr size_t k = 2; + const uint16_t max_finite = 0x7bff; // 65504 + for (uint16_t bad : + {uint16_t(0x7e00), uint16_t(0x7c00), uint16_t(0xfc00)}) { + auto encoded = encode_fp16(make_training_data(n, d)); + encoded[0] = max_finite; + encoded[n * d - 1] = bad; + faiss::Clustering clus(d, k, small_clustering_params(2)); + faiss::IndexFlatL2 index(d); + EXPECT_THROW( + clus.train_ex( + n, encoded.data(), faiss::NumericType::Float16, index), + faiss::FaissException) + << "bits=" << bad; + } + + auto encoded = encode_fp16(make_training_data(n, d)); + encoded[0] = max_finite; + faiss::Clustering clus(d, k, small_clustering_params(2)); + faiss::IndexFlatL2 index(d); + EXPECT_NO_THROW(clus.train_ex( + n, encoded.data(), faiss::NumericType::Float16, index)); +} + +TEST(Clustering, train_ex_rejects_unsupported_numeric_type) { + constexpr size_t n = 8; + constexpr size_t d = 4; + std::vector x(n * d, 1); + faiss::Clustering clus(d, 2, small_clustering_params(2)); + faiss::IndexFlatL2 index(d); + EXPECT_THROW( + clus.train_ex(n, x.data(), faiss::NumericType::UInt8, index), + faiss::FaissException); +} + +TEST(IVF, train_float16_with_flat_assigner_matches_float32) { + constexpr int d = 10; + constexpr int n = 120; + constexpr int nlist = 4; + auto encoded = encode_fp16(make_training_data(n, d)); + auto rounded = decode_fp16(encoded); + + faiss::IndexFlatL2 float_quantizer(d); + faiss::IndexIVFFlat float_index(&float_quantizer, d, nlist); + float_index.quantizer_trains_alone = 2; + float_index.cp = small_clustering_params(4); + float_index.train(n, rounded.data()); + + faiss::IndexFlatL2 half_quantizer(d); + faiss::IndexIVFFlat half_index(&half_quantizer, d, nlist); + half_index.quantizer_trains_alone = 2; + half_index.cp = small_clustering_params(4); + half_index.train_ex(n, encoded.data(), faiss::NumericType::Float16); + + ASSERT_TRUE(half_index.is_trained); + std::vector float_centroids(nlist * d); + std::vector half_centroids(nlist * d); + float_quantizer.reconstruct_n(0, nlist, float_centroids.data()); + half_quantizer.reconstruct_n(0, nlist, half_centroids.data()); + EXPECT_EQ(half_centroids, float_centroids); +} + +TEST(IVF, train_float16_clustering_index_sees_bounded_batches) { + constexpr int d = 8; + constexpr int n = 60; + constexpr int nlist = 3; + auto encoded = encode_fp16(make_training_data(n, d)); + + faiss::IndexFlatL2 quantizer(d); + TrackingFlatL2 clustering_index(d); + faiss::IndexIVFScalarQuantizer index( + &quantizer, d, nlist, faiss::ScalarQuantizer::QT_8bit); + index.clustering_index = &clustering_index; + index.cp = small_clustering_params(3); + index.cp.decode_block_size = 5; + index.train_ex(n, encoded.data(), faiss::NumericType::Float16); + + EXPECT_TRUE(index.is_trained); + EXPECT_EQ(quantizer.ntotal, nlist); + EXPECT_GT(clustering_index.max_search_rows, 0); + EXPECT_LE(clustering_index.max_search_rows, index.cp.decode_block_size); +} + +TEST(IVF, train_float16_rejected_by_dedup) { + constexpr int d = 4; + constexpr int n = 16; + auto encoded = encode_fp16(make_training_data(n, d)); + faiss::IndexFlatL2 quantizer(d); + faiss::IndexIVFFlatDedup index(&quantizer, d, 2); + EXPECT_THROW( + index.train_ex(n, encoded.data(), faiss::NumericType::Float16), + faiss::FaissException); +} + +TEST(IVF, train_ex_float32_preserves_deduplication) { + constexpr int d = 1; + constexpr int n = 4; + constexpr int nlist = 1; + std::vector input{0.0f, 0.0f, 0.0f, 10.0f}; + + faiss::IndexFlatL2 direct_quantizer(d); + faiss::IndexIVFFlatDedup direct(&direct_quantizer, d, nlist); + direct.cp = small_clustering_params(1); + direct.train(n, input.data()); + + faiss::IndexFlatL2 train_ex_quantizer(d); + faiss::IndexIVFFlatDedup train_ex(&train_ex_quantizer, d, nlist); + train_ex.cp = small_clustering_params(1); + train_ex.train_ex(n, input.data(), faiss::NumericType::Float32); + + float direct_centroid; + float train_ex_centroid; + direct_quantizer.reconstruct(0, &direct_centroid); + train_ex_quantizer.reconstruct(0, &train_ex_centroid); + EXPECT_FLOAT_EQ(train_ex_centroid, direct_centroid); + EXPECT_FLOAT_EQ(train_ex_centroid, 5.0f); +} + +namespace { + +/// forces the BLAS kernels with small query tiles, restored on scope exit +struct ScopedBlasTiles { + int threshold = faiss::distance_compute_blas_threshold; + int query_bs = faiss::distance_compute_blas_query_bs; + + explicit ScopedBlasTiles(int query_bs_in) { + faiss::distance_compute_blas_threshold = 0; + faiss::distance_compute_blas_query_bs = query_bs_in; + } + + ~ScopedBlasTiles() { + faiss::distance_compute_blas_threshold = threshold; + faiss::distance_compute_blas_query_bs = query_bs; + } +}; + +std::vector make_queries(size_t n, size_t d) { + std::vector x(n * d); + for (size_t i = 0; i < x.size(); ++i) { + x[i] = static_cast((i * 37) % 113) / 19.0f - 3.0f; + } + return x; +} + +} // namespace + +TEST(IndexFlat, search_ex_float16_native_matches_float32) { + constexpr size_t nb = 500; + constexpr size_t nq = 300; + ScopedBlasTiles tiles(64); // several query tiles per search + for (faiss::MetricType metric : + {faiss::METRIC_L2, faiss::METRIC_INNER_PRODUCT}) { + for (size_t d : {37, 64}) { + for (size_t k : {1, 5}) { + auto base = make_training_data(nb, d); + auto queries = encode_fp16(make_queries(nq, d)); + auto rounded = decode_fp16(queries); + faiss::IndexFlat index(d, metric); + index.add(nb, base.data()); + + std::vector fd(nq * k), hd(nq * k); + std::vector fi(nq * k), hi(nq * k); + index.search(nq, rounded.data(), k, fd.data(), fi.data()); + index.search_ex( + nq, + queries.data(), + faiss::NumericType::Float16, + k, + hd.data(), + hi.data()); + EXPECT_EQ(hi, fi) + << "metric=" << metric << " d=" << d << " k=" << k; + EXPECT_EQ(hd, fd) + << "metric=" << metric << " d=" << d << " k=" << k; + } + } + } +} + +TEST(IndexFlat, search_ex_float16_with_selector_matches_float32) { + constexpr size_t d = 48; + constexpr size_t nb = 400; + constexpr size_t nq = 50; + constexpr size_t k = 3; + ScopedBlasTiles tiles(16); + auto base = make_training_data(nb, d); + auto queries = encode_fp16(make_queries(nq, d)); + auto rounded = decode_fp16(queries); + faiss::IndexFlatL2 index(d); + index.add(nb, base.data()); + + faiss::IDSelectorRange sel(100, 250); + faiss::SearchParameters params; + params.sel = &sel; + std::vector fd(nq * k), hd(nq * k); + std::vector fi(nq * k), hi(nq * k); + index.search(nq, rounded.data(), k, fd.data(), fi.data(), ¶ms); + index.search_ex( + nq, + queries.data(), + faiss::NumericType::Float16, + k, + hd.data(), + hi.data(), + ¶ms); + EXPECT_EQ(hi, fi); + EXPECT_EQ(hd, fd); + for (auto id : hi) { + EXPECT_TRUE(id >= 100 && id < 250); + } +} + +TEST(IndexFlat, search_ex_float16_keeps_flat1d_search) { + constexpr size_t nb = 100; + constexpr size_t nq = 20; + constexpr size_t k = 4; + auto base = make_training_data(nb, 1); + auto queries = encode_fp16(make_queries(nq, 1)); + auto rounded = decode_fp16(queries); + faiss::IndexFlat1D index; + index.add(nb, base.data()); + + std::vector fd(nq * k), hd(nq * k); + std::vector fi(nq * k), hi(nq * k); + index.search(nq, rounded.data(), k, fd.data(), fi.data()); + index.search_ex( + nq, + queries.data(), + faiss::NumericType::Float16, + k, + hd.data(), + hi.data()); + EXPECT_EQ(hi, fi); + EXPECT_EQ(hd, fd); // L1 distances from IndexFlat1D::search +} + +TEST(Clustering, train_ex_float16_native_flat_assignment_matches_float32) { + constexpr size_t n = 3000; + constexpr size_t d = 48; + constexpr size_t k = 8; + ScopedBlasTiles tiles(256); + auto encoded = encode_fp16(make_training_data(n, d)); + auto rounded = decode_fp16(encoded); + auto cp = small_clustering_params(5); + + faiss::Clustering full(d, k, cp); + faiss::IndexFlatL2 full_index(d); + full.train(n, rounded.data(), full_index); + + faiss::Clustering half(d, k, cp); + faiss::IndexFlatL2 half_index(d); + half.train_ex(n, encoded.data(), faiss::NumericType::Float16, half_index); + + EXPECT_EQ(half.centroids, full.centroids); + ASSERT_EQ(half.iteration_stats.size(), full.iteration_stats.size()); + for (size_t i = 0; i < full.iteration_stats.size(); ++i) { + EXPECT_EQ(half.iteration_stats[i].obj, full.iteration_stats[i].obj); + } +}