mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-12 02:13:36 +00:00
Jaeyoon Jeong
430a41ee4c
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
53 lines
1.2 KiB
C++
53 lines
1.2 KiB
C++
/*
|
|
* 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
|