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

Add AVX-512 fp16 conversion kernels (#5694)

Browse Source 
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
This commit is contained in:
Jaeyoon Jeong
2026-10-02 18:24:52 -07:00
committed by meta-codesync[bot]
parent ae4466d57e
commit e7c44eb000
6 changed files with 248 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files
faiss/impl/ClusteringHelpers.cpp
+16 -20
View File
@@ -227,26 +227,22 @@ void compute_centroids_fp16(
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);
});
});
with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_BASE>([&]<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) {
faiss/impl/ClusteringInitialization.cpp
+28 -29
View File
@@ -485,35 +485,34 @@ void init_centroids_fp16(
n_existing_centroids,
existing_centroids);
with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_AVX2_NEON>(
[&]<SIMDLevel SL>() {
switch (initializer.method) {
case ClusteringInitMethod::KMEANS_PLUS_PLUS:
init_kmeans_plus_plus_impl(
initializer,
n,
Fp16Rows<SL>{initializer.d, x},
centroids,
n_existing_centroids,
existing_centroids);
break;
case ClusteringInitMethod::AFK_MC2:
init_afkmc2_impl(
initializer,
n,
Fp16Rows<SL>{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<AVAILABLE_SIMD_LEVELS_BASE>([&]<SIMDLevel SL>() {
switch (initializer.method) {
case ClusteringInitMethod::KMEANS_PLUS_PLUS:
init_kmeans_plus_plus_impl(
initializer,
n,
Fp16Rows<SL>{initializer.d, x},
centroids,
n_existing_centroids,
existing_centroids);
break;
case ClusteringInitMethod::AFK_MC2:
init_afkmc2_impl(
initializer,
n,
Fp16Rows<SL>{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
faiss/utils/simd_impl/distances_avx512.cpp
+41
View File
@@ -15,6 +15,7 @@
#include <faiss/impl/ResultHandler.h>
#include <faiss/utils/distances_fused/distances_fused.h>
#include <faiss/utils/simd_impl/exhaustive_L2sqr_blas_cmax.h>
#include <faiss/utils/simd_impl/fp16_kernels.h>
#ifndef FINTEGER
#define FINTEGER long
@@ -48,6 +49,46 @@ int sgemm_(
namespace faiss {
namespace detail {
template <>
void fp16_madd<SIMDLevel::AVX512>(
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<const __m256i*>(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<SIMDLevel::AVX512>(
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<const __m256i*>(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<SIMDLevel::AVX512>(
const size_t n,
faiss/utils/simd_impl/fp16_kernels.h
+15
View File
@@ -42,6 +42,21 @@ void fp16_to_fp32_kernel<SIMDLevel::AVX2>(
float* out);
#endif
#ifdef COMPILE_SIMD_AVX512
template <>
void fp16_madd<SIMDLevel::AVX512>(
size_t d,
const uint16_t* x,
float w,
float* c);
template <>
void fp16_to_fp32_kernel<SIMDLevel::AVX512>(
size_t n,
const uint16_t* x,
float* out);
#endif
#ifdef COMPILE_SIMD_ARM_NEON
template <>
void fp16_madd<SIMDLevel::ARM_NEON>(
faiss/utils/utils.cpp
+12 -14
View File
@@ -644,21 +644,19 @@ 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);
with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_BASE>([&]<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);
}
});
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) {
tests/test_distances_dispatch.cpp
+136
View File
@@ -34,11 +34,15 @@
#include <algorithm>
#include <cmath>
#include <cstring>
#include <random>
#include <vector>
#include <faiss/impl/ClusteringHelpers.h>
#include <faiss/utils/distances.h>
#include <faiss/utils/fp16.h>
#include <faiss/utils/simd_levels.h>
#include <faiss/utils/utils.h>
using namespace faiss;
@@ -74,6 +78,26 @@ std::vector<SIMDLevel> available_levels() {
return out;
}
std::vector<SIMDLevel> available_fp16_levels() {
static const std::vector<SIMDLevel> all = {
SIMDLevel::NONE,
SIMDLevel::AVX2,
SIMDLevel::AVX512,
SIMDLevel::AVX512_SPR,
SIMDLevel::AVX512_VPOPCNT,
SIMDLevel::ARM_NEON,
SIMDLevel::ARM_SVE,
SIMDLevel::RISCV_RVV,
};
std::vector<SIMDLevel> out;
for (auto lv : all) {
if (SIMDConfig::is_simd_level_available(lv)) {
out.push_back(lv);
}
}
return out;
}
std::vector<float> rand_vec(size_t n, uint32_t seed) {
std::mt19937 rng(seed);
std::uniform_real_distribution<float> u(-1.f, 1.f);
@@ -144,6 +168,118 @@ void check_index_at_levels(Fn fn, const std::vector<SIMDLevel>& 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<size_t> 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<uint16_t> input(n);
std::vector<float> 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<float> 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<int>(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<uint16_t> 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<float> expected_hassign(k, 0.0f);
std::vector<float> expected(k * d, 0.0f);
for (size_t i = 0; i < n; ++i) {
const size_t ci = static_cast<size_t>(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<float> hassign(k, 0.0f);
std::vector<float> 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<int>(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<int>(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();