Skip to content

[MOD-17916] Reduce native-FP16 SIMD lanes in FP32 #1021

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
dor-forer wants to merge 1 commit into from
Open

[MOD-17916] Reduce native-FP16 SIMD lanes in FP32 #1021

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 11 additions & 2 deletions src/VecSim/spaces/IP/IP_AVX512FP16_VL_FP16.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,13 +6,16 @@
* (RSALv2); or (b) the Server Side Public License v1 (SSPLv1); or (c) the
* GNU Affero General Public License v3 (AGPLv3).
*/
#pragma once
#include <cstdint>
#include "VecSim/spaces/space_includes.h"
#include "VecSim/types/float16.h"
#include <cstring>

using float16 = vecsim_types::float16;

// Keep the throughput-critical loop in native fp16. Only the final horizontal reduction widens to
// fp32, so finite lane accumulators cannot overflow merely while being combined across lanes.
static void InnerProductStep(float16 *&pVect1, float16 *&pVect2, __m512h &sum) {
__m512h v1 = _mm512_loadu_ph(pVect1);
__m512h v2 = _mm512_loadu_ph(pVect2);
Expand All @@ -22,6 +25,13 @@ static void InnerProductStep(float16 *&pVect1, float16 *&pVect2, __m512h &sum) {
pVect2 += 32;
}

static inline float ReduceFP16AccumulatorToFP32(__m512h sum) {
__m512 lo = _mm512_cvtxph_ps(_mm512_castph512_ph256(sum));
__m512 hi = _mm512_cvtxph_ps(
_mm256_castsi256_ph(_mm512_extracti64x4_epi64(_mm512_castph_si512(sum), 1)));
return _mm512_reduce_add_ps(_mm512_add_ps(lo, hi));
}

template <unsigned short residual> // 0..31
float FP16_InnerProductSIMD32_AVX512FP16_VL(const void *pVect1v, const void *pVect2v,
size_t dimension) {
Expand All @@ -46,6 +56,5 @@ float FP16_InnerProductSIMD32_AVX512FP16_VL(const void *pVect1v, const void *pVe
InnerProductStep(pVect1, pVect2, sum);
} while (pVect1 < pEnd1);

_Float16 res = _mm512_reduce_add_ph(sum);
return _Float16(1) - res;
return 1.0f - ReduceFP16AccumulatorToFP32(sum);
}
25 changes: 12 additions & 13 deletions src/VecSim/spaces/IP/IP_NEON_FP16.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@
*/
#include <arm_neon.h>

// Keep the throughput-critical loop in native fp16. The accumulators are widened separately before
// the final horizontal sum, avoiding fp16 overflow while combining otherwise-finite lanes.
inline void InnerProduct_Step(const float16_t *&vec1, const float16_t *&vec2, float16x8_t &acc) {
// Load half-precision vectors
float16x8_t v1 = vld1q_f16(vec1);
Expand Down Expand Up @@ -77,19 +79,16 @@ float FP16_InnerProduct_NEON_HP(const void *pVect1v, const void *pVect2v, size_t
InnerProduct_Step(vec1, vec2, acc4);
}

// Accumulate accumulators
acc1 = vpaddq_f16(acc1, acc3);
acc2 = vpaddq_f16(acc2, acc4);
acc1 = vpaddq_f16(acc1, acc2);

// Horizontal sum of the accumulated values
float32x4_t sum_f32 = vcvt_f32_f16(vget_low_f16(acc1));
sum_f32 = vaddq_f32(sum_f32, vcvt_f32_f16(vget_high_f16(acc1)));

// Pairwise add to get horizontal sum
float32x2_t sum_2 = vadd_f32(vget_low_f32(sum_f32), vget_high_f32(sum_f32));
sum_2 = vpadd_f32(sum_2, sum_2);
// Widen every accumulator before combining them; no horizontal addition is performed in fp16.
float32x4_t sum_f32 =
vaddq_f32(vcvt_f32_f16(vget_low_f16(acc1)), vcvt_f32_f16(vget_high_f16(acc1)));
sum_f32 = vaddq_f32(sum_f32, vcvt_f32_f16(vget_low_f16(acc2)));
sum_f32 = vaddq_f32(sum_f32, vcvt_f32_f16(vget_high_f16(acc2)));
sum_f32 = vaddq_f32(sum_f32, vcvt_f32_f16(vget_low_f16(acc3)));
sum_f32 = vaddq_f32(sum_f32, vcvt_f32_f16(vget_high_f16(acc3)));
sum_f32 = vaddq_f32(sum_f32, vcvt_f32_f16(vget_low_f16(acc4)));
sum_f32 = vaddq_f32(sum_f32, vcvt_f32_f16(vget_high_f16(acc4)));

// Extract result
return 1.0f - vget_lane_f32(sum_2, 0);
return 1.0f - vaddvq_f32(sum_f32);
}
33 changes: 24 additions & 9 deletions src/VecSim/spaces/IP/IP_SVE_FP16.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@
*/
#include <arm_sve.h>

// Keep the throughput-critical loop in native fp16. The accumulators are widened separately before
// the final horizontal sum, avoiding fp16 overflow while combining otherwise-finite lanes.
inline void InnerProduct_Step(const float16_t *vec1, const float16_t *vec2, svfloat16_t &acc,
size_t &offset, const size_t chunk) {
svbool_t all = svptrue_b16();
Expand All @@ -22,12 +24,20 @@ inline void InnerProduct_Step(const float16_t *vec1, const float16_t *vec2, svfl
offset += chunk;
}

inline void WidenIPAccumulator(svfloat16_t value, svfloat32_t &lo, svfloat32_t &hi) {
// ZIP1/ZIP2 interleave the lower/upper contiguous halves with zero, placing each half value in
// the low 16 bits of a 32-bit lane for the widening conversion.
const svfloat16_t zero = svdup_f16(0.0f);
const svbool_t all32 = svptrue_b32();
lo = svcvt_f32_f16_x(all32, svzip1_f16(value, zero));
hi = svcvt_f32_f16_x(all32, svzip2_f16(value, zero));
}

template <bool partial_chunk, unsigned char additional_steps> // [t/f, 0..3]
float FP16_InnerProduct_SVE(const void *pVect1v, const void *pVect2v, size_t dimension) {
const auto *vec1 = static_cast<const float16_t *>(pVect1v);
const auto *vec2 = static_cast<const float16_t *>(pVect2v);
const size_t chunk = svcnth(); // number of 16-bit elements in a register
svbool_t all = svptrue_b16();
svfloat16_t acc1 = svdup_f16(0.0f);
svfloat16_t acc2 = svdup_f16(0.0f);
svfloat16_t acc3 = svdup_f16(0.0f);
Expand All @@ -54,7 +64,6 @@ float FP16_InnerProduct_SVE(const void *pVect1v, const void *pVect2v, size_t dim
// Handle the tail with the residual predicate
if constexpr (partial_chunk) {
svbool_t pg = svwhilelt_b16_u64(offset, dimension);

// Load half-precision vectors.
svfloat16_t v1 = svld1_f16(pg, vec1 + offset);
svfloat16_t v2 = svld1_f16(pg, vec2 + offset);
Expand All @@ -63,12 +72,18 @@ float FP16_InnerProduct_SVE(const void *pVect1v, const void *pVect2v, size_t dim
acc4 = svmla_f16_m(pg, acc4, v1, v2);
}

// Accumulate accumulators
acc1 = svadd_f16_x(all, acc1, acc3);
acc2 = svadd_f16_x(all, acc2, acc4);
acc1 = svadd_f16_x(all, acc1, acc2);
// Widen every accumulator before combining them; no horizontal addition is performed in fp16.
svfloat32_t acc1_lo, acc1_hi, acc2_lo, acc2_hi;
svfloat32_t acc3_lo, acc3_hi, acc4_lo, acc4_hi;
WidenIPAccumulator(acc1, acc1_lo, acc1_hi);
WidenIPAccumulator(acc2, acc2_lo, acc2_hi);
WidenIPAccumulator(acc3, acc3_lo, acc3_hi);
WidenIPAccumulator(acc4, acc4_lo, acc4_hi);

// Reduce the accumulated sum.
float result = svaddv_f16(all, acc1);
return 1.0f - result;
const svbool_t all32 = svptrue_b32();
svfloat32_t sum = svadd_f32_x(all32, svadd_f32_x(all32, acc1_lo, acc1_hi),
svadd_f32_x(all32, acc2_lo, acc2_hi));
sum = svadd_f32_x(all32, sum, svadd_f32_x(all32, acc3_lo, acc3_hi));
sum = svadd_f32_x(all32, sum, svadd_f32_x(all32, acc4_lo, acc4_hi));
return 1.0f - svaddv_f32(all32, sum);
}
15 changes: 11 additions & 4 deletions src/VecSim/spaces/L2/L2_AVX512FP16_VL_FP16.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,24 +6,33 @@
* (RSALv2); or (b) the Server Side Public License v1 (SSPLv1); or (c) the
* GNU Affero General Public License v3 (AGPLv3).
*/
#pragma once
#include <cstdint>
#include "VecSim/spaces/space_includes.h"
#include "VecSim/types/float16.h"
#include <cstring>

using float16 = vecsim_types::float16;

// Keep the throughput-critical loop in native fp16. Only the final horizontal reduction widens to
// fp32, so finite lane accumulators cannot overflow merely while being combined across lanes.
static inline void L2SqrStep(float16 *&pVect1, float16 *&pVect2, __m512h &sum) {
__m512h v1 = _mm512_loadu_ph(pVect1);
__m512h v2 = _mm512_loadu_ph(pVect2);

__m512h diff = _mm512_sub_ph(v1, v2);

sum = _mm512_fmadd_ph(diff, diff, sum);
pVect1 += 32;
pVect2 += 32;
}

static inline float ReduceL2FP16AccumulatorToFP32(__m512h sum) {
__m512 lo = _mm512_cvtxph_ps(_mm512_castph512_ph256(sum));
__m512 hi = _mm512_cvtxph_ps(
_mm256_castsi256_ph(_mm512_extracti64x4_epi64(_mm512_castph_si512(sum), 1)));
return _mm512_reduce_add_ps(_mm512_add_ps(lo, hi));
}

template <unsigned short residual> // 0..31
float FP16_L2SqrSIMD32_AVX512FP16_VL(const void *pVect1v, const void *pVect2v, size_t dimension) {
auto *pVect1 = (float16 *)pVect1v;
Expand All @@ -40,7 +49,6 @@ float FP16_L2SqrSIMD32_AVX512FP16_VL(const void *pVect1v, const void *pVect2v, s
__m512h v2 = _mm512_loadu_ph(pVect2);
pVect2 += residual;
__m512h diff = _mm512_maskz_sub_ph(mask, v1, v2);

sum = _mm512_mul_ph(diff, diff);
}

Expand All @@ -49,6 +57,5 @@ float FP16_L2SqrSIMD32_AVX512FP16_VL(const void *pVect1v, const void *pVect2v, s
L2SqrStep(pVect1, pVect2, sum);
} while (pVect1 < pEnd1);

_Float16 res = _mm512_reduce_add_ph(sum);
return res;
return ReduceL2FP16AccumulatorToFP32(sum);
}
25 changes: 12 additions & 13 deletions src/VecSim/spaces/L2/L2_NEON_FP16.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@
*/
#include <arm_neon.h>

// Keep the throughput-critical loop in native fp16. The accumulators are widened separately before
// the final horizontal sum, avoiding fp16 overflow while combining otherwise-finite lanes.
inline void L2Sqr_Step(const float16_t *&vec1, const float16_t *&vec2, float16x8_t &acc) {
// Load half-precision vectors
float16x8_t v1 = vld1q_f16(vec1);
Expand Down Expand Up @@ -81,19 +83,16 @@ float FP16_L2Sqr_NEON_HP(const void *pVect1v, const void *pVect2v, size_t dimens
L2Sqr_Step(vec1, vec2, acc4);
}

// Accumulate accumulators
acc1 = vpaddq_f16(acc1, acc3);
acc2 = vpaddq_f16(acc2, acc4);
acc1 = vpaddq_f16(acc1, acc2);

// Horizontal sum of the accumulated values
float32x4_t sum_f32 = vcvt_f32_f16(vget_low_f16(acc1));
sum_f32 = vaddq_f32(sum_f32, vcvt_f32_f16(vget_high_f16(acc1)));

// Pairwise add to get horizontal sum
float32x2_t sum_2 = vadd_f32(vget_low_f32(sum_f32), vget_high_f32(sum_f32));
sum_2 = vpadd_f32(sum_2, sum_2);
// Widen every accumulator before combining them; no horizontal addition is performed in fp16.
float32x4_t sum_f32 =
vaddq_f32(vcvt_f32_f16(vget_low_f16(acc1)), vcvt_f32_f16(vget_high_f16(acc1)));
sum_f32 = vaddq_f32(sum_f32, vcvt_f32_f16(vget_low_f16(acc2)));
sum_f32 = vaddq_f32(sum_f32, vcvt_f32_f16(vget_high_f16(acc2)));
sum_f32 = vaddq_f32(sum_f32, vcvt_f32_f16(vget_low_f16(acc3)));
sum_f32 = vaddq_f32(sum_f32, vcvt_f32_f16(vget_high_f16(acc3)));
sum_f32 = vaddq_f32(sum_f32, vcvt_f32_f16(vget_low_f16(acc4)));
sum_f32 = vaddq_f32(sum_f32, vcvt_f32_f16(vget_high_f16(acc4)));

// Extract result
return vget_lane_f32(sum_2, 0);
return vaddvq_f32(sum_f32);
}
32 changes: 24 additions & 8 deletions src/VecSim/spaces/L2/L2_SVE_FP16.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@
*/
#include <arm_sve.h>

// Keep the throughput-critical loop in native fp16. The accumulators are widened separately before
// the final horizontal sum, avoiding fp16 overflow while combining otherwise-finite lanes.
inline void L2Sqr_Step(const float16_t *vec1, const float16_t *vec2, svfloat16_t &acc,
size_t &offset, const size_t chunk) {
svbool_t all = svptrue_b16();
Expand All @@ -21,12 +23,20 @@ inline void L2Sqr_Step(const float16_t *vec1, const float16_t *vec2, svfloat16_t
offset += chunk;
}

inline void WidenL2Accumulator(svfloat16_t value, svfloat32_t &lo, svfloat32_t &hi) {
// ZIP1/ZIP2 interleave the lower/upper contiguous halves with zero, placing each half value in
// the low 16 bits of a 32-bit lane for the widening conversion.
const svfloat16_t zero = svdup_f16(0.0f);
const svbool_t all32 = svptrue_b32();
lo = svcvt_f32_f16_x(all32, svzip1_f16(value, zero));
hi = svcvt_f32_f16_x(all32, svzip2_f16(value, zero));
}

template <bool partial_chunk, unsigned char additional_steps> // [t/f, 0..3]
float FP16_L2Sqr_SVE(const void *pVect1v, const void *pVect2v, size_t dimension) {
const auto *vec1 = static_cast<const float16_t *>(pVect1v);
const auto *vec2 = static_cast<const float16_t *>(pVect2v);
const size_t chunk = svcnth(); // number of 16-bit elements in a register
svbool_t all = svptrue_b16();
svfloat16_t acc1 = svdup_f16(0.0f);
svfloat16_t acc2 = svdup_f16(0.0f);
svfloat16_t acc3 = svdup_f16(0.0f);
Expand Down Expand Up @@ -64,12 +74,18 @@ float FP16_L2Sqr_SVE(const void *pVect1v, const void *pVect2v, size_t dimension)
acc4 = svmla_f16_m(pg, acc4, diff, diff);
}

// Accumulate accumulators
acc1 = svadd_f16_x(all, acc1, acc3);
acc2 = svadd_f16_x(all, acc2, acc4);
acc1 = svadd_f16_x(all, acc1, acc2);
// Widen every accumulator before combining them; no horizontal addition is performed in fp16.
svfloat32_t acc1_lo, acc1_hi, acc2_lo, acc2_hi;
svfloat32_t acc3_lo, acc3_hi, acc4_lo, acc4_hi;
WidenL2Accumulator(acc1, acc1_lo, acc1_hi);
WidenL2Accumulator(acc2, acc2_lo, acc2_hi);
WidenL2Accumulator(acc3, acc3_lo, acc3_hi);
WidenL2Accumulator(acc4, acc4_lo, acc4_hi);

// Reduce the accumulated sum.
float result = svaddv_f16(all, acc1);
return result;
const svbool_t all32 = svptrue_b32();
svfloat32_t sum = svadd_f32_x(all32, svadd_f32_x(all32, acc1_lo, acc1_hi),
svadd_f32_x(all32, acc2_lo, acc2_hi));
sum = svadd_f32_x(all32, sum, svadd_f32_x(all32, acc3_lo, acc3_hi));
sum = svadd_f32_x(all32, sum, svadd_f32_x(all32, acc4_lo, acc4_hi));
return svaddv_f32(all32, sum);
}
58 changes: 48 additions & 10 deletions tests/unit/test_spaces.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
#include <random>
#include <cmath>
#include <limits>
#include <vector>

#include "gtest/gtest.h"
#include "VecSim/spaces/space_includes.h"
Expand Down Expand Up @@ -1427,22 +1428,19 @@ TEST_P(FP16SpacesOptimizationTestAdvanced, FP16InnerProductTestAdv) {
std::mt19937 gen(42);
std::uniform_real_distribution<> dis(-0.99, 0.99);

#if defined(CPU_FEATURES_ARCH_AARCH64) && defined(__GNUC__) && (__GNUC__ < 13)
// https://github.com/pytorch/executorch/issues/6844
__fp16 baseline = 0;
#else
_Float16 baseline = 0;
#endif
// Use the scalar fp32 contract as the reference. Native-fp16 tiers deliberately retain half
// precision in their hot loops for throughput, and the tolerance below bounds that trade-off.
float baseline = 0;

for (size_t i = 0; i < dim; i++) {
float val1 = (dis(gen));
float val2 = (dis(gen));
v1[i] = vecsim_types::FP32_to_FP16((val1));
v2[i] = vecsim_types::FP32_to_FP16((val2));

baseline += static_cast<decltype(baseline)>(val1) * static_cast<decltype(baseline)>(val2);
baseline += vecsim_types::FP16_to_FP32(v1[i]) * vecsim_types::FP16_to_FP32(v2[i]);
}
baseline = decltype(baseline)(1) - baseline;
baseline = 1.0f - baseline;

auto expected_alignment = [](size_t reg_bit_size, size_t dim) {
size_t elements_in_reg = reg_bit_size / sizeof(float16) / 8;
Expand Down Expand Up @@ -1532,14 +1530,15 @@ TEST_P(FP16SpacesOptimizationTestAdvanced, FP16L2SqrTestAdv) {
std::mt19937 gen(42);
std::uniform_real_distribution<float> dis(-0.99f, 0.99f);

_Float16 baseline = 0;
// Use the scalar fp32 contract as the reference; see the inner-product note above.
float baseline = 0;
for (size_t i = 0; i < dim; i++) {
float val1 = (dis(gen));
float val2 = (dis(gen));
v1[i] = vecsim_types::FP32_to_FP16((val1));
v2[i] = vecsim_types::FP32_to_FP16((val2));

_Float16 diff = static_cast<_Float16>(val1) - static_cast<_Float16>(val2);
float diff = vecsim_types::FP16_to_FP32(v1[i]) - vecsim_types::FP16_to_FP32(v2[i]);
baseline += diff * diff;
}

Expand Down Expand Up @@ -1570,6 +1569,45 @@ INSTANTIATE_TEST_SUITE_P(, FP16SpacesOptimizationTestAdvanced,

#endif // defined(OPT_AVX512_FP16_VL) || defined(CPU_FEATURES_ARCH_AARCH64)

// Regression test for the final reduction width of the native-fp16 SIMD kernels.
//
// The hot loop intentionally accumulates in fp16 for throughput, so an individual product or lane
// accumulator can still overflow for sufficiently large input values. These inputs isolate the
// narrower guarantee made here: every lane accumulator stays finite (at most 40,000), but reducing
// all lanes in fp16 would exceed 65,504 and return infinity. Widening only for the final horizontal
// sum preserves the native-fp16 loop while returning the finite total.
//
// The public choosers are called without a feature override, so whichever tier this CPU selects is
// the one under test, and the scalar path is covered on machines with no SIMD tier at all.
TEST(FP16SpacesTest, FinalReductionDoesNotOverflowFiniteAccumulators) {
// Exactly `dim` elements, no padding: the residual paths issue a full-width load at the front,
// mask it to the residual, then continue at that residual offset. At dim=35 the x86 native-half
// kernel masks elements 0..2 and then processes 3..34. Tight buffers keep ASan meaningful.
for (size_t dim : {32UL, 35UL, 40UL, 64UL, 128UL}) {
std::vector<float16> v1(dim, vecsim_types::FP32_to_FP16(0.0f));
std::vector<float16> v2(dim, vecsim_types::FP32_to_FP16(0.0f));
for (size_t i = 0; i < dim; i++) {
v1[i] = vecsim_types::FP32_to_FP16(100.0f);
}

// L2: v2 is all zeros, so the distance is dim * 100^2.
const float expected_l2 = static_cast<float>(dim) * 100.0f * 100.0f;
const float l2 = L2_FP16_GetDistFunc(dim)(v1.data(), v2.data(), dim);
ASSERT_TRUE(std::isfinite(l2)) << "L2 final reduction overflowed at dim " << dim;
ASSERT_NEAR(l2, expected_l2, expected_l2 * 1e-6f) << "L2 at dim " << dim;

// Inner product: both sides at 100, so the raw product sums to dim * 100^2 and the
// returned distance is 1 - that.
for (size_t i = 0; i < dim; i++) {
v2[i] = vecsim_types::FP32_to_FP16(100.0f);
}
const float expected_ip = 1.0f - static_cast<float>(dim) * 100.0f * 100.0f;
const float ip = IP_FP16_GetDistFunc(dim)(v1.data(), v2.data(), dim);
ASSERT_TRUE(std::isfinite(ip)) << "IP final reduction overflowed at dim " << dim;
ASSERT_NEAR(ip, expected_ip, std::abs(expected_ip) * 1e-6f) << "IP at dim " << dim;
}
}

class INT8SpacesOptimizationTest : public testing::TestWithParam<size_t> {};

TEST_P(INT8SpacesOptimizationTest, INT8L2SqrTest) {
Expand Down
Loading