From ac4d89268f7ef47e7ce02291c7f6e60e4bcecf25 Mon Sep 17 00:00:00 2001 From: Dor Forer Date: Sun, 4 Oct 2026 10:47:31 +0300 Subject: [PATCH 1/8] MOD-19169 Add native FP16 arithmetic for SQ8 FP16 inner product --- cmake/x86_64InstructionFlags.cmake | 4 + src/VecSim/spaces/CMakeLists.txt | 6 + .../spaces/IP/IP_AVX512FP16_BW_VL_SQ8_FP16.h | 72 +++++ src/VecSim/spaces/IP_space.cpp | 19 +- .../spaces/functions/AVX512FP16_BW_VL.cpp | 29 ++ .../spaces/functions/AVX512FP16_BW_VL.h | 18 ++ tests/unit/test_spaces.cpp | 278 +++++++++++++++++- 7 files changed, 422 insertions(+), 4 deletions(-) create mode 100644 src/VecSim/spaces/IP/IP_AVX512FP16_BW_VL_SQ8_FP16.h create mode 100644 src/VecSim/spaces/functions/AVX512FP16_BW_VL.cpp create mode 100644 src/VecSim/spaces/functions/AVX512FP16_BW_VL.h diff --git a/cmake/x86_64InstructionFlags.cmake b/cmake/x86_64InstructionFlags.cmake index ff0e43e97..6eba175f7 100644 --- a/cmake/x86_64InstructionFlags.cmake +++ b/cmake/x86_64InstructionFlags.cmake @@ -61,6 +61,10 @@ if(CXX_AVX512VL AND CXX_AVX512FP16) add_compile_definitions(OPT_AVX512_FP16_VL) endif() +if(CXX_AVX512F AND CXX_AVX512FP16 AND CXX_AVX512BW AND CXX_AVX512VL) + add_compile_definitions(OPT_AVX512_FP16_BW_VL) +endif() + if(CXX_AVX512F) add_compile_definitions(OPT_AVX512F) endif() diff --git a/src/VecSim/spaces/CMakeLists.txt b/src/VecSim/spaces/CMakeLists.txt index eb8a15fdd..722964c23 100644 --- a/src/VecSim/spaces/CMakeLists.txt +++ b/src/VecSim/spaces/CMakeLists.txt @@ -32,6 +32,12 @@ if(CMAKE_SYSTEM_PROCESSOR MATCHES "(x86_64)|(AMD64|amd64)|(^i.86$)") list(APPEND OPTIMIZATIONS functions/AVX512FP16_VL.cpp) endif() + if(CXX_AVX512F AND CXX_AVX512FP16 AND CXX_AVX512BW AND CXX_AVX512VL) + message("Building with AVX512F, AVX512FP16, AVX512BW and AVX512VL") + set_source_files_properties(functions/AVX512FP16_BW_VL.cpp PROPERTIES COMPILE_FLAGS "-mavx512f -mavx512fp16 -mavx512bw -mavx512vl") + list(APPEND OPTIMIZATIONS functions/AVX512FP16_BW_VL.cpp) + endif() + if(CXX_AVX512BW AND CXX_AVX512VBMI2) message("Building with AVX512BW and AVX512VBMI2") set_source_files_properties(functions/AVX512BW_VBMI2.cpp PROPERTIES COMPILE_FLAGS "-mavx512bw -mavx512vbmi2") diff --git a/src/VecSim/spaces/IP/IP_AVX512FP16_BW_VL_SQ8_FP16.h b/src/VecSim/spaces/IP/IP_AVX512FP16_BW_VL_SQ8_FP16.h new file mode 100644 index 000000000..c3456fde6 --- /dev/null +++ b/src/VecSim/spaces/IP/IP_AVX512FP16_BW_VL_SQ8_FP16.h @@ -0,0 +1,72 @@ +/* + * Copyright (c) 2006-Present, Redis Ltd. + * All rights reserved. + * + * Licensed under your choice of the Redis Source Available License 2.0 + * (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/spaces/functions/AVX512F.h" +#include "VecSim/types/sq8.h" +#include "VecSim/types/float16.h" +#include "VecSim/utils/alignment.h" + +static inline __m512h SQ8_FP16_LoadCodes_AVX512FP16(const uint8_t *codes) { + const __m256i bytes = _mm256_loadu_si256(reinterpret_cast(codes)); + return _mm512_cvtepu16_ph(_mm512_cvtepu8_epi16(bytes)); +} + +static inline void SQ8_FP16_InnerProductStep_AVX512FP16(const uint8_t *&codes, + const vecsim_types::float16 *&query, + __m512h &sum) { + const __m512h values = SQ8_FP16_LoadCodes_AVX512FP16(codes); + const __m512h query_values = _mm512_loadu_ph(query); + sum = _mm512_fmadd_ph(values, query_values, sum); + codes += 32; + query += 32; +} + +// dim >= 32 keeps the residual's full loads within the vector payloads. +template +float SQ8_FP16_InnerProductSIMD32_AVX512FP16_BW_VL(const void *storage, const void *query_blob, + size_t dimension) { + using sq8 = vecsim_types::sq8; + using float16 = vecsim_types::float16; + // Directed rounding can saturate half overflow, hiding it from the finite check below. + if ((_mm_getcsr() & _MM_ROUND_MASK) != _MM_ROUND_NEAREST) + return spaces::Choose_SQ8_FP16_IP_implementation_AVX512F(dimension)(storage, query_blob, + dimension); + const auto *codes = static_cast(storage); + const auto *query = static_cast(query_blob); + const auto *end = codes + dimension; + __m512h sum = _mm512_setzero_ph(); + if constexpr (residual) { + constexpr __mmask32 mask = (1U << residual) - 1; + const __m512h values = SQ8_FP16_LoadCodes_AVX512FP16(codes); + const __m512h query_values = _mm512_loadu_ph(query); + sum = _mm512_maskz_mul_ph(mask, values, query_values); + codes += residual; + query += residual; + } + do { + SQ8_FP16_InnerProductStep_AVX512FP16(codes, query, sum); + } while (codes < end); + const _Float16 reduced = _mm512_reduce_add_ph(sum); + const float dot = static_cast(reduced); + // Nearest rounding leaves intermediate and reduction overflow nonfinite. + if (!std::isfinite(dot)) + return spaces::Choose_SQ8_FP16_IP_implementation_AVX512F(dimension)(storage, query_blob, + dimension); + + const auto *storage_meta = static_cast(storage) + dimension; + const float min_val = load_unaligned(storage_meta + sq8::MIN_VAL * sizeof(float)); + const float delta = load_unaligned(storage_meta + sq8::DELTA * sizeof(float)); + const auto *query_meta = + reinterpret_cast(static_cast(query_blob) + dimension); + const float query_sum = load_unaligned(query_meta + sq8::SUM_QUERY * sizeof(float)); + return 1.0f - (min_val * query_sum + delta * dot); +} diff --git a/src/VecSim/spaces/IP_space.cpp b/src/VecSim/spaces/IP_space.cpp index 7ae4da228..3bac383ea 100644 --- a/src/VecSim/spaces/IP_space.cpp +++ b/src/VecSim/spaces/IP_space.cpp @@ -17,6 +17,7 @@ #include "VecSim/spaces/functions/SSE.h" #include "VecSim/spaces/functions/AVX512BW_VBMI2.h" #include "VecSim/spaces/functions/AVX512FP16_VL.h" +#include "VecSim/spaces/functions/AVX512FP16_BW_VL.h" #include "VecSim/spaces/functions/AVX512BF16_VL.h" #include "VecSim/spaces/functions/AVX512F_BW_VL_VNNI.h" #include "VecSim/spaces/functions/AVX2.h" @@ -193,7 +194,15 @@ dist_func_t IP_SQ8_FP16_GetDistFunc(size_t dim, unsigned char *alignment, return ret_dist_func; } // Alignment hints below refer to the SQ8 (first) operand per the GetDistFunc contract. - // AVX-512 tier only needs AVX-512F (cvtph_ps is part of AVX-512F, no VNNI/BW/VL required). +#ifdef OPT_AVX512_FP16_BW_VL + if (dim >= 32 && dim <= spaces::FP16_MAX_UNIT_IP_SIMD_DIM && features.avx512f && + features.avx512_fp16 && features.avx512bw && features.avx512vl) { + if (dim % 32 == 0) + *alignment = 32 * sizeof(uint8_t); + return Choose_SQ8_FP16_IP_implementation_AVX512FP16_BW_VL(dim); + } +#endif + // The mixed-arithmetic tier needs only AVX-512F; cvtph_ps is part of that ISA. #ifdef OPT_AVX512F if (features.avx512f) { if (dim % 16 == 0) // SQ8 chunk = 16 bytes @@ -271,6 +280,14 @@ dist_func_t Cosine_SQ8_FP16_GetDistFunc(size_t dim, unsigned char *alignm if (dim < 16) { return ret_dist_func; } +#ifdef OPT_AVX512_FP16_BW_VL + if (dim >= 32 && dim <= spaces::FP16_MAX_UNIT_IP_SIMD_DIM && features.avx512f && + features.avx512_fp16 && features.avx512bw && features.avx512vl) { + if (dim % 32 == 0) + *alignment = 32 * sizeof(uint8_t); + return Choose_SQ8_FP16_Cosine_implementation_AVX512FP16_BW_VL(dim); + } +#endif #ifdef OPT_AVX512F if (features.avx512f) { if (dim % 16 == 0) diff --git a/src/VecSim/spaces/functions/AVX512FP16_BW_VL.cpp b/src/VecSim/spaces/functions/AVX512FP16_BW_VL.cpp new file mode 100644 index 000000000..99f17bf0f --- /dev/null +++ b/src/VecSim/spaces/functions/AVX512FP16_BW_VL.cpp @@ -0,0 +1,29 @@ +/* + * Copyright (c) 2006-Present, Redis Ltd. + * All rights reserved. + * + * Licensed under your choice of the Redis Source Available License 2.0 + * (RSALv2); or (b) the Server Side Public License v1 (SSPLv1); or (c) the + * GNU Affero General Public License v3 (AGPLv3). + */ +#include "AVX512FP16_BW_VL.h" + +#include "VecSim/spaces/IP/IP_AVX512FP16_BW_VL_SQ8_FP16.h" + +namespace spaces { + +#include "implementation_chooser.h" + +dist_func_t Choose_SQ8_FP16_IP_implementation_AVX512FP16_BW_VL(size_t dim) { + dist_func_t ret_dist_func; + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 32, SQ8_FP16_InnerProductSIMD32_AVX512FP16_BW_VL); + return ret_dist_func; +} + +dist_func_t Choose_SQ8_FP16_Cosine_implementation_AVX512FP16_BW_VL(size_t dim) { + return Choose_SQ8_FP16_IP_implementation_AVX512FP16_BW_VL(dim); +} + +#include "implementation_chooser_cleanup.h" + +} // namespace spaces diff --git a/src/VecSim/spaces/functions/AVX512FP16_BW_VL.h b/src/VecSim/spaces/functions/AVX512FP16_BW_VL.h new file mode 100644 index 000000000..a44cac3b5 --- /dev/null +++ b/src/VecSim/spaces/functions/AVX512FP16_BW_VL.h @@ -0,0 +1,18 @@ +/* + * Copyright (c) 2006-Present, Redis Ltd. + * All rights reserved. + * + * Licensed under your choice of the Redis Source Available License 2.0 + * (RSALv2); or (b) the Server Side Public License v1 (SSPLv1); or (c) the + * GNU Affero General Public License v3 (AGPLv3). + */ +#pragma once + +#include "VecSim/spaces/spaces.h" + +namespace spaces { + +dist_func_t Choose_SQ8_FP16_IP_implementation_AVX512FP16_BW_VL(size_t dim); +dist_func_t Choose_SQ8_FP16_Cosine_implementation_AVX512FP16_BW_VL(size_t dim); + +} // namespace spaces diff --git a/tests/unit/test_spaces.cpp b/tests/unit/test_spaces.cpp index 923bb7c1b..9ed24ff35 100644 --- a/tests/unit/test_spaces.cpp +++ b/tests/unit/test_spaces.cpp @@ -33,6 +33,7 @@ #include "VecSim/spaces/functions/AVX512BW_VBMI2.h" #include "VecSim/spaces/functions/AVX512BF16_VL.h" #include "VecSim/spaces/functions/AVX512FP16_VL.h" +#include "VecSim/spaces/functions/AVX512FP16_BW_VL.h" #include "VecSim/spaces/functions/AVX512F_BW_VL_VNNI.h" #include "VecSim/spaces/functions/AVX2.h" #include "VecSim/spaces/functions/AVX2_F16C.h" @@ -3604,9 +3605,51 @@ INSTANTIATE_TEST_SUITE_P(SQ8_FP16_NoOpt, SQ8_FP16_NoOptimizationSpacesTest, /* ======================== SQ8_FP16 SIMD optimisation tests ========================= */ -// Walks down the x86 ISA tiers (AVX-512 → AVX2+FMA → AVX2 → SSE4 → scalar) and asserts -// that {IP,Cosine,L2}_SQ8_FP16_GetDistFunc returns the expected Choose_* symbol and that -// its output matches the scalar baseline within 0.01. +#if defined(CPU_FEATURES_ARCH_X86_64) && defined(OPT_AVX512_FP16_BW_VL) +namespace { + +bool HasSQ8FP16NativeISA() { + const auto features = getCpuOptimizationFeatures(); + return features.avx512f && features.avx512bw && features.avx512vl && features.avx512_fp16; +} + +float SQ8FP16NativeReference(const uint8_t *storage, const float16 *query, size_t dim) { + std::vector codes(dim); + for (size_t i = 0; i < dim; ++i) + codes[i] = vecsim_types::FP32_to_FP16(storage[i]); + const float dense_distance = + Choose_FP16_IP_implementation_AVX512FP16_VL(dim)(codes.data(), query, dim); + if (!std::isfinite(dense_distance)) + return Choose_SQ8_FP16_IP_implementation_AVX512F(dim)(storage, query, dim); + const float min = load_unaligned(storage + dim + sq8::MIN_VAL * sizeof(float)); + const float delta = load_unaligned(storage + dim + sq8::DELTA * sizeof(float)); + const float sum = load_unaligned(reinterpret_cast(query + dim)); + return 1.0f - (min * sum + delta * (1.0f - dense_distance)); +} + +struct SQ8FP16NativeCase { + std::vector storage; + std::vector query; + + SQ8FP16NativeCase(size_t dim, float min = 0.0f, float delta = 1.0f) + : storage(dim + sq8::storage_metadata_count() * sizeof(float)), + query(dim + sizeof(float) / sizeof(float16)) { + std::memcpy(storage.data() + dim + sq8::MIN_VAL * sizeof(float), &min, sizeof(min)); + std::memcpy(storage.data() + dim + sq8::DELTA * sizeof(float), &delta, sizeof(delta)); + } + + void setQuerySum(size_t dim) { + float sum = 0.0f; + for (size_t i = 0; i < dim; ++i) + sum += vecsim_types::FP16_to_FP32(query[i]); + std::memcpy(query.data() + dim, &sum, sizeof(sum)); + } +}; + +} // namespace +#endif + +// Walks down the ISA tiers and checks the selected symbol and each tier's arithmetic. class SQ8_FP16_SpacesOptimizationTest : public testing::TestWithParam {}; TEST_P(SQ8_FP16_SpacesOptimizationTest, SQ8_FP16_L2SqrTest) { @@ -3752,6 +3795,15 @@ TEST_P(SQ8_FP16_SpacesOptimizationTest, SQ8_FP16_InnerProductTest) { dist_func_t arch_opt_func; float baseline = SQ8_FP16_InnerProduct(v2_compressed.data(), v1_query.data(), dim); +#ifdef OPT_AVX512_FP16_BW_VL + if (dim >= 32 && HasSQ8FP16NativeISA()) { + arch_opt_func = IP_SQ8_FP16_GetDistFunc(dim, nullptr, &optimization); + ASSERT_EQ(arch_opt_func, Choose_SQ8_FP16_IP_implementation_AVX512FP16_BW_VL(dim)); + EXPECT_FLOAT_EQ(arch_opt_func(v2_compressed.data(), v1_query.data(), dim), + SQ8FP16NativeReference(v2_compressed.data(), v1_query.data(), dim)); + optimization.avx512_fp16 = 0; + } +#endif #ifdef OPT_AVX512F if (optimization.avx512f) { unsigned char alignment = 0; @@ -3878,6 +3930,15 @@ TEST_P(SQ8_FP16_SpacesOptimizationTest, SQ8_FP16_CosineTest) { dist_func_t arch_opt_func; float baseline = SQ8_FP16_Cosine(v2_compressed.data(), v1_query.data(), dim); +#ifdef OPT_AVX512_FP16_BW_VL + if (dim >= 32 && HasSQ8FP16NativeISA()) { + arch_opt_func = Cosine_SQ8_FP16_GetDistFunc(dim, nullptr, &optimization); + ASSERT_EQ(arch_opt_func, Choose_SQ8_FP16_Cosine_implementation_AVX512FP16_BW_VL(dim)); + EXPECT_FLOAT_EQ(arch_opt_func(v2_compressed.data(), v1_query.data(), dim), + SQ8FP16NativeReference(v2_compressed.data(), v1_query.data(), dim)); + optimization.avx512_fp16 = 0; + } +#endif #ifdef OPT_AVX512F if (optimization.avx512f) { unsigned char alignment = 0; @@ -5408,3 +5469,214 @@ TEST_F(SpacesTest, SQ8_SQ8_L2_self_distance_is_near_zero) { << "dispatched kernel, dim " << dim; } } + +#if defined(CPU_FEATURES_ARCH_X86_64) && defined(OPT_AVX512_FP16_BW_VL) +TEST(SQ8FP16NativeIPTest, DispatchRequiresEveryFeatureAndDimension32) { + auto features = getCpuOptimizationFeatures(); + features.avx512f = features.avx512bw = features.avx512vl = features.avx512_fp16 = 1; + for (size_t dim : {32UL, 63UL, 64UL, 513UL}) { + SCOPED_TRACE(dim); + const auto native_ip = Choose_SQ8_FP16_IP_implementation_AVX512FP16_BW_VL(dim); + const auto native_cosine = Choose_SQ8_FP16_Cosine_implementation_AVX512FP16_BW_VL(dim); + EXPECT_EQ(IP_SQ8_FP16_GetDistFunc(dim, nullptr, &features), native_ip); + EXPECT_EQ(Cosine_SQ8_FP16_GetDistFunc(dim, nullptr, &features), native_cosine); + EXPECT_EQ(L2_SQ8_FP16_GetDistFunc(dim, nullptr, &features), + Choose_SQ8_FP16_L2_implementation_AVX512F(dim)); + for (int feature = 0; feature < 4; ++feature) { + auto without = features; + switch (feature) { + case 0: + without.avx512f = 0; + break; + case 1: + without.avx512bw = 0; + break; + case 2: + without.avx512vl = 0; + break; + case 3: + without.avx512_fp16 = 0; + break; + } + EXPECT_NE(IP_SQ8_FP16_GetDistFunc(dim, nullptr, &without), native_ip); + EXPECT_NE(Cosine_SQ8_FP16_GetDistFunc(dim, nullptr, &without), native_cosine); + } + } + const size_t limit = spaces::FP16_MAX_UNIT_IP_SIMD_DIM; + EXPECT_EQ(IP_SQ8_FP16_GetDistFunc(limit, nullptr, &features), + Choose_SQ8_FP16_IP_implementation_AVX512FP16_BW_VL(limit)); + EXPECT_EQ(Cosine_SQ8_FP16_GetDistFunc(limit, nullptr, &features), + Choose_SQ8_FP16_Cosine_implementation_AVX512FP16_BW_VL(limit)); + EXPECT_EQ(IP_SQ8_FP16_GetDistFunc(limit + 1, nullptr, &features), + Choose_SQ8_FP16_IP_implementation_AVX512F(limit + 1)); + EXPECT_EQ(Cosine_SQ8_FP16_GetDistFunc(limit + 1, nullptr, &features), + Choose_SQ8_FP16_Cosine_implementation_AVX512F(limit + 1)); + for (size_t dim = 1; dim < 32; ++dim) { + EXPECT_EQ(IP_SQ8_FP16_GetDistFunc(dim, nullptr, &features), + dim < 16 ? SQ8_FP16_InnerProduct + : Choose_SQ8_FP16_IP_implementation_AVX512F(dim)); + EXPECT_EQ(Cosine_SQ8_FP16_GetDistFunc(dim, nullptr, &features), + dim < 16 ? SQ8_FP16_Cosine : Choose_SQ8_FP16_Cosine_implementation_AVX512F(dim)); + } +} + +TEST(SQ8FP16NativeIPTest, DenseFP16ArithmeticAcrossResidualsAndBlocks) { + if (!HasSQ8FP16NativeISA()) + GTEST_SKIP() << "AVX512F/BW/VL/FP16 is unavailable"; + std::vector dimensions{32, 511, 512, 513, 768, 1024, 4096}; + for (size_t residual = 0; residual < 32; ++residual) + dimensions.push_back(64 + residual); + for (size_t dim : dimensions) { + SCOPED_TRACE(dim); + SQ8FP16NativeCase input(dim); + std::vector codes(dim); + for (size_t i = 0; i < dim; ++i) { + input.storage[i] = static_cast((i * 17 + 29) % 256); + codes[i] = vecsim_types::FP32_to_FP16(input.storage[i]); + input.query[i] = vecsim_types::FP32_to_FP16(float(i % 7 + 1) / 1024.0f); + } + input.setQuerySum(dim); + const auto original = input.query; + const float expected = + Choose_FP16_IP_implementation_AVX512FP16_VL(dim)(codes.data(), input.query.data(), dim); + ASSERT_TRUE(std::isfinite(expected)); + EXPECT_FLOAT_EQ(Choose_SQ8_FP16_IP_implementation_AVX512FP16_BW_VL(dim)( + input.storage.data(), input.query.data(), dim), + expected); + EXPECT_FLOAT_EQ(Choose_SQ8_FP16_Cosine_implementation_AVX512FP16_BW_VL(dim)( + input.storage.data(), input.query.data(), dim), + expected); + EXPECT_EQ( + std::memcmp(input.query.data(), original.data(), original.size() * sizeof(float16)), 0); + } +} + +TEST(SQ8FP16NativeIPTest, ScaleAndOffsetCorrectionRemainFP32) { + if (!HasSQ8FP16NativeISA()) + GTEST_SKIP() << "AVX512F/BW/VL/FP16 is unavailable"; + constexpr size_t dim = 65; + SQ8FP16NativeCase input(dim, -0.12345678f, 0.0012345678f); + for (size_t i = 0; i < dim; ++i) { + input.storage[i] = static_cast((i * 17 + 29) % 256); + input.query[i] = vecsim_types::FP32_to_FP16(float(i % 7 + 1) / 1024.0f); + } + input.setQuerySum(dim); + const float expected = SQ8FP16NativeReference(input.storage.data(), input.query.data(), dim); + ASSERT_NE(expected, vecsim_types::FP16_to_FP32(vecsim_types::FP32_to_FP16(expected))); + EXPECT_FLOAT_EQ(Choose_SQ8_FP16_IP_implementation_AVX512FP16_BW_VL(dim)( + input.storage.data(), input.query.data(), dim), + expected); + EXPECT_FLOAT_EQ(Choose_SQ8_FP16_Cosine_implementation_AVX512FP16_BW_VL(dim)( + input.storage.data(), input.query.data(), dim), + expected); +} + +TEST(SQ8FP16NativeIPTest, MinimumSmallAndLargestFiniteFP16Values) { + if (!HasSQ8FP16NativeISA()) + GTEST_SKIP() << "AVX512F/BW/VL/FP16 is unavailable"; + constexpr size_t dim = 32; + for (float value : + {0.0f, 0x1p-24f, -0x1p-24f, 37.0f * 0x1p-24f, 0x1p-14f, -0x1p-14f, 65504.0f, -65504.0f}) { + SCOPED_TRACE(value); + SQ8FP16NativeCase input(dim); + input.storage[0] = 1; + input.query[0] = vecsim_types::FP32_to_FP16(value); + input.setQuerySum(dim); + ASSERT_FLOAT_EQ(vecsim_types::FP16_to_FP32(input.query[0]), value); + // Compare distances directly: subtracting a distance near one can erase a tiny dot. + const float expected = 1.0f - value; + EXPECT_FLOAT_EQ(Choose_SQ8_FP16_IP_implementation_AVX512FP16_BW_VL(dim)( + input.storage.data(), input.query.data(), dim), + expected); + EXPECT_FLOAT_EQ(Choose_SQ8_FP16_Cosine_implementation_AVX512FP16_BW_VL(dim)( + input.storage.data(), input.query.data(), dim), + expected); + } +} + +TEST(SQ8FP16NativeIPTest, CenteredCancellationRetainsNativeCodeDotRounding) { + if (!HasSQ8FP16NativeISA()) + GTEST_SKIP() << "AVX512F/BW/VL/FP16 is unavailable"; + constexpr size_t dim = 64; + SQ8FP16NativeCase input(dim, -100.0f); + for (size_t i = 0; i < dim; ++i) { + input.storage[i] = i < 32 ? 100 : (i % 2 == 0 ? 0 : 255); + if (i < 32) + input.query[i] = vecsim_types::FP32_to_FP16(1.0009765625f); + } + input.setQuerySum(dim); + const float expected = SQ8FP16NativeReference(input.storage.data(), input.query.data(), dim); + // The code dot rounds in FP16 before the FP32 offset cancels it; decoded-vector IP is one. + ASSERT_FLOAT_EQ(expected, 0.125f); + EXPECT_FLOAT_EQ(Choose_SQ8_FP16_IP_implementation_AVX512FP16_BW_VL(dim)( + input.storage.data(), input.query.data(), dim), + expected); + EXPECT_FLOAT_EQ(Choose_SQ8_FP16_Cosine_implementation_AVX512FP16_BW_VL(dim)( + input.storage.data(), input.query.data(), dim), + expected); +} + +TEST(SQ8FP16NativeIPTest, ProductLaneAndReductionOverflowFallBack) { + if (!HasSQ8FP16NativeISA()) + GTEST_SKIP() << "AVX512F/BW/VL/FP16 is unavailable"; + struct OverflowCase { + size_t dim; + float value; + size_t negative_from; + }; + for (const auto [dim, value, negative_from] : {OverflowCase{32, 1000.0f, 32}, + {96, 200.0f, 64}, + {32, 32.0f, 32}, + {64, 1000.0f, 32}, + {32, 1000.0f, 16}}) { + SCOPED_TRACE(::testing::Message() + << "dim=" << dim << " value=" << value << " negative_from=" << negative_from); + SQ8FP16NativeCase input(dim, -0.125f, 1.0f / 256.0f); + std::vector codes(dim, vecsim_types::FP32_to_FP16(255.0f)); + for (size_t i = 0; i < dim; ++i) { + input.storage[i] = 255; + input.query[i] = vecsim_types::FP32_to_FP16(i < negative_from ? value : -value); + } + input.setQuerySum(dim); + EXPECT_FALSE(std::isfinite(Choose_FP16_IP_implementation_AVX512FP16_VL(dim)( + codes.data(), input.query.data(), dim))); + const float expected = Choose_SQ8_FP16_IP_implementation_AVX512F(dim)( + input.storage.data(), input.query.data(), dim); + ASSERT_TRUE(std::isfinite(expected)); + EXPECT_FLOAT_EQ(Choose_SQ8_FP16_IP_implementation_AVX512FP16_BW_VL(dim)( + input.storage.data(), input.query.data(), dim), + expected); + EXPECT_FLOAT_EQ(Choose_SQ8_FP16_Cosine_implementation_AVX512FP16_BW_VL(dim)( + input.storage.data(), input.query.data(), dim), + expected); + } +} + +TEST(SQ8FP16NativeIPTest, NonNearestMXCSRRoundModesFallBack) { + if (!HasSQ8FP16NativeISA()) + GTEST_SKIP() << "AVX512F/BW/VL/FP16 is unavailable"; + constexpr size_t dim = 97; + SQ8FP16NativeCase input(dim, -0.12345678f, 0.0012345678f); + for (size_t i = 0; i < dim; ++i) { + input.storage[i] = static_cast((i * 17 + 29) % 256); + input.query[i] = vecsim_types::FP32_to_FP16(1.0009765625f); + } + input.setQuerySum(dim); + struct RestoreMXCSR { + unsigned int value = _mm_getcsr(); + ~RestoreMXCSR() { _mm_setcsr(value); } + } restore; + const auto native = Choose_SQ8_FP16_IP_implementation_AVX512FP16_BW_VL(dim); + const auto cosine = Choose_SQ8_FP16_Cosine_implementation_AVX512FP16_BW_VL(dim); + const auto fallback = Choose_SQ8_FP16_IP_implementation_AVX512F(dim); + for (unsigned int rounding : {_MM_ROUND_DOWN, _MM_ROUND_UP, _MM_ROUND_TOWARD_ZERO}) { + _mm_setcsr((restore.value & ~_MM_ROUND_MASK) | rounding); + const float expected = fallback(input.storage.data(), input.query.data(), dim); + const float actual = native(input.storage.data(), input.query.data(), dim); + const float actual_cosine = cosine(input.storage.data(), input.query.data(), dim); + _mm_setcsr(restore.value); + EXPECT_FLOAT_EQ(actual, expected) << "rounding=" << rounding; + EXPECT_FLOAT_EQ(actual_cosine, expected) << "rounding=" << rounding; + } +} +#endif From ef7311a07a08a565e3769b3c61b2a659ae43cf0f Mon Sep 17 00:00:00 2001 From: Dor Forer Date: Sun, 4 Oct 2026 19:26:55 +0300 Subject: [PATCH 2/8] Use four SQ8 FP16 accumulators with FP32 final reduction --- .../spaces/IP/IP_AVX512FP16_BW_VL_SQ8_FP16.h | 47 +++- src/VecSim/spaces/IP_space.cpp | 4 +- .../spaces/functions/AVX512FP16_BW_VL.cpp | 7 +- tests/unit/test_spaces.cpp | 262 +++++++++++++----- 4 files changed, 243 insertions(+), 77 deletions(-) diff --git a/src/VecSim/spaces/IP/IP_AVX512FP16_BW_VL_SQ8_FP16.h b/src/VecSim/spaces/IP/IP_AVX512FP16_BW_VL_SQ8_FP16.h index c3456fde6..d979304df 100644 --- a/src/VecSim/spaces/IP/IP_AVX512FP16_BW_VL_SQ8_FP16.h +++ b/src/VecSim/spaces/IP/IP_AVX512FP16_BW_VL_SQ8_FP16.h @@ -31,7 +31,7 @@ static inline void SQ8_FP16_InnerProductStep_AVX512FP16(const uint8_t *&codes, } // dim >= 32 keeps the residual's full loads within the vector payloads. -template +template float SQ8_FP16_InnerProductSIMD32_AVX512FP16_BW_VL(const void *storage, const void *query_blob, size_t dimension) { using sq8 = vecsim_types::sq8; @@ -52,11 +52,38 @@ float SQ8_FP16_InnerProductSIMD32_AVX512FP16_BW_VL(const void *storage, const vo codes += residual; query += residual; } - do { - SQ8_FP16_InnerProductStep_AVX512FP16(codes, query, sum); - } while (codes < end); - const _Float16 reduced = _mm512_reduce_add_ph(sum); - const float dot = static_cast(reduced); + if constexpr (four_sums) { + __m512h sum1 = _mm512_setzero_ph(); + __m512h sum2 = _mm512_setzero_ph(); + __m512h sum3 = _mm512_setzero_ph(); + while (end - codes >= 128) { + SQ8_FP16_InnerProductStep_AVX512FP16(codes, query, sum); + SQ8_FP16_InnerProductStep_AVX512FP16(codes, query, sum1); + SQ8_FP16_InnerProductStep_AVX512FP16(codes, query, sum2); + SQ8_FP16_InnerProductStep_AVX512FP16(codes, query, sum3); + } + if (codes < end) + SQ8_FP16_InnerProductStep_AVX512FP16(codes, query, sum); + if (codes < end) + SQ8_FP16_InnerProductStep_AVX512FP16(codes, query, sum1); + if (codes < end) + SQ8_FP16_InnerProductStep_AVX512FP16(codes, query, sum2); + sum = _mm512_add_ph(_mm512_add_ph(sum, sum1), _mm512_add_ph(sum2, sum3)); + } else { + do { + SQ8_FP16_InnerProductStep_AVX512FP16(codes, query, sum); + } while (codes < end); + } + float dot; + if constexpr (four_sums) { + const __m512i bits = _mm512_castph_si512(sum); + const __m512 low = _mm512_cvtph_ps(_mm512_castsi512_si256(bits)); + const __m512 high = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(bits, 1)); + dot = _mm512_reduce_add_ps(_mm512_add_ps(low, high)); + } else { + const _Float16 reduced = _mm512_reduce_add_ph(sum); + dot = static_cast(reduced); + } // Nearest rounding leaves intermediate and reduction overflow nonfinite. if (!std::isfinite(dot)) return spaces::Choose_SQ8_FP16_IP_implementation_AVX512F(dimension)(storage, query_blob, @@ -70,3 +97,11 @@ float SQ8_FP16_InnerProductSIMD32_AVX512FP16_BW_VL(const void *storage, const vo const float query_sum = load_unaligned(query_meta + sq8::SUM_QUERY * sizeof(float)); return 1.0f - (min_val * query_sum + delta * dot); } + +template +float SQ8_FP16_InnerProductSIMD32_FourSums_AVX512FP16_BW_VL(const void *storage, + const void *query_blob, + size_t dimension) { + return SQ8_FP16_InnerProductSIMD32_AVX512FP16_BW_VL(storage, query_blob, + dimension); +} diff --git a/src/VecSim/spaces/IP_space.cpp b/src/VecSim/spaces/IP_space.cpp index 3bac383ea..36d55bcd5 100644 --- a/src/VecSim/spaces/IP_space.cpp +++ b/src/VecSim/spaces/IP_space.cpp @@ -195,7 +195,7 @@ dist_func_t IP_SQ8_FP16_GetDistFunc(size_t dim, unsigned char *alignment, } // Alignment hints below refer to the SQ8 (first) operand per the GetDistFunc contract. #ifdef OPT_AVX512_FP16_BW_VL - if (dim >= 32 && dim <= spaces::FP16_MAX_UNIT_IP_SIMD_DIM && features.avx512f && + if (dim >= 128 && dim <= spaces::FP16_MAX_UNIT_IP_SIMD_DIM && features.avx512f && features.avx512_fp16 && features.avx512bw && features.avx512vl) { if (dim % 32 == 0) *alignment = 32 * sizeof(uint8_t); @@ -281,7 +281,7 @@ dist_func_t Cosine_SQ8_FP16_GetDistFunc(size_t dim, unsigned char *alignm return ret_dist_func; } #ifdef OPT_AVX512_FP16_BW_VL - if (dim >= 32 && dim <= spaces::FP16_MAX_UNIT_IP_SIMD_DIM && features.avx512f && + if (dim >= 128 && dim <= spaces::FP16_MAX_UNIT_IP_SIMD_DIM && features.avx512f && features.avx512_fp16 && features.avx512bw && features.avx512vl) { if (dim % 32 == 0) *alignment = 32 * sizeof(uint8_t); diff --git a/src/VecSim/spaces/functions/AVX512FP16_BW_VL.cpp b/src/VecSim/spaces/functions/AVX512FP16_BW_VL.cpp index 99f17bf0f..d9cfa162c 100644 --- a/src/VecSim/spaces/functions/AVX512FP16_BW_VL.cpp +++ b/src/VecSim/spaces/functions/AVX512FP16_BW_VL.cpp @@ -16,7 +16,12 @@ namespace spaces { dist_func_t Choose_SQ8_FP16_IP_implementation_AVX512FP16_BW_VL(size_t dim) { dist_func_t ret_dist_func; - CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 32, SQ8_FP16_InnerProductSIMD32_AVX512FP16_BW_VL); + if (dim >= 256) { + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 32, + SQ8_FP16_InnerProductSIMD32_FourSums_AVX512FP16_BW_VL); + } else { + CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 32, SQ8_FP16_InnerProductSIMD32_AVX512FP16_BW_VL); + } return ret_dist_func; } diff --git a/tests/unit/test_spaces.cpp b/tests/unit/test_spaces.cpp index 9ed24ff35..c96f8be92 100644 --- a/tests/unit/test_spaces.cpp +++ b/tests/unit/test_spaces.cpp @@ -7,6 +7,7 @@ * GNU Affero General Public License v3 (AGPLv3). */ +#include #include #include #include @@ -3613,18 +3614,62 @@ bool HasSQ8FP16NativeISA() { return features.avx512f && features.avx512bw && features.avx512vl && features.avx512_fp16; } +// Round exact half-input arithmetic directly to binary16, including ties-to-even. +float SQ8FP16RoundHalf(double value) { + if (!std::isfinite(value) || value == 0.0) + return static_cast(value); + const double magnitude = std::abs(value); + if (magnitude >= 65520.0) + return std::copysign(std::numeric_limits::infinity(), value); + int exponent; + std::frexp(magnitude, &exponent); + const int step_exponent = std::max(exponent - 11, -24); + const double scaled = std::ldexp(magnitude, -step_exponent); + const double lower = std::floor(scaled); + const double fraction = scaled - lower; + const bool round_up = fraction > 0.5 || (fraction == 0.5 && (static_cast(lower) & 1)); + return std::copysign(static_cast(std::ldexp(lower + round_up, step_exponent)), value); +} + +constexpr size_t sq8_fp16_native_min_dim = 128; + +float SQ8FP16NativeDotReference(const uint8_t *storage, const float16 *query, size_t dim, + bool four_sums) { + std::array sums[4]{}; + const size_t residual = dim % 32; + for (size_t lane = 0; lane < residual; ++lane) + sums[0][lane] = + SQ8FP16RoundHalf(double(storage[lane]) * vecsim_types::FP16_to_FP32(query[lane])); + for (size_t block = 0, offset = residual; offset < dim; ++block, offset += 32) { + auto &sum = sums[four_sums ? block % 4 : 0]; + for (size_t lane = 0; lane < 32; ++lane) + sum[lane] = SQ8FP16RoundHalf(std::fma( + double(storage[offset + lane]), + double(vecsim_types::FP16_to_FP32(query[offset + lane])), double(sum[lane]))); + } + if (four_sums) + for (size_t lane = 0; lane < 32; ++lane) { + const float first = SQ8FP16RoundHalf(double(sums[0][lane]) + sums[1][lane]); + const float second = SQ8FP16RoundHalf(double(sums[2][lane]) + sums[3][lane]); + sums[0][lane] = SQ8FP16RoundHalf(double(first) + second); + } + // Both reduction trees combine opposite halves; the four-sum path widens first. + for (size_t stride = 16; stride; stride /= 2) + for (size_t lane = 0; lane < stride; ++lane) { + const double sum = double(sums[0][lane]) + sums[0][lane + stride]; + sums[0][lane] = four_sums ? static_cast(sum) : SQ8FP16RoundHalf(sum); + } + return sums[0][0]; +} + float SQ8FP16NativeReference(const uint8_t *storage, const float16 *query, size_t dim) { - std::vector codes(dim); - for (size_t i = 0; i < dim; ++i) - codes[i] = vecsim_types::FP32_to_FP16(storage[i]); - const float dense_distance = - Choose_FP16_IP_implementation_AVX512FP16_VL(dim)(codes.data(), query, dim); - if (!std::isfinite(dense_distance)) + const float dot = SQ8FP16NativeDotReference(storage, query, dim, dim >= 256); + if (!std::isfinite(dot)) return Choose_SQ8_FP16_IP_implementation_AVX512F(dim)(storage, query, dim); const float min = load_unaligned(storage + dim + sq8::MIN_VAL * sizeof(float)); const float delta = load_unaligned(storage + dim + sq8::DELTA * sizeof(float)); const float sum = load_unaligned(reinterpret_cast(query + dim)); - return 1.0f - (min * sum + delta * (1.0f - dense_distance)); + return 1.0f - (min * sum + delta * dot); } struct SQ8FP16NativeCase { @@ -3796,7 +3841,7 @@ TEST_P(SQ8_FP16_SpacesOptimizationTest, SQ8_FP16_InnerProductTest) { float baseline = SQ8_FP16_InnerProduct(v2_compressed.data(), v1_query.data(), dim); #ifdef OPT_AVX512_FP16_BW_VL - if (dim >= 32 && HasSQ8FP16NativeISA()) { + if (dim >= sq8_fp16_native_min_dim && HasSQ8FP16NativeISA()) { arch_opt_func = IP_SQ8_FP16_GetDistFunc(dim, nullptr, &optimization); ASSERT_EQ(arch_opt_func, Choose_SQ8_FP16_IP_implementation_AVX512FP16_BW_VL(dim)); EXPECT_FLOAT_EQ(arch_opt_func(v2_compressed.data(), v1_query.data(), dim), @@ -3931,7 +3976,7 @@ TEST_P(SQ8_FP16_SpacesOptimizationTest, SQ8_FP16_CosineTest) { float baseline = SQ8_FP16_Cosine(v2_compressed.data(), v1_query.data(), dim); #ifdef OPT_AVX512_FP16_BW_VL - if (dim >= 32 && HasSQ8FP16NativeISA()) { + if (dim >= sq8_fp16_native_min_dim && HasSQ8FP16NativeISA()) { arch_opt_func = Cosine_SQ8_FP16_GetDistFunc(dim, nullptr, &optimization); ASSERT_EQ(arch_opt_func, Choose_SQ8_FP16_Cosine_implementation_AVX512FP16_BW_VL(dim)); EXPECT_FLOAT_EQ(arch_opt_func(v2_compressed.data(), v1_query.data(), dim), @@ -5471,10 +5516,10 @@ TEST_F(SpacesTest, SQ8_SQ8_L2_self_distance_is_near_zero) { } #if defined(CPU_FEATURES_ARCH_X86_64) && defined(OPT_AVX512_FP16_BW_VL) -TEST(SQ8FP16NativeIPTest, DispatchRequiresEveryFeatureAndDimension32) { +TEST(SQ8FP16NativeIPTest, DispatchRequiresEveryFeatureAndDimension128) { auto features = getCpuOptimizationFeatures(); features.avx512f = features.avx512bw = features.avx512vl = features.avx512_fp16 = 1; - for (size_t dim : {32UL, 63UL, 64UL, 513UL}) { + for (size_t dim : {128UL, 255UL, 256UL, 513UL}) { SCOPED_TRACE(dim); const auto native_ip = Choose_SQ8_FP16_IP_implementation_AVX512FP16_BW_VL(dim); const auto native_cosine = Choose_SQ8_FP16_Cosine_implementation_AVX512FP16_BW_VL(dim); @@ -5511,7 +5556,7 @@ TEST(SQ8FP16NativeIPTest, DispatchRequiresEveryFeatureAndDimension32) { Choose_SQ8_FP16_IP_implementation_AVX512F(limit + 1)); EXPECT_EQ(Cosine_SQ8_FP16_GetDistFunc(limit + 1, nullptr, &features), Choose_SQ8_FP16_Cosine_implementation_AVX512F(limit + 1)); - for (size_t dim = 1; dim < 32; ++dim) { + for (size_t dim = 1; dim < sq8_fp16_native_min_dim; ++dim) { EXPECT_EQ(IP_SQ8_FP16_GetDistFunc(dim, nullptr, &features), dim < 16 ? SQ8_FP16_InnerProduct : Choose_SQ8_FP16_IP_implementation_AVX512F(dim)); @@ -5520,55 +5565,93 @@ TEST(SQ8FP16NativeIPTest, DispatchRequiresEveryFeatureAndDimension32) { } } -TEST(SQ8FP16NativeIPTest, DenseFP16ArithmeticAcrossResidualsAndBlocks) { +TEST(SQ8FP16NativeIPTest, HalfArithmeticAcrossResidualsAndBlocks) { if (!HasSQ8FP16NativeISA()) GTEST_SKIP() << "AVX512F/BW/VL/FP16 is unavailable"; - std::vector dimensions{32, 511, 512, 513, 768, 1024, 4096}; - for (size_t residual = 0; residual < 32; ++residual) - dimensions.push_back(64 + residual); + std::vector dimensions{255, 511, 512, 513, 768, 1024, 4096}; + for (size_t dim = 32; dim < 160; ++dim) + dimensions.push_back(dim); + for (size_t base : {256UL, 288UL, 320UL, 352UL, 384UL, 512UL, 768UL}) + for (size_t residual = 0; residual < 32; ++residual) + dimensions.push_back(base + residual); for (size_t dim : dimensions) { SCOPED_TRACE(dim); - SQ8FP16NativeCase input(dim); - std::vector codes(dim); + for (bool mixed_signs : {false, true}) { + SCOPED_TRACE(mixed_signs); + SQ8FP16NativeCase input(dim); + std::vector codes(dim); + for (size_t i = 0; i < dim; ++i) { + input.storage[i] = static_cast((i * 17 + 29) % 256); + codes[i] = vecsim_types::FP32_to_FP16(input.storage[i]); + const float sign = mixed_signs && i % 3 ? -1.0f : 1.0f; + input.query[i] = vecsim_types::FP32_to_FP16(sign * float(i % 7 + 1) / 1024.0f); + } + input.setQuerySum(dim); + const auto original = input.query; + const float one_sum = 1.0f - SQ8FP16NativeDotReference(input.storage.data(), + input.query.data(), dim, false); + EXPECT_FLOAT_EQ(Choose_FP16_IP_implementation_AVX512FP16_VL(dim)( + codes.data(), input.query.data(), dim), + one_sum); + const float expected = + SQ8FP16NativeReference(input.storage.data(), input.query.data(), dim); + ASSERT_TRUE(std::isfinite(expected)); + EXPECT_FLOAT_EQ(Choose_SQ8_FP16_IP_implementation_AVX512FP16_BW_VL(dim)( + input.storage.data(), input.query.data(), dim), + expected); + EXPECT_FLOAT_EQ(Choose_SQ8_FP16_Cosine_implementation_AVX512FP16_BW_VL(dim)( + input.storage.data(), input.query.data(), dim), + expected); + EXPECT_EQ( + std::memcmp(input.query.data(), original.data(), original.size() * sizeof(float16)), + 0); + } + } + constexpr size_t dim = 256; + SQ8FP16NativeCase input(dim); + for (size_t offset : {0UL, 32UL, 128UL}) + input.storage[offset] = 1; + input.query[0] = vecsim_types::FP32_to_FP16(2048.0f); + input.query[32] = vecsim_types::FP32_to_FP16(1.0f); + input.query[128] = vecsim_types::FP32_to_FP16(-2048.0f); + input.setQuerySum(dim); + ASSERT_FLOAT_EQ(SQ8FP16NativeDotReference(input.storage.data(), input.query.data(), dim, false), + 0.0f); + ASSERT_FLOAT_EQ(SQ8FP16NativeDotReference(input.storage.data(), input.query.data(), dim, true), + 1.0f); + const auto original = input.query; + const float expected = 0.0f; + EXPECT_FLOAT_EQ(Choose_SQ8_FP16_IP_implementation_AVX512FP16_BW_VL(dim)( + input.storage.data(), input.query.data(), dim), + expected); + EXPECT_FLOAT_EQ(Choose_SQ8_FP16_Cosine_implementation_AVX512FP16_BW_VL(dim)( + input.storage.data(), input.query.data(), dim), + expected); + EXPECT_EQ(std::memcmp(input.query.data(), original.data(), original.size() * sizeof(float16)), + 0); +} + +TEST(SQ8FP16NativeIPTest, ScaleAndOffsetCorrectionRemainFP32) { + if (!HasSQ8FP16NativeISA()) + GTEST_SKIP() << "AVX512F/BW/VL/FP16 is unavailable"; + for (size_t dim : {65UL, 128UL, 255UL, 256UL, 257UL, 512UL, 768UL}) { + SCOPED_TRACE(dim); + SQ8FP16NativeCase input(dim, -0.12345678f, 0.0012345678f); for (size_t i = 0; i < dim; ++i) { input.storage[i] = static_cast((i * 17 + 29) % 256); - codes[i] = vecsim_types::FP32_to_FP16(input.storage[i]); input.query[i] = vecsim_types::FP32_to_FP16(float(i % 7 + 1) / 1024.0f); } input.setQuerySum(dim); - const auto original = input.query; const float expected = - Choose_FP16_IP_implementation_AVX512FP16_VL(dim)(codes.data(), input.query.data(), dim); - ASSERT_TRUE(std::isfinite(expected)); + SQ8FP16NativeReference(input.storage.data(), input.query.data(), dim); + ASSERT_NE(expected, vecsim_types::FP16_to_FP32(vecsim_types::FP32_to_FP16(expected))); EXPECT_FLOAT_EQ(Choose_SQ8_FP16_IP_implementation_AVX512FP16_BW_VL(dim)( input.storage.data(), input.query.data(), dim), expected); EXPECT_FLOAT_EQ(Choose_SQ8_FP16_Cosine_implementation_AVX512FP16_BW_VL(dim)( input.storage.data(), input.query.data(), dim), expected); - EXPECT_EQ( - std::memcmp(input.query.data(), original.data(), original.size() * sizeof(float16)), 0); - } -} - -TEST(SQ8FP16NativeIPTest, ScaleAndOffsetCorrectionRemainFP32) { - if (!HasSQ8FP16NativeISA()) - GTEST_SKIP() << "AVX512F/BW/VL/FP16 is unavailable"; - constexpr size_t dim = 65; - SQ8FP16NativeCase input(dim, -0.12345678f, 0.0012345678f); - for (size_t i = 0; i < dim; ++i) { - input.storage[i] = static_cast((i * 17 + 29) % 256); - input.query[i] = vecsim_types::FP32_to_FP16(float(i % 7 + 1) / 1024.0f); } - input.setQuerySum(dim); - const float expected = SQ8FP16NativeReference(input.storage.data(), input.query.data(), dim); - ASSERT_NE(expected, vecsim_types::FP16_to_FP32(vecsim_types::FP32_to_FP16(expected))); - EXPECT_FLOAT_EQ(Choose_SQ8_FP16_IP_implementation_AVX512FP16_BW_VL(dim)( - input.storage.data(), input.query.data(), dim), - expected); - EXPECT_FLOAT_EQ(Choose_SQ8_FP16_Cosine_implementation_AVX512FP16_BW_VL(dim)( - input.storage.data(), input.query.data(), dim), - expected); } TEST(SQ8FP16NativeIPTest, MinimumSmallAndLargestFiniteFP16Values) { @@ -5616,32 +5699,14 @@ TEST(SQ8FP16NativeIPTest, CenteredCancellationRetainsNativeCodeDotRounding) { expected); } -TEST(SQ8FP16NativeIPTest, ProductLaneAndReductionOverflowFallBack) { +TEST(SQ8FP16NativeIPTest, ProductLaneMergeAndReductionOverflowHandling) { if (!HasSQ8FP16NativeISA()) GTEST_SKIP() << "AVX512F/BW/VL/FP16 is unavailable"; - struct OverflowCase { - size_t dim; - float value; - size_t negative_from; - }; - for (const auto [dim, value, negative_from] : {OverflowCase{32, 1000.0f, 32}, - {96, 200.0f, 64}, - {32, 32.0f, 32}, - {64, 1000.0f, 32}, - {32, 1000.0f, 16}}) { - SCOPED_TRACE(::testing::Message() - << "dim=" << dim << " value=" << value << " negative_from=" << negative_from); - SQ8FP16NativeCase input(dim, -0.125f, 1.0f / 256.0f); - std::vector codes(dim, vecsim_types::FP32_to_FP16(255.0f)); - for (size_t i = 0; i < dim; ++i) { - input.storage[i] = 255; - input.query[i] = vecsim_types::FP32_to_FP16(i < negative_from ? value : -value); - } + auto check = [&](SQ8FP16NativeCase &input, size_t dim) { input.setQuerySum(dim); - EXPECT_FALSE(std::isfinite(Choose_FP16_IP_implementation_AVX512FP16_VL(dim)( - codes.data(), input.query.data(), dim))); - const float expected = Choose_SQ8_FP16_IP_implementation_AVX512F(dim)( - input.storage.data(), input.query.data(), dim); + const auto original = input.query; + const float expected = + SQ8FP16NativeReference(input.storage.data(), input.query.data(), dim); ASSERT_TRUE(std::isfinite(expected)); EXPECT_FLOAT_EQ(Choose_SQ8_FP16_IP_implementation_AVX512FP16_BW_VL(dim)( input.storage.data(), input.query.data(), dim), @@ -5649,13 +5714,71 @@ TEST(SQ8FP16NativeIPTest, ProductLaneAndReductionOverflowFallBack) { EXPECT_FLOAT_EQ(Choose_SQ8_FP16_Cosine_implementation_AVX512FP16_BW_VL(dim)( input.storage.data(), input.query.data(), dim), expected); + EXPECT_EQ( + std::memcmp(input.query.data(), original.data(), original.size() * sizeof(float16)), 0); + }; + for (size_t dim : {32UL, 128UL, 256UL, 257UL}) { + SCOPED_TRACE(dim); + SQ8FP16NativeCase input(dim, -0.125f, 1.0f / 256.0f); + input.storage[0] = 255; + input.query[0] = vecsim_types::FP32_to_FP16(1000.0f); + ASSERT_FALSE(std::isfinite( + SQ8FP16NativeDotReference(input.storage.data(), input.query.data(), dim, dim >= 256))); + check(input, dim); + } + { + SCOPED_TRACE("half lane overflow followed by cancellation"); + constexpr size_t dim = 640; + SQ8FP16NativeCase input(dim, -0.125f, 1.0f / 256.0f); + for (size_t offset : {0UL, 128UL, 256UL}) { + input.storage[offset] = 255; + input.query[offset] = vecsim_types::FP32_to_FP16(offset == 256 ? -200.0f : 200.0f); + } + ASSERT_FALSE(std::isfinite( + SQ8FP16NativeDotReference(input.storage.data(), input.query.data(), dim, dim >= 256))); + check(input, dim); + } + { + SCOPED_TRACE("half merge overflow"); + constexpr size_t dim = 256; + SQ8FP16NativeCase input(dim, -0.125f, 1.0f / 256.0f); + for (size_t offset : {0UL, 32UL}) { + input.storage[offset] = 255; + input.query[offset] = vecsim_types::FP32_to_FP16(200.0f); + } + ASSERT_TRUE(std::isfinite(SQ8FP16RoundHalf(255.0 * 200.0))); + ASSERT_FALSE(std::isfinite( + SQ8FP16NativeDotReference(input.storage.data(), input.query.data(), dim, dim >= 256))); + check(input, dim); + } + { + SCOPED_TRACE("FP32 reduction retains a finite native result above half range"); + constexpr size_t dim = 256; + SQ8FP16NativeCase input(dim, -0.125f, 1.0f / 256.0f); + for (size_t block = 0; block < 8; ++block) + for (size_t lane : {0UL, 1UL}) { + input.storage[block * 32 + lane] = 255; + input.query[block * 32 + lane] = vecsim_types::FP32_to_FP16(32.03125f); + } + input.setQuerySum(dim); + ASSERT_FALSE(std::isfinite( + SQ8FP16NativeDotReference(input.storage.data(), input.query.data(), dim, false))); + ASSERT_FLOAT_EQ( + SQ8FP16NativeDotReference(input.storage.data(), input.query.data(), dim, true), + 130688.0f); + const float expected = + SQ8FP16NativeReference(input.storage.data(), input.query.data(), dim); + const float fallback = Choose_SQ8_FP16_IP_implementation_AVX512F(dim)( + input.storage.data(), input.query.data(), dim); + ASSERT_NE(expected, fallback); + check(input, dim); } } TEST(SQ8FP16NativeIPTest, NonNearestMXCSRRoundModesFallBack) { if (!HasSQ8FP16NativeISA()) GTEST_SKIP() << "AVX512F/BW/VL/FP16 is unavailable"; - constexpr size_t dim = 97; + constexpr size_t dim = 513; SQ8FP16NativeCase input(dim, -0.12345678f, 0.0012345678f); for (size_t i = 0; i < dim; ++i) { input.storage[i] = static_cast((i * 17 + 29) % 256); @@ -5669,6 +5792,7 @@ TEST(SQ8FP16NativeIPTest, NonNearestMXCSRRoundModesFallBack) { const auto native = Choose_SQ8_FP16_IP_implementation_AVX512FP16_BW_VL(dim); const auto cosine = Choose_SQ8_FP16_Cosine_implementation_AVX512FP16_BW_VL(dim); const auto fallback = Choose_SQ8_FP16_IP_implementation_AVX512F(dim); + const auto original = input.query; for (unsigned int rounding : {_MM_ROUND_DOWN, _MM_ROUND_UP, _MM_ROUND_TOWARD_ZERO}) { _mm_setcsr((restore.value & ~_MM_ROUND_MASK) | rounding); const float expected = fallback(input.storage.data(), input.query.data(), dim); @@ -5678,5 +5802,7 @@ TEST(SQ8FP16NativeIPTest, NonNearestMXCSRRoundModesFallBack) { EXPECT_FLOAT_EQ(actual, expected) << "rounding=" << rounding; EXPECT_FLOAT_EQ(actual_cosine, expected) << "rounding=" << rounding; } + EXPECT_EQ(std::memcmp(input.query.data(), original.data(), original.size() * sizeof(float16)), + 0); } #endif From 25968215c721c25847588e14e3d1503a0b794d38 Mon Sep 17 00:00:00 2001 From: Dor Forer Date: Mon, 5 Oct 2026 09:49:07 +0300 Subject: [PATCH 3/8] Name SQ8 FP16 dispatch thresholds and reduction policy --- .../spaces/IP/IP_AVX512FP16_BW_VL_SQ8_FP16.h | 6 ++-- src/VecSim/spaces/IP_space.cpp | 8 +++--- .../spaces/functions/AVX512FP16_BW_VL.cpp | 2 +- src/VecSim/spaces/spaces.h | 3 ++ tests/unit/test_spaces.cpp | 28 +++++++++++-------- 5 files changed, 27 insertions(+), 20 deletions(-) diff --git a/src/VecSim/spaces/IP/IP_AVX512FP16_BW_VL_SQ8_FP16.h b/src/VecSim/spaces/IP/IP_AVX512FP16_BW_VL_SQ8_FP16.h index d979304df..d52827d34 100644 --- a/src/VecSim/spaces/IP/IP_AVX512FP16_BW_VL_SQ8_FP16.h +++ b/src/VecSim/spaces/IP/IP_AVX512FP16_BW_VL_SQ8_FP16.h @@ -31,7 +31,7 @@ static inline void SQ8_FP16_InnerProductStep_AVX512FP16(const uint8_t *&codes, } // dim >= 32 keeps the residual's full loads within the vector payloads. -template +template float SQ8_FP16_InnerProductSIMD32_AVX512FP16_BW_VL(const void *storage, const void *query_blob, size_t dimension) { using sq8 = vecsim_types::sq8; @@ -52,7 +52,7 @@ float SQ8_FP16_InnerProductSIMD32_AVX512FP16_BW_VL(const void *storage, const vo codes += residual; query += residual; } - if constexpr (four_sums) { + if constexpr (four_sums_fp32_reduce) { __m512h sum1 = _mm512_setzero_ph(); __m512h sum2 = _mm512_setzero_ph(); __m512h sum3 = _mm512_setzero_ph(); @@ -75,7 +75,7 @@ float SQ8_FP16_InnerProductSIMD32_AVX512FP16_BW_VL(const void *storage, const vo } while (codes < end); } float dot; - if constexpr (four_sums) { + if constexpr (four_sums_fp32_reduce) { const __m512i bits = _mm512_castph_si512(sum); const __m512 low = _mm512_cvtph_ps(_mm512_castsi512_si256(bits)); const __m512 high = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(bits, 1)); diff --git a/src/VecSim/spaces/IP_space.cpp b/src/VecSim/spaces/IP_space.cpp index 36d55bcd5..5bffee8d1 100644 --- a/src/VecSim/spaces/IP_space.cpp +++ b/src/VecSim/spaces/IP_space.cpp @@ -195,8 +195,8 @@ dist_func_t IP_SQ8_FP16_GetDistFunc(size_t dim, unsigned char *alignment, } // Alignment hints below refer to the SQ8 (first) operand per the GetDistFunc contract. #ifdef OPT_AVX512_FP16_BW_VL - if (dim >= 128 && dim <= spaces::FP16_MAX_UNIT_IP_SIMD_DIM && features.avx512f && - features.avx512_fp16 && features.avx512bw && features.avx512vl) { + if (dim >= spaces::SQ8_FP16_NATIVE_MIN_DIM && dim <= spaces::FP16_MAX_UNIT_IP_SIMD_DIM && + features.avx512f && features.avx512_fp16 && features.avx512bw && features.avx512vl) { if (dim % 32 == 0) *alignment = 32 * sizeof(uint8_t); return Choose_SQ8_FP16_IP_implementation_AVX512FP16_BW_VL(dim); @@ -281,8 +281,8 @@ dist_func_t Cosine_SQ8_FP16_GetDistFunc(size_t dim, unsigned char *alignm return ret_dist_func; } #ifdef OPT_AVX512_FP16_BW_VL - if (dim >= 128 && dim <= spaces::FP16_MAX_UNIT_IP_SIMD_DIM && features.avx512f && - features.avx512_fp16 && features.avx512bw && features.avx512vl) { + if (dim >= spaces::SQ8_FP16_NATIVE_MIN_DIM && dim <= spaces::FP16_MAX_UNIT_IP_SIMD_DIM && + features.avx512f && features.avx512_fp16 && features.avx512bw && features.avx512vl) { if (dim % 32 == 0) *alignment = 32 * sizeof(uint8_t); return Choose_SQ8_FP16_Cosine_implementation_AVX512FP16_BW_VL(dim); diff --git a/src/VecSim/spaces/functions/AVX512FP16_BW_VL.cpp b/src/VecSim/spaces/functions/AVX512FP16_BW_VL.cpp index d9cfa162c..c4caef596 100644 --- a/src/VecSim/spaces/functions/AVX512FP16_BW_VL.cpp +++ b/src/VecSim/spaces/functions/AVX512FP16_BW_VL.cpp @@ -16,7 +16,7 @@ namespace spaces { dist_func_t Choose_SQ8_FP16_IP_implementation_AVX512FP16_BW_VL(size_t dim) { dist_func_t ret_dist_func; - if (dim >= 256) { + if (dim >= spaces::SQ8_FP16_FOUR_SUMS_MIN_DIM) { CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 32, SQ8_FP16_InnerProductSIMD32_FourSums_AVX512FP16_BW_VL); } else { diff --git a/src/VecSim/spaces/spaces.h b/src/VecSim/spaces/spaces.h index f6de7823e..146b82f23 100644 --- a/src/VecSim/spaces/spaces.h +++ b/src/VecSim/spaces/spaces.h @@ -75,6 +75,9 @@ static constexpr size_t UINT8_MAX_EXACT_SIMD_DIM = static constexpr size_t FP16_MAX_UNIT_IP_SIMD_DIM = 65504; static constexpr size_t FP16_MAX_UNIT_L2_SIMD_DIM = FP16_MAX_UNIT_IP_SIMD_DIM / 4; +static constexpr size_t SQ8_FP16_NATIVE_MIN_DIM = 128; +static constexpr size_t SQ8_FP16_FOUR_SUMS_MIN_DIM = 256; + static inline auto getCpuOptimizationFeatures(const void *arch_opt = nullptr) { #if defined(CPU_FEATURES_ARCH_AARCH64) diff --git a/tests/unit/test_spaces.cpp b/tests/unit/test_spaces.cpp index c96f8be92..91dfdf46d 100644 --- a/tests/unit/test_spaces.cpp +++ b/tests/unit/test_spaces.cpp @@ -3631,8 +3631,6 @@ float SQ8FP16RoundHalf(double value) { return std::copysign(static_cast(std::ldexp(lower + round_up, step_exponent)), value); } -constexpr size_t sq8_fp16_native_min_dim = 128; - float SQ8FP16NativeDotReference(const uint8_t *storage, const float16 *query, size_t dim, bool four_sums) { std::array sums[4]{}; @@ -3663,7 +3661,8 @@ float SQ8FP16NativeDotReference(const uint8_t *storage, const float16 *query, si } float SQ8FP16NativeReference(const uint8_t *storage, const float16 *query, size_t dim) { - const float dot = SQ8FP16NativeDotReference(storage, query, dim, dim >= 256); + const float dot = + SQ8FP16NativeDotReference(storage, query, dim, dim >= spaces::SQ8_FP16_FOUR_SUMS_MIN_DIM); if (!std::isfinite(dot)) return Choose_SQ8_FP16_IP_implementation_AVX512F(dim)(storage, query, dim); const float min = load_unaligned(storage + dim + sq8::MIN_VAL * sizeof(float)); @@ -3841,7 +3840,7 @@ TEST_P(SQ8_FP16_SpacesOptimizationTest, SQ8_FP16_InnerProductTest) { float baseline = SQ8_FP16_InnerProduct(v2_compressed.data(), v1_query.data(), dim); #ifdef OPT_AVX512_FP16_BW_VL - if (dim >= sq8_fp16_native_min_dim && HasSQ8FP16NativeISA()) { + if (dim >= spaces::SQ8_FP16_NATIVE_MIN_DIM && HasSQ8FP16NativeISA()) { arch_opt_func = IP_SQ8_FP16_GetDistFunc(dim, nullptr, &optimization); ASSERT_EQ(arch_opt_func, Choose_SQ8_FP16_IP_implementation_AVX512FP16_BW_VL(dim)); EXPECT_FLOAT_EQ(arch_opt_func(v2_compressed.data(), v1_query.data(), dim), @@ -3976,7 +3975,7 @@ TEST_P(SQ8_FP16_SpacesOptimizationTest, SQ8_FP16_CosineTest) { float baseline = SQ8_FP16_Cosine(v2_compressed.data(), v1_query.data(), dim); #ifdef OPT_AVX512_FP16_BW_VL - if (dim >= sq8_fp16_native_min_dim && HasSQ8FP16NativeISA()) { + if (dim >= spaces::SQ8_FP16_NATIVE_MIN_DIM && HasSQ8FP16NativeISA()) { arch_opt_func = Cosine_SQ8_FP16_GetDistFunc(dim, nullptr, &optimization); ASSERT_EQ(arch_opt_func, Choose_SQ8_FP16_Cosine_implementation_AVX512FP16_BW_VL(dim)); EXPECT_FLOAT_EQ(arch_opt_func(v2_compressed.data(), v1_query.data(), dim), @@ -5517,6 +5516,8 @@ TEST_F(SpacesTest, SQ8_SQ8_L2_self_distance_is_near_zero) { #if defined(CPU_FEATURES_ARCH_X86_64) && defined(OPT_AVX512_FP16_BW_VL) TEST(SQ8FP16NativeIPTest, DispatchRequiresEveryFeatureAndDimension128) { + EXPECT_EQ(spaces::SQ8_FP16_NATIVE_MIN_DIM, 128UL); + EXPECT_EQ(spaces::SQ8_FP16_FOUR_SUMS_MIN_DIM, 256UL); auto features = getCpuOptimizationFeatures(); features.avx512f = features.avx512bw = features.avx512vl = features.avx512_fp16 = 1; for (size_t dim : {128UL, 255UL, 256UL, 513UL}) { @@ -5556,7 +5557,7 @@ TEST(SQ8FP16NativeIPTest, DispatchRequiresEveryFeatureAndDimension128) { Choose_SQ8_FP16_IP_implementation_AVX512F(limit + 1)); EXPECT_EQ(Cosine_SQ8_FP16_GetDistFunc(limit + 1, nullptr, &features), Choose_SQ8_FP16_Cosine_implementation_AVX512F(limit + 1)); - for (size_t dim = 1; dim < sq8_fp16_native_min_dim; ++dim) { + for (size_t dim = 1; dim < spaces::SQ8_FP16_NATIVE_MIN_DIM; ++dim) { EXPECT_EQ(IP_SQ8_FP16_GetDistFunc(dim, nullptr, &features), dim < 16 ? SQ8_FP16_InnerProduct : Choose_SQ8_FP16_IP_implementation_AVX512F(dim)); @@ -5722,8 +5723,9 @@ TEST(SQ8FP16NativeIPTest, ProductLaneMergeAndReductionOverflowHandling) { SQ8FP16NativeCase input(dim, -0.125f, 1.0f / 256.0f); input.storage[0] = 255; input.query[0] = vecsim_types::FP32_to_FP16(1000.0f); - ASSERT_FALSE(std::isfinite( - SQ8FP16NativeDotReference(input.storage.data(), input.query.data(), dim, dim >= 256))); + ASSERT_FALSE( + std::isfinite(SQ8FP16NativeDotReference(input.storage.data(), input.query.data(), dim, + dim >= spaces::SQ8_FP16_FOUR_SUMS_MIN_DIM))); check(input, dim); } { @@ -5734,8 +5736,9 @@ TEST(SQ8FP16NativeIPTest, ProductLaneMergeAndReductionOverflowHandling) { input.storage[offset] = 255; input.query[offset] = vecsim_types::FP32_to_FP16(offset == 256 ? -200.0f : 200.0f); } - ASSERT_FALSE(std::isfinite( - SQ8FP16NativeDotReference(input.storage.data(), input.query.data(), dim, dim >= 256))); + ASSERT_FALSE( + std::isfinite(SQ8FP16NativeDotReference(input.storage.data(), input.query.data(), dim, + dim >= spaces::SQ8_FP16_FOUR_SUMS_MIN_DIM))); check(input, dim); } { @@ -5747,8 +5750,9 @@ TEST(SQ8FP16NativeIPTest, ProductLaneMergeAndReductionOverflowHandling) { input.query[offset] = vecsim_types::FP32_to_FP16(200.0f); } ASSERT_TRUE(std::isfinite(SQ8FP16RoundHalf(255.0 * 200.0))); - ASSERT_FALSE(std::isfinite( - SQ8FP16NativeDotReference(input.storage.data(), input.query.data(), dim, dim >= 256))); + ASSERT_FALSE( + std::isfinite(SQ8FP16NativeDotReference(input.storage.data(), input.query.data(), dim, + dim >= spaces::SQ8_FP16_FOUR_SUMS_MIN_DIM))); check(input, dim); } { From 48f05a3a0257507274729a6e10c37dd612a8398a Mon Sep 17 00:00:00 2001 From: Dor Forer Date: Mon, 5 Oct 2026 12:25:37 +0300 Subject: [PATCH 4/8] Rename SQ8 FP16 helper to LoadQuantizedValues --- src/VecSim/spaces/IP/IP_AVX512FP16_BW_VL_SQ8_FP16.h | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/VecSim/spaces/IP/IP_AVX512FP16_BW_VL_SQ8_FP16.h b/src/VecSim/spaces/IP/IP_AVX512FP16_BW_VL_SQ8_FP16.h index d52827d34..4d8e88060 100644 --- a/src/VecSim/spaces/IP/IP_AVX512FP16_BW_VL_SQ8_FP16.h +++ b/src/VecSim/spaces/IP/IP_AVX512FP16_BW_VL_SQ8_FP16.h @@ -15,7 +15,7 @@ #include "VecSim/types/float16.h" #include "VecSim/utils/alignment.h" -static inline __m512h SQ8_FP16_LoadCodes_AVX512FP16(const uint8_t *codes) { +static inline __m512h SQ8_FP16_LoadQuantizedValues_AVX512FP16(const uint8_t *codes) { const __m256i bytes = _mm256_loadu_si256(reinterpret_cast(codes)); return _mm512_cvtepu16_ph(_mm512_cvtepu8_epi16(bytes)); } @@ -23,7 +23,7 @@ static inline __m512h SQ8_FP16_LoadCodes_AVX512FP16(const uint8_t *codes) { static inline void SQ8_FP16_InnerProductStep_AVX512FP16(const uint8_t *&codes, const vecsim_types::float16 *&query, __m512h &sum) { - const __m512h values = SQ8_FP16_LoadCodes_AVX512FP16(codes); + const __m512h values = SQ8_FP16_LoadQuantizedValues_AVX512FP16(codes); const __m512h query_values = _mm512_loadu_ph(query); sum = _mm512_fmadd_ph(values, query_values, sum); codes += 32; @@ -46,7 +46,7 @@ float SQ8_FP16_InnerProductSIMD32_AVX512FP16_BW_VL(const void *storage, const vo __m512h sum = _mm512_setzero_ph(); if constexpr (residual) { constexpr __mmask32 mask = (1U << residual) - 1; - const __m512h values = SQ8_FP16_LoadCodes_AVX512FP16(codes); + const __m512h values = SQ8_FP16_LoadQuantizedValues_AVX512FP16(codes); const __m512h query_values = _mm512_loadu_ph(query); sum = _mm512_maskz_mul_ph(mask, values, query_values); codes += residual; From 490a38476ffb9b37dfa5d83d06ce9ac1abd2727c Mon Sep 17 00:00:00 2001 From: Dor Forer Date: Mon, 5 Oct 2026 12:37:20 +0300 Subject: [PATCH 5/8] Split native SQ8 FP16 inner product into focused helpers --- .../spaces/IP/IP_AVX512FP16_BW_VL_SQ8_FP16.h | 49 +++++++++++++------ 1 file changed, 34 insertions(+), 15 deletions(-) diff --git a/src/VecSim/spaces/IP/IP_AVX512FP16_BW_VL_SQ8_FP16.h b/src/VecSim/spaces/IP/IP_AVX512FP16_BW_VL_SQ8_FP16.h index 4d8e88060..f264a0418 100644 --- a/src/VecSim/spaces/IP/IP_AVX512FP16_BW_VL_SQ8_FP16.h +++ b/src/VecSim/spaces/IP/IP_AVX512FP16_BW_VL_SQ8_FP16.h @@ -31,15 +31,11 @@ static inline void SQ8_FP16_InnerProductStep_AVX512FP16(const uint8_t *&codes, } // dim >= 32 keeps the residual's full loads within the vector payloads. -template -float SQ8_FP16_InnerProductSIMD32_AVX512FP16_BW_VL(const void *storage, const void *query_blob, - size_t dimension) { - using sq8 = vecsim_types::sq8; +template +static inline __m512h SQ8_FP16_AccumulateInnerProduct_AVX512FP16(const void *storage, + const void *query_blob, + size_t dimension) { using float16 = vecsim_types::float16; - // Directed rounding can saturate half overflow, hiding it from the finite check below. - if ((_mm_getcsr() & _MM_ROUND_MASK) != _MM_ROUND_NEAREST) - return spaces::Choose_SQ8_FP16_IP_implementation_AVX512F(dimension)(storage, query_blob, - dimension); const auto *codes = static_cast(storage); const auto *query = static_cast(query_blob); const auto *end = codes + dimension; @@ -74,21 +70,27 @@ float SQ8_FP16_InnerProductSIMD32_AVX512FP16_BW_VL(const void *storage, const vo SQ8_FP16_InnerProductStep_AVX512FP16(codes, query, sum); } while (codes < end); } - float dot; + return sum; +} + +template +static inline float SQ8_FP16_ReduceInnerProduct_AVX512FP16(__m512h sum) { if constexpr (four_sums_fp32_reduce) { const __m512i bits = _mm512_castph_si512(sum); const __m512 low = _mm512_cvtph_ps(_mm512_castsi512_si256(bits)); const __m512 high = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(bits, 1)); - dot = _mm512_reduce_add_ps(_mm512_add_ps(low, high)); + return _mm512_reduce_add_ps(_mm512_add_ps(low, high)); } else { const _Float16 reduced = _mm512_reduce_add_ph(sum); - dot = static_cast(reduced); + return static_cast(reduced); } - // Nearest rounding leaves intermediate and reduction overflow nonfinite. - if (!std::isfinite(dot)) - return spaces::Choose_SQ8_FP16_IP_implementation_AVX512F(dimension)(storage, query_blob, - dimension); +} +static inline float SQ8_FP16_ApplyInnerProductCorrection_AVX512FP16(const void *storage, + const void *query_blob, + size_t dimension, float dot) { + using sq8 = vecsim_types::sq8; + using float16 = vecsim_types::float16; const auto *storage_meta = static_cast(storage) + dimension; const float min_val = load_unaligned(storage_meta + sq8::MIN_VAL * sizeof(float)); const float delta = load_unaligned(storage_meta + sq8::DELTA * sizeof(float)); @@ -98,6 +100,23 @@ float SQ8_FP16_InnerProductSIMD32_AVX512FP16_BW_VL(const void *storage, const vo return 1.0f - (min_val * query_sum + delta * dot); } +template +float SQ8_FP16_InnerProductSIMD32_AVX512FP16_BW_VL(const void *storage, const void *query_blob, + size_t dimension) { + // Directed rounding can saturate half overflow, hiding it from the finite check below. + if ((_mm_getcsr() & _MM_ROUND_MASK) != _MM_ROUND_NEAREST) + return spaces::Choose_SQ8_FP16_IP_implementation_AVX512F(dimension)(storage, query_blob, + dimension); + const __m512h sum = SQ8_FP16_AccumulateInnerProduct_AVX512FP16( + storage, query_blob, dimension); + const float dot = SQ8_FP16_ReduceInnerProduct_AVX512FP16(sum); + // Nearest rounding leaves intermediate and reduction overflow nonfinite. + if (!std::isfinite(dot)) + return spaces::Choose_SQ8_FP16_IP_implementation_AVX512F(dimension)(storage, query_blob, + dimension); + return SQ8_FP16_ApplyInnerProductCorrection_AVX512FP16(storage, query_blob, dimension, dot); +} + template float SQ8_FP16_InnerProductSIMD32_FourSums_AVX512FP16_BW_VL(const void *storage, const void *query_blob, From 7897ea8d98537a257e85d8f43838d2559db747c0 Mon Sep 17 00:00:00 2001 From: Dor Forer Date: Mon, 5 Oct 2026 12:52:27 +0300 Subject: [PATCH 6/8] Rename SQ8 byte pointers to quantized_values --- .../spaces/IP/IP_AVX512FP16_BW_VL_SQ8_FP16.h | 44 +++++++++---------- 1 file changed, 22 insertions(+), 22 deletions(-) diff --git a/src/VecSim/spaces/IP/IP_AVX512FP16_BW_VL_SQ8_FP16.h b/src/VecSim/spaces/IP/IP_AVX512FP16_BW_VL_SQ8_FP16.h index f264a0418..6d1ef747f 100644 --- a/src/VecSim/spaces/IP/IP_AVX512FP16_BW_VL_SQ8_FP16.h +++ b/src/VecSim/spaces/IP/IP_AVX512FP16_BW_VL_SQ8_FP16.h @@ -15,18 +15,18 @@ #include "VecSim/types/float16.h" #include "VecSim/utils/alignment.h" -static inline __m512h SQ8_FP16_LoadQuantizedValues_AVX512FP16(const uint8_t *codes) { - const __m256i bytes = _mm256_loadu_si256(reinterpret_cast(codes)); +static inline __m512h SQ8_FP16_LoadQuantizedValues_AVX512FP16(const uint8_t *quantized_values) { + const __m256i bytes = _mm256_loadu_si256(reinterpret_cast(quantized_values)); return _mm512_cvtepu16_ph(_mm512_cvtepu8_epi16(bytes)); } -static inline void SQ8_FP16_InnerProductStep_AVX512FP16(const uint8_t *&codes, +static inline void SQ8_FP16_InnerProductStep_AVX512FP16(const uint8_t *&quantized_values, const vecsim_types::float16 *&query, __m512h &sum) { - const __m512h values = SQ8_FP16_LoadQuantizedValues_AVX512FP16(codes); + const __m512h values = SQ8_FP16_LoadQuantizedValues_AVX512FP16(quantized_values); const __m512h query_values = _mm512_loadu_ph(query); sum = _mm512_fmadd_ph(values, query_values, sum); - codes += 32; + quantized_values += 32; query += 32; } @@ -36,39 +36,39 @@ static inline __m512h SQ8_FP16_AccumulateInnerProduct_AVX512FP16(const void *sto const void *query_blob, size_t dimension) { using float16 = vecsim_types::float16; - const auto *codes = static_cast(storage); + const auto *quantized_values = static_cast(storage); const auto *query = static_cast(query_blob); - const auto *end = codes + dimension; + const auto *end = quantized_values + dimension; __m512h sum = _mm512_setzero_ph(); if constexpr (residual) { constexpr __mmask32 mask = (1U << residual) - 1; - const __m512h values = SQ8_FP16_LoadQuantizedValues_AVX512FP16(codes); + const __m512h values = SQ8_FP16_LoadQuantizedValues_AVX512FP16(quantized_values); const __m512h query_values = _mm512_loadu_ph(query); sum = _mm512_maskz_mul_ph(mask, values, query_values); - codes += residual; + quantized_values += residual; query += residual; } if constexpr (four_sums_fp32_reduce) { __m512h sum1 = _mm512_setzero_ph(); __m512h sum2 = _mm512_setzero_ph(); __m512h sum3 = _mm512_setzero_ph(); - while (end - codes >= 128) { - SQ8_FP16_InnerProductStep_AVX512FP16(codes, query, sum); - SQ8_FP16_InnerProductStep_AVX512FP16(codes, query, sum1); - SQ8_FP16_InnerProductStep_AVX512FP16(codes, query, sum2); - SQ8_FP16_InnerProductStep_AVX512FP16(codes, query, sum3); + while (end - quantized_values >= 128) { + SQ8_FP16_InnerProductStep_AVX512FP16(quantized_values, query, sum); + SQ8_FP16_InnerProductStep_AVX512FP16(quantized_values, query, sum1); + SQ8_FP16_InnerProductStep_AVX512FP16(quantized_values, query, sum2); + SQ8_FP16_InnerProductStep_AVX512FP16(quantized_values, query, sum3); } - if (codes < end) - SQ8_FP16_InnerProductStep_AVX512FP16(codes, query, sum); - if (codes < end) - SQ8_FP16_InnerProductStep_AVX512FP16(codes, query, sum1); - if (codes < end) - SQ8_FP16_InnerProductStep_AVX512FP16(codes, query, sum2); + if (quantized_values < end) + SQ8_FP16_InnerProductStep_AVX512FP16(quantized_values, query, sum); + if (quantized_values < end) + SQ8_FP16_InnerProductStep_AVX512FP16(quantized_values, query, sum1); + if (quantized_values < end) + SQ8_FP16_InnerProductStep_AVX512FP16(quantized_values, query, sum2); sum = _mm512_add_ph(_mm512_add_ph(sum, sum1), _mm512_add_ph(sum2, sum3)); } else { do { - SQ8_FP16_InnerProductStep_AVX512FP16(codes, query, sum); - } while (codes < end); + SQ8_FP16_InnerProductStep_AVX512FP16(quantized_values, query, sum); + } while (quantized_values < end); } return sum; } From 742193c4daef41666c9c74e30e06fcea0f420aac Mon Sep 17 00:00:00 2001 From: Dor Forer Date: Mon, 5 Oct 2026 15:35:20 +0300 Subject: [PATCH 7/8] Make SQ8 FP16 tail block counts explicit --- src/VecSim/spaces/IP/IP_AVX512FP16_BW_VL_SQ8_FP16.h | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/VecSim/spaces/IP/IP_AVX512FP16_BW_VL_SQ8_FP16.h b/src/VecSim/spaces/IP/IP_AVX512FP16_BW_VL_SQ8_FP16.h index 6d1ef747f..3cc0e5582 100644 --- a/src/VecSim/spaces/IP/IP_AVX512FP16_BW_VL_SQ8_FP16.h +++ b/src/VecSim/spaces/IP/IP_AVX512FP16_BW_VL_SQ8_FP16.h @@ -58,11 +58,12 @@ static inline __m512h SQ8_FP16_AccumulateInnerProduct_AVX512FP16(const void *sto SQ8_FP16_InnerProductStep_AVX512FP16(quantized_values, query, sum2); SQ8_FP16_InnerProductStep_AVX512FP16(quantized_values, query, sum3); } - if (quantized_values < end) + const size_t remaining = end - quantized_values; + if (remaining >= 32) SQ8_FP16_InnerProductStep_AVX512FP16(quantized_values, query, sum); - if (quantized_values < end) + if (remaining >= 64) SQ8_FP16_InnerProductStep_AVX512FP16(quantized_values, query, sum1); - if (quantized_values < end) + if (remaining >= 96) SQ8_FP16_InnerProductStep_AVX512FP16(quantized_values, query, sum2); sum = _mm512_add_ph(_mm512_add_ph(sum, sum1), _mm512_add_ph(sum2, sum3)); } else { From e56df68c4f2b433bfaa7a134fb0043d115780a1d Mon Sep 17 00:00:00 2001 From: Dor Forer Date: Wed, 21 Oct 2026 11:54:06 +0300 Subject: [PATCH 8/8] Test minimum FP16 subnormals without distance cancellation --- tests/unit/test_spaces.cpp | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/unit/test_spaces.cpp b/tests/unit/test_spaces.cpp index 91dfdf46d..36129d017 100644 --- a/tests/unit/test_spaces.cpp +++ b/tests/unit/test_spaces.cpp @@ -5662,13 +5662,14 @@ TEST(SQ8FP16NativeIPTest, MinimumSmallAndLargestFiniteFP16Values) { for (float value : {0.0f, 0x1p-24f, -0x1p-24f, 37.0f * 0x1p-24f, 0x1p-14f, -0x1p-14f, 65504.0f, -65504.0f}) { SCOPED_TRACE(value); - SQ8FP16NativeCase input(dim); + // Amplify subnormal dots so subtracting from one cannot hide their loss. + const float delta = std::abs(value) < 0x1p-14f ? 0x1p24f : 1.0f; + SQ8FP16NativeCase input(dim, 0.0f, delta); input.storage[0] = 1; input.query[0] = vecsim_types::FP32_to_FP16(value); input.setQuerySum(dim); ASSERT_FLOAT_EQ(vecsim_types::FP16_to_FP32(input.query[0]), value); - // Compare distances directly: subtracting a distance near one can erase a tiny dot. - const float expected = 1.0f - value; + const float expected = 1.0f - delta * value; EXPECT_FLOAT_EQ(Choose_SQ8_FP16_IP_implementation_AVX512FP16_BW_VL(dim)( input.storage.data(), input.query.data(), dim), expected);