From 930bc73b900ea390be0c1e46939fae299cfcbccb Mon Sep 17 00:00:00 2001 From: Dor Forer Date: Thu, 20 Aug 2026 17:46:57 +0300 Subject: [PATCH] [MOD-17916] Reduce native FP16 accumulators in FP32 --- src/VecSim/spaces/IP/IP_AVX512FP16_VL_FP16.h | 13 ++++- src/VecSim/spaces/IP/IP_NEON_FP16.h | 25 ++++----- src/VecSim/spaces/IP/IP_SVE_FP16.h | 33 ++++++++--- src/VecSim/spaces/L2/L2_AVX512FP16_VL_FP16.h | 15 +++-- src/VecSim/spaces/L2/L2_NEON_FP16.h | 25 ++++----- src/VecSim/spaces/L2/L2_SVE_FP16.h | 32 ++++++++--- tests/unit/test_spaces.cpp | 58 ++++++++++++++++---- 7 files changed, 142 insertions(+), 59 deletions(-) diff --git a/src/VecSim/spaces/IP/IP_AVX512FP16_VL_FP16.h b/src/VecSim/spaces/IP/IP_AVX512FP16_VL_FP16.h index 130ff2c7a..0d3c8816d 100644 --- a/src/VecSim/spaces/IP/IP_AVX512FP16_VL_FP16.h +++ b/src/VecSim/spaces/IP/IP_AVX512FP16_VL_FP16.h @@ -6,6 +6,7 @@ * (RSALv2); or (b) the Server Side Public License v1 (SSPLv1); or (c) the * GNU Affero General Public License v3 (AGPLv3). */ +#pragma once #include #include "VecSim/spaces/space_includes.h" #include "VecSim/types/float16.h" @@ -13,6 +14,8 @@ 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); @@ -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 // 0..31 float FP16_InnerProductSIMD32_AVX512FP16_VL(const void *pVect1v, const void *pVect2v, size_t dimension) { @@ -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); } diff --git a/src/VecSim/spaces/IP/IP_NEON_FP16.h b/src/VecSim/spaces/IP/IP_NEON_FP16.h index fd547c457..a0af33a90 100644 --- a/src/VecSim/spaces/IP/IP_NEON_FP16.h +++ b/src/VecSim/spaces/IP/IP_NEON_FP16.h @@ -8,6 +8,8 @@ */ #include +// 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); @@ -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); } diff --git a/src/VecSim/spaces/IP/IP_SVE_FP16.h b/src/VecSim/spaces/IP/IP_SVE_FP16.h index ac464977e..374e77f37 100644 --- a/src/VecSim/spaces/IP/IP_SVE_FP16.h +++ b/src/VecSim/spaces/IP/IP_SVE_FP16.h @@ -8,6 +8,8 @@ */ #include +// 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(); @@ -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 // [t/f, 0..3] float FP16_InnerProduct_SVE(const void *pVect1v, const void *pVect2v, size_t dimension) { const auto *vec1 = static_cast(pVect1v); const auto *vec2 = static_cast(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); @@ -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); @@ -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); } diff --git a/src/VecSim/spaces/L2/L2_AVX512FP16_VL_FP16.h b/src/VecSim/spaces/L2/L2_AVX512FP16_VL_FP16.h index 27e909a30..00fd002f3 100644 --- a/src/VecSim/spaces/L2/L2_AVX512FP16_VL_FP16.h +++ b/src/VecSim/spaces/L2/L2_AVX512FP16_VL_FP16.h @@ -6,6 +6,7 @@ * (RSALv2); or (b) the Server Side Public License v1 (SSPLv1); or (c) the * GNU Affero General Public License v3 (AGPLv3). */ +#pragma once #include #include "VecSim/spaces/space_includes.h" #include "VecSim/types/float16.h" @@ -13,17 +14,25 @@ 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 // 0..31 float FP16_L2SqrSIMD32_AVX512FP16_VL(const void *pVect1v, const void *pVect2v, size_t dimension) { auto *pVect1 = (float16 *)pVect1v; @@ -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); } @@ -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); } diff --git a/src/VecSim/spaces/L2/L2_NEON_FP16.h b/src/VecSim/spaces/L2/L2_NEON_FP16.h index e2786aa7a..dcb54e9a5 100644 --- a/src/VecSim/spaces/L2/L2_NEON_FP16.h +++ b/src/VecSim/spaces/L2/L2_NEON_FP16.h @@ -8,6 +8,8 @@ */ #include +// 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); @@ -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); } diff --git a/src/VecSim/spaces/L2/L2_SVE_FP16.h b/src/VecSim/spaces/L2/L2_SVE_FP16.h index 24b5ee2df..228cb50cf 100644 --- a/src/VecSim/spaces/L2/L2_SVE_FP16.h +++ b/src/VecSim/spaces/L2/L2_SVE_FP16.h @@ -8,6 +8,8 @@ */ #include +// 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(); @@ -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 // [t/f, 0..3] float FP16_L2Sqr_SVE(const void *pVect1v, const void *pVect2v, size_t dimension) { const auto *vec1 = static_cast(pVect1v); const auto *vec2 = static_cast(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); @@ -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); } diff --git a/tests/unit/test_spaces.cpp b/tests/unit/test_spaces.cpp index f9ffa2c3e..49599554f 100644 --- a/tests/unit/test_spaces.cpp +++ b/tests/unit/test_spaces.cpp @@ -14,6 +14,7 @@ #include #include #include +#include #include "gtest/gtest.h" #include "VecSim/spaces/space_includes.h" @@ -1427,12 +1428,9 @@ 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)); @@ -1440,9 +1438,9 @@ TEST_P(FP16SpacesOptimizationTestAdvanced, FP16InnerProductTestAdv) { v1[i] = vecsim_types::FP32_to_FP16((val1)); v2[i] = vecsim_types::FP32_to_FP16((val2)); - baseline += static_cast(val1) * static_cast(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; @@ -1532,14 +1530,15 @@ TEST_P(FP16SpacesOptimizationTestAdvanced, FP16L2SqrTestAdv) { std::mt19937 gen(42); std::uniform_real_distribution 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; } @@ -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 v1(dim, vecsim_types::FP32_to_FP16(0.0f)); + std::vector 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(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(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 {}; TEST_P(INT8SpacesOptimizationTest, INT8L2SqrTest) {