facebookresearch/faiss
1
0
Fork 0
mirror of https://github.com/facebookresearch/faiss.git synced 2026-10-11 22:50:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

Native fp16 k-means training for Faiss IVF (#5671)

Browse Source 
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
This commit is contained in:
Jaeyoon Jeong
2026-10-01 04:06:09 -07:00
committed by meta-codesync[bot]
parent 88a28bcd80
commit 430a41ee4c
30 changed files with 2394 additions and 150 deletions
Download Patch File Download Diff File Expand all files Collapse all files
faiss/CMakeLists.txt
+3
View File
@@ -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
faiss/Clustering.cpp
+124 -43
View File
@@ -11,6 +11,7 @@
#include <faiss/VectorTransform.h>
#include <faiss/impl/AuxIndexStructures.h>
#include <algorithm>
#include <cinttypes>
#include <cmath>
#include <cstdio>
@@ -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<const float*>(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<const uint8_t*>(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<idx_t>(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<size_t>(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<const float*>(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<const uint16_t*>(x_in)),
"input contains NaN's or Inf's");
} else {
const float* x = reinterpret_cast<const float*>(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<uint8_t[]> del1;
std::unique_ptr<float[]> 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<const uint16_t*>(rows), out);
} else {
memcpy(out, rows, n * line_size);
}
};
if (static_cast<size_t>(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<float> decode_buffer(codec ? d * decode_block_size : 0);
// temporary buffer to decode codec-encoded vectors during the optimization
std::vector<float> decode_buffer(
codec ? d * std::min(static_cast<size_t>(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<int> 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<float> x_decoded;
if (!codec) {
if (!encoded) {
x_float = reinterpret_cast<const float*>(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<const float*>(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<size_t>(nx);
i0 += decode_block_size) {
size_t i1 = i0 + decode_block_size;
if (i1 > static_cast<size_t>(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<float> 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<const uint16_t*>(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());
faiss/Clustering.h
+29
View File
@@ -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
faiss/Index.cpp
+36
View File
@@ -13,8 +13,11 @@
#include <faiss/impl/DistanceComputer.h>
#include <faiss/impl/FaissAssert.h>
#include <faiss/utils/distances.h>
#include <faiss/utils/utils.h>
#include <algorithm>
#include <cstring>
#include <vector>
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<const float*>(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<idx_t>(
1, idx_t(64) * 1024 * 1024 / (sizeof(float) * idx_t(d)));
const uint16_t* x16 = static_cast<const uint16_t*>(x);
std::vector<float> 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");
}
faiss/Index.h
+7 -12
View File
@@ -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<const float*>(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(
faiss/IndexFlat.cpp
+56
View File
@@ -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<const uint16_t*>(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,
faiss/IndexFlat.h
+31
View File
@@ -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
faiss/IndexIVF.cpp
+201 -11
View File
@@ -18,6 +18,7 @@
#include <cstdio>
#include <limits>
#include <faiss/utils/random.h>
#include <faiss/utils/utils.h>
#include <faiss/IndexFlat.h>
@@ -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<size_t>(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<size_t>(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<size_t>(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<ClusteringParameters&>(super_cp) = cp;
SuperKMeans clus(
static_cast<int>(d), static_cast<int>(nlist), super_cp);
clus.train(n, x);
clus.train(n, reinterpret_cast<const float*>(x));
quantizer->add(nlist, clus.centroids.data());
} else {
Clustering clus(static_cast<int>(d), static_cast<int>(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<int>(d), static_cast<int>(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<const uint8_t*>(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 <class DecodeRows>
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<float> decoded(static_cast<size_t>(max_nt) * d);
if (max_nt == n) {
decode_rows(n, x, decoded.data());
} else {
FAISS_THROW_IF_NOT_FMT(
n <= static_cast<idx_t>(std::numeric_limits<int>::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<int> subset(static_cast<size_t>(n));
rand_perm(subset.data(), n, 1234);
for (idx_t i = 0; i < max_nt; ++i) {
decode_rows(
1,
x + static_cast<size_t>(subset[i]) * row_size,
decoded.data() + static_cast<size_t>(i) * d);
}
}
n = max_nt;
if (ivf.by_residual) {
FAISS_THROW_IF_NOT_MSG(ivf.quantizer, "IVF quantizer must not be null");
std::vector<idx_t> assign(n);
ivf.quantizer->assign(n, decoded.data(), assign.data());
std::vector<float> residuals(static_cast<size_t>(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<const uint8_t*>(x),
sizeof(uint16_t) * dim,
[dim](idx_t ni, const uint8_t* rows, float* out) {
fp16_to_fp32(
ni * dim, reinterpret_cast<const uint16_t*>(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");
faiss/IndexIVF.h
+45 -1
View File
@@ -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;
faiss/IndexIVFFlat.cpp
+33
View File
@@ -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<const float*>(x));
}
void IndexIVFFlatDedup::train(idx_t n, const float* x) {
std::unordered_map<uint64_t, idx_t> map;
std::unique_ptr<float[]> x2(new float[n * d]);
faiss/IndexIVFFlat.h
+16
View File
@@ -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;
faiss/gpu/GpuIndex.cu
+28
View File
@@ -15,6 +15,7 @@
#include <faiss/gpu/utils/DeviceUtils.h>
#include <faiss/gpu/utils/StaticUtils.h>
#include <faiss/impl/FaissAssert.h>
#include <faiss/gpu/utils/ConversionOperators.cuh>
#include <faiss/gpu/utils/CopyUtils.cuh>
#include <faiss/gpu/utils/Float16.cuh>
@@ -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<const float*>(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<half, 2, true> queries(
const_cast<half*>(static_cast<const half*>(x)), {n, this->d});
auto queriesFloat = convertTensorTemporary<half, float, 2>(
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,
faiss/gpu/GpuIndex.h
+4 -13
View File
@@ -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<const float*>(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
faiss/gpu/GpuIndexFlat.cu
+35
View File
@@ -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<half, 2, true> queries(
const_cast<half*>(static_cast<const half*>(x)), {n, this->d});
Tensor<float, 2, true> outDistances(distances, {n, k});
Tensor<idx_t, 2, true> 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);
faiss/gpu/GpuIndexFlat.h
+11
View File
@@ -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_;
faiss/gpu/test/TestGpuIndexFlat.cpp
+176
View File
@@ -5,14 +5,19 @@
* LICENSE file in the root directory of this source tree.
*/
#include <faiss/Clustering.h>
#include <faiss/IndexFlat.h>
#include <faiss/gpu/GpuIndexFlat.h>
#include <faiss/gpu/GpuIndexIVFFlat.h>
#include <faiss/gpu/StandardGpuResources.h>
#include <faiss/gpu/impl/IndexUtils.h>
#include <faiss/gpu/test/TestUtils.h>
#include <faiss/gpu/utils/DeviceUtils.h>
#include <faiss/impl/IDSelector.h>
#include <faiss/utils/fp16.h>
#include <gtest/gtest.h>
#include <algorithm>
#include <cstdint>
#include <sstream>
#include <unordered_map>
#include <vector>
@@ -807,6 +812,177 @@ TEST(TestCuvsGpuIndexFlat, IDSelector_IP) {
}
#endif
// Packs float vectors as IEEE binary16, the Float16 input of search_ex
std::vector<uint16_t> toFp16(const std::vector<float>& v) {
std::vector<uint16_t> 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<float> refDistances(numQuery * k);
std::vector<faiss::idx_t> refLabels(numQuery * k);
index.search(
numQuery,
roundedQueries.data(),
k,
refDistances.data(),
refLabels.data());
auto expectSearchExMatches = [&](const void* x, const std::string& where) {
std::vector<float> distances(numQuery * k, 0);
std::vector<faiss::idx_t> 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 <typename GpuIndexFlatT>
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<faiss::gpu::GpuIndexFlatL2>();
}
TEST(TestGpuIndexFlat, SearchExFloat16_IP) {
testFlatSearchExFloat16<faiss::gpu::GpuIndexFlatIP>();
}
// 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);
faiss/impl/ClusteringHelpers.cpp
+136 -45
View File
@@ -7,6 +7,8 @@
#include <faiss/impl/ClusteringHelpers.h>
#include <algorithm>
#include <atomic>
#include <cassert>
#include <chrono>
#include <cinttypes>
@@ -20,11 +22,86 @@
#include <faiss/Index.h>
#include <faiss/impl/FaissAssert.h>
#include <faiss/impl/simd_dispatch.h>
#include <faiss/utils/random.h>
#include <faiss/utils/simd_impl/fp16_kernels.h>
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 <class AddRow>
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<bool> invalid_assignment{false};
const int64_t num_centroids = static_cast<int64_t>(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<float> 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<int64_t>(c0) &&
ci < static_cast<int64_t>(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<idx_t>(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<float> 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<int64_t>(c0) &&
ci < static_cast<int64_t>(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<const float*>(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<idx_t>(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<AVAILABLE_SIMD_LEVELS_AVX2_NEON>(
[&]<SIMDLevel SL>() {
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<SL>(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<bool> all_finite{true};
#pragma omp parallel for if (n > fp16_parallel_threshold)
for (int64_t i = 0; i < static_cast<int64_t>(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
faiss/impl/ClusteringHelpers.h
+27
View File
@@ -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
faiss/utils/distances.cpp
+222 -10
View File
@@ -26,6 +26,7 @@
#include <faiss/utils/distances_dispatch.h>
#include <faiss/utils/distances_fused/distances_fused.h>
#include <faiss/utils/simd_impl/exhaustive_L2sqr_blas_cmax.h>
#include <faiss/utils/utils.h>
#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 <class BlockResultHandler>
void exhaustive_inner_product_blas(
const float* x,
/** Find the nearest neighbors for nx queries (see Fp32QueryTiles) in a set
* of ny vectors */
template <class BlockResultHandler, class QueryTiles>
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 <class BlockResultHandler>
void exhaustive_L2sqr_blas_default_impl(
const float* x,
template <class BlockResultHandler, class QueryTiles>
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<float[]> x_norms(new float[nx]);
std::unique_ptr<float[]> 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 <class BlockResultHandler>
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 <class BlockResultHandler>
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 <class BlockResultHandler>
void exhaustive_L2sqr_blas(
const float* x,
@@ -556,6 +587,44 @@ void exhaustive_L2sqr_blas<Top1BlockResultHandler<CMax<float, int64_t>>>(
});
}
/// exhaustive_L2sqr_blas for packed IEEE binary16 queries
template <class BlockResultHandler>
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<Top1BlockResultHandler<CMax<float, int64_t>>>(
const uint16_t* x,
const float* y,
size_t d,
size_t nx,
size_t ny,
Top1BlockResultHandler<CMax<float, int64_t>>& res,
const float* y_norms) {
with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_BASE_WITH_SVE>(
[&]<SIMDLevel SL>() {
if constexpr (
SL == SIMDLevel::AVX2 || SL == SIMDLevel::AVX512 ||
SL == SIMDLevel::ARM_SVE) {
exhaustive_L2sqr_blas_cmax_fp16<SL>(
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 <class BlockResultHandler>
@@ -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 <class BlockResultHandler>
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 <class BlockResultHandler>
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<size_t>(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 <class Fp32Search>
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<size_t>(1, (size_t(8) << 20) / (bs_x * d * sizeof(float)));
std::vector<float> 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
***************************************************************************/
faiss/utils/distances.h
+30
View File
@@ -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
*
faiss/utils/simd_impl/distances_arm_sve.cpp
+37 -5
View File
@@ -758,9 +758,11 @@ size_t fvec_L2sqr_ny_nearest_y_transposed<SIMDLevel::ARM_SVE>(
return current_min_idx;
}
template <>
void exhaustive_L2sqr_blas_cmax<SIMDLevel::ARM_SVE>(
const float* x,
namespace {
template <class QueryTiles>
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<SIMDLevel::ARM_SVE>(
std::unique_ptr<float[]> x_norms(new float[nx]);
std::unique_ptr<float[]> 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<SIMDLevel::ARM_SVE>(
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<SIMDLevel::ARM_SVE>(
&one,
y + j0 * d,
&di,
x + i0 * d,
xt,
&di,
&zero,
ip_block.get(),
@@ -938,4 +942,32 @@ void exhaustive_L2sqr_blas_cmax<SIMDLevel::ARM_SVE>(
}
}
} // namespace
template <>
void exhaustive_L2sqr_blas_cmax<SIMDLevel::ARM_SVE>(
const float* x,
const float* y,
size_t d,
size_t nx,
size_t ny,
Top1BlockResultHandler<CMax<float, int64_t>>& 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<SIMDLevel::ARM_SVE>(
const uint16_t* x,
const float* y,
size_t d,
size_t nx,
size_t ny,
Top1BlockResultHandler<CMax<float, int64_t>>& 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
faiss/utils/simd_impl/distances_avx2.cpp
+37 -5
View File
@@ -1213,9 +1213,11 @@ int fvec_madd_and_argmin<SIMDLevel::AVX2>(
return fvec_madd_and_argmin_sse(n, a, bf, b, c);
}
template <>
void exhaustive_L2sqr_blas_cmax<SIMDLevel::AVX2>(
const float* x,
namespace {
template <class QueryTiles>
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<SIMDLevel::AVX2>(
std::unique_ptr<float[]> x_norms(new float[nx]);
std::unique_ptr<float[]> 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<SIMDLevel::AVX2>(
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<SIMDLevel::AVX2>(
&one,
y + j0 * d,
&di,
x + i0 * d,
xt,
&di,
&zero,
ip_block.get(),
@@ -1427,4 +1431,32 @@ void exhaustive_L2sqr_blas_cmax<SIMDLevel::AVX2>(
}
}
} // namespace
template <>
void exhaustive_L2sqr_blas_cmax<SIMDLevel::AVX2>(
const float* x,
const float* y,
size_t d,
size_t nx,
size_t ny,
Top1BlockResultHandler<CMax<float, int64_t>>& 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<SIMDLevel::AVX2>(
const uint16_t* x,
const float* y,
size_t d,
size_t nx,
size_t ny,
Top1BlockResultHandler<CMax<float, int64_t>>& 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
faiss/utils/simd_impl/distances_avx512.cpp
+37 -5
View File
@@ -1133,9 +1133,11 @@ int fvec_madd_and_argmin<SIMDLevel::AVX512>(
return fvec_madd_and_argmin_sse(n, a, bf, b, c);
}
template <>
void exhaustive_L2sqr_blas_cmax<SIMDLevel::AVX512>(
const float* x,
namespace {
template <class QueryTiles>
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<SIMDLevel::AVX512>(
std::unique_ptr<float[]> x_norms(new float[nx]);
std::unique_ptr<float[]> 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<SIMDLevel::AVX512>(
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<SIMDLevel::AVX512>(
&one,
y + j0 * d,
&di,
x + i0 * d,
xt,
&di,
&zero,
ip_block.get(),
@@ -1340,4 +1344,32 @@ void exhaustive_L2sqr_blas_cmax<SIMDLevel::AVX512>(
}
}
} // namespace
template <>
void exhaustive_L2sqr_blas_cmax<SIMDLevel::AVX512>(
const float* x,
const float* y,
size_t d,
size_t nx,
size_t ny,
Top1BlockResultHandler<CMax<float, int64_t>>& 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<SIMDLevel::AVX512>(
const uint16_t* x,
const float* y,
size_t d,
size_t nx,
size_t ny,
Top1BlockResultHandler<CMax<float, int64_t>>& 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
faiss/utils/simd_impl/exhaustive_L2sqr_blas_cmax.h
+93
View File
@@ -7,11 +7,93 @@
#pragma once
#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <memory>
#include <faiss/impl/ResultHandler.h>
#include <faiss/utils/distances.h>
#include <faiss/utils/simd_levels.h>
#include <faiss/utils/utils.h>
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<float[]> 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<size_t>(
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<int64_t>(chunk_end - chunk_begin);
#pragma omp parallel for if (n * static_cast<int64_t>(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 <SIMDLevel>
@@ -24,4 +106,15 @@ void exhaustive_L2sqr_blas_cmax(
Top1BlockResultHandler<CMax<float, int64_t>>& res,
const float* y_norms);
/// Same as exhaustive_L2sqr_blas_cmax for packed IEEE binary16 queries.
template <SIMDLevel>
void exhaustive_L2sqr_blas_cmax_fp16(
const uint16_t* x,
const float* y,
size_t d,
size_t nx,
size_t ny,
Top1BlockResultHandler<CMax<float, int64_t>>& res,
const float* y_norms);
} // namespace faiss
faiss/utils/simd_impl/fp16_kernels.h
+61
View File
@@ -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 <cstddef>
#include <cstdint>
#include <faiss/utils/fp16.h>
#include <faiss/utils/simd_levels.h>
namespace faiss {
namespace detail {
// c[j] += fp32(x[j]) * w for j < d
template <SIMDLevel Level>
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 <SIMDLevel Level>
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<SIMDLevel::AVX2>(size_t d, const uint16_t* x, float w, float* c);
template <>
void fp16_to_fp32_kernel<SIMDLevel::AVX2>(
size_t n,
const uint16_t* x,
float* out);
#endif
#ifdef COMPILE_SIMD_ARM_NEON
template <>
void fp16_madd<SIMDLevel::ARM_NEON>(
size_t d,
const uint16_t* x,
float w,
float* c);
template <>
void fp16_to_fp32_kernel<SIMDLevel::ARM_NEON>(
size_t n,
const uint16_t* x,
float* out);
#endif
} // namespace detail
} // namespace faiss
faiss/utils/simd_impl/fp16_kernels_avx2.cpp
+56
View File
@@ -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 <faiss/utils/simd_impl/fp16_kernels.h>
#include <immintrin.h>
namespace faiss {
namespace detail {
template <>
void fp16_madd<SIMDLevel::AVX2>(
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<const __m128i*>(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<SIMDLevel::AVX2>(
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<const __m128i*>(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
faiss/utils/simd_impl/fp16_kernels_neon.cpp
+52
View File
@@ -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 <faiss/utils/simd_impl/fp16_kernels.h>
#include <arm_neon.h>
namespace faiss {
namespace detail {
template <>
void fp16_madd<SIMDLevel::ARM_NEON>(
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<SIMDLevel::ARM_NEON>(
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
faiss/utils/utils.cpp
+22
View File
@@ -36,7 +36,9 @@
#include <faiss/impl/AuxIndexStructures.h>
#include <faiss/impl/FaissAssert.h>
#include <faiss/impl/simd_dispatch.h>
#include <faiss/utils/random.h>
#include <faiss/utils/simd_impl/fp16_kernels.h>
#ifndef FINTEGER
#define FINTEGER long
@@ -639,6 +641,26 @@ void CombinerRangeKNN<T>::write_result(T* D_res, int64_t* I_res) {
template struct CombinerRangeKNN<float>;
template struct CombinerRangeKNN<int16_t>;
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<AVAILABLE_SIMD_LEVELS_AVX2_NEON>(
[&]<SIMDLevel SL>() {
if (n <= block) {
detail::fp16_to_fp32_kernel<SL>(n, x, out);
return;
}
const int64_t nblock =
static_cast<int64_t>((n + block - 1) / block);
#pragma omp parallel for
for (int64_t b = 0; b < nblock; b++) {
const size_t j0 = static_cast<size_t>(b) * block;
detail::fp16_to_fp32_kernel<SL>(
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(
faiss/utils/utils.h
+8
View File
@@ -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);
tests/test_ivf_index.cpp
+741
View File
@@ -7,6 +7,7 @@
#include <omp.h>
#include <algorithm>
#include <atomic>
#include <cstddef>
#include <limits>
#include <map>
@@ -17,9 +18,13 @@
#include <faiss/IndexFlat.h>
#include <faiss/IndexIVFFlat.h>
#include <faiss/IndexScalarQuantizer.h>
#include <faiss/impl/AuxIndexStructures.h>
#include <faiss/impl/FaissAssert.h>
#include <faiss/impl/IDSelector.h>
#include <faiss/impl/ResultHandler.h>
#include <faiss/utils/distances.h>
#include <faiss/utils/fp16.h>
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<float> encoder_input;
std::vector<faiss::idx_t> 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<size_t>(n));
faiss::IndexScalarQuantizer::sa_decode(n, bytes, x);
}
mutable std::atomic<size_t> decode_calls{0};
mutable size_t max_decode_rows = 0;
};
std::vector<uint16_t> encode_fp16(const std::vector<float>& values) {
std::vector<uint16_t> encoded(values.size());
for (size_t i = 0; i < values.size(); ++i) {
encoded[i] = faiss::encode_fp16(values[i]);
}
return encoded;
}
std::vector<float> decode_fp16(const std::vector<uint16_t>& values) {
std::vector<float> 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<float> input(n * d);
for (size_t i = 0; i < input.size(); ++i) {
input[i] = static_cast<float>((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> float_centroids(nlist * d);
std::vector<float> 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<float>::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<float> input(n * d);
for (size_t i = 0; i < input.size(); ++i) {
input[i] = static_cast<float>((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<const uint8_t*>(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<const uint8_t*>(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<float> input(n * d);
for (size_t i = 0; i < input.size(); ++i) {
input[i] = static_cast<float>((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<uint16_t> encoded(n * (d + 1), faiss::encode_fp16(1.0f));
faiss::IndexFlatL2 quantizer(d);
std::vector<float> 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<const uint8_t*>(encoded.data()),
&wrong_dimension_codec),
faiss::FaissException);
EXPECT_THROW(
index.train_encoded(
n,
reinterpret_cast<const uint8_t*>(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<float> input(n * d);
for (size_t i = 0; i < input.size(); ++i) {
input[i] = static_cast<float>((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<size_t>(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<size_t>(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<size_t>(n) * static_cast<size_t>(k);
std::fill_n(distances, result_size, 0.0f);
std::fill_n(labels, result_size, faiss::idx_t{-1});
}
};
std::vector<float> make_training_data(size_t n, size_t d) {
std::vector<float> x(n * d);
for (size_t i = 0; i < x.size(); ++i) {
x[i] = static_cast<float>((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<faiss::idx_t>(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<float> weights(n);
for (size_t i = 0; i < n; ++i) {
weights[i] = 0.25f + static_cast<float>(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<float> queries_in(nq * d);
for (size_t i = 0; i < queries_in.size(); ++i) {
queries_in[i] = static_cast<float>((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> float_distances(nq * k), half_distances(nq * k);
std::vector<faiss::idx_t> 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<uint8_t> 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> float_centroids(nlist * d);
std::vector<float> 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<float> 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<float> make_queries(size_t n, size_t d) {
std::vector<float> x(n * d);
for (size_t i = 0; i < x.size(); ++i) {
x[i] = static_cast<float>((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<float> fd(nq * k), hd(nq * k);
std::vector<faiss::idx_t> 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<float> fd(nq * k), hd(nq * k);
std::vector<faiss::idx_t> fi(nq * k), hi(nq * k);
index.search(nq, rounded.data(), k, fd.data(), fi.data(), &params);
index.search_ex(
nq,
queries.data(),
faiss::NumericType::Float16,
k,
hd.data(),
hi.data(),
&params);
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<float> fd(nq * k), hd(nq * k);
std::vector<faiss::idx_t> 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);
}
}