From e7c44eb000bebb16f84be38a115caa5d333a8229 Mon Sep 17 00:00:00 2001 From: Jaeyoon Jeong Date: Fri, 2 Oct 2026 18:24:52 -0700 Subject: [PATCH] Add AVX-512 fp16 conversion kernels (#5694) Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5694 Add exact AVX-512 widening and fp32 accumulation specializations for packed fp16 data. Reuse the existing AVX-512 SIMD translation unit, expand runtime dispatch to the AVX-512 level, and cover conversion tails, large blocks, weighted and unweighted centroid updates, and architecture fallbacks. Reviewed By: mnorris11 Differential Revision: D122781509 fbshipit-source-id: 67cbc4546149f069706e80dd25cae25d3588b0cc --- faiss/impl/ClusteringHelpers.cpp | 36 +++--- faiss/impl/ClusteringInitialization.cpp | 57 +++++---- faiss/utils/simd_impl/distances_avx512.cpp | 41 +++++++ faiss/utils/simd_impl/fp16_kernels.h | 15 +++ faiss/utils/utils.cpp | 26 ++-- tests/test_distances_dispatch.cpp | 136 +++++++++++++++++++++ 6 files changed, 248 insertions(+), 63 deletions(-) diff --git a/faiss/impl/ClusteringHelpers.cpp b/faiss/impl/ClusteringHelpers.cpp index a4accebd5..121aebfeb 100644 --- a/faiss/impl/ClusteringHelpers.cpp +++ b/faiss/impl/ClusteringHelpers.cpp @@ -227,26 +227,22 @@ void compute_centroids_fp16( const float* weights, float* hassign, float* centroids) { - with_selected_simd_levels( - [&]() { - compute_centroids_impl( - d, - k, - n, - k_frozen, - assign, - weights, - hassign, - centroids, - [&](size_t i, - const float* w, - float* c, - float* /*scratch*/) { - // x * 1.0f is exact, so unweighted sums match the - // fp32 path. - fp16_madd(d, x + i * d, w ? *w : 1.0f, c); - }); - }); + with_selected_simd_levels([&]() { + compute_centroids_impl( + d, + k, + n, + k_frozen, + assign, + weights, + hassign, + centroids, + [&](size_t i, const float* w, float* c, float* /*scratch*/) { + // x * 1.0f is exact, so unweighted sums match the + // fp32 path. + fp16_madd(d, x + i * d, w ? *w : 1.0f, c); + }); + }); } bool fp16_all_finite(size_t n, const uint16_t* x) { diff --git a/faiss/impl/ClusteringInitialization.cpp b/faiss/impl/ClusteringInitialization.cpp index fbd2160eb..b6693a7ee 100644 --- a/faiss/impl/ClusteringInitialization.cpp +++ b/faiss/impl/ClusteringInitialization.cpp @@ -485,35 +485,34 @@ void init_centroids_fp16( n_existing_centroids, existing_centroids); - with_selected_simd_levels( - [&]() { - switch (initializer.method) { - case ClusteringInitMethod::KMEANS_PLUS_PLUS: - init_kmeans_plus_plus_impl( - initializer, - n, - Fp16Rows{initializer.d, x}, - centroids, - n_existing_centroids, - existing_centroids); - break; - case ClusteringInitMethod::AFK_MC2: - init_afkmc2_impl( - initializer, - n, - Fp16Rows{initializer.d, x}, - centroids, - n_existing_centroids, - existing_centroids); - break; - case ClusteringInitMethod::RANDOM: - FAISS_THROW_MSG( - "fp16 centroid helper requires non-random " - "initialization"); - default: - FAISS_THROW_MSG("Unknown initialization method"); - } - }); + with_selected_simd_levels([&]() { + switch (initializer.method) { + case ClusteringInitMethod::KMEANS_PLUS_PLUS: + init_kmeans_plus_plus_impl( + initializer, + n, + Fp16Rows{initializer.d, x}, + centroids, + n_existing_centroids, + existing_centroids); + break; + case ClusteringInitMethod::AFK_MC2: + init_afkmc2_impl( + initializer, + n, + Fp16Rows{initializer.d, x}, + centroids, + n_existing_centroids, + existing_centroids); + break; + case ClusteringInitMethod::RANDOM: + FAISS_THROW_MSG( + "fp16 centroid helper requires non-random " + "initialization"); + default: + FAISS_THROW_MSG("Unknown initialization method"); + } + }); } } // namespace detail diff --git a/faiss/utils/simd_impl/distances_avx512.cpp b/faiss/utils/simd_impl/distances_avx512.cpp index 284a43c51..a13205da2 100644 --- a/faiss/utils/simd_impl/distances_avx512.cpp +++ b/faiss/utils/simd_impl/distances_avx512.cpp @@ -15,6 +15,7 @@ #include #include #include +#include #ifndef FINTEGER #define FINTEGER long @@ -48,6 +49,46 @@ int sgemm_( namespace faiss { +namespace detail { + +template <> +void fp16_madd( + size_t d, + const uint16_t* x, + float w, + float* c) { + const __m512 wv = _mm512_set1_ps(w); + size_t j = 0; + for (; j + 16 <= d; j += 16) { + const __m512 xv = _mm512_cvtph_ps( + _mm256_loadu_si256(reinterpret_cast(x + j))); + _mm512_storeu_ps( + c + j, + _mm512_add_ps(_mm512_loadu_ps(c + j), _mm512_mul_ps(xv, wv))); + } + for (; j < d; ++j) { + c[j] += decode_fp16(x[j]) * w; + } +} + +template <> +void fp16_to_fp32_kernel( + size_t n, + const uint16_t* x, + float* out) { + size_t j = 0; + for (; j + 16 <= n; j += 16) { + const __m256i h = + _mm256_loadu_si256(reinterpret_cast(x + j)); + _mm512_storeu_ps(out + j, _mm512_cvtph_ps(h)); + } + for (; j < n; ++j) { + out[j] = decode_fp16(x[j]); + } +} + +} // namespace detail + template <> void fvec_madd( const size_t n, diff --git a/faiss/utils/simd_impl/fp16_kernels.h b/faiss/utils/simd_impl/fp16_kernels.h index 4b7f707e6..b813ea166 100644 --- a/faiss/utils/simd_impl/fp16_kernels.h +++ b/faiss/utils/simd_impl/fp16_kernels.h @@ -42,6 +42,21 @@ void fp16_to_fp32_kernel( float* out); #endif +#ifdef COMPILE_SIMD_AVX512 +template <> +void fp16_madd( + size_t d, + const uint16_t* x, + float w, + float* c); + +template <> +void fp16_to_fp32_kernel( + size_t n, + const uint16_t* x, + float* out); +#endif + #ifdef COMPILE_SIMD_ARM_NEON template <> void fp16_madd( diff --git a/faiss/utils/utils.cpp b/faiss/utils/utils.cpp index e39a7dda5..dec1552e3 100644 --- a/faiss/utils/utils.cpp +++ b/faiss/utils/utils.cpp @@ -644,21 +644,19 @@ template struct CombinerRangeKNN; void fp16_to_fp32(size_t n, const uint16_t* x, float* out) { // below this many values, stay on the calling thread constexpr size_t block = size_t(1) << 16; - with_selected_simd_levels( - [&]() { - if (n <= block) { - detail::fp16_to_fp32_kernel(n, x, out); - return; - } - const int64_t nblock = - static_cast((n + block - 1) / block); + with_selected_simd_levels([&]() { + if (n <= block) { + detail::fp16_to_fp32_kernel(n, x, out); + return; + } + const int64_t nblock = static_cast((n + block - 1) / block); #pragma omp parallel for - for (int64_t b = 0; b < nblock; b++) { - const size_t j0 = static_cast(b) * block; - detail::fp16_to_fp32_kernel( - std::min(block, n - j0), x + j0, out + j0); - } - }); + for (int64_t b = 0; b < nblock; b++) { + const size_t j0 = static_cast(b) * block; + detail::fp16_to_fp32_kernel( + std::min(block, n - j0), x + j0, out + j0); + } + }); } void CodeSet::insert(size_t n, const uint8_t* codes, bool* inserted) { diff --git a/tests/test_distances_dispatch.cpp b/tests/test_distances_dispatch.cpp index 91e4911ab..90aab4452 100644 --- a/tests/test_distances_dispatch.cpp +++ b/tests/test_distances_dispatch.cpp @@ -34,11 +34,15 @@ #include #include +#include #include #include +#include #include +#include #include +#include using namespace faiss; @@ -74,6 +78,26 @@ std::vector available_levels() { return out; } +std::vector available_fp16_levels() { + static const std::vector all = { + SIMDLevel::NONE, + SIMDLevel::AVX2, + SIMDLevel::AVX512, + SIMDLevel::AVX512_SPR, + SIMDLevel::AVX512_VPOPCNT, + SIMDLevel::ARM_NEON, + SIMDLevel::ARM_SVE, + SIMDLevel::RISCV_RVV, + }; + std::vector out; + for (auto lv : all) { + if (SIMDConfig::is_simd_level_available(lv)) { + out.push_back(lv); + } + } + return out; +} + std::vector rand_vec(size_t n, uint32_t seed) { std::mt19937 rng(seed); std::uniform_real_distribution u(-1.f, 1.f); @@ -144,6 +168,118 @@ void check_index_at_levels(Fn fn, const std::vector& levels) { } // namespace +TEST(DistancesDispatch, Fp16ConversionTails_AllLevels) { + SIMDLevelGuard guard; + auto levels = available_fp16_levels(); + constexpr uint16_t kValues[] = { + 0x0000, + 0x8000, + 0x3c00, + 0xc000, + 0x3555, + 0x0400, + 0x03ff, + 0x7bff, + 0x0001, + }; + std::vector sizes; + for (size_t n = 1; n <= 31; ++n) { + sizes.push_back(n); + } + sizes.push_back((size_t(1) << 16) + 17); + + for (size_t n : sizes) { + std::vector input(n); + std::vector expected(n); + for (size_t i = 0; i < n; ++i) { + input[i] = kValues[i % (sizeof(kValues) / sizeof(kValues[0]))]; + expected[i] = decode_fp16(input[i]); + } + for (auto lv : levels) { + SIMDConfig::set_level(lv); + std::vector actual(n); + fp16_to_fp32(n, input.data(), actual.data()); + EXPECT_EQ( + 0, + std::memcmp( + expected.data(), actual.data(), n * sizeof(float))) + << "diverged at level " << static_cast(lv) << " size " + << n; + } + } +} + +TEST(DistancesDispatch, Fp16CentroidMaddTails_AllLevels) { + SIMDLevelGuard guard; + auto levels = available_fp16_levels(); + constexpr size_t n = 3; + constexpr int64_t kUnweightedAssign[n] = {0, 1, 0}; + constexpr int64_t kWeightedAssign[n] = {0, 0, 0}; + constexpr float kWeights[n] = {0.1f, 0.3f, 0.7f}; + constexpr float kValues[] = {0.0f, 0.5f, -1.0f, 2.0f, -0.25f, 4.0f, 0.125f}; + + for (size_t d = 1; d <= 31; ++d) { + std::vector input(n * d); + for (size_t i = 0; i < input.size(); ++i) { + input[i] = encode_fp16( + kValues[i % (sizeof(kValues) / sizeof(kValues[0]))]); + } + + auto check_case = [&](size_t k, + const int64_t* assign, + const float* weights, + const char* case_name) { + std::vector expected_hassign(k, 0.0f); + std::vector expected(k * d, 0.0f); + for (size_t i = 0; i < n; ++i) { + const size_t ci = static_cast(assign[i]); + const float weight = weights ? weights[i] : 1.0f; + expected_hassign[ci] += weight; + for (size_t j = 0; j < d; ++j) { + expected[ci * d + j] += + decode_fp16(input[i * d + j]) * weight; + } + } + for (size_t ci = 0; ci < k; ++ci) { + for (size_t j = 0; j < d; ++j) { + expected[ci * d + j] /= expected_hassign[ci]; + } + } + + for (auto lv : levels) { + SIMDConfig::set_level(lv); + std::vector hassign(k, 0.0f); + std::vector actual(k * d); + detail::compute_centroids_fp16( + d, + k, + n, + 0, + input.data(), + assign, + weights, + hassign.data(), + actual.data()); + for (size_t ci = 0; ci < k; ++ci) { + EXPECT_NEAR(expected_hassign[ci], hassign[ci], tol_for_d(d)) + << case_name << " histogram diverged at level " + << static_cast(lv) << " dimension " << d + << " centroid " << ci; + } + for (size_t i = 0; i < actual.size(); ++i) { + EXPECT_NEAR(expected[i], actual[i], tol_for_d(d)) + << case_name << " centroid diverged at level " + << static_cast(lv) << " dimension " << d + << " index " << i; + } + } + }; + + check_case(2, kUnweightedAssign, nullptr, "unweighted"); + check_case(1, kWeightedAssign, kWeights, "weighted"); + } +} + TEST(DistancesDispatch, FvecL1_AllLevels) { SIMDLevelGuard guard; auto levels = available_levels();