diff --git a/src/VecSim/spaces/AVX_utils.h b/src/VecSim/spaces/AVX_utils.h index 2fe0b904e..b631ee5d6 100644 --- a/src/VecSim/spaces/AVX_utils.h +++ b/src/VecSim/spaces/AVX_utils.h @@ -35,3 +35,14 @@ static inline float my_mm256_reduce_add_ps(__m256 x) { return TmpRes[0] + TmpRes[1] + TmpRes[2] + TmpRes[3] + TmpRes[4] + TmpRes[5] + TmpRes[6] + TmpRes[7]; } + +// As my_mm256_reduce_add_ps, but folded in-register in three steps instead of spilling 8 floats +// and summing them with 7 dependent scalar adds. That fixed cost dominates at small dimensions. +// Reassociating changes the low bits, so this is added alongside the original rather than +// replacing it -- the original's other callers would each need their own benchmarking. +static inline float my_mm256_reduce_add_ps_tree(__m256 x) { + __m128 sum128 = _mm_add_ps(_mm256_castps256_ps128(x), _mm256_extractf128_ps(x, 1)); + sum128 = _mm_add_ps(sum128, _mm_movehl_ps(sum128, sum128)); + sum128 = _mm_add_ss(sum128, _mm_shuffle_ps(sum128, sum128, 0x1)); + return _mm_cvtss_f32(sum128); +} diff --git a/src/VecSim/spaces/L2/L2.cpp b/src/VecSim/spaces/L2/L2.cpp index 253041fb6..3ba75fb4c 100644 --- a/src/VecSim/spaces/L2/L2.cpp +++ b/src/VecSim/spaces/L2/L2.cpp @@ -19,29 +19,38 @@ using float16 = vecsim_types::float16; using sq8 = vecsim_types::sq8; /* - * Optimized asymmetric SQ8-FP32 L2 squared distance using algebraic identity: - * ||x - y||² = Σx_i² - 2*IP(x, y) + Σy_i² - * = x_sum_squares - 2 * IP(x, y) + y_sum_squares - * where IP(x, y) = min * y_sum + delta * Σ(q_i * y_i) + * Asymmetric SQ8-FP32 L2 squared distance computed via direct residual accumulation: + * ||x - y||² = Σ(dequant(x_i) - y_i)² + * where dequant(x_i) = min_val + delta * q_i + * + * This avoids the algebraic-identity/cancellation approach (||x||² + ||y||² - 2*IP(x, y)), + * which catastrophically cancels in FP32 when x and y share a large common offset relative to + * their spread. + * + * The operand order below relies on FP addition NOT being reassociated, so `-ffast-math` / + * `-Ofast` would reinstate the bug. The repo's -O3 builds are safe. * * pVect1 is storage (SQ8): [uint8_t values (dim)] [min_val] [delta] [x_sum] [x_sum_squares] * pVect2 is query (FP32): [float values (dim)] [y_sum] [y_sum_squares] */ float SQ8_FP32_L2Sqr(const void *pVect1v, const void *pVect2v, size_t dimension) { - // Get the raw inner product using the common implementation - const float ip = SQ8_FP32_InnerProduct_Impl(pVect1v, pVect2v, dimension); - // Storage metadata follows a byte payload and is not necessarily float-aligned. const auto *pVect1 = static_cast(pVect1v); - const float x_sum_sq = - load_unaligned(pVect1 + dimension + sq8::SUM_SQUARES * sizeof(float)); - - // Get precomputed sum of squares from query blob (pVect2 is FP32) const auto *pVect2 = static_cast(pVect2v); - const float y_sum_sq = pVect2[dimension + sq8::SUM_SQUARES_QUERY]; - // L2² = ||x||² + ||y||² - 2*IP(x, y) - return x_sum_sq + y_sum_sq - 2.0f * ip; + const auto *params1 = pVect1 + dimension; + const float min_val = load_unaligned(params1 + sq8::MIN_VAL * sizeof(float)); + const float delta = load_unaligned(params1 + sq8::DELTA * sizeof(float)); + + float res = 0; + for (size_t i = 0; i < dimension; i++) { + // Order matters: min_val - y_i is exact (Sterbenz) since both are large and close, so + // adding the small delta*q_i correction afterward preserves the residual. Computing + // (min_val + delta*q_i) - y_i instead rounds it away at the large offset's precision. + float diff = delta * static_cast(pVect1[i]) + (min_val - pVect2[i]); + res += diff * diff; + } + return res; } /* diff --git a/src/VecSim/spaces/L2/L2_AVX2_FMA_SQ8_FP32.h b/src/VecSim/spaces/L2/L2_AVX2_FMA_SQ8_FP32.h index 0f9d5bde9..285bb2ab7 100644 --- a/src/VecSim/spaces/L2/L2_AVX2_FMA_SQ8_FP32.h +++ b/src/VecSim/spaces/L2/L2_AVX2_FMA_SQ8_FP32.h @@ -9,38 +9,119 @@ #pragma once #include "VecSim/spaces/space_includes.h" #include "VecSim/spaces/AVX_utils.h" -#include "VecSim/spaces/IP/IP_AVX2_FMA_SQ8_FP32.h" #include "VecSim/types/sq8.h" using sq8 = vecsim_types::sq8; /* - * Optimized asymmetric SQ8 L2 squared distance using algebraic identity: + * Asymmetric SQ8 L2 squared distance computed via direct residual accumulation: * - * ||x - y||² = Σx_i² - 2*IP(x, y) + Σy_i² - * = x_sum_squares - 2 * IP(x, y) + y_sum_squares + * ||x - y||² = Σ(dequant(x_i) - y_i)² + * where dequant(x_i) = min_val + delta * q_i * - * where: - * - IP(x, y) = min * y_sum + delta * Σ(q_i * y_i) (computed via SQ8_FP32_InnerProductImp_FMA) - * - x_sum_squares and y_sum_squares are precomputed + * This avoids the algebraic-identity/cancellation approach, which catastrophically cancels in + * FP32 when x and y share a large common offset relative to their spread. * - * This avoids dequantization in the hot loop. + * This version uses FMA instructions. Critically, the subtract is fused into the FMA + * (diff = fma(delta, q, min - y)) rather than computed separately, which matters for + * performance. */ +// Helper: compute Σ(diff_i²) for 8 elements, where diff_i = dequant(x_i) - y_i. +// pVect1 = SQ8 storage (quantized values), pVect2 = FP32 query. +// min_val_vec/delta_vec are broadcast scalars from the stored vector's metadata. +static inline void L2StepSQ8_FP32_FMA(const uint8_t *&pVect1, const float *&pVect2, __m256 &sum, + __m256 min_val_vec, __m256 delta_vec) { + // Load 8 uint8 elements and convert to float + __m128i v1_128 = _mm_loadl_epi64(reinterpret_cast(pVect1)); + pVect1 += 8; + __m256i v1_256 = _mm256_cvtepu8_epi32(v1_128); + __m256 v1_f = _mm256_cvtepi32_ps(v1_256); + + // Load 8 float elements from query + __m256 v2 = _mm256_loadu_ps(pVect2); + pVect2 += 8; + + // min - y computed once per lane, then fuse the dequantize-and-subtract into a single FMA: + // diff = delta*q + (min - y). + __m256 min_minus_y = _mm256_sub_ps(min_val_vec, v2); + __m256 diff = _mm256_fmadd_ps(delta_vec, v1_f, min_minus_y); + + sum = _mm256_fmadd_ps(diff, diff, sum); +} + // pVect1v = SQ8 storage, pVect2v = FP32 query template // 0..31 float SQ8_FP32_L2SqrSIMD16_AVX2_FMA(const void *pVect1v, const void *pVect2v, size_t dimension) { - // Get the raw inner product using the common SIMD implementation - const float ip = SQ8_FP32_InnerProductImp_FMA(pVect1v, pVect2v, dimension); + const uint8_t *pVect1 = static_cast(pVect1v); // SQ8 storage + const float *pVect2 = static_cast(pVect2v); // FP32 query + const uint8_t *pEnd1 = pVect1 + dimension; + + // Get quantization parameters from stored vector (after quantized data) + const uint8_t *pVect1Base = static_cast(pVect1v); + const auto *params1 = pVect1Base + dimension; + const float min_val_scalar = load_unaligned(params1 + sq8::MIN_VAL * sizeof(float)); + const float delta_scalar = load_unaligned(params1 + sq8::DELTA * sizeof(float)); + const __m256 min_val_vec = _mm256_set1_ps(min_val_scalar); + const __m256 delta_vec = _mm256_set1_ps(delta_scalar); + + // Initialize sum accumulators. Four accumulators break the FMA dependency chain, letting + // more FMAs be in flight at once. + __m256 sum0 = _mm256_setzero_ps(); + __m256 sum1 = _mm256_setzero_ps(); + __m256 sum2 = _mm256_setzero_ps(); + __m256 sum3 = _mm256_setzero_ps(); + + // Handle residual elements first (0-7 elements). The full-width query load is safe because + // `dim` is at least 8, so the query spans at least 8 floats. + if constexpr (residual % 8) { + __mmask8 constexpr mask = (1 << (residual % 8)) - 1; + + // Load uint8 elements and convert to float + __m128i v1_128 = _mm_loadl_epi64(reinterpret_cast(pVect1)); + pVect1 += residual % 8; + + __m256i v1_256 = _mm256_cvtepu8_epi32(v1_128); + __m256 v1_f = _mm256_cvtepi32_ps(v1_256); + + // Load masked float elements from query + __m256 v2 = my_mm256_maskz_loadu_ps(pVect2); + pVect2 += residual % 8; + + // min - y, then dequantize-and-subtract + __m256 min_minus_y = _mm256_sub_ps(min_val_vec, v2); + __m256 diff = _mm256_fmadd_ps(delta_vec, v1_f, min_minus_y); + + // Masked-out lanes carry garbage (v2 was zeroed, not set to min_val), so blend the + // squared diff with zero for those lanes before accumulating. + __m256 diff_sq = _mm256_mul_ps(diff, diff); + sum0 = _mm256_blend_ps(_mm256_setzero_ps(), diff_sq, mask); + } - // Get precomputed sum of squares from storage blob (pVect1v is SQ8 storage) - const uint8_t *pVect1 = static_cast(pVect1v); - const float x_sum_sq = - load_unaligned(pVect1 + dimension + sq8::SUM_SQUARES * sizeof(float)); + // Handle the remaining full 8-element blocks of the residual (compile-time resolved). + if constexpr (residual >= 8) { + L2StepSQ8_FP32_FMA(pVect1, pVect2, sum1, min_val_vec, delta_vec); + } + if constexpr (residual >= 16) { + L2StepSQ8_FP32_FMA(pVect1, pVect2, sum2, min_val_vec, delta_vec); + } + if constexpr (residual >= 24) { + L2StepSQ8_FP32_FMA(pVect1, pVect2, sum3, min_val_vec, delta_vec); + } - // Get precomputed sum of squares from query blob (pVect2v is FP32 query) - const float y_sum_sq = static_cast(pVect2v)[dimension + sq8::SUM_SQUARES_QUERY]; + // We dealt with the residual part. We are left with some multiple of 32 elements. + // In each iteration we calculate 32 elements = 4 chunks of 8. The loop may run zero times + // (dim can be as small as 8). + while (pVect1 < pEnd1) { + L2StepSQ8_FP32_FMA(pVect1, pVect2, sum0, min_val_vec, delta_vec); + L2StepSQ8_FP32_FMA(pVect1, pVect2, sum1, min_val_vec, delta_vec); + L2StepSQ8_FP32_FMA(pVect1, pVect2, sum2, min_val_vec, delta_vec); + L2StepSQ8_FP32_FMA(pVect1, pVect2, sum3, min_val_vec, delta_vec); + } - // L2² = ||x||² + ||y||² - 2*IP(x, y) - return x_sum_sq + y_sum_sq - 2.0f * ip; + // Reduce to get Σ(diff_i²) + // Tree reduction: the fixed cost of the stack-based fold is a large share of + // total time at small dimensions. See my_mm256_reduce_add_ps_tree in AVX_utils.h. + return my_mm256_reduce_add_ps_tree( + _mm256_add_ps(_mm256_add_ps(sum0, sum1), _mm256_add_ps(sum2, sum3))); } diff --git a/src/VecSim/spaces/L2/L2_AVX2_SQ8_FP32.h b/src/VecSim/spaces/L2/L2_AVX2_SQ8_FP32.h index d08474f71..f3f1a31f8 100644 --- a/src/VecSim/spaces/L2/L2_AVX2_SQ8_FP32.h +++ b/src/VecSim/spaces/L2/L2_AVX2_SQ8_FP32.h @@ -9,38 +9,115 @@ #pragma once #include "VecSim/spaces/space_includes.h" #include "VecSim/spaces/AVX_utils.h" -#include "VecSim/spaces/IP/IP_AVX2_SQ8_FP32.h" #include "VecSim/types/sq8.h" using sq8 = vecsim_types::sq8; /* - * Optimized asymmetric SQ8 L2 squared distance using algebraic identity: + * Asymmetric SQ8 L2 squared distance computed via direct residual accumulation: * - * ||x - y||² = Σx_i² - 2*IP(x, y) + Σy_i² - * = x_sum_squares - 2 * IP(x, y) + y_sum_squares + * ||x - y||² = Σ(dequant(x_i) - y_i)² + * where dequant(x_i) = min_val + delta * q_i * - * where: - * - IP(x, y) = min * y_sum + delta * Σ(q_i * y_i) (computed via SQ8_FP32_InnerProductImp_AVX2) - * - x_sum_squares and y_sum_squares are precomputed - * - * This avoids dequantization in the hot loop. + * This avoids the algebraic-identity/cancellation approach, which catastrophically cancels in + * FP32 when x and y share a large common offset relative to their spread. */ +// Helper: compute Σ(diff_i²) for 8 elements, where diff_i = dequant(x_i) - y_i. +// pVect1 = SQ8 storage (quantized values), pVect2 = FP32 query. +// min_val_vec/delta_vec are broadcast scalars from the stored vector's metadata. +static inline void L2StepSQ8_FP32_AVX2(const uint8_t *&pVect1, const float *&pVect2, __m256 &sum, + __m256 min_val_vec, __m256 delta_vec) { + // Load 8 uint8 elements and convert to float + __m128i v1_128 = _mm_loadl_epi64(reinterpret_cast(pVect1)); + pVect1 += 8; + __m256i v1_256 = _mm256_cvtepu8_epi32(v1_128); + __m256 v1_f = _mm256_cvtepi32_ps(v1_256); + + // Load 8 float elements from query + __m256 v2 = _mm256_loadu_ps(pVect2); + pVect2 += 8; + + // min - y computed once per lane, then dequantize-and-subtract: diff = delta*q + (min - y). + // No FMA in this variant, so mul + add. + __m256 min_minus_y = _mm256_sub_ps(min_val_vec, v2); + __m256 diff = _mm256_add_ps(_mm256_mul_ps(delta_vec, v1_f), min_minus_y); + + sum = _mm256_add_ps(sum, _mm256_mul_ps(diff, diff)); +} + // pVect1v = SQ8 storage, pVect2v = FP32 query template // 0..31 float SQ8_FP32_L2SqrSIMD16_AVX2(const void *pVect1v, const void *pVect2v, size_t dimension) { - // Get the raw inner product using the common SIMD implementation - const float ip = SQ8_FP32_InnerProductImp_AVX2(pVect1v, pVect2v, dimension); + const uint8_t *pVect1 = static_cast(pVect1v); // SQ8 storage + const float *pVect2 = static_cast(pVect2v); // FP32 query + const uint8_t *pEnd1 = pVect1 + dimension; + + // Get quantization parameters from stored vector (after quantized data) + const uint8_t *pVect1Base = static_cast(pVect1v); + const auto *params1 = pVect1Base + dimension; + const float min_val_scalar = load_unaligned(params1 + sq8::MIN_VAL * sizeof(float)); + const float delta_scalar = load_unaligned(params1 + sq8::DELTA * sizeof(float)); + const __m256 min_val_vec = _mm256_set1_ps(min_val_scalar); + const __m256 delta_vec = _mm256_set1_ps(delta_scalar); + + // Initialize sum accumulators. Four accumulators break the dependency chain, letting more + // ops be in flight at once. + __m256 sum0 = _mm256_setzero_ps(); + __m256 sum1 = _mm256_setzero_ps(); + __m256 sum2 = _mm256_setzero_ps(); + __m256 sum3 = _mm256_setzero_ps(); + + // Handle residual elements first (0-7 elements). The full-width query load is safe because + // `dim` is at least 8, so the query spans at least 8 floats. + if constexpr (residual % 8) { + __mmask8 constexpr mask = (1 << (residual % 8)) - 1; + + // Load uint8 elements and convert to float + __m128i v1_128 = _mm_loadl_epi64(reinterpret_cast(pVect1)); + pVect1 += residual % 8; + + __m256i v1_256 = _mm256_cvtepu8_epi32(v1_128); + __m256 v1_f = _mm256_cvtepi32_ps(v1_256); + + // Load masked float elements from query + __m256 v2 = my_mm256_maskz_loadu_ps(pVect2); + pVect2 += residual % 8; + + // min - y, then dequantize-and-subtract + __m256 min_minus_y = _mm256_sub_ps(min_val_vec, v2); + __m256 diff = _mm256_add_ps(_mm256_mul_ps(delta_vec, v1_f), min_minus_y); + + // Masked-out lanes carry garbage (v2 was zeroed, not set to min_val), so blend the + // squared diff with zero for those lanes before accumulating. + __m256 diff_sq = _mm256_mul_ps(diff, diff); + sum0 = _mm256_blend_ps(_mm256_setzero_ps(), diff_sq, mask); + } - // Get precomputed sum of squares from storage blob (pVect1v is SQ8 storage) - const uint8_t *pVect1 = static_cast(pVect1v); - const float x_sum_sq = - load_unaligned(pVect1 + dimension + sq8::SUM_SQUARES * sizeof(float)); + // Handle the remaining full 8-element blocks of the residual (compile-time resolved). + if constexpr (residual >= 8) { + L2StepSQ8_FP32_AVX2(pVect1, pVect2, sum1, min_val_vec, delta_vec); + } + if constexpr (residual >= 16) { + L2StepSQ8_FP32_AVX2(pVect1, pVect2, sum2, min_val_vec, delta_vec); + } + if constexpr (residual >= 24) { + L2StepSQ8_FP32_AVX2(pVect1, pVect2, sum3, min_val_vec, delta_vec); + } - // Get precomputed sum of squares from query blob (pVect2v is FP32 query) - const float y_sum_sq = static_cast(pVect2v)[dimension + sq8::SUM_SQUARES_QUERY]; + // We dealt with the residual part. We are left with some multiple of 32 elements. + // In each iteration we calculate 32 elements = 4 chunks of 8. The loop may run zero times + // (dim can be as small as 8). + while (pVect1 < pEnd1) { + L2StepSQ8_FP32_AVX2(pVect1, pVect2, sum0, min_val_vec, delta_vec); + L2StepSQ8_FP32_AVX2(pVect1, pVect2, sum1, min_val_vec, delta_vec); + L2StepSQ8_FP32_AVX2(pVect1, pVect2, sum2, min_val_vec, delta_vec); + L2StepSQ8_FP32_AVX2(pVect1, pVect2, sum3, min_val_vec, delta_vec); + } - // L2² = ||x||² + ||y||² - 2*IP(x, y) - return x_sum_sq + y_sum_sq - 2.0f * ip; + // Reduce to get Σ(diff_i²) + // Tree reduction: the fixed cost of the stack-based fold is a large share of + // total time at small dimensions. See my_mm256_reduce_add_ps_tree in AVX_utils.h. + return my_mm256_reduce_add_ps_tree( + _mm256_add_ps(_mm256_add_ps(sum0, sum1), _mm256_add_ps(sum2, sum3))); } diff --git a/src/VecSim/spaces/L2/L2_AVX512F_BW_VL_VNNI_SQ8_FP32.h b/src/VecSim/spaces/L2/L2_AVX512F_BW_VL_VNNI_SQ8_FP32.h index f1c86d689..0c321a228 100644 --- a/src/VecSim/spaces/L2/L2_AVX512F_BW_VL_VNNI_SQ8_FP32.h +++ b/src/VecSim/spaces/L2/L2_AVX512F_BW_VL_VNNI_SQ8_FP32.h @@ -8,39 +8,109 @@ */ #pragma once #include "VecSim/spaces/space_includes.h" -#include "VecSim/spaces/IP/IP_AVX512F_BW_VL_VNNI_SQ8_FP32.h" #include "VecSim/types/sq8.h" +#include using sq8 = vecsim_types::sq8; /* - * Optimized asymmetric SQ8 L2 squared distance using algebraic identity: + * Asymmetric SQ8 L2 squared distance computed via direct residual accumulation: * - * ||x - y||² = Σx_i² - 2*IP(x, y) + Σy_i² - * = x_sum_squares - 2 * IP(x, y) + y_sum_squares + * ||x - y||² = Σ(dequant(x_i) - y_i)² + * where dequant(x_i) = min_val + delta * q_i * - * where: - * - IP(x, y) = min * y_sum + delta * Σ(q_i * y_i) (computed via SQ8_FP32_InnerProductImp_AVX512) - * - x_sum_squares and y_sum_squares are precomputed + * This avoids the algebraic-identity/cancellation approach, which catastrophically cancels in + * FP32 when x and y share a large common offset relative to their spread. * - * This avoids dequantization in the hot loop. + * The subtract is fused into the FMA (diff = fma(delta, q, min - y)) rather than computed + * separately, which matters for performance. */ -// pVect1v = SQ8 storage, pVect2v = FP32 query +// Helper: compute Σ(diff_i²) for 16 elements, where diff_i = dequant(x_i) - y_i. +// pVec1 = SQ8 storage (quantized values), pVec2 = FP32 query. +// min_val_vec/delta_vec are broadcast scalars from the stored vector's metadata. +static inline void L2StepSQ8_FP32_AVX512(const uint8_t *&pVec1, const float *&pVec2, __m512 &sum, + __m512 min_val_vec, __m512 delta_vec) { + // Load 16 uint8 elements from quantized vector and convert to float + __m128i v1_128 = _mm_loadu_si128(reinterpret_cast(pVec1)); + __m512i v1_512 = _mm512_cvtepu8_epi32(v1_128); + __m512 v1_f = _mm512_cvtepi32_ps(v1_512); + + // Load 16 float elements from query (pVec2) + __m512 v2 = _mm512_loadu_ps(pVec2); + + // min - y computed once per lane, then fuse the dequantize-and-subtract into a single FMA: + // diff = delta*q + (min - y). + __m512 min_minus_y = _mm512_sub_ps(min_val_vec, v2); + __m512 diff = _mm512_fmadd_ps(delta_vec, v1_f, min_minus_y); + + sum = _mm512_fmadd_ps(diff, diff, sum); + + pVec1 += 16; + pVec2 += 16; +} + +// pVec1v = SQ8 storage, pVec2v = FP32 query template // 0..31 -float SQ8_FP32_L2SqrSIMD16_AVX512F_BW_VL_VNNI(const void *pVect1v, const void *pVect2v, +float SQ8_FP32_L2SqrSIMD16_AVX512F_BW_VL_VNNI(const void *pVec1v, const void *pVec2v, size_t dimension) { - // Get the raw inner product using the common SIMD implementation - const float ip = SQ8_FP32_InnerProductImp_AVX512(pVect1v, pVect2v, dimension); + const uint8_t *pVec1 = static_cast(pVec1v); // SQ8 storage + const float *pVec2 = static_cast(pVec2v); // FP32 query + const uint8_t *pEnd1 = pVec1 + dimension; + + // Get quantization parameters from stored vector (after quantized data) + const uint8_t *pVec1Base = static_cast(pVec1v); + const auto *params1 = pVec1Base + dimension; + const float min_val_scalar = load_unaligned(params1 + sq8::MIN_VAL * sizeof(float)); + const float delta_scalar = load_unaligned(params1 + sq8::DELTA * sizeof(float)); + const __m512 min_val_vec = _mm512_set1_ps(min_val_scalar); + const __m512 delta_vec = _mm512_set1_ps(delta_scalar); + + // Initialize sum accumulators for Σ(diff_i²). Two accumulators break the FMA dependency + // chain, letting more FMAs be in flight at once. + __m512 sum0 = _mm512_setzero_ps(); + __m512 sum1 = _mm512_setzero_ps(); + + // Handle the sub-16 residual elements first + if constexpr (residual % 16) { + __mmask16 constexpr mask = (1U << (residual % 16)) - 1; + + // Load uint8 elements (safe to load 16 bytes due to the metadata padding after the + // quantized values). The query load is masked, which suppresses faults on masked-out + // lanes, so both loads are safe for any dimension. + __m128i v1_128 = _mm_loadu_si128(reinterpret_cast(pVec1)); + __m512i v1_512 = _mm512_cvtepu8_epi32(v1_128); + __m512 v1_f = _mm512_cvtepi32_ps(v1_512); + + // Load masked float elements from query + __m512 v2 = _mm512_maskz_loadu_ps(mask, pVec2); + + // min - y, then dequantize-and-subtract + __m512 min_minus_y = _mm512_sub_ps(min_val_vec, v2); + __m512 diff = _mm512_fmadd_ps(delta_vec, v1_f, min_minus_y); + + // Masked-out lanes carry garbage (v2 was zeroed, not set to min_val), so mask the + // squared diff to zero for those lanes before accumulating. + sum0 = _mm512_maskz_mul_ps(mask, diff, diff); + + pVec1 += residual % 16; + pVec2 += residual % 16; + } - // Get precomputed sum of squares from storage blob (pVect1v is SQ8 storage) - const uint8_t *pVect1 = static_cast(pVect1v); - const float x_sum_sq = - load_unaligned(pVect1 + dimension + sq8::SUM_SQUARES * sizeof(float)); + // Handle the remaining full 16-element block of the residual (compile-time resolved). + if constexpr (residual >= 16) { + L2StepSQ8_FP32_AVX512(pVec1, pVec2, sum1, min_val_vec, delta_vec); + } - // Get precomputed sum of squares from query blob (pVect2v is FP32 query) - const float y_sum_sq = static_cast(pVect2v)[dimension + sq8::SUM_SQUARES_QUERY]; + // We dealt with the residual part. We are left with some multiple of 32 elements. + // In each iteration we calculate 32 elements = 2 chunks of 16. The loop may run zero times + // (dim can be as small as 8). + while (pVec1 < pEnd1) { + L2StepSQ8_FP32_AVX512(pVec1, pVec2, sum0, min_val_vec, delta_vec); + L2StepSQ8_FP32_AVX512(pVec1, pVec2, sum1, min_val_vec, delta_vec); + } - // L2² = ||x||² + ||y||² - 2*IP(x, y) - return x_sum_sq + y_sum_sq - 2.0f * ip; + // Reduce to get Σ(diff_i²) + __m512 sum = _mm512_add_ps(sum0, sum1); + return _mm512_reduce_add_ps(sum); } diff --git a/src/VecSim/spaces/L2/L2_NEON_SQ8_FP32.h b/src/VecSim/spaces/L2/L2_NEON_SQ8_FP32.h index f6f7a6bc0..5de7f6a34 100644 --- a/src/VecSim/spaces/L2/L2_NEON_SQ8_FP32.h +++ b/src/VecSim/spaces/L2/L2_NEON_SQ8_FP32.h @@ -8,40 +8,146 @@ */ #pragma once #include "VecSim/spaces/space_includes.h" -#include "VecSim/spaces/IP/IP_NEON_SQ8_FP32.h" #include "VecSim/types/sq8.h" #include using sq8 = vecsim_types::sq8; /* - * Optimized asymmetric SQ8 L2 squared distance using algebraic identity: + * Asymmetric SQ8-FP32 L2 squared distance via direct residual accumulation: * - * ||x - y||² = Σx_i² - 2*IP(x, y) + Σy_i² - * = x_sum_squares - 2 * IP(x, y) + y_sum_squares + * ||x - y||² = Σ(dequant(x_i) - y_i)², where dequant(x_i) = min_val + delta * q_i * - * where: - * - IP(x, y) = min * y_sum + delta * Σ(q_i * y_i) (computed via - * SQ8_FP32_InnerProductSIMD16_NEON_IMP) - * - x_sum_squares and y_sum_squares are precomputed - * - * This avoids dequantization in the hot loop. + * Not the ||x||² + ||y||² - 2*IP identity, which cancels catastrophically in FP32 when x and y + * share a large common offset relative to their spread (MOD-17526). */ +// 4 elements of Σ(diff_i²). Used for the tail; the main loop uses the 16-element form below. +static inline void L2StepSQ8_FP32_NEON(const uint8_t *&pVect1, const float *&pVect2, + float32x4_t &sum, float32x4_t min_val_vec, + float32x4_t delta_vec) { + uint8x8_t v1_u8 = vld1_u8(pVect1); + pVect1 += 4; + + uint32x4_t v1_u32 = vmovl_u16(vget_low_u16(vmovl_u8(v1_u8))); + float32x4_t v1_f = vcvtq_f32_u32(v1_u32); + + float32x4_t v2 = vld1q_f32(pVect2); + pVect2 += 4; + + // Explicit vfmaq_f32, not vmlaq_f32: keeping min - y first is what preserves the residual. + float32x4_t min_minus_y = vsubq_f32(min_val_vec, v2); + float32x4_t diff = vfmaq_f32(min_minus_y, delta_vec, v1_f); + + sum = vfmaq_f32(sum, diff, diff); +} + +// 16 elements off a single 16-byte load. The 4-element form reads 8 bytes and uses only 4, so +// back-to-back calls re-read half of every load. Per-lane arithmetic and accumulator mapping are +// unchanged, so results match the 4-element path exactly. +static inline void L2Step16SQ8_FP32_NEON(const uint8_t *&pVect1, const float *&pVect2, + float32x4_t &sum0, float32x4_t &sum1, float32x4_t &sum2, + float32x4_t &sum3, float32x4_t min_val_vec, + float32x4_t delta_vec) { + uint8x16_t v1_u8 = vld1q_u8(pVect1); + pVect1 += 16; + + const uint16x8_t wide_lo = vmovl_u8(vget_low_u8(v1_u8)); + const uint16x8_t wide_hi = vmovl_u8(vget_high_u8(v1_u8)); + + const float32x4_t q0 = vcvtq_f32_u32(vmovl_u16(vget_low_u16(wide_lo))); + const float32x4_t q1 = vcvtq_f32_u32(vmovl_u16(vget_high_u16(wide_lo))); + const float32x4_t q2 = vcvtq_f32_u32(vmovl_u16(vget_low_u16(wide_hi))); + const float32x4_t q3 = vcvtq_f32_u32(vmovl_u16(vget_high_u16(wide_hi))); + + const float32x4_t y0 = vld1q_f32(pVect2); + const float32x4_t y1 = vld1q_f32(pVect2 + 4); + const float32x4_t y2 = vld1q_f32(pVect2 + 8); + const float32x4_t y3 = vld1q_f32(pVect2 + 12); + pVect2 += 16; + + const float32x4_t d0 = vfmaq_f32(vsubq_f32(min_val_vec, y0), delta_vec, q0); + const float32x4_t d1 = vfmaq_f32(vsubq_f32(min_val_vec, y1), delta_vec, q1); + const float32x4_t d2 = vfmaq_f32(vsubq_f32(min_val_vec, y2), delta_vec, q2); + const float32x4_t d3 = vfmaq_f32(vsubq_f32(min_val_vec, y3), delta_vec, q3); + + sum0 = vfmaq_f32(sum0, d0, d0); + sum1 = vfmaq_f32(sum1, d1, d1); + sum2 = vfmaq_f32(sum2, d2, d2); + sum3 = vfmaq_f32(sum3, d3, d3); +} + // pVect1v = SQ8 storage, pVect2v = FP32 query template // 0..15 float SQ8_FP32_L2SqrSIMD16_NEON(const void *pVect1v, const void *pVect2v, size_t dimension) { - // Get the raw inner product using the common SIMD implementation - const float ip = SQ8_FP32_InnerProductSIMD16_NEON_IMP(pVect1v, pVect2v, dimension); + const uint8_t *pVect1 = static_cast(pVect1v); // SQ8 storage + const float *pVect2 = static_cast(pVect2v); // FP32 query + + // Get quantization parameters from stored vector (after quantized data) + const uint8_t *pVect1Base = static_cast(pVect1v); + const auto *params1 = pVect1Base + dimension; + const float min_val_scalar = load_unaligned(params1 + sq8::MIN_VAL * sizeof(float)); + const float delta_scalar = load_unaligned(params1 + sq8::DELTA * sizeof(float)); + const float32x4_t min_val_vec = vdupq_n_f32(min_val_scalar); + const float32x4_t delta_vec = vdupq_n_f32(delta_scalar); + + // Multiple accumulators for ILP + float32x4_t sum0 = vdupq_n_f32(0.0f); + float32x4_t sum1 = vdupq_n_f32(0.0f); + float32x4_t sum2 = vdupq_n_f32(0.0f); + float32x4_t sum3 = vdupq_n_f32(0.0f); + + const size_t num_of_chunks = dimension / 16; + + // Process 16 elements at a time in the main loop, one 16-byte load per iteration. + for (size_t i = 0; i < num_of_chunks; i++) { + L2Step16SQ8_FP32_NEON(pVect1, pVect2, sum0, sum1, sum2, sum3, min_val_vec, delta_vec); + } + + // Handle remaining complete 4-element blocks within residual + if constexpr (residual >= 4) { + L2StepSQ8_FP32_NEON(pVect1, pVect2, sum0, min_val_vec, delta_vec); + } + if constexpr (residual >= 8) { + L2StepSQ8_FP32_NEON(pVect1, pVect2, sum1, min_val_vec, delta_vec); + } + if constexpr (residual >= 12) { + L2StepSQ8_FP32_NEON(pVect1, pVect2, sum2, min_val_vec, delta_vec); + } + + // Handle final residual elements (0-3 elements) + constexpr size_t final_residual = residual % 4; + if constexpr (final_residual > 0) { + // Padding lanes get q=0, y=min_val_scalar, so diff = delta*0 + (min - min) = 0. + float32x4_t v1_f = vdupq_n_f32(0.0f); + float32x4_t v2 = vdupq_n_f32(min_val_scalar); + + if constexpr (final_residual >= 1) { + float q0 = static_cast(pVect1[0]); + v1_f = vld1q_lane_f32(&q0, v1_f, 0); + v2 = vld1q_lane_f32(pVect2, v2, 0); + } + if constexpr (final_residual >= 2) { + float q1 = static_cast(pVect1[1]); + v1_f = vld1q_lane_f32(&q1, v1_f, 1); + v2 = vld1q_lane_f32(pVect2 + 1, v2, 1); + } + if constexpr (final_residual >= 3) { + float q2 = static_cast(pVect1[2]); + v1_f = vld1q_lane_f32(&q2, v1_f, 2); + v2 = vld1q_lane_f32(pVect2 + 2, v2, 2); + } - // Get precomputed sum of squares from storage blob (pVect1v is SQ8 storage) - const uint8_t *pVect1 = static_cast(pVect1v); - const float x_sum_sq = - load_unaligned(pVect1 + dimension + sq8::SUM_SQUARES * sizeof(float)); + float32x4_t min_minus_y = vsubq_f32(min_val_vec, v2); + float32x4_t diff = vfmaq_f32(min_minus_y, delta_vec, v1_f); + sum3 = vfmaq_f32(sum3, diff, diff); + } - // Get precomputed sum of squares from query blob (pVect2v is FP32 query) - const float y_sum_sq = static_cast(pVect2v)[dimension + sq8::SUM_SQUARES_QUERY]; + // Combine all four sum accumulators + float32x4_t sum_combined = vaddq_f32(vaddq_f32(sum0, sum1), vaddq_f32(sum2, sum3)); - // L2² = ||x||² + ||y||² - 2*IP(x, y) - return x_sum_sq + y_sum_sq - 2.0f * ip; + // Horizontal sum to get Σ(diff_i²) + float32x2_t sum_halves = vadd_f32(vget_low_f32(sum_combined), vget_high_f32(sum_combined)); + float32x2_t summed = vpadd_f32(sum_halves, sum_halves); + return vget_lane_f32(summed, 0); } diff --git a/src/VecSim/spaces/L2/L2_SSE4_SQ8_FP32.h b/src/VecSim/spaces/L2/L2_SSE4_SQ8_FP32.h index 3a3a4d12d..2cf5cec6c 100644 --- a/src/VecSim/spaces/L2/L2_SSE4_SQ8_FP32.h +++ b/src/VecSim/spaces/L2/L2_SSE4_SQ8_FP32.h @@ -8,39 +8,111 @@ */ #pragma once #include "VecSim/spaces/space_includes.h" -#include "VecSim/spaces/IP/IP_SSE4_SQ8_FP32.h" #include "VecSim/types/sq8.h" using sq8 = vecsim_types::sq8; /* - * Optimized asymmetric SQ8 L2 squared distance using algebraic identity: + * Asymmetric SQ8 L2 squared distance computed via direct residual accumulation: * - * ||x - y||² = Σx_i² - 2*IP(x, y) + Σy_i² - * = x_sum_squares - 2 * IP(x, y) + y_sum_squares + * ||x - y||² = Σ(dequant(x_i) - y_i)² + * where dequant(x_i) = min_val + delta * q_i * - * where: - * - IP(x, y) = min * y_sum + delta * Σ(q_i * y_i) (computed via - * SQ8_FP32_InnerProductSIMD16_SSE4_IMP) - * - x_sum_squares and y_sum_squares are precomputed - * - * This avoids dequantization in the hot loop. + * This avoids the algebraic-identity/cancellation approach, which catastrophically cancels in + * FP32 when x and y share a large common offset relative to their spread. */ +// Helper: compute Σ(diff_i²) for 4 elements, where diff_i = dequant(x_i) - y_i. +// pVect1 = SQ8 storage (quantized values), pVect2 = FP32 query. +// min_val/delta are broadcast scalars from the stored vector's metadata. +static inline void L2StepSQ8_FP32_SSE4(const uint8_t *&pVect1, const float *&pVect2, __m128 &sum, + __m128 min_val, __m128 delta) { + // Load 4 uint8 elements and convert to float + __m128i v1_i = _mm_cvtepu8_epi32(_mm_cvtsi32_si128(load_unaligned(pVect1))); + pVect1 += 4; + __m128 v1_f = _mm_cvtepi32_ps(v1_i); + + // Load 4 float elements from query + __m128 v2 = _mm_loadu_ps(pVect2); + pVect2 += 4; + + // min - y computed once per lane, then fuse the dequantize-and-subtract: diff = delta*q + + // (min - y). SSE has no FMA, so this is mul + add. + __m128 min_minus_y = _mm_sub_ps(min_val, v2); + __m128 diff = _mm_add_ps(_mm_mul_ps(delta, v1_f), min_minus_y); + + sum = _mm_add_ps(sum, _mm_mul_ps(diff, diff)); +} + // pVect1v = SQ8 storage, pVect2v = FP32 query template // 0..15 float SQ8_FP32_L2SqrSIMD16_SSE4(const void *pVect1v, const void *pVect2v, size_t dimension) { - // Get the raw inner product using the common SIMD implementation - const float ip = SQ8_FP32_InnerProductSIMD16_SSE4_IMP(pVect1v, pVect2v, dimension); + const uint8_t *pVect1 = static_cast(pVect1v); // SQ8 storage + const float *pVect2 = static_cast(pVect2v); // FP32 query + const uint8_t *pEnd1 = pVect1 + dimension; + + // Get quantization parameters from stored vector (after quantized data) + const uint8_t *pVect1Base = static_cast(pVect1v); + const auto *params1 = pVect1Base + dimension; + const float min_val_scalar = load_unaligned(params1 + sq8::MIN_VAL * sizeof(float)); + const float delta_scalar = load_unaligned(params1 + sq8::DELTA * sizeof(float)); + const __m128 min_val = _mm_set1_ps(min_val_scalar); + const __m128 delta = _mm_set1_ps(delta_scalar); + + // Initialize sum accumulators. Four accumulators break the dependency chain, letting more + // ops be in flight at once. + __m128 sum0 = _mm_setzero_ps(); + __m128 sum1 = _mm_setzero_ps(); + __m128 sum2 = _mm_setzero_ps(); + __m128 sum3 = _mm_setzero_ps(); + + // Residual elements (1-3), loaded at full width with the lanes past the residual masked + // off. Staging them through stack arrays instead costs a store-to-load stall. In bounds + // because the x86 chooser never reaches this kernel below dim 8. + if constexpr (residual % 4) { + constexpr unsigned char r = residual % 4; + + __m128i v1_i = _mm_cvtepu8_epi32(_mm_cvtsi32_si128(load_unaligned(pVect1))); + __m128 v1_f = _mm_cvtepi32_ps(v1_i); + __m128 v2 = _mm_loadu_ps(pVect2); + + __m128 min_minus_y = _mm_sub_ps(min_val, v2); + __m128 diff = _mm_add_ps(_mm_mul_ps(delta, v1_f), min_minus_y); + + // Lanes >= r hold elements the main loop will process; zero their contribution. + const __m128 lane_mask = _mm_castsi128_ps( + _mm_set_epi32(r > 3 ? -1 : 0, r > 2 ? -1 : 0, r > 1 ? -1 : 0, r > 0 ? -1 : 0)); + diff = _mm_and_ps(diff, lane_mask); + + pVect1 += r; + pVect2 += r; + + sum0 = _mm_mul_ps(diff, diff); + } - // Get precomputed sum of squares from storage blob (pVect1v is SQ8 storage) - const uint8_t *pVect1 = static_cast(pVect1v); - const float x_sum_sq = - load_unaligned(pVect1 + dimension + sq8::SUM_SQUARES * sizeof(float)); + // Handle remaining residual in chunks of 4 (for residual 4-15) + if constexpr (residual >= 4) { + L2StepSQ8_FP32_SSE4(pVect1, pVect2, sum1, min_val, delta); + } + if constexpr (residual >= 8) { + L2StepSQ8_FP32_SSE4(pVect1, pVect2, sum2, min_val, delta); + } + if constexpr (residual >= 12) { + L2StepSQ8_FP32_SSE4(pVect1, pVect2, sum3, min_val, delta); + } - // Get precomputed sum of squares from query blob (pVect2v is FP32 query) - const float y_sum_sq = static_cast(pVect2v)[dimension + sq8::SUM_SQUARES_QUERY]; + // Process remaining full chunks of 16 elements (4x4). The loop may run zero times + // (dim can be as small as 8). + while (pVect1 < pEnd1) { + L2StepSQ8_FP32_SSE4(pVect1, pVect2, sum0, min_val, delta); + L2StepSQ8_FP32_SSE4(pVect1, pVect2, sum1, min_val, delta); + L2StepSQ8_FP32_SSE4(pVect1, pVect2, sum2, min_val, delta); + L2StepSQ8_FP32_SSE4(pVect1, pVect2, sum3, min_val, delta); + } - // L2² = ||x||² + ||y||² - 2*IP(x, y) - return x_sum_sq + y_sum_sq - 2.0f * ip; + // Horizontal sum to get Σ(diff_i²) + __m128 sum = _mm_add_ps(_mm_add_ps(sum0, sum1), _mm_add_ps(sum2, sum3)); + float PORTABLE_ALIGN16 TmpRes[4]; + _mm_store_ps(TmpRes, sum); + return TmpRes[0] + TmpRes[1] + TmpRes[2] + TmpRes[3]; } diff --git a/src/VecSim/spaces/L2/L2_SVE_SQ8_FP32.h b/src/VecSim/spaces/L2/L2_SVE_SQ8_FP32.h index a1a800727..2b891adae 100644 --- a/src/VecSim/spaces/L2/L2_SVE_SQ8_FP32.h +++ b/src/VecSim/spaces/L2/L2_SVE_SQ8_FP32.h @@ -8,7 +8,6 @@ */ #pragma once #include "VecSim/spaces/space_includes.h" -#include "VecSim/spaces/IP/IP_SVE_SQ8_FP32.h" #include "VecSim/types/sq8.h" #include @@ -21,35 +20,108 @@ namespace { using sq8 = vecsim_types::sq8; /* - * Optimized asymmetric SQ8-FP32 L2 squared distance using algebraic identity: + * Asymmetric SQ8-FP32 L2 squared distance via direct residual accumulation: * - * ||x - y||² = Σx_i² - 2*IP(x, y) + Σy_i² - * = x_sum_squares - 2 * IP(x, y) + y_sum_squares + * ||x - y||² = Σ(dequant(x_i) - y_i)², where dequant(x_i) = min_val + delta * q_i * - * where: - * - IP(x, y) = min * y_sum + delta * Σ(q_i * y_i) (computed via - * SQ8_FP32_InnerProductSIMD_SVE_IMP) - * - x_sum_squares and y_sum_squares are precomputed - * - * This avoids dequantization in the hot loop. + * Not the ||x||² + ||y||² - 2*IP identity, which cancels catastrophically in FP32 when x and y + * share a large common offset relative to their spread (MOD-17526). */ +// One SVE vector width of Σ(diff_i²). pVect1 = SQ8 storage, pVect2 = FP32 query. +static inline void L2StepSQ8_FP32_SVE(const uint8_t *pVect1, const float *pVect2, size_t offset, + svfloat32_t &sum, svfloat32_t min_val_vec, + svfloat32_t delta_vec) { + svbool_t pg = svptrue_b32(); + + svuint32_t v1_u32 = svld1ub_u32(pg, pVect1 + offset); + svfloat32_t v1_f = svcvt_f32_u32_x(pg, v1_u32); + svfloat32_t v2 = svld1_f32(pg, pVect2 + offset); + + // Subtract fused into the multiply-add; keeping min - y first is what preserves the residual. + svfloat32_t min_minus_y = svsub_f32_x(pg, min_val_vec, v2); + svfloat32_t diff = svmla_f32_x(pg, min_minus_y, delta_vec, v1_f); + + sum = svmla_f32_x(pg, sum, diff, diff); +} + // pVect1v = SQ8 storage, pVect2v = FP32 query template float SQ8_FP32_L2SqrSIMD_SVE(const void *pVect1v, const void *pVect2v, size_t dimension) { - // Get the raw inner product using the common SIMD implementation - const float ip = SQ8_FP32_InnerProductSIMD_SVE_IMP( - pVect1v, pVect2v, dimension); + const uint8_t *pVect1 = static_cast(pVect1v); // SQ8 storage + const float *pVect2 = static_cast(pVect2v); // FP32 query + size_t offset = 0; + + svbool_t pg = svptrue_b32(); + + // Get the number of 32-bit elements per vector at runtime + uint64_t chunk = svcntw(); + + // Get quantization parameters from stored vector (after quantized data) + const auto *params1 = pVect1 + dimension; + const float min_val_scalar = load_unaligned(params1 + sq8::MIN_VAL * sizeof(float)); + const float delta_scalar = load_unaligned(params1 + sq8::DELTA * sizeof(float)); + const svfloat32_t min_val_vec = svdup_f32(min_val_scalar); + const svfloat32_t delta_vec = svdup_f32(delta_scalar); + + // Multiple accumulators for ILP + svfloat32_t sum0 = svdup_f32(0.0f); + svfloat32_t sum1 = svdup_f32(0.0f); + svfloat32_t sum2 = svdup_f32(0.0f); + svfloat32_t sum3 = svdup_f32(0.0f); + + // Full vectors first, predicated tail last, bounds compared rather than divided: `chunk` is + // runtime (svcntw), so `dimension % chunk` would emit a real udiv, and the chooser already + // divided once to produce `partial_chunk`/`additional_steps`. Deliberately unlike the + // prefix-first sibling SQ8 SVE kernels; see MOD-17526 for the measured trade-off. + // A bitmask cannot replace the modulo here: SVE vector length is a multiple of 128 bits, + // not necessarily a power of two. + const size_t chunk_size = 4 * chunk; + for (; offset + chunk_size <= dimension; offset += chunk_size) { + L2StepSQ8_FP32_SVE(pVect1, pVect2, offset, sum0, min_val_vec, delta_vec); + L2StepSQ8_FP32_SVE(pVect1, pVect2, offset + chunk, sum1, min_val_vec, delta_vec); + L2StepSQ8_FP32_SVE(pVect1, pVect2, offset + 2 * chunk, sum2, min_val_vec, delta_vec); + L2StepSQ8_FP32_SVE(pVect1, pVect2, offset + 3 * chunk, sum3, min_val_vec, delta_vec); + } + + // Handle remaining full-width steps (0-3), resolved at compile time. + if constexpr (additional_steps > 0) { + L2StepSQ8_FP32_SVE(pVect1, pVect2, offset, sum0, min_val_vec, delta_vec); + offset += chunk; + } + if constexpr (additional_steps > 1) { + L2StepSQ8_FP32_SVE(pVect1, pVect2, offset, sum1, min_val_vec, delta_vec); + offset += chunk; + } + if constexpr (additional_steps > 2) { + L2StepSQ8_FP32_SVE(pVect1, pVect2, offset, sum2, min_val_vec, delta_vec); + offset += chunk; + } + + // Predicated tail for the final partial vector. + if constexpr (partial_chunk) { + svbool_t pg_tail = + svwhilelt_b32(static_cast(offset), static_cast(dimension)); + + svuint32_t v1_u32 = svld1ub_u32(pg_tail, pVect1 + offset); + svfloat32_t v1_f = svcvt_f32_u32_z(pg_tail, v1_u32); + svfloat32_t v2 = svld1_f32(pg_tail, pVect2 + offset); + + svfloat32_t min_minus_y = svsub_f32_z(pg_tail, min_val_vec, v2); + svfloat32_t diff = svmla_f32_z(pg_tail, min_minus_y, delta_vec, v1_f); - // Get precomputed sum of squares from storage blob (pVect1v is SQ8 storage) - const uint8_t *pVect1 = static_cast(pVect1v); - const float x_sum_sq = - load_unaligned(pVect1 + dimension + sq8::SUM_SQUARES * sizeof(float)); + // Merging (_m), not zeroing: sum3 already holds partial sums from the main loop, and _z + // would zero every lane outside pg_tail. Only the accumulator needs this; the + // temporaries above are fresh values. + sum3 = svmla_f32_m(pg_tail, sum3, diff, diff); + } - // Get precomputed sum of squares from query blob (pVect2v is FP32 query) - const float y_sum_sq = static_cast(pVect2v)[dimension + sq8::SUM_SQUARES_QUERY]; + // Combine the accumulators + svfloat32_t sum = svadd_f32_z(pg, sum0, sum1); + sum = svadd_f32_z(pg, sum, sum2); + sum = svadd_f32_z(pg, sum, sum3); - // L2² = ||x||² + ||y||² - 2*IP(x, y) - return x_sum_sq + y_sum_sq - 2.0f * ip; + // Horizontal sum to get Σ(diff_i²) + return svaddv_f32(pg, sum); } } // namespace diff --git a/tests/unit/test_spaces.cpp b/tests/unit/test_spaces.cpp index fe0138246..66d82fe17 100644 --- a/tests/unit/test_spaces.cpp +++ b/tests/unit/test_spaces.cpp @@ -14,6 +14,7 @@ #include #include #include +#include #include "gtest/gtest.h" #include "VecSim/spaces/space_includes.h" @@ -409,6 +410,238 @@ TEST_F(SpacesTest, SQ8_FP32_odd_dim_unaligned_metadata_test) { } } +/* ==================== MOD-17526 regression tests ==================== + * The SQ8-FP32 L2 kernels used to compute L2^2 via the ||x||^2 + ||y||^2 - 2*IP identity, which + * cancels catastrophically in FP32 when both vectors share a large common offset. + * + * These check against an independent double-precision reference rather than another kernel + * (pre-fix, scalar and SIMD were wrong together), and use relative rather than absolute + * tolerances, since a ~0.01 absolute bound is what let the bug hide. + */ + +namespace { + +// Ground truth: dequantizes the storage by hand and accumulates in double. Independent of the +// production kernels and of test_utils::SQ8_FP32_NotOptimized_L2Sqr, which accumulates in float. +double SQ8_FP32_L2Sqr_DoubleReference(const uint8_t *storage, const float *query, size_t dim) { + const float min_val = load_unaligned(storage + dim + sq8::MIN_VAL * sizeof(float)); + const float delta = load_unaligned(storage + dim + sq8::DELTA * sizeof(float)); + double res = 0.0; + for (size_t i = 0; i < dim; i++) { + const double dequantized = + static_cast(min_val) + static_cast(delta) * storage[i]; + const double diff = static_cast(query[i]) - dequantized; + res += diff * diff; + } + return res; +} + +// Checks the scalar kernel, the dispatched kernel, and every SIMD variant in this binary against +// `expected`, within `rel_tol` scaled by max(expected, 1.0). +void ExpectSQ8_FP32_L2SqrNear(const uint8_t *storage, const float *query, size_t dim, + double expected, double rel_tol, const std::string &context) { + const double scale = std::max(expected, 1.0); + auto check = [&](double actual, const char *label) { + EXPECT_NEAR(actual, expected, rel_tol * scale) + << context << " [" << label << "]: actual=" << actual << ", expected=" << expected; + }; + + check(SQ8_FP32_L2Sqr(storage, query, dim), "scalar"); + check(L2_SQ8_FP32_GetDistFunc(dim, nullptr)(storage, query, dim), "dispatched"); + + auto optimization = getCpuOptimizationFeatures(); + + // The chooser only reaches the x86 kernels at dim >= 8, and the AVX2 residual path reads a + // full 32-byte vector (my_mm256_maskz_loadu_ps is an unconditional load), which overreads a + // smaller query blob. aarch64 has no such floor, so those tiers stay ungated to match + // production. + const bool x86_simd_reachable = dim >= 8; +#ifdef OPT_AVX512_F_BW_VL_VNNI + if (x86_simd_reachable && optimization.avx512f && optimization.avx512bw && + optimization.avx512vl && optimization.avx512vnni) { + check(Choose_SQ8_FP32_L2_implementation_AVX512F_BW_VL_VNNI(dim)(storage, query, dim), + "AVX512F_BW_VL_VNNI"); + } +#endif +#ifdef OPT_AVX2_FMA + if (x86_simd_reachable && optimization.avx2 && optimization.fma3) { + check(Choose_SQ8_FP32_L2_implementation_AVX2_FMA(dim)(storage, query, dim), "AVX2_FMA"); + } +#endif +#ifdef OPT_AVX2 + if (x86_simd_reachable && optimization.avx2) { + check(Choose_SQ8_FP32_L2_implementation_AVX2(dim)(storage, query, dim), "AVX2"); + } +#endif +#ifdef OPT_SSE4 + if (x86_simd_reachable && optimization.sse4_1) { + check(Choose_SQ8_FP32_L2_implementation_SSE4(dim)(storage, query, dim), "SSE4"); + } +#endif +#ifdef OPT_SVE2 + if (optimization.sve2) { + check(Choose_SQ8_FP32_L2_implementation_SVE2(dim)(storage, query, dim), "SVE2"); + } +#endif +#ifdef OPT_SVE + if (optimization.sve) { + check(Choose_SQ8_FP32_L2_implementation_SVE(dim)(storage, query, dim), "SVE"); + } +#endif +#ifdef OPT_NEON + if (optimization.asimd) { + check(Choose_SQ8_FP32_L2_implementation_NEON(dim)(storage, query, dim), "NEON"); + } +#endif +} + +// Builds the blobs from explicit float vectors, kept separate from generation so a caller can +// derive a shifted pair by adding a constant to the same values. Re-sampling +// std::uniform_real_distribution over a shifted [min,max] with the same seed does NOT give "the +// same vector plus a constant", which is what a translation-invariance test needs. +struct SQ8_FP32_L2_TestVectors { + std::vector query; // [float values (dim)] [sum] [sum_squares] + std::vector storage; // [uint8_t values (dim)] [min] [delta] [sum] [sum_squares] +}; + +SQ8_FP32_L2_TestVectors +BuildSQ8_FP32_L2_TestVectorsFromValues(const std::vector &query_values, + const std::vector &storage_values) { + const size_t dim = query_values.size(); + SQ8_FP32_L2_TestVectors v; + v.query.resize(dim + sq8::query_metadata_count()); + std::copy(query_values.begin(), query_values.end(), v.query.begin()); + test_utils::preprocess_sq8_fp32_query(v.query.data(), dim); + + v.storage.resize(dim * sizeof(uint8_t) + + sq8::storage_metadata_count() * sizeof(float)); + test_utils::quantize_float_vec_to_sq8_with_metadata(storage_values.data(), dim, + v.storage.data()); + return v; +} + +// One base (query, storage) pair, plus the same pair with `offset` added to every element. +std::pair +BuildSQ8_FP32_L2_ShiftedPair(size_t dim, float spread, float offset, int seed) { + std::vector query_values(dim), storage_values(dim); + test_utils::populate_float_vec(query_values.data(), dim, seed, -spread, spread); + test_utils::populate_float_vec(storage_values.data(), dim, seed + 1, -spread, spread); + + std::vector shifted_query(dim), shifted_storage(dim); + for (size_t i = 0; i < dim; i++) { + shifted_query[i] = query_values[i] + offset; + shifted_storage[i] = storage_values[i] + offset; + } + + return {BuildSQ8_FP32_L2_TestVectorsFromValues(query_values, storage_values), + BuildSQ8_FP32_L2_TestVectorsFromValues(shifted_query, shifted_storage)}; +} + +} // namespace + +TEST_F(SpacesTest, SQ8_FP32_L2Sqr_MOD17526_TranslationInvariance) { + // L2 must satisfy L2(x, y) == L2(x + C, y + C) for any constant offset C. The old + // identity-based kernel violated this for large C relative to the vectors' spread; the + // direct-residual kernel must satisfy it (up to FP32 rounding, hence the relative tolerance). + const float spread = 1.0f; + const float large_offset = 100000.0f; // offset/spread == 100000, far past the ~4000 threshold + // where the old identity started to catastrophically + // cancel in FP32. + // Mixes exact multiples of the SIMD chunk widths with awkward remainders, so the + // large-offset case also runs through the residual/masked-lane tails. + for (const size_t dim : {128UL, 129UL, 145UL, 512UL, 513UL, 527UL, 768UL}) { + auto [no_offset, with_offset] = + BuildSQ8_FP32_L2_ShiftedPair(dim, spread, large_offset, 1234); + + const double expected_no_offset = + SQ8_FP32_L2Sqr_DoubleReference(no_offset.storage.data(), no_offset.query.data(), dim); + const double expected_with_offset = SQ8_FP32_L2Sqr_DoubleReference( + with_offset.storage.data(), with_offset.query.data(), dim); + + // The reference itself must be translation-invariant. Not exact equality: float32 + // spacing near 1e5 is coarser than near 0, which can shift a value to an adjacent + // quantization code. + ASSERT_NEAR(expected_no_offset, expected_with_offset, + 1e-2 * std::max(expected_no_offset, 1.0)) + << "dim " << dim << ": independent reference is not translation-invariant"; + + ExpectSQ8_FP32_L2SqrNear(no_offset.storage.data(), no_offset.query.data(), dim, + expected_no_offset, 1e-3, "dim " + std::to_string(dim) + ", C=0"); + + // `min - y` is exact in FP32 by Sterbenz (both large, within a factor of 2), so no + // large-scale cancellation remains and the tolerance stays as tight as the C=0 case. + ExpectSQ8_FP32_L2SqrNear( + with_offset.storage.data(), with_offset.query.data(), dim, expected_with_offset, 1e-3, + "dim " + std::to_string(dim) + ", C=" + std::to_string(large_offset)); + } +} + +TEST_F(SpacesTest, SQ8_FP32_L2Sqr_MOD17526_TicketReproShape) { + // The ticket's literal repro: x=[100000,100008], y=[100000,100000]. True L2^2 = (100000 - + // 100000)^2 + (100008 - 100000)^2 = 64. The old ||x||^2+||y||^2-2*IP identity computed this + // as a huge value minus a nearly-equal huge value in FP32 and could return 0, a negative + // number, or other unbounded garbage instead of 64. + const size_t dim = 2; + const float storage_values[dim] = {100000.0f, 100008.0f}; + const float query_values[dim] = {100000.0f, 100000.0f}; + + std::vector query(dim + sq8::query_metadata_count()); + std::copy(query_values, query_values + dim, query.begin()); + test_utils::preprocess_sq8_fp32_query(query.data(), dim); + + std::vector storage(dim * sizeof(uint8_t) + + sq8::storage_metadata_count() * sizeof(float)); + test_utils::quantize_float_vec_to_sq8_with_metadata(storage_values, dim, storage.data()); + + ExpectSQ8_FP32_L2SqrNear(storage.data(), query.data(), dim, /*expected=*/64.0, 1e-4, + "MOD-17526 ticket repro shape"); +} + +TEST_F(SpacesTest, SQ8_FP32_L2Sqr_MOD17526_ScalarAssociativityRegression) { + // Pins a scalar-only ordering bug: `(min_val + delta*q) - y` rounds the dequantized value at + // the large offset's precision before subtracting, discarding the residual, while + // `delta*q + (min_val - y)` keeps it. With min=100000, delta=8/255, q=1 the dequantized value + // is 100000.031372549..., which FP32 rounds to exactly y=100000.03125, so the buggy order + // yields diff=0 instead of ~0.0001225. + const size_t dim = 8; + // storage: min=100000 (index 0), max=100008 (index 7) => delta = 8/255, so index 1 + // quantizes to q=1 (dequantizes to 100000 + 8/255 = 100000.031372549...). + const float storage_values[dim] = { + 100000.0f, 100000.0f + 8.0f / 255.0f, 100000.0f, 100000.0f, 100000.0f, 100000.0f, 100000.0f, + 100008.0f}; + // query matches storage exactly everywhere except index 1, so every other element + // contributes exactly 0 and the whole result isolates the associativity bug at index 1. + const float query_values[dim] = {100000.0f, 100000.03125f, 100000.0f, 100000.0f, + 100000.0f, 100000.0f, 100000.0f, 100008.0f}; + + std::vector query(dim + sq8::query_metadata_count()); + std::copy(query_values, query_values + dim, query.begin()); + test_utils::preprocess_sq8_fp32_query(query.data(), dim); + + std::vector storage(dim * sizeof(uint8_t) + + sq8::storage_metadata_count() * sizeof(float)); + test_utils::quantize_float_vec_to_sq8_with_metadata(storage_values, dim, storage.data()); + + const double expected = SQ8_FP32_L2Sqr_DoubleReference(storage.data(), query.data(), dim); + ASSERT_GT(expected, 0.0) << "test construction should produce a nonzero true distance"; + // Absolute, not relative: the true value is ~1.5e-8 and the bug returns exactly 0. + EXPECT_NEAR(SQ8_FP32_L2Sqr(storage.data(), query.data(), dim), expected, 1e-9) + << "scalar kernel: expected=" << expected; +} + +TEST_F(SpacesTest, SQ8_FP32_L2Sqr_MOD17526_RegularVectorSanity) { + // Sanity check that the fix didn't regress the common case: ordinary small-magnitude + // vectors, no large shared offset, realistic embedding-sized dims. + for (const size_t dim : {128UL, 768UL}) { + auto [v, _unused] = + BuildSQ8_FP32_L2_ShiftedPair(dim, /*spread=*/1.0f, /*offset=*/0.0f, 4242); + const double expected = + SQ8_FP32_L2Sqr_DoubleReference(v.storage.data(), v.query.data(), dim); + ExpectSQ8_FP32_L2SqrNear(v.storage.data(), v.query.data(), dim, expected, 1e-3, + "dim " + std::to_string(dim) + ", regular vectors"); + } +} + /* ======================== Tests SQ8-FP16 ========================= */ TEST_F(SpacesTest, SQ8_FP16_ip_no_optimization_norm_func_test) {