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..3cc0e5582 --- /dev/null +++ b/src/VecSim/spaces/IP/IP_AVX512FP16_BW_VL_SQ8_FP16.h @@ -0,0 +1,127 @@ +/* + * 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_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 *&quantized_values, + const vecsim_types::float16 *&query, + __m512h &sum) { + 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); + quantized_values += 32; + query += 32; +} + +// dim >= 32 keeps the residual's full loads within the vector payloads. +template +static inline __m512h SQ8_FP16_AccumulateInnerProduct_AVX512FP16(const void *storage, + const void *query_blob, + size_t dimension) { + using float16 = vecsim_types::float16; + const auto *quantized_values = static_cast(storage); + const auto *query = static_cast(query_blob); + 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(quantized_values); + const __m512h query_values = _mm512_loadu_ph(query); + sum = _mm512_maskz_mul_ph(mask, values, query_values); + 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 - 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); + } + const size_t remaining = end - quantized_values; + if (remaining >= 32) + SQ8_FP16_InnerProductStep_AVX512FP16(quantized_values, query, sum); + if (remaining >= 64) + SQ8_FP16_InnerProductStep_AVX512FP16(quantized_values, query, sum1); + 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 { + do { + SQ8_FP16_InnerProductStep_AVX512FP16(quantized_values, query, sum); + } while (quantized_values < end); + } + 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)); + return _mm512_reduce_add_ps(_mm512_add_ps(low, high)); + } else { + const _Float16 reduced = _mm512_reduce_add_ph(sum); + return static_cast(reduced); + } +} + +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)); + 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); +} + +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, + 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 7ae4da228..5bffee8d1 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 >= 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); + } +#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 >= 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); + } +#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..c4caef596 --- /dev/null +++ b/src/VecSim/spaces/functions/AVX512FP16_BW_VL.cpp @@ -0,0 +1,34 @@ +/* + * 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; + if (dim >= spaces::SQ8_FP16_FOUR_SUMS_MIN_DIM) { + 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; +} + +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/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 923bb7c1b..36129d017 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 @@ -33,6 +34,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 +3606,94 @@ 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; +} + +// 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); +} + +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) { + 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)); + 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 * dot); +} + +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 +3839,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 >= 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), + 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 +3974,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 >= 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), + 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 +5513,301 @@ 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, 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}) { + 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 < 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)); + EXPECT_EQ(Cosine_SQ8_FP16_GetDistFunc(dim, nullptr, &features), + dim < 16 ? SQ8_FP16_Cosine : Choose_SQ8_FP16_Cosine_implementation_AVX512F(dim)); + } +} + +TEST(SQ8FP16NativeIPTest, HalfArithmeticAcrossResidualsAndBlocks) { + if (!HasSQ8FP16NativeISA()) + GTEST_SKIP() << "AVX512F/BW/VL/FP16 is unavailable"; + 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); + 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); + 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); + // 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); + 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); + 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, ProductLaneMergeAndReductionOverflowHandling) { + if (!HasSQ8FP16NativeISA()) + GTEST_SKIP() << "AVX512F/BW/VL/FP16 is unavailable"; + auto check = [&](SQ8FP16NativeCase &input, size_t dim) { + input.setQuerySum(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), + 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); + }; + 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 >= spaces::SQ8_FP16_FOUR_SUMS_MIN_DIM))); + 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 >= spaces::SQ8_FP16_FOUR_SUMS_MIN_DIM))); + 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 >= spaces::SQ8_FP16_FOUR_SUMS_MIN_DIM))); + 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 = 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); + 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); + 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); + 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; + } + EXPECT_EQ(std::memcmp(input.query.data(), original.data(), original.size() * sizeof(float16)), + 0); +} +#endif